lambdarec for singly-recursive lambda's

6 files changed, 57 insertions, 10 deletions

diff --git a/AST.hs b/AST.hs
index f8b5e5b..a096227 100644
--- a/AST.hs
+++ b/AST.hs
@@ -15,6 +15,7 @@ data Value
     | VQuoted Value
     | VDefine Name Value
     | VLambda [Name] Value
+    | VLambdaRec Name [Name] Value
     | VLet [(Name, Value)] Value
     | VBuiltin String
     | VEllipsis
@@ -32,6 +33,7 @@ instance Show Value where
     show (VQuoted e) = '\'' : show e
     show (VDefine n v) = "(define " ++ n ++ " " ++ show v ++ ")"
     show (VLambda as v) = "(lambda (" ++ intercalate " " as ++ ") " ++ show v ++ ")"
+    show (VLambdaRec rn as v) = "(lambdarec " ++ rn ++ " (" ++ intercalate " " as ++ ") " ++ show v ++ ")"
     show (VLet ps v) = "(let (" ++ intercalate " " ["(" ++ n ++ " " ++ show w ++ ")" | (n, w) <- ps] ++ ") " ++ show v ++ ")"
     show (VBuiltin str) = "[[builtin " ++ str ++ "]]"
     show VEllipsis = "..."
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
diff --git a/Intermediate.hs b/Intermediate.hs
index c72e81c..0e181a4 100644
--- a/Intermediate.hs
+++ b/Intermediate.hs
@@ -29,7 +29,7 @@ data Ref
 
 data InsCode
     = IAssign Ref
-    | IParam Int
+    | IParam Int  -- first param is self-recurse link
     | IClosure Int
     | IData Int
     | ICallC Ref [Ref]
diff --git a/VM.hs b/VM.hs
index 1250c1e..7762b75 100644
--- a/VM.hs
+++ b/VM.hs
@@ -81,7 +81,8 @@ vmRunInstr info@(Info bbmap gfds datas) state@(State { sTempMap = tmap, sArgs =
                         (rv, _) <- vmRunBB info (state { sArgs = map (findRef tmap) as, sCloVals = clvals }) bb
                         return (assignRef state dest rv)
                 Nothing -> do
-                    (rv, state') <- vmRunBuiltin state clname (map (findRef tmap) as)
+                    -- Take 'tail as' to skip the first self-link argument
+                    (rv, state') <- vmRunBuiltin state clname (map (findRef tmap) (tail as))
                     return (assignRef state' dest rv)
             obj -> error $ "VM: Cannot call non-closure object: " ++ show obj
     IAllocClo name clrefs -> return (assignRef state dest (RVClosure name (map (findRef tmap) clrefs)))
diff --git a/tests/lambdarec.lisp b/tests/lambdarec.lisp
index 69b830c..a6caa55 100644
--- a/tests/lambdarec.lisp
+++ b/tests/lambdarec.lisp
@@ -1,6 +1,10 @@
+#include "stdlib.lisp"
+
 (define fibo (n)
     (let ((helper (lambdarec rec (m a b)
-                        (if (= m n) b (rec m b (+ a b))))))
+                        (if (= m n) b (rec (+ m 1) b (+ a b))))))
         (if (<= n 0) 0
          (if (<= n 2) 1
           (helper 2 1 1)))))
+
+(for 1 25 (lambda (i) (print (fibo i))))
diff --git a/tests/lambdarec.out b/tests/lambdarec.out
new file mode 100644
index 0000000..14950e3
--- /dev/null
+++ b/tests/lambdarec.out
@@ -0,0 +1,25 @@
+1
+1
+2
+3
+5
+8
+13
+21
+34
+55
+89
+144
+233
+377
+610
+987
+1597
+2584
+4181
+6765
+10946
+17711
+28657
+46368
+75025