Split off macro expansion from analyseValue

2 files changed, 60 insertions, 42 deletions

diff --git a/AST.hs b/AST.hs
index df946dd..f8b5e5b 100644
--- a/AST.hs
+++ b/AST.hs
@@ -13,6 +13,7 @@ data Value
     | VString String
     | VName Name
     | VQuoted Value
+    | VDefine Name Value
     | VLambda [Name] Value
     | VLet [(Name, Value)] Value
     | VBuiltin String
@@ -29,6 +30,7 @@ instance Show Value where
     show (VString s) = show s
     show (VName n) = n
     show (VQuoted e) = '\'' : show e
+    show (VDefine n v) = "(define " ++ n ++ " " ++ show v ++ ")"
     show (VLambda as v) = "(lambda (" ++ intercalate " " as ++ ") " ++ show v ++ ")"
     show (VLet ps v) = "(let (" ++ intercalate " " ["(" ++ n ++ " " ++ show w ++ ")" | (n, w) <- ps] ++ ") " ++ show v ++ ")"
     show (VBuiltin str) = "[[builtin " ++ str ++ "]]"
diff --git a/Compiler.hs b/Compiler.hs
index 0272950..b221694 100644
--- a/Compiler.hs
+++ b/Compiler.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving, LambdaCase #-}
 module Compiler(IRProgram, compileProgram) where
 
 import Control.Monad.Except
@@ -24,51 +24,65 @@ data TaggedValue
     | TVEllipsis
   deriving Show
 
--- also does some preprocessing, like parsing lambda's and defines
+preprocess :: Value -> Value
+preprocess (VList [VName "define", VName name, VList args, body])
+    | Just names <- mapM fromVName args = preprocess (VList [VName "define", VName name, VLambda names body])
+    | otherwise = error "Invalid 'define' shorthand syntax: Invalid argument list"
+preprocess (VList [VName "define", VName name, value]) = VDefine name (preprocess value)
+preprocess (VList (VName "define" : _)) = error "Invalid 'define' syntax"
+
+preprocess (VList [VName "lambda", VList args, body])
+    | Just names <- mapM fromVName args = preprocess (VLambda names body)
+    | otherwise = error "Invalid 'lambda' syntax: Invalid argument list"
+preprocess (VList (VName "lambda" : _)) = error "Invalid 'lambda' 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 ->
+--             preprocess (VList (VList [VName "lambda", VList (map fst pairs), body] : map snd pairs))
+--         Nothing ->
+--             error "Invalid 'let' syntax: Invalid variable list (not all pairs)"
+preprocess (VList [VName "let", VList args, body]) =
+    case sequence [case p of { [VName n, v] -> Just (n, v); _ -> Nothing } | VList p <- args] of
+        Just pairs ->
+            preprocess (VLet pairs body)
+        Nothing ->
+            error "Invalid 'let' syntax: Invalid variable list (not all pairs)"
+preprocess (VList (VName "let" : _)) = error "Invalid 'let' syntax: Invalid argument list"
+
+preprocess (VList values) = VList (map preprocess values)
+
+preprocess (VDefine name body) = VDefine name (preprocess body)
+preprocess (VLambda args body) = VLambda args (preprocess body)
+preprocess (VLet args body) = VLet [(name, preprocess value) | (name, value) <- args] (preprocess body)
+preprocess v@(VNum _) = v
+preprocess v@(VString _) = v
+preprocess v@(VName _) = v
+preprocess v@(VQuoted _) = v
+preprocess v@(VBuiltin _) = v
+preprocess v@VEllipsis = v
+
 analyseValue :: Value -> TaggedValue
 analyseValue = go (Map.empty, 0)
   where
     -- 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 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 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 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 env (VList values) = TVList (map (go env) values)
-    go _ (VNum n) = TVNum n
-    go _ (VString s) = TVString s
-    go (envd, _) (VName name) = TVName name (Map.lookup name envd)
-    go _ (VQuoted value) = TVQuoted value
-    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 env@(envd, depth) = \case
+        VList values -> TVList (map (go env) values)
+        VNum n -> TVNum n
+        VString s -> TVString s
+        VName name -> TVName name (Map.lookup name envd)
+        VQuoted value -> TVQuoted value
+        VDefine name value -> TVDefine name (go env value)
+        VLambda args body ->
+            let t = go (foldr (flip Map.insert depth) envd args, depth + 1) body
+            in TVLambda 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))
+        VLet [] body -> go env body
+        VBuiltin _ -> undefined
+        VEllipsis -> TVEllipsis
 
     collectEscapes :: Int -> TaggedValue -> Set.Set Name
     collectEscapes depth (TVList values) = Set.unions (map (collectEscapes depth) values)
@@ -186,7 +200,9 @@ compileProgram (Program values) = runCM $ do
     bstart <- newBlockSwitch
     forM_ values $ \value -> do
         bnext <- newBlock
-        ref <- genTValue (analyseValue value) bnext
+        let value' = preprocess value
+            tvalue = analyseValue value'
+        ref <- genTValue tvalue bnext
         switchBlock bnext
         addIns (RNone, IDiscard ref)
     setTerm IExit