From 1e962be827b7191663ead3bd57de297d5165a3a1 Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Tue, 28 Jan 2025 17:51:47 -0800
Subject: [PATCH] fix issue #643 with wrong return type

---
 src/Compile/TypeCheck.hs |  3 ++-
 src/Type/Infer.hs        | 17 +++++++++--------
 test/cgen/return6.kk     | 19 +++++++++++++++++++
 test/cgen/return6.kk.out |  0
 4 files changed, 30 insertions(+), 9 deletions(-)
 create mode 100644 test/cgen/return6.kk
 create mode 100644 test/cgen/return6.kk.out

diff --git a/src/Compile/TypeCheck.hs b/src/Compile/TypeCheck.hs
index 34b47cf8c..6a8eb0335 100644
--- a/src/Compile/TypeCheck.hs
+++ b/src/Compile/TypeCheck.hs
@@ -113,7 +113,8 @@ typeCheck flags defs coreImports program0
         unreturn penv
 
         -- checkCoreDefs "unreturn"
-        let borrowed = borrowedExtendICore (coreProgram{ Core.coreProgDefs = coreDefs }) (defsBorrowed defs)
+        coreDefs1 <- Core.getCoreDefs
+        let borrowed = borrowedExtendICore (coreProgram{ Core.coreProgDefs = coreDefs1 }) (defsBorrowed defs)
         checkFBIP penv (platform flags) newtypes borrowed gamma
 
         -- initial simplify
diff --git a/src/Type/Infer.hs b/src/Type/Infer.hs
index bba7936be..e057e6948 100644
--- a/src/Type/Infer.hs
+++ b/src/Type/Infer.hs
@@ -710,18 +710,19 @@ inferExpr propagated expect (App (Var name _ nameRng) [(_,expr)] rng)  | name ==
   = do allowed <- isReturnAllowed
        if (False && not allowed)
         then infError rng (text "illegal expression context for a return statement")
-        else  do (tp,eff,core) <- inferExpr propagated expect expr
-                 mbTp <- lookupInfName nameReturn -- (unqualify nameReturn)
+        else  do mbTp <- lookupInfName nameReturn -- (unqualify nameReturn)
                  case mbTp of
                    Nothing
                     -> do infError rng (text "illegal context for a return statement")
+                          inferExpr propagated expect expr
                    Just (_,retTp)
-                    -> do inferUnify (checkReturn rng) (getRange expr) retTp tp
-                 resTp <- Op.freshStar
-                 let typeReturn = typeFun [(nameNil,tp)] typeTotal resTp
-                 addRangeInfo nameRng (RM.Id (newName "return") (RM.NIValue "expr" tp "" False) [] False)
-                 return (resTp, eff, Core.App (Core.Var (Core.TName nameReturn typeReturn)
-                                      (Core.InfoExternal [(Default,"return #1")])) [core])
+                    -> do (tp,eff,core) <- inferExpr (Just (retTp,nameRng)) expect expr
+                          inferUnify (checkReturn rng) (getRange expr) retTp tp
+                          resTp <- Op.freshStar
+                          let typeReturn = typeFun [(nameNil,tp)] typeTotal resTp
+                          addRangeInfo nameRng (RM.Id (newName "return") (RM.NIValue "expr" tp "" False) [] False)
+                          return (resTp, eff, Core.App (Core.Var (Core.TName nameReturn typeReturn)
+                                                (Core.InfoExternal [(Default,"return #1")])) [core])
 -- | Assign expression
 inferExpr propagated expect (App assign@(Var name _ arng) [lhs@(_,lval),rhs@(_,rexpr)] rng) | name == nameAssign
   = case lval of
diff --git a/test/cgen/return6.kk b/test/cgen/return6.kk
new file mode 100644
index 000000000..eee06d9c3
--- /dev/null
+++ b/test/cgen/return6.kk
@@ -0,0 +1,19 @@
+// issue #643
+fun test(c: char): maybe<char>
+  val c1 =
+    match f(c)
+      Just(c1) ->
+        c1
+      Nothing ->
+        return Nothing // or `return test(c)`
+
+  Just(c1)
+
+fun f(c: char): maybe<char>
+  val x = match Just(c)
+    Just(_cx) -> return Nothing
+    Nothing  -> Just(c)
+  x
+
+fun main()
+  ()
\ No newline at end of file
diff --git a/test/cgen/return6.kk.out b/test/cgen/return6.kk.out
new file mode 100644
index 000000000..e69de29bb