Refactor analyseValue, fix Let

2 files changed, 50 insertions, 22 deletions

diff --git a/Compiler.hs b/Compiler.hs
index 913d568..0272950 100644
--- a/Compiler.hs
+++ b/Compiler.hs
@@ -16,50 +16,71 @@ data TaggedValue
     = TVList [TaggedValue]
     | TVNum Int
     | TVString String
-    | TVName Name (Maybe Int)  -- Nothing: unknown, Just n: defined n lambdas up (0 = current lambda arg)
+    | 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)
+    | TVLet (Name, TaggedValue) TaggedValue
     | TVEllipsis
   deriving Show
 
 -- also does some preprocessing, like parsing lambda's and defines
 analyseValue :: Value -> TaggedValue
-analyseValue = go []
+analyseValue = go (Map.empty, 0)
   where
-    go :: [Set.Set Name] -> Value -> TaggedValue
-    go scopes (VList [VName "define", VName name, VList args, body])
-        | Just names <- mapM fromVName args = go scopes (VList [VName "define", VName name, VLambda names body])
+    -- env: (name -> lambda depth at which name is defined, current lambda depth)
+    go :: (Map.Map Name Int, Int) -> Value -> TaggedValue
+    go env (VList [VName "define", VName name, VList args, body])
+        | Just names <- mapM fromVName args = go env (VList [VName "define", VName name, VLambda names body])
         | otherwise = error "Invalid 'define' shorthand syntax: Invalid argument list"
-    go scopes (VList [VName "define", VName name, value]) = TVDefine name (go scopes value)
-    go scopes (VList [VName "lambda", VList args, body])
-        | Just names <- mapM fromVName args = go scopes (VLambda names body)
+    go env (VList [VName "define", VName name, value]) = TVDefine name (go env value)
+    go _ (VList (VName "define" : _)) = error "Invalid 'define' syntax"
+
+    go (envd, depth) (VList [VName "lambda", VList args, body])
+        | Just names <- mapM fromVName args = go (envd, depth + 1) (VLambda names body)
         | otherwise = error "Invalid 'lambda' syntax: Invalid argument list"
     go _ (VList (VName "lambda" : _)) = error "Invalid 'lambda' syntax"
-    go scopes (VList [VName "let", VList args, body]) =
-        case sequence [if length p == 2 then Just (p !! 0, p !! 1) else Nothing | VList p <- args] of
+
+    -- go env (VList [VName "let", VList args, body]) =
+    --     case sequence [if length p == 2 then Just (p !! 0, p !! 1) else Nothing | VList p <- args] of
+    --         Just pairs ->
+    --             go env (VList (VList [VName "lambda", VList (map fst pairs), body] : map snd pairs))
+    --         Nothing ->
+    --             error "Invalid 'let' syntax: Invalid variable list (not all pairs)"
+    go env (VList [VName "let", VList args, body]) =
+        case sequence [case p of { [VName n, v] -> Just (n, v); _ -> Nothing } | VList p <- args] of
             Just pairs ->
-                go scopes (VList (VList [VName "lambda", VList (map fst pairs), body] : map snd pairs))
+                go env (VLet pairs body)
             Nothing ->
                 error "Invalid 'let' syntax: Invalid variable list (not all pairs)"
     go _ (VList (VName "let" : _)) = error "Invalid 'let' syntax: Invalid argument list"
-    go scopes (VList values) = TVList (map (go scopes) values)
+
+    go env (VList values) = TVList (map (go env) values)
     go _ (VNum n) = TVNum n
     go _ (VString s) = TVString s
-    go scopes (VName name) = TVName name (findIndex id (map (Set.member name) scopes))
+    go (envd, _) (VName name) = TVName name (Map.lookup name envd)
     go _ (VQuoted value) = TVQuoted value
-    go scopes (VLambda args body) =
-        let t = go (Set.fromList args : scopes) body
-        in TVLambda args t (Set.toList (collectEscapes 0 t))
+    go (envd, depth) (VLambda args body) =
+        let t = go (foldr (flip Map.insert depth) envd args, depth) body
+        in TVLambda args t (Set.toList (collectEscapes depth t))
+    go (envd, depth) (VLet ((name, value) : args) body) =
+        TVLet (name, go (envd, depth) value)
+              (go (Map.insert name depth envd, depth) (VLet args body))
+    go env (VLet [] body) = go env body
     go _ (VBuiltin _) = undefined
     go _ VEllipsis = TVEllipsis
-    go _ (VLet _ _) = error "'let' should not be present yet at this point"
 
     collectEscapes :: Int -> TaggedValue -> Set.Set Name
-    collectEscapes limit (TVList values) = Set.unions (map (collectEscapes limit) values)
-    collectEscapes limit (TVName name (Just n)) | n > limit = Set.singleton name
-    collectEscapes limit (TVLambda _ body _) = collectEscapes (limit + 1) body
-    collectEscapes _ _ = Set.empty
+    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 (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
+    collectEscapes _ (TVString _) = Set.empty
+    collectEscapes _ (TVQuoted _) = Set.empty
+    collectEscapes _ TVEllipsis = Set.empty
 
 
 data CompState = CompState
@@ -73,7 +94,7 @@ data CompState = CompState
     , csDatas :: [Value] }
   deriving Show
 
-data ScopeItem = SIParam Int | SIClosure Int | SIGlobal
+data ScopeItem = SIParam Int | SIClosure Int | SIGlobal | SIRef Ref
   deriving Show
 
 newtype CM a = CM {unCM :: StateT CompState (Except String) a}
@@ -272,10 +293,16 @@ genTValue (TVLambda args body closure) nextnext = do
             return r
     setTerm $ IJmp nextnext
     return resref
+genTValue (TVLet (name, value) body) nextnext = do
+    b <- newBlock
+    r <- genTValue value b
+    switchBlock b
+    withScope (Map.singleton name (SIRef r)) (genTValue body nextnext)
 genTValue (TVName name _) nextnext = do
     r <- genTemp
     lookupVar name >>= \si -> case si of
         Right ref -> addIns (r, IAssign ref)
+        Left (SIRef ref) -> addIns (r, IAssign ref)
         Left (SIParam n) -> addIns (r, IParam n)
         Left (SIClosure n) -> addIns (r, IClosure n)
         Left SIGlobal -> do
diff --git a/VM.hs b/VM.hs
index a272829..1250c1e 100644
--- a/VM.hs
+++ b/VM.hs
@@ -111,6 +111,7 @@ vmRunBuiltin state "print" as = putStrLn (intercalate ", " (map printshow as)) >
 vmRunBuiltin state "=" [a, b] = return (if equalOp a b then RVNum 1 else RVNum 0, state)
 vmRunBuiltin state "<=" [RVNum a, RVNum b] = return (RVNum (fromEnum (a <= b)), state)
 vmRunBuiltin state "+" [RVNum a, RVNum b] = return (RVNum (a + b), state)
+vmRunBuiltin state "+" [RVString a, RVString b] = return (RVString (a ++ b), state)
 vmRunBuiltin state "-" [RVNum a, RVNum b] = return (RVNum (a - b), state)
 vmRunBuiltin state "null?" [v] = return (RVNum (case v of { RVList [] -> 1; _ -> 0 }), state)
 vmRunBuiltin state "car" [RVList l] = case l of