fix(tc): fixed return check for functions without specified return-type

src/Plume/TypeChecker/Checker.hs

@@ -15,6 +15,7 @@ import Plume.TypeChecker.Checker.Interface (synthInterface)
 import Plume.TypeChecker.Checker.Native
 import Plume.TypeChecker.Checker.Switch
 import Plume.TypeChecker.Constraints.Solver
+import Plume.TypeChecker.Constraints.Unification
 import Plume.TypeChecker.Constraints.Typeclass
 import Plume.TypeChecker.Monad
 import Plume.TypeChecker.TLIR qualified as Post

src/Plume/TypeChecker/Constraints/Unification.hs

@@ -3,6 +3,8 @@ module Plume.TypeChecker.Constraints.Unification where
 
 import Plume.TypeChecker.Constraints.Definition
 import Plume.TypeChecker.Monad
+import Plume.TypeChecker.TLIR qualified as Typed
+import System.IO.Pretty
 
 -- infix 4 `unifiesTo`
 
@@ -133,4 +135,49 @@ compressPaths (TypeApp t ts) = do
   t' <- compressPaths t
   ts' <- traverse compressPaths ts
   pure (TypeApp t' ts')
-compressPaths t = pure t
+compressPaths t = pure t
+
+-- | Lift a block of expressions to check if any return is present
+-- | and if the return type matches the expected return type.
+liftBlock :: 
+  Placeholder Typed.Expression ->
+  [PlumeType] -> 
+  PlumeType -> 
+  Placeholder Typed.Expression
+liftBlock block _ t = do
+  f <- ask
+  res <- liftIO $ runReaderT block f
+  ty <- liftIO $ compressPaths t
+
+  case res of
+    Typed.EBlock exprs 
+      | any Typed.containsReturn exprs 
+        || ty == TUnit -> do
+          pure $ Typed.EBlock exprs
+    Typed.EBlock exprs
+      | not (any Typed.containsReturn exprs)
+        && isTVar ty -> case ty of
+          TypeVar ref -> do
+            writeIORef ref (Link TUnit)
+            pure $ Typed.EBlock exprs
+          _ -> error "Not a type variable"
+
+    Typed.EBlock _ -> liftIO $ do
+      pos <- fetchPositionIO
+      printErrorFromString 
+        mempty 
+        ( "No return found in the expression for type " <> show ty,
+          Just (hintMsg ty),
+          pos
+        )
+        "while performing typechecking"
+      exitFailure
+    _ -> error "Not a block"
+
+  where 
+    hintMsg ty' = case ty' of
+      TypeVar _ -> "Did you perhaps forget to specify unit? Every function must return a value"
+      _ -> "Every function must have a return in its body"
+
+    isTVar (TypeVar _) = True
+    isTVar _ = False

src/Plume/TypeChecker/Monad.hs

@@ -7,7 +7,7 @@ module Plume.TypeChecker.Monad
     Placeholder,
     Substitution,
     liftPlaceholders,
-    liftBlock,
+    fetchPositionIO,
     throw,
     throwRaw,
     insertWith,
@@ -285,39 +285,6 @@ fetchPositionIO = do
 
 type Substitution = Map Text PlumeType
 
--- | Lift a block of expressions to check if any return is present
--- | and if the return type matches the expected return type.
-liftBlock :: 
-  Placeholder Typed.Expression ->
-  [PlumeType] -> 
-  PlumeType -> 
-  Placeholder Typed.Expression
-liftBlock block _ ty = do
-  f <- ask
-  res <- liftIO $ runReaderT block f
-
-  case res of
-    Typed.EBlock exprs 
-      | any Typed.containsReturn exprs 
-        || ty == TUnit -> pure $ Typed.EBlock exprs
-    Typed.EBlock _ -> liftIO $ do
-      pos <- fetchPositionIO
-      printErrorFromString 
-        mempty 
-        ( "No return found in the expression for type " <> showTy ty,
-          Just (hintMsg ty),
-          pos
-        )
-        "while performing typechecking"
-      exitFailure
-    _ -> error "Not a block"
-
-  where 
-    hintMsg ty' = case ty' of
-      TypeVar _ -> "Did you perhaps forget to specify unit? Every function must return a value"
-      _ -> "Every function must have a return in its body"
-
-
 liftPlaceholders ::
   Text ->
   PlumeType ->