fold constant can return errors

1 files changed, 28 insertions, 16 deletions

diff --git a/src/Micro/Compiler.hs b/src/Micro/Compiler.hs
index 40d9c1a..f305893 100644
--- a/src/Micro/Compiler.hs
+++ b/src/Micro/Compiler.hs
@@ -179,24 +179,33 @@ compileOne x = do
         Just Sym {symType = t} -> pure $ Just t
         Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos
 
-foldConstant :: A.Expr -> A.Expr
+foldConstant :: A.Expr -> Either [Error] A.Expr
 foldConstant x =
   case x of
     -- FIXME: overflow, invalid, etc
-    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) -> A.Num (a + b) pos
-    (A.BinOp A.Minus pos (A.Num a _) (A.Num b _)) -> A.Num (a - b) pos
-    (A.BinOp A.Mul pos (A.Num a _) (A.Num b _)) -> A.Num (a * b) pos
-    (A.BinOp A.Div pos (A.Num a _) (A.Num b _)) -> A.Num (a `div` b) pos
+    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a + b) pos
+    (A.BinOp A.Minus pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a - b) pos
+    (A.BinOp A.Mul pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a * b) pos
+    (A.BinOp A.Div pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a `div` b) pos
     (A.BinOp op pos a b) -> do
-      let newOp = A.BinOp op pos (foldConstant a) (foldConstant b)
-      if newOp /= x then foldConstant newOp else newOp
-    (A.Func ident params ret body priv anon pos) ->
-      A.Func ident params ret (map foldConstant body) priv anon pos
-    (A.Call ident args pos) ->
-      A.Call (foldConstant ident) (map foldConstant args) pos
-    (A.Var ident typ val priv pos) -> A.Var ident typ (foldConstant val) priv pos
-    (A.Return value pos) -> A.Return (fmap foldConstant value) pos
-    _ -> x
+      fa <- foldConstant a
+      fb <- foldConstant b
+      let newOp = A.BinOp op pos fa fb
+      if newOp /= x then foldConstant newOp else Right $ newOp
+    (A.Func ident params ret body priv anon pos) -> do
+      fbody <- traverse foldConstant body
+      Right $ A.Func ident params ret fbody priv anon pos
+    (A.Call ident args pos) -> do
+      fid <- foldConstant ident
+      fargs <- traverse foldConstant args
+      Right $ A.Call fid fargs pos
+    (A.Var ident typ val priv pos) -> do
+      fv <- foldConstant val
+      Right $ A.Var ident typ fv priv pos
+    (A.Return value pos) -> do
+      fv <- traverse foldConstant value
+      Right $ A.Return fv pos
+    _ -> Right x
 
 compileAll :: [A.Expr] -> State CompState (Either [Error] SymMap)
 compileAll ast = do
@@ -209,9 +218,12 @@ compileAll ast = do
     _ -> return $ Left errs
 
 compileToAst :: [A.Expr] -> Either [Error] [A.Expr]
-compileToAst ast = fmap (\_ -> map foldConstant ast) (evalState (compileAll ast) startState)
+compileToAst ast = do
+  _ <- evalState (compileAll ast) startState
+  traverse foldConstant ast
 
 compile :: [A.Expr] -> Either [Error] String
 compile ast = do
   sym <- evalState (compileAll ast) startState
-  return $ generate version sym $ map foldConstant ast
+  fast <- traverse foldConstant ast
+  return $ generate version sym fast