Fixed constant folding in function body and return

1 files changed, 16 insertions, 17 deletions

diff --git a/src/Micro/Compiler.hs b/src/Micro/Compiler.hs
index 01456d2..71d9400 100644
--- a/src/Micro/Compiler.hs
+++ b/src/Micro/Compiler.hs
@@ -6,7 +6,7 @@ module Micro.Compiler
 where
 
 import Control.Monad.State
-import Data.Maybe (catMaybes, fromMaybe)
+import Data.Maybe (catMaybes, fromMaybe, isNothing)
 import Micro.Asm.Sdcc (generate)
 import qualified Micro.Ast as A
 import Micro.Env
@@ -87,16 +87,6 @@ typecheck expected found
   | expected == found = Nothing
   | otherwise = Just $ "type mismatch\n     found: " ++ showMaybet found ++ "\n  expected: " ++ showMaybet expected
 
--- | @typecheckReturn value fret@ resolves @value@ and compares it with @fret@,
--- returning a string decribing an error or Nothing in case of a type match.
-typecheckReturn :: Maybe A.Expr -> Maybe A.Type -> State CompState (Maybe String)
-typecheckReturn Nothing Nothing = return $ Nothing
-typecheckReturn Nothing fret = return $ Just $ "invalid return value\n     found: ()\n  expected: " ++ showMaybet fret
-typecheckReturn (Just value) fret = do
-  -- FIXME: this break constant folding
-  r <- compileOne value
-  return $ typecheck (crType r) fret
-
 -- built-in types
 types :: [String]
 types = ["bool", "u8", "s8", "u16", "s16"]
@@ -161,11 +151,14 @@ compileOne x = do
       -- helper for return
       nev <- return $ addSym nev $ Sym "$fn$" ftype True True pos
       put st {stEnv = nev, stErr = errs}
-      _ <- compileAll body
+      r <- compileAll body
       st <- get
       -- store updated errors and the env with the function
       put st {stEnv = ev}
-      return $ typeResult (Just ftype) x
+      rbody <- case r of
+        Right (_, xs) -> pure $ xs
+        Left _ -> pure $ body
+      return $ typeResult (Just ftype) (A.Func ident params ret rbody priv anon pos)
       where
         ftype = A.toFuncType params ret
     (A.Call ident args pos) -> do
@@ -194,10 +187,16 @@ compileOne x = do
       st <- get
       case getSyml (stEnv st) "$fn$" of
         Just Sym {symType = A.FuncType _ rtyp} -> do
-          r <- typecheckReturn value rtyp
-          case r of
-            Just err -> addError $ Error TypeError err pos
-            Nothing -> return $ typeResult rtyp x
+          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
+            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)
         _ -> addError $ Error UnexpectedReturn "return without function call" pos
     (A.Variable ident pos) -> do
       st <- get