Typechecking return

1 files changed, 29 insertions, 4 deletions

diff --git a/src/Compiler.hs b/src/Compiler.hs
index 3d3f0ac..62be364 100644
--- a/src/Compiler.hs
+++ b/src/Compiler.hs
@@ -38,7 +38,7 @@ addError e = do
   return $ Right Nothing
 
 -- | @typecheckCall args params@ resolves @args@ and compares it with @params@,
--- returning a a string describing an error or Nothing in case of type match.
+-- returning a string describing an error or Nothing in case of type match.
 typecheckCall :: [A.Expr] -> [A.Type] -> State CompState (Maybe String)
 typecheckCall args params
   | length args /= length params = return $ Just "invalid number of arguments in function call"
@@ -59,6 +59,23 @@ typecheckCall args params
           -- there was an error in on argument
           return $ Nothing
 
+-- | @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 (Just value) fret = do
+  r <- compile value
+  case r of
+    Right r ->
+      if r == fret
+        then return $ Nothing
+        else return $ Just $ "invalid return value:\n  unexpected " ++ showMaybet r ++ "\n   expecting " ++ showMaybet fret
+    Left _ -> return $ Nothing -- error resolving return value
+  where
+    showMaybet :: Maybe A.Type -> String
+    showMaybet Nothing = "()"
+    showMaybet (Just t) = show t
+
 compile :: A.Expr -> State CompState CompResult
 compile x = do
   case x of
@@ -78,6 +95,8 @@ compile x = do
           Right ev -> (ev, errs)
       -- with parameters
       (nev, errs) <- return $ foldlEither addSymUniq (addEnv ev, errs) params
+      -- helper for return
+      nev <- return $ addSym nev ("$fn$", ftype, pos)
       put (nev, errs)
       r <- compileAll body
       (_, errs) <- get
@@ -94,9 +113,15 @@ compile x = do
           case r of
             Just err -> addError $ Error err pos
             Nothing -> return $ Right rtyp
-    (A.Return value pos) -> case value of
-      Just v -> compile v
-      Nothing -> return $ Right Nothing
+    (A.Return value pos) -> do
+      (ev, errs) <- get
+      case getSyml ev "$fn$" of
+        Nothing -> addError $ Error "return without function call" pos
+        Just (_, A.FuncType _ rtyp, _) -> do
+          r <- typecheckReturn value rtyp
+          case r of
+            Just err -> addError $ Error err pos
+            Nothing -> return $ Right rtyp
     (A.Var ident pos) -> do
       (ev, errs) <- get
       case getSym ev ident of