a bit more flexible VarI processing

3 files changed, 55 insertions, 26 deletions

diff --git a/app/Compiler.hs b/app/Compiler.hs
index 3c98d70..6ea6cbc 100644
--- a/app/Compiler.hs
+++ b/app/Compiler.hs
@@ -2,10 +2,9 @@ module Compiler where
 
 import Data.Char (isUpper)
 import Data.Containers.ListUtils (nubOrd)
-import Data.List (elemIndex)
 import qualified Data.Map as M
 
-import Code (Code, Datum(..), Instr(..))
+import Code (Code, Datum(..), Heap, Instr(..), heapStruct)
 import IR (Id(..), PrlgInt(..), StrTable(..))
 
 varname :: String -> Bool
@@ -13,38 +12,48 @@ varname ('_':_) = True
 varname (c:_) = isUpper c
 varname _ = False
 
-varIds :: StrTable -> PrlgInt -> [Int]
-varIds st (CallI _ xs) = nubOrd $ concatMap (varIds st) xs
-varIds (StrTable _ _ st) (LiteralI x)
-  | Just s <- st M.!? x
-  , varname s = [x]
-  | otherwise = []
-
-variablizePrlg :: Int -> [Int] -> PrlgInt -> PrlgInt
-variablizePrlg void vs (CallI id ps) =
-  CallI id $ map (variablizePrlg void vs) ps
-variablizePrlg void vs (LiteralI i)
-  | i == void = VoidI
-  | Just idx <- elemIndex i vs = VarI idx i
-  | otherwise = LiteralI i
+varIds :: PrlgInt -> [Int]
+varIds (CallI _ xs) = nubOrd $ concatMap varIds xs
+varIds (VarI idx _) = [idx]
+varIds _ = []
+
+variablizePrlg :: Int -> StrTable -> PrlgInt -> PrlgInt
+variablizePrlg void (StrTable _ _ itos) = go
+  where
+    go (CallI id ps) = CallI id $ map go ps
+    go (LiteralI i)
+      | i == void = VoidI
+      | varname (itos M.! i) = VarI i i
+      | otherwise = LiteralI i
+
+renumVars :: [(Int, Int)] -> 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
+    go x = x
+
+squashVars :: PrlgInt -> PrlgInt
+squashVars x = renumVars (zip (varIds x) [1 ..]) x
 
 compileGoals :: Id -> Int -> PrlgInt -> [Code]
-compileGoals andop cut = go
+compileGoals andop cut = go'
   where
+    go' = go . struct2goal
     go p@(CallI x args)
-      | x == andop = concatMap go args
-    go p@(LiteralI x)
+      | x == andop = concatMap go' args
+    go p@(CallI (Id x 0) [])
       | x == cut = [[Cut]]
     go x = [compileGoal x]
 
 compileGoal :: PrlgInt -> Code
-compileGoal (LiteralI x) = [U (Struct $ Id x 0)]
-compileGoal x = compileArg x
+compileGoal = compileArg
 
 compileArg :: PrlgInt -> Code
-compileArg (CallI x args) = U (Struct x) : concatMap compileArg args
-compileArg (LiteralI x) = [U (Atom x)]
-compileArg (VarI x i) = [U (LocalRef x i)]
+compileArg (CallI s args) = U (Struct s) : concatMap compileArg args
+compileArg (LiteralI s) = [U (Atom s)]
+compileArg (VarI x s) = [U (LocalRef x s)]
 compileArg (VoidI) = [U VoidRef]
 
 seqGoals :: [Code] -> Code
@@ -54,3 +63,23 @@ seqGoals [x] = [Goal] ++ x ++ [LastCall]
 seqGoals [x, [Cut]] = [Goal] ++ x ++ [LastCall, Cut]
 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
+  where
+    atom (Atom s) = pure $ LiteralI s
+    atom VoidRef = pure $ VoidI
+    struct (Struct s) args = pure $ CallI s args
+    hrec (HeapRef r) ref
+      | r == ref = pure $ VarI r 0
+      | otherwise = heaperr
+
+-- TODO check if this is used
+goal2struct :: PrlgInt -> PrlgInt
+goal2struct (CallI (Id s 0) []) = LiteralI s
+goal2struct x = x
+
+struct2goal :: PrlgInt -> PrlgInt
+struct2goal (LiteralI s) = CallI (Id s 0) []
+struct2goal x = x
diff --git a/app/Frontend.hs b/app/Frontend.hs
index a5fa71f..39689f2 100644
--- a/app/Frontend.hs
+++ b/app/Frontend.hs
@@ -49,7 +49,7 @@ interpret = (>> return True) . lex
       underscore <- findAtom "_"
       prlgv <-
         withStrTable $ \st ->
-          (st, C.variablizePrlg underscore (C.varIds st prlgi) prlgi)
+          (st, C.squashVars $ C.variablizePrlg underscore st prlgi)
       compile prlgv
     compile prlgv = do
       commaId <- findStruct "," 2
diff --git a/app/IR.hs b/app/IR.hs
index ac929fa..8507a3e 100644
--- a/app/IR.hs
+++ b/app/IR.hs
@@ -18,7 +18,7 @@ data Id =
 data PrlgInt
   = CallI Id [PrlgInt]
   | LiteralI Int
-  | VarI Int Int
+  | VarI Int Int -- VarI localIndex strTableString
   | VoidI
   deriving (Show)