From 5d186de9c8483b4de749459fda9c507c68f8fa73 Mon Sep 17 00:00:00 2001
From: Mirek Kratochvil <exa.exa@gmail.com>
Date: Wed, 9 Nov 2022 19:37:43 +0100
Subject: [PATCH] hopefully vars

---
 app/Interpreter.hs | 65 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/app/Interpreter.hs b/app/Interpreter.hs
index ffc812c..6a80232 100644
--- a/app/Interpreter.hs
+++ b/app/Interpreter.hs
@@ -137,13 +137,22 @@ proveStep c f i = go i
             Just x' -> BoundRef x x'
             _ -> NoRef
         writeHeap addr x (Heap nxt m) = Heap nxt (M.adjust (const x) addr m)
-        newHeapVar (Heap nxt m) =
-          (nxt, Heap (nxt + 1) (M.insert nxt (HeapRef nxt) m))
-        allocLocal scope reg cont
+        newHeapVar h = head <$> newHeapVars 1 h
+        newHeapVars n (Heap nxt m) =
+          let addrs = [nxt + i - 1 | i <- [1 .. n]]
+           in ( Heap (nxt + n) $
+                foldr (uncurry M.insert) m $ zip addrs (map HeapRef addrs)
+              , addrs)
+        allocLocal reg scope cont
           | Just addr <- scope M.!? reg = cont scope heap addr
-          | (addr, heap') <- newHeapVar heap =
+          | (heap', addr) <- newHeapVar heap =
             cont (M.insert reg addr scope) heap' addr
-        newHeapStruct addr s@(Struct Id {arity = arity}) cont = undefined
+        newHeapStruct addr s@(Struct Id {arity = arity}) cont =
+          let (Heap nxt' m', addrs) = newHeapVars (arity + 1) heap
+              m'' =
+                M.insert addr (HeapRef $ head addrs) . M.insert (head addrs) s $
+                m'
+           in cont (map HeapRef $ tail addrs) (Heap nxt' m'')
         {- simple cases first -}
         unify VoidVar VoidVar = uok
         unify (Atom a) (Atom b)
@@ -161,9 +170,35 @@ proveStep c f i = go i
         unify (LocalRef _) VoidVar = uok
         unify VoidVar (LocalRef _) = uok
         {- allocate heap for LocalRefs and retry with HeapRefs -}
-        unify (LocalRef hv) (LocalRef gv) = undefined -- avoid allocating 2 things
-        unify (LocalRef hv) _ = undefined
-        unify _ (LocalRef gv) = undefined
+        unify (LocalRef hv) (LocalRef gv) =
+          allocLocal gv (gvar cur) $ \gvar' heap' addr ->
+            c
+              i
+                { cur =
+                    cur
+                      { hed = U (HeapRef addr) : hs
+                      , hvar = M.insert hv addr (hvar cur)
+                      , gol = U (HeapRef addr) : gs
+                      , gvar = gvar'
+                      , heap = heap'
+                      }
+                }
+        unify (LocalRef hv) _ =
+          allocLocal hv (hvar cur) $ \hvar' heap' addr ->
+            c
+              i
+                { cur =
+                    cur
+                      {hed = U (HeapRef addr) : hs, hvar = hvar', heap = heap'}
+                }
+        unify _ (LocalRef gv) =
+          allocLocal gv (gvar cur) $ \gvar' heap' addr ->
+            c
+              i
+                { cur =
+                    cur
+                      {gol = U (HeapRef addr) : gs, gvar = gvar', heap = heap'}
+                }
         {- handle heap refs; first ignore their combination with voids again -}
         unify (HeapRef _) VoidVar = uok
         unify VoidVar (HeapRef _) = uok
@@ -178,7 +213,12 @@ proveStep c f i = go i
                     hr
                     s
                     (\nhs nheap ->
-                       c i {cur = cur {hed = nhs ++ hs, gol = gs, heap = nheap}})
+                       c
+                         i
+                           { cur =
+                               cur
+                                 {hed = map U nhs ++ hs, gol = gs, heap = nheap}
+                           })
                 HeapRef gr' ->
                   case deref gr' of
                     FreeRef gr -> setHeap hr (HeapRef gr)
@@ -207,7 +247,12 @@ proveStep c f i = go i
                     gr
                     s
                     (\ngs nheap ->
-                       c i {cur = cur {hed = hs, gol = ngs ++ gs, heap = nheap}})
+                       c
+                         i
+                           { cur =
+                               cur
+                                 {hed = hs, gol = map U ngs ++ gs, heap = nheap}
+                           })
             BoundRef _ atom@(Atom b) -> unify h atom
             BoundRef addr struct@(Struct Id {arity = arity}) ->
               c