lambdarec for singly-recursive lambda's

1 files changed, 22 insertions, 7 deletions

diff --git a/Compiler.hs b/Compiler.hs
index 922528c..37f3d0d 100644
--- a/Compiler.hs
+++ b/Compiler.hs
@@ -19,7 +19,7 @@ data TaggedValue
     | TVName Name (Maybe Int)  -- Nothing: unknown (external/builtin), Just n: defined n lambdas up (0 = current lambda arg)
     | TVQuoted Value
     | TVDefine Name TaggedValue
-    | TVLambda [Name] TaggedValue [Name]  -- (args) (body) (closure slot names)
+    | TVLambda (Maybe Name) [Name] TaggedValue [Name]  -- (recname) (args) (body) (closure slot names)
     | TVLet (Name, TaggedValue) TaggedValue
     | TVEllipsis
   deriving Show
@@ -36,6 +36,11 @@ preprocess (VList [VName "lambda", VList args, body])
     | otherwise = error "Invalid 'lambda' syntax: Invalid argument list"
 preprocess (VList (VName "lambda" : _)) = error "Invalid 'lambda' syntax"
 
+preprocess (VList [VName "lambdarec", VName recname, VList args, body])
+    | Just names <- mapM fromVName args = preprocess (VLambdaRec recname names body)
+    | otherwise = error "Invalid 'lambdarec' syntax: Invalid argument list"
+preprocess (VList (VName "lambdarec" : _)) = error "Invalid 'lambdarec' syntax"
+
 -- preprocess (VList [VName "let", VList args, body]) =
 --     case sequence [case p of { [n, v] -> Just (n, v); _ -> Nothing } | VList p <- args] of
 --         Just pairs ->
@@ -54,6 +59,7 @@ preprocess (VList values) = VList (map preprocess values)
 
 preprocess (VDefine name body) = VDefine name (preprocess body)
 preprocess (VLambda args body) = VLambda args (preprocess body)
+preprocess (VLambdaRec recname args body) = VLambdaRec recname args (preprocess body)
 preprocess (VLet args body) = VLet [(name, preprocess value) | (name, value) <- args] (preprocess body)
 preprocess v@(VNum _) = v
 preprocess v@(VString _) = v
@@ -76,8 +82,14 @@ analyseValue = go (Map.empty, 0)
         VDefine name value -> TVDefine name (go env value)
         VLambda args body ->
             let depth' = depth + 1
-                t = go (foldr (flip Map.insert depth') envd args, depth') body
-            in TVLambda args t (Set.toList (collectEscapes depth' t))
+                envd' = foldr (flip Map.insert depth') envd args
+                t = go (envd', depth') body
+            in TVLambda Nothing args t (Set.toList (collectEscapes depth' t))
+        VLambdaRec recname args body ->
+            let depth' = depth + 1
+                envd' = Map.insert recname depth' (foldr (flip Map.insert depth') envd args)
+                t = go (envd', depth') body
+            in TVLambda (Just recname) args t (Set.toList (collectEscapes depth' t))
         VLet ((name, value) : args) body ->
             TVLet (name, go (envd, depth) value)
                   (go (Map.insert name depth envd, depth) (VLet args body))
@@ -89,7 +101,7 @@ analyseValue = go (Map.empty, 0)
     collectEscapes depth (TVList values) = Set.unions (map (collectEscapes depth) values)
     collectEscapes depth (TVName name (Just d)) = if d < depth then Set.singleton name else Set.empty
     collectEscapes depth (TVDefine _ value) = collectEscapes depth value
-    collectEscapes depth (TVLambda _ body _) = collectEscapes depth body
+    collectEscapes depth (TVLambda _ _ body _) = collectEscapes depth body
     collectEscapes depth (TVLet (_, value) body) = collectEscapes depth value <> collectEscapes depth body
     collectEscapes _ (TVName _ Nothing) = Set.empty  -- external/builtin names do not need to be captured
     collectEscapes _ (TVNum _) = Set.empty
@@ -260,7 +272,7 @@ genTValue (TVList (funcexpr : args)) nextnext = do
     funcref <- genTValue funcexpr b
     switchBlock b
     resref <- genTemp
-    addIns (resref, ICallC funcref refs)
+    addIns (resref, ICallC funcref (funcref : refs))
     setTerm $ IJmp nextnext
     return resref
 genTValue (TVNum n) nextnext = do
@@ -285,9 +297,12 @@ genTValue (TVDefine name value) nextnext = do
     -- traceM $ "Defining '" ++ name ++ "', ref " ++ show dref ++ ", with value " ++ show vref
     addIns (dref, IAssign vref)
     return RNone
-genTValue (TVLambda args body closure) nextnext = do
+genTValue (TVLambda mrecname args body closure) nextnext = do
+    let bindpairs = maybe [] (\n -> [(n, SIParam 0)]) mrecname ++
+                    zip args (map SIParam [1..]) ++
+                    zip closure (map SIClosure [0..])
     startb <- rememberBlock $
-        withScope (Map.fromList (zip args (map SIParam [0..]) ++ zip closure (map SIClosure [0..]))) $ do
+        withScope (Map.fromList bindpairs) $ do
             b <- newBlockSwitch
             b2 <- newBlock
             ref <- genTValue body b2