Better job splitting constnt folding

1 files changed, 52 insertions, 72 deletions

diff --git a/src/Micro/Compiler.hs b/src/Micro/Compiler.hs
index 9cfaddf..3899ead 100644
--- a/src/Micro/Compiler.hs
+++ b/src/Micro/Compiler.hs
@@ -22,15 +22,6 @@ data CompState = CompState
   }
   deriving (Show)
 
-data CompResult = Result
-  { -- last resolved type
-    crType :: Maybe A.Type,
-    crExpr :: A.Expr
-  }
-
-typeResult :: Maybe A.Type -> A.Expr -> CompResult
-typeResult t e = Result t e
-
 startState :: CompState
 startState = CompState emptyEnv []
 
@@ -49,11 +40,11 @@ foldlEither fn init xs =
 
 -- | @addError error@ adds @error@ to the state and returns no type to allow
 -- the compilation to continue.
-addError :: Error -> State CompState CompResult
+addError :: Error -> State CompState (Maybe A.Type)
 addError e = do
   st <- get
   put st {stErr = e : stErr st}
-  pure $ typeResult Nothing A.Nop
+  pure $ Nothing
 
 -- | @typecheckCall args params@ resolves @args@ and compares it with @params@,
 -- returning a string describing an error or Nothing in case of type match.
@@ -63,9 +54,7 @@ typecheckCall args params
   | length params == 0 = pure Nothing
   | otherwise = do
       -- resolve all args types
-      -- FIXME: this break constant folding
-      rargs <- traverse compileOne args
-      let targs = map crType rargs
+      targs <- traverse compileOne args
       case sequence targs of
         Just t
           -- after resolution, we could have less "rights"
@@ -78,14 +67,16 @@ typecheckCall args params
           -- there was an error in on argument
           pure Nothing
 
-showMaybet :: Maybe A.Type -> String
-showMaybet Nothing = "()"
-showMaybet (Just t) = show t
+-- | @showMaybeType t@ is a helper to show nicely a @Maybe A.Type@.
+showMaybeType :: Maybe A.Type -> String
+showMaybeType Nothing = "()"
+showMaybeType (Just t) = show t
 
-typecheck :: Maybe A.Type -> Maybe A.Type -> Maybe String
-typecheck expected found
-  | expected == found = Nothing
-  | otherwise = Just $ "type mismatch\n     found: " ++ showMaybet found ++ "\n  expected: " ++ showMaybet expected
+-- | @typecheck expected found pos@ compares expected and found, returning either expected of adding an type error to the state.
+typecheck :: Maybe A.Type -> Maybe A.Type -> SourcePos -> State CompState (Maybe A.Type)
+typecheck expected found pos
+  | expected == found = pure expected
+  | otherwise = addError $ Error TypeError ("type mismatch\n     found: " ++ showMaybeType found ++ "\n  expected: " ++ showMaybeType expected) pos
 
 -- built-in types
 types :: [String]
@@ -111,38 +102,25 @@ verifyFuncType ident params ret pos = do
       then [Error UndefinedType ("undefined return type in \"" ++ ident ++ "\"") pos]
       else []
 
-compileOne :: A.Expr -> State CompState CompResult
+compileOne :: A.Expr -> State CompState (Maybe A.Type)
 compileOne x = do
   case x of
-    (A.Module _ _) -> return $ typeResult Nothing x
-    (A.Num _ _) -> return $ typeResult (Just $ A.Type "u8") x -- TODO: placeholder
-    (A.Bool' _ _) -> return $ typeResult (Just $ A.Type "bool") x
-    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) ->
-      -- TODO: overflow check
-      return $ typeResult (Just $ A.Type "u8") (A.Num (a + b) pos)
-    orig@(A.BinOp op pos a b) -> do
-      ra <- compileOne a
-      let (ta, ea) = (crType ra, crExpr ra)
-      rb <- compileOne b
-      let (tb, eb) = (crType rb, crExpr rb)
+    (A.Module _ _) -> pure Nothing
+    (A.Num _ _) -> pure $ Just $ A.Type "u8" -- TODO: placeholder
+    (A.Bool' _ _) -> pure $ Just $ A.Type "bool"
+    (A.BinOp op pos a b) -> do
+      ta <- compileOne a
+      tb <- compileOne b
       case op of
         A.Assign -> case a of
-          (A.Variable _ _) -> case typecheck ta tb of
-            Nothing -> return $ typeResult ta (A.BinOp A.Assign pos ea eb)
-            Just err -> addError $ Error TypeError err pos
+          (A.Variable _ _) -> typecheck ta tb pos
           _ -> addError $ Error InvalidTarget "invalid assignment target" pos
-        _ -> case typecheck ta tb of
-          Nothing -> do
-            let new = (A.BinOp op pos ea eb)
-            if orig == new
-              then return $ typeResult ta orig
-              else compileOne new
-          Just err -> addError $ Error TypeError err pos
+        _ -> typecheck ta tb pos
     (A.Func ident params ret body priv anon pos) -> do
       -- current env
       st <- get
       -- check for undefined types
-      (ev, errs) <- return $ (stEnv st, (verifyFuncType ident params ret pos) ++ stErr st)
+      (ev, errs) <- return (stEnv st, (verifyFuncType ident params ret pos) ++ stErr st)
       -- updated with the function
       (ev, errs) <-
         return $ case addSymUniq ev (Sym ident ftype priv False pos) of
@@ -155,24 +133,21 @@ compileOne x = do
       -- helper for return
       nev <- return $ addSym nev $ Sym "$fn$" ftype True True pos
       put st {stEnv = nev, stErr = errs}
-      r <- compileAll body
+      _ <- compileAll body
       st <- get
       -- store updated errors and the env with the function
       put st {stEnv = ev}
-      rbody <- case r of
-        Right (_, xs) -> pure $ xs
-        Left _ -> pure $ body -- there was an error, so just keep the old body
-      return $ typeResult (Just ftype) (A.Func ident params ret rbody priv anon pos)
+      pure Nothing
       where
         ftype = A.toFuncType params ret
     (A.Call ident args pos) -> do
       r <- compileOne ident
-      case crType r of
+      case r of
         Just (A.FuncType params rtyp) -> do
-          r <- typecheckCall args params
-          case r of
+          e <- typecheckCall args params
+          case e of
             Just err -> addError $ Error TypeError err pos
-            Nothing -> return $ typeResult rtyp x
+            Nothing -> pure rtyp
         _ -> addError $ Error NonCallable "non callable value in function call" pos
     (A.Var ident typ val priv pos) -> do
       st <- get
@@ -183,10 +158,8 @@ compileOne x = do
             then Error UndefinedType ("undefined type in declaration \"" ++ ident ++ "\"") pos : errs
             else errs
       put st {stEnv = ev, stErr = errs}
-      r <- compileOne val
-      case typecheck (Just typ) (crType r) of
-        Just err -> addError $ Error TypeError err pos
-        Nothing -> return $ typeResult (Just typ) x
+      t <- compileOne val
+      typecheck (Just typ) t pos
     (A.Return value pos) -> do
       st <- get
       case getSyml (stEnv st) "$fn$" of
@@ -194,37 +167,44 @@ compileOne x = do
           case value of
             Nothing ->
               if isNothing rtyp
-                then return (typeResult Nothing x)
-                else addError $ Error TypeError ("invalid return value\n     found: ()\n  expected: " ++ showMaybet rtyp) pos
+                then pure $ Nothing
+                else addError $ Error TypeError ("invalid return value\n     found: ()\n  expected: " ++ showMaybeType rtyp) pos
             Just v -> do
               r <- compileOne v
-              case typecheck rtyp (crType r) of
-                Just err -> addError $ Error TypeError err pos
-                Nothing -> return $ typeResult rtyp (A.Return (Just (crExpr r)) pos)
+              typecheck rtyp r pos
         _ -> addError $ Error UnexpectedReturn "return without function call" pos
     (A.Variable ident pos) -> do
       st <- get
       case getSym (stEnv st) ident of
-        Just Sym {symType = t} -> return $ typeResult (Just t) x
+        Just Sym {symType = t} -> pure $ Just t
         Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos
-    (A.Nop) -> return $ typeResult Nothing x
 
-compileAll :: [A.Expr] -> State CompState (Either [Error] (SymMap, [A.Expr]))
+foldConstant :: A.Expr -> A.Expr
+foldConstant x =
+  case x of
+    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) -> A.Num (a + 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.Return value pos) -> A.Return (fmap foldConstant value) pos
+    _ -> x
+
+compileAll :: [A.Expr] -> State CompState (Either [Error] SymMap)
 compileAll ast = do
-  result <- traverse compileOne ast
+  _ <- traverse compileOne ast
   st <- get
   let (Env sym _) = stEnv st
   let errs = stErr st
   case errs of
-    [] -> pure $ Right $ (sym, (map (\r -> crExpr r) result))
+    [] -> pure $ Right sym
     _ -> return $ Left errs
 
 compileToAst :: [A.Expr] -> Either [Error] [A.Expr]
-compileToAst ast = do
-  (_, expr) <- evalState (compileAll ast) startState
-  pure $ expr
+compileToAst ast = fmap (\_ -> map foldConstant ast) (evalState (compileAll ast) startState)
 
 compile :: [A.Expr] -> Either [Error] String
 compile ast = do
-  (sym, expr) <- evalState (compileAll ast) startState
-  return $ generate version sym expr
+  sym <- evalState (compileAll ast) startState
+  return $ generate version sym $ map foldConstant ast