Fix return scrutinee

src/Compile/TypeCheck.hs

@@ -111,9 +111,9 @@ typeCheck flags defs coreImports program0
 
         -- remove return statements
         unreturn penv
-
+        coreDefs' <- Core.getCoreDefs
         -- checkCoreDefs "unreturn"
-        let borrowed = borrowedExtendICore (coreProgram{ Core.coreProgDefs = coreDefs }) (defsBorrowed defs)
+        let borrowed = borrowedExtendICore (coreProgram{ Core.coreProgDefs = coreDefs' }) (defsBorrowed defs)
         checkFBIP penv (platform flags) newtypes borrowed gamma
 
         -- initial simplify
@@ -146,7 +146,7 @@ typeCheck flags defs coreImports program0
             currentImports   = S.fromList (map Core.importName coreImports)
             typeImports      = [Core.Import name "" Core.ImportTypes Private "" | name <- typeDeps, not (S.member name currentImports) && not (name == progName)]
             coreFinal        = coreUnique{ Core.coreProgImports = Core.coreProgImports coreUnique ++ typeImports }
-
+        
         return (coreFinal,mbRangeMap)
 
   where

src/Core/UnReturn.hs

@@ -44,14 +44,13 @@ import Core.Pretty
 import Core.CoreVar
 
 trace s x =
-   -- Lib.Trace.trace s
+    -- Lib.Trace.trace s
     x
 
 unreturn :: Pretty.Env -> CorePhase b ()
 unreturn penv
   = liftCorePhaseUniq $ \uniq defs -> runUR penv uniq (urTopDefGroups defs)
 
-
 {--------------------------------------------------------------------------
   transform definition groups
 --------------------------------------------------------------------------}
@@ -116,7 +115,15 @@ urExpr expr
 
       -- case: scrutinee cannot contain return due to grammar
       Case scruts branches
-        -> urCase expr scruts branches
+        -> if any containsReturn scruts then do
+            names <- mapM (\scrut -> (uniqueName "scrut")) scruts
+            let tnscrut = zipWith (\nm scrut -> (TName nm (typeOf scrut), scrut)) names scruts
+            let addDefs :: [(TName, Expr)] -> Expr -> Expr
+                addDefs dfs = Let (map (\(tn, dexpr) -> DefNonRec (Def (getName tn) (typeOf tn) dexpr Private DefVal InlineNever rangeNull "")) dfs)
+            let case' = addDefs tnscrut (Case (map (\(nm, scrut) -> Var nm InfoNone) tnscrut) branches)
+            urExpr case'
+           else
+            urCase expr scruts branches
 
       -- return
       App ret@(Var v _) [arg] | getName v == nameReturn
@@ -130,6 +137,20 @@ urExpr expr
               return (I (App f' args'))
       _ -> return (U expr)
 
+containsReturn :: Expr -> Bool
+containsReturn expr
+  = case expr of
+      Lam pars eff body -> containsReturn body
+      Var _ _ -> False
+      App ret@(Var v _) [arg] | getName v == nameReturn -> True
+      App f args -> any containsReturn (f:args)
+      TypeLam tvars body -> containsReturn body
+      TypeApp body targs -> containsReturn body
+      Lit _ -> False
+      Con _ _ -> False
+      Let defgs body -> any (containsReturn . defExpr) (flattenDefGroups defgs)
+      Case scruts branches -> any containsReturn scruts || any (any (containsReturn . guardExpr) . branchGuards) branches
+
 urPure :: Expr -> UR Expr
 urPure expr
   = do kexpr <- urExpr expr
@@ -156,7 +177,7 @@ urLet org defgroups kbody
     fold (Right (makeDefGroup,kdefexpr) : kdefgs) kexpr
       = fold kdefgs (bind Nothing combine kdefexpr kexpr)
       where
-        combine e1 e2 = trace ("combine: " ++ show (e1,e2)) $
+        combine e1 e2 = trace ("combine: " ++ show (e1, e2) ++ "\n") $
                         addDef (makeDefGroup e1) e2
 
     addDef :: DefGroup -> Expr -> Expr
@@ -242,7 +263,11 @@ data KExpr  = U Expr   -- unchanged expression
             | R Expr   -- returned expression
             | F ((Expr -> Expr) -> Expr)  -- cps transformed, needs continuation
 
-
+showKExpr :: KExpr -> String
+showKExpr (U e) = "U(" ++ show e ++ ")"
+showKExpr (I e) = "I(" ++ show e ++ ")"
+showKExpr (R e) = "R(" ++ show e ++ ")"
+showKExpr (F _) = "F"
 
 isU (U _) = True
 isU _     = False
@@ -284,7 +309,7 @@ emapK mbOrg g kexpr
 
 
 bind :: Maybe Expr -> (Expr -> Expr -> Expr) -> KExpr -> KExpr -> KExpr
-bind mbOrg combine  ke1 ke2
+bind mbOrg combine ke1 ke2
   = case (ke1,ke2) of
       (R r, _) -> R r
       (U a, k) -> case k of

test/cgen/returnscrut.kk

@@ -0,0 +1,29 @@
+// Issue: https://github.com/koka-lang/koka/issues/644
+fun main()
+  println("Compiled")
+
+type token
+  Int
+
+value struct token-iter<h>
+  chars: vector<char>
+  idx: ref<h, int>
+
+fun token-iter/next(token-iter: token-iter<h>): <div, read<h>, write<h>> maybe<(int, token)>
+  // Strange code below to work around issues with early returns in Koka 3.1.2.
+  // (https://github.com/koka-lang/koka/discussions/643)
+
+  // Skip whitespace
+  val next-char0 =
+    match Nothing
+      Just((idx, c)) ->
+        if True then
+          return token-iter.next
+        else
+          Nothing
+      Nothing -> Nothing
+  val tmp = match next-char0
+          Nothing -> return Nothing
+          Just(c) -> c
+  val (char-idx, char) = tmp
+  Nothing

test/cgen/returnscrut.kk.out

@@ -0,0 +1,5 @@
+Compiled
+test/cgen/returnscrut@kk(18, 7): type warning: pattern variable c is unused (or a wrongly spelled constructor?)
+ hint: prepend an underscore to make it a wildcard pattern
+test/cgen/returnscrut@kk(27, 7): type warning: pattern variable char is unused (or a wrongly spelled constructor?)
+ hint: prepend an underscore to make it a wildcard pattern