Folding constants in addition

3 files changed, 37 insertions, 3 deletions

diff --git a/src/Micro/Compiler.hs b/src/Micro/Compiler.hs
index 629ea45..901a1e3 100644
--- a/src/Micro/Compiler.hs
+++ b/src/Micro/Compiler.hs
@@ -1,6 +1,7 @@
 module Micro.Compiler
   ( version,
     compile,
+    compileToAst,
   )
 where
 
@@ -142,6 +143,9 @@ compileOne x = do
       typecheckBinOp x a b pos
     (A.BinOp A.Assign pos _ _) ->
       addError $ Error InvalidTarget "invalid assignment target" pos
+    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) ->
+      -- TODO: overflow check, actual type
+      return $ typeResult (Just $ A.Type "u8") (A.Num (a + b) pos)
     (A.BinOp _ pos a b) ->
       -- TODO: types and invalid operators
       typecheckBinOp x a b pos
@@ -207,15 +211,22 @@ compileOne x = do
         Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos
     (A.Nop) -> return $ typeResult Nothing x
 
-compileAll :: [A.Expr] -> State CompState (Either [Error] String)
+compileAll :: [A.Expr] -> State CompState (Either [Error] (SymMap, [A.Expr]))
 compileAll ast = do
   result <- traverse compileOne ast
   st <- get
   let (Env sym _) = stEnv st
   let errs = stErr st
   case errs of
-    [] -> pure $ Right $ generate version sym (map (\r -> crExpr r) result)
+    [] -> pure $ Right $ (sym, (map (\r -> crExpr r) result))
     _ -> return $ Left errs
 
+compileToAst :: [A.Expr] -> Either [Error] [A.Expr]
+compileToAst ast = do
+  (_, expr) <- evalState (compileAll ast) startState
+  pure $ expr
+
 compile :: [A.Expr] -> Either [Error] String
-compile ast = evalState (compileAll ast) startState
+compile ast = do
+  (sym, expr) <- evalState (compileAll ast) startState
+  return $ generate version sym expr
diff --git a/src/Micro/Parser.hs b/src/Micro/Parser.hs
index 45ba44d..3e2b0a3 100644
--- a/src/Micro/Parser.hs
+++ b/src/Micro/Parser.hs
@@ -18,6 +18,7 @@ binary s f assoc =
   E.Infix
     ( reservedOp s
         >> do
+          -- FIXME: this is the second operand
           pos <- getPosition
           return $ BinOp f pos
     )
diff --git a/test/Language.hs b/test/Language.hs
index 83736b2..604b06b 100644
--- a/test/Language.hs
+++ b/test/Language.hs
@@ -9,6 +9,17 @@ import Micro.Parser (parseFromString)
 import Test.HUnit
 import Text.Parsec.Pos (newPos)
 
+assertCompileAst :: String -> [A.Expr] -> Assertion
+assertCompileAst input expected = do
+  r <- return $ parseFromString input
+  case r of
+    Left e -> assertFailure $ show e
+    Right ast -> do
+      res <- return $ compileToAst ast
+      case res of
+        Left e -> assertFailure $ show e
+        Right ast -> assertEqual "" expected ast
+
 assertAst :: String -> [A.Expr] -> Assertion
 assertAst input expected = do
   r <- return $ parseFromString input
@@ -238,6 +249,16 @@ testCase17 =
           A.BinOp A.Assign (newPos "-" 3 5) (A.Variable "a" $ newPos "-" 3 1) (A.Num 10 $ newPos "-" 3 5)
         ]
 
+testCase18 =
+  TestLabel "fold constant addition" $
+    TestCase $
+      assertCompileAst
+        "module main\n\
+        \1 + 2;"
+        [ A.Module "main" $ newPos "-" 1 1,
+          A.Num 3 $ newPos "-" 2 5
+        ]
+
 -- test errors
 
 testCaseE1 =
@@ -402,6 +423,7 @@ language =
     testCase15,
     testCase16,
     testCase17,
+    testCase18,
     -- errors
     testCaseE1,
     testCaseE2,