fix minor stuff

1 files changed, 14 insertions, 10 deletions

diff --git a/app/Compiler.hs b/app/Compiler.hs
index efa641c..67ca20f 100644
--- a/app/Compiler.hs
+++ b/app/Compiler.hs
@@ -1,7 +1,6 @@
 module Compiler where
 
 import Data.Char (isUpper)
-import Data.Containers.ListUtils (nubOrd)
 import qualified Data.Map as M
 
 import Code (Code, Datum(..), Heap, Instr(..), heapStruct)
@@ -12,10 +11,10 @@ varname ('_':_) = True
 varname (c:_) = isUpper c
 varname _ = False
 
-varIds :: PrlgInt -> [Int]
-varIds (CallI _ xs) = nubOrd $ concatMap varIds xs
-varIds (VarI idx _) = [idx]
-varIds _ = []
+varOccurs :: PrlgInt -> M.Map Int Int
+varOccurs (CallI _ xs) = M.unionsWith (+) $ map varOccurs xs
+varOccurs (VarI idx _) = M.singleton idx 1
+varOccurs _ = M.empty
 
 variablizePrlg :: Int -> StrTable -> PrlgInt -> PrlgInt
 variablizePrlg void (StrTable _ _ itos) = go
@@ -26,16 +25,22 @@ variablizePrlg void (StrTable _ _ itos) = go
       | varname (itos M.! i) = VarI i i
       | otherwise = LiteralI i
 
-renumVars :: [(Int, Int)] -> PrlgInt -> PrlgInt
+renumVars :: (Int -> Maybe PrlgInt) -> PrlgInt -> PrlgInt
 renumVars rename = go
   where
     go (CallI id ps) = CallI id $ map go ps
     go (VarI idx i)
-      | Just idx' <- lookup idx rename = VarI idx' i
+      | Just new <- rename idx = new
     go x = x
 
 squashVars :: PrlgInt -> PrlgInt
-squashVars x = renumVars (zip (varIds x) [1 ..]) x
+squashVars x =
+  let occurs = M.assocs (varOccurs x)
+      m' =
+        M.fromList $
+        [(idx, VoidI) | (idx, n) <- occurs, n <= 1] ++
+        [(idx, VarI idx' 0) | ((idx, n), idx') <- zip occurs [1 ..], n > 1]
+   in renumVars (m' M.!?) x
 
 compileGoals :: Id -> Int -> PrlgInt -> [Code]
 compileGoals andop cut = go'
@@ -65,8 +70,7 @@ seqGoals ([Cut]:xs) = [Cut] ++ seqGoals xs
 seqGoals (x:xs) = [Goal] ++ x ++ [Call] ++ seqGoals xs
 
 heapStructPrlgInt :: Monad m => m PrlgInt -> Heap -> Int -> m PrlgInt
-heapStructPrlgInt heaperr heap ref =
-  squashVars <$> heapStruct atom struct hrec heap ref
+heapStructPrlgInt heaperr heap ref = heapStruct atom struct hrec heap ref
   where
     atom (Atom s) = pure $ LiteralI s
     atom VoidRef = pure $ VoidI