1 files changed, 50 insertions, 1 deletions
diff --git a/app/Code.hs b/app/Code.hs
index df3401e..b1c474b 100644
--- a/app/Code.hs
+++ b/app/Code.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE TupleSections #-}
+
 module Code where
 
 import Control.Monad.Trans.State.Lazy
@@ -32,7 +34,7 @@ data Heap =
   Heap Int (M.Map Int Datum)
   deriving (Show)
 
-emptyHeap = Heap 0 M.empty
+emptyHeap = Heap 1 M.empty
 
 type Scope = M.Map Int (Int, Int)
 
@@ -70,3 +72,50 @@ data Builtin =
 
 instance Show Builtin where
   show _ = "Builtin _"
+
+codeStruct ::
+     Monad m
+  => (Datum -> m a)
+  -> (Datum -> [a] -> m a)
+  -> (Datum -> m (Either Int a))
+  -> (Datum -> Int -> m a)
+  -> m a
+  -> Heap
+  -> Code
+  -> m (Code, a)
+codeStruct atom struct local rec end heap = go
+  where
+    go [] = ([], ) <$> end
+    go (U lr@(LocalRef _ _):cs) = do
+      x <- local lr
+      case x of
+        Left ref -> (cs, ) <$> heapStruct atom struct rec heap ref
+        Right a -> pure (cs, a)
+    go (U s@(Struct (IR.Id _ arity)):cs) = eat arity cs >>= traverse (struct s)
+    go (U x:cs) = (cs, ) <$> atom x
+    go cs = (cs, ) <$> end
+    eat n cs
+      | n <= 0 = pure (cs, [])
+      | otherwise = do
+        (rest, a) <- go cs
+        fmap (a :) <$> eat (n - 1) rest
+
+heapStruct ::
+     Monad m
+  => (Datum -> m a)
+  -> (Datum -> [a] -> m a)
+  -> (Datum -> Int -> m a)
+  -> Heap
+  -> Int
+  -> m a
+heapStruct atom struct rec (Heap _ heap) hr = go [hr] hr
+  where
+    go visited ref
+      | rr@(HeapRef r) <- heap M.! ref =
+        if r == ref || r `elem` visited
+          then rec rr ref
+          else go (r : visited) r
+      | s@(Struct (IR.Id _ arity)) <- heap M.! ref =
+        sequence [go (ref + i : visited) (ref + i) | i <- [1 .. arity]] >>=
+        struct s
+      | x <- heap M.! ref = atom x