Fix while loop rewriting (hkust-taco#403)

Co-authored-by: Lionel Parreaux <lionel.parreaux@gmail.com>

Author: AnsonYeung

hkmc2/shared/src/main/scala/hkmc2/semantics/ucs/Normalization.scala

@@ -454,7 +454,6 @@ class Normalization(lowering: Lowering)(using tl: TL)(using Raise, Ctx, State) e
             val loopEnd: Path =
               Select(Value.Ref(State.runtimeSymbol), Tree.Ident("LoopEnd"))(S(State.loopEndSymbol))
             val blk = blockBuilder
-              .assign(l, Value.Lit(Tree.UnitLit(false)))
               .define(FunDefn(N, f, tSym, PlainParamList(Nil) :: Nil, Begin(body, Return(loopEnd, false)))(forceTailRec = false))
               .assign(loopResult, Call(Value.Ref(f, S(tSym)), Nil)(true, true, false))
             if summon[LoweringCtx].mayRet then
@@ -463,9 +462,9 @@ class Normalization(lowering: Lowering)(using tl: TL)(using Raise, Ctx, State) e
                   loopResult.asPath.asArg :: loopEnd.asArg :: Nil)(true, false, false))
                 .ifthen(Value.Ref(isReturned), Case.Lit(Tree.BoolLit(true)),
                   Return(Value.Ref(loopResult), false),
-                  S(rest)
+                  N
                 )
-                .end
+                .rest(rest)
             else
               blk.rest(rest)
           else

hkmc2/shared/src/test/mlscript-compile/Runtime.mls

@@ -10,6 +10,9 @@ module Runtime with ...
 object Unit with
   fun toString() = "()"
 
+// This object is used as a marker for the end of a while loop in the compiled codegen.
+// It is possible for the loop body to contain a return statement,
+// in which case we want to distinguish between returning from the loop body and loop termination.
 object LoopEnd
 val short_and = RuntimeJS.short_and
 val short_or = RuntimeJS.short_or

hkmc2/shared/src/test/mlscript/codegen/ScopedBlocks.mls

@@ -169,8 +169,7 @@ fun f(x) =
 //│ JS (unsanitized):
 //│ let f7;
 //│ f7 = function f(x1) {
-//│   let while1, tmp2, tmp3, tmp4;
-//│   tmp4 = undefined;
+//│   let while1, tmp2, tmp3;
 //│   while1 = (undefined, function () {
 //│     let y1;
 //│     if (x1 === true) {
@@ -182,7 +181,8 @@ fun f(x) =
 //│   });
 //│   tmp2 = while1();
 //│   tmp3 = tmp2 !== runtime.LoopEnd;
-//│   if (tmp3 === true) { return tmp2 } else { return tmp4 }
+//│   if (tmp3 === true) { return tmp2 }
+//│   return runtime.Unit
 //│ };
 
 

hkmc2/shared/src/test/mlscript/codegen/While.mls

@@ -359,11 +359,10 @@ while
 
 
 // * https://github.com/hkust-taco/mlscript/issues/401
-:fixme
 // :sjs
 :rewriteWhile
 fun f = while false do ()
 print(f)
-//│ ═══[RUNTIME ERROR] Error: MLscript call unexpectedly returned `undefined`, the forbidden value.
+//│ > ()