Non local return (#288)

* Remove old implementation of non local return and add detection for non local return
* Add OuterCtx and removed some dummy code
* Add tests for non local return

---------

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

Author: AnsonYeung

hkmc2/shared/src/main/scala/hkmc2/MLsCompiler.scala

@@ -72,11 +72,11 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
 
     val elab = Elaborator(etl, wd, Ctx.empty)
 
-    val initState = State.init.nest(N)
+    val initState = State.init.nestLocal
 
     val (pblk, newCtx) = elab.importFrom(preludeParse.resultBlk)(using initState)
 
-    newCtx.nest(N).givenIn:
+    newCtx.nestLocal.givenIn:
       val elab = Elaborator(etl, wd, newCtx)
       val parsed = mainParse.resultBlk
       val (blk0, _) = elab.importFrom(parsed)

hkmc2/shared/src/main/scala/hkmc2/codegen/Block.scala

@@ -42,12 +42,11 @@ sealed abstract class Block extends Product with AutoLocated:
       val rest = rst.definedVars
       if defn.isOwned then rest else rest + defn.sym
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) => bod.definedVars ++ rst.definedVars + lhs
-    case HandleBlockReturn(_) => Set.empty
     case TryBlock(sub, fin, rst) => sub.definedVars ++ fin.definedVars ++ rst.definedVars
     case Label(lbl, bod, rst) => bod.definedVars ++ rst.definedVars
 
   lazy val size: Int = this match
-    case _: Return | _: Throw | _: End | _: Break | _: Continue | _: HandleBlockReturn => 1
+    case _: Return | _: Throw | _: End | _: Break | _: Continue => 1
     case Begin(sub, rst) => sub.size + rst.size
     case Assign(_, _, rst) => 1 + rst.size
     case AssignField(_, _, _, rst) => 1 + rst.size
@@ -97,7 +96,6 @@ sealed abstract class Block extends Product with AutoLocated:
     case Define(defn, rest) => defn.freeVars ++ rest.freeVars
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) =>
       (bod.freeVars - lhs) ++ rst.freeVars ++ hdr.flatMap(_.freeVars)
-    case HandleBlockReturn(res) => res.freeVars
     case End(msg) => Set.empty
 
   // TODO: freeVarsLLIR skips `fun` and `cls` in `Call` and `Instantiate` respectively, which is needed in some
@@ -121,7 +119,6 @@ sealed abstract class Block extends Product with AutoLocated:
     case Define(defn, rest) => defn.freeVarsLLIR ++ rest.freeVarsLLIR
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) =>
       (bod.freeVarsLLIR - lhs) ++ rst.freeVarsLLIR ++ hdr.flatMap(_.freeVars)
-    case HandleBlockReturn(res) => res.freeVarsLLIR
     case End(msg) => Set.empty
 
   lazy val subBlocks: Ls[Block] = this match
@@ -137,10 +134,9 @@ sealed abstract class Block extends Product with AutoLocated:
 
     // TODO rm Lam from values and thus the need for these cases
     case Return(r, _) => r.subBlocks
-    case HandleBlockReturn(r) => r.subBlocks
     case Throw(r) => r.subBlocks
 
-    case _: Return | _: Throw | _: Break | _: Continue | _: End | _: HandleBlockReturn => Nil
+    case _: Return | _: Throw | _: Break | _: Continue | _: End => Nil
 
   // Moves definitions in a block to the top. Only scans the top-level definitions of the block;
   // i.e, definitions inside other definitions are not moved out. Definitions inside `match`/`if`
@@ -295,8 +291,6 @@ case class HandleBlock(
     rest: Block
 ) extends Block with ProductWithTail
 
-case class HandleBlockReturn(res: Result) extends BlockTail
-
 sealed abstract class Defn:
   val innerSym: Opt[MemberSymbol[?]]
   val sym: BlockMemberSymbol

hkmc2/shared/src/main/scala/hkmc2/codegen/BlockTransformer.scala

@@ -29,9 +29,6 @@ class BlockTransformer(subst: SymbolSubst):
     case Throw(exc) =>
       applyResult2(exc): exc2 =>
         if exc2 is exc then b else Throw(exc2)
-    case HandleBlockReturn(res) =>
-      applyResult2(res): res2 =>
-        if res2 is res then b else HandleBlockReturn(res2)
     case Match(scrut, arms, dflt, rst) =>
       val scrut2 = applyPath(scrut)
       val arms2 = arms.mapConserve: arm =>

hkmc2/shared/src/main/scala/hkmc2/codegen/BlockTraverser.scala

@@ -22,7 +22,6 @@ class BlockTraverser(subst: SymbolSubst):
     case Continue(lbl) => applyLocal(lbl)
     case Return(res, implct) => applyResult(res)
     case Throw(exc) => applyResult(exc)
-    case HandleBlockReturn(res) => applyResult(res)
     case Match(scrut, arms, dflt, rst) =>
       val scrut2 = applyPath(scrut)
       arms.foreach: arm =>

hkmc2/shared/src/main/scala/hkmc2/codegen/HandlerLowering.scala

@@ -29,21 +29,17 @@ object HandlerLowering:
 
   private case class LinkState(res: Path, cls: Path, uid: StateId)
 
-  // shouldUnwrapRet: whether the current block should unwrap the runtime.Return if it encounter one
   // isTopLevel:
   // whether the current block is the top level block, as we do not emit code for continuation class on the top level
   // since we cannot return an effect signature on the top level (we are not in a function so return statement are invalid)
   // and we do not have any `return` statement in the top level block so we do not need the `runtime.Return` workarounds.
-  // isHandler: whether the current block is the body of handler method
   // contName: the name of the continuation class
   // ctorThis: the path to `this` in the constructor, this is used to insert `return this;` at the end of constructor.
   // linkAndHandle:
   // a function that takes a LinkState and returns a block that links the continuation class and handles the effect
   // this is a convenience function which initializes the continuation class in function context or throw an error in top level
   private case class HandlerCtx(
-    shouldUnwrapRet: Bool,
     isTopLevel: Bool,
-    isHandlerMtd: Bool,
     isHandlerBody: Bool,
     contName: Str,
     ctorThis: Option[Path],
@@ -59,27 +55,25 @@ class HandlerPaths(using Elaborator.State):
   val effectSigPath: Path = runtimePath.selN(Tree.Ident("EffectSig")).selN(Tree.Ident("class"))
   val effectSigSym: ClassSymbol = State.effectSigSymbol
   val contClsPath: Path = runtimePath.selN(Tree.Ident("FunctionContFrame")).selN(Tree.Ident("class"))
-  val retClsPath: Path = runtimePath.selN(Tree.Ident("Return")).selN(Tree.Ident("class"))
-  val retClsSym: ClassSymbol = State.returnClsSymbol
   val mkEffectPath: Path = runtimePath.selN(Tree.Ident("mkEffect"))
   val handleBlockImplPath: Path = runtimePath.selN(Tree.Ident("handleBlockImpl"))
   val stackDelayClsPath: Path = runtimePath.selN(Tree.Ident("StackDelay"))
 
   def isHandlerClsPath(p: Path) =
-    (p eq contClsPath)  || (p eq stackDelayClsPath) || (p eq effectSigPath) || (p eq retClsPath)
+    (p eq contClsPath)  || (p eq stackDelayClsPath) || (p eq effectSigPath)
 
 class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, Elaborator.Ctx):
 
   private def funcLikeHandlerCtx(ctorThis: Option[Path], isHandlerMtd: Bool, nme: Str) =
-    HandlerCtx(!isHandlerMtd, false, isHandlerMtd, false, nme, ctorThis, state =>
+    HandlerCtx(false, false, nme, ctorThis, state =>
       blockBuilder
         .assignFieldN(state.res.contTrace.last, nextIdent, Instantiate(
           state.cls.selN(Tree.Ident("class")),
           Value.Lit(Tree.IntLit(state.uid)) :: Nil))
         .assignFieldN(state.res.contTrace, lastIdent, state.res.contTrace.last.next)
         .ret(state.res))
   private def functionHandlerCtx(nme: Str) = funcLikeHandlerCtx(N, false, nme)
-  private def topLevelCtx(nme: Str) = HandlerCtx(false, true, false, false, nme, N, _ => rtThrowMsg("Unhandled effects"))
+  private def topLevelCtx(nme: Str) = HandlerCtx(true, false, nme, N, _ => rtThrowMsg("Unhandled effects"))
   private def ctorCtx(ctorThis: Path, nme: Str) = funcLikeHandlerCtx(S(ctorThis), false, nme)
   private def handlerMtdCtx(nme: Str) = funcLikeHandlerCtx(N, true, nme)
   private def handlerCtx(using HandlerCtx): HandlerCtx = summon
@@ -116,25 +110,25 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
     def apply(res: Local, uid: StateId) =
       Assign(res, PureCall(Value.Ref(returnContSymbol), List(Value.Lit(Tree.IntLit(uid)))), End(""))
     def unapply(blk: Block) = blk match
-      case Assign(res, PureCall(Value.Ref(`returnContSymbol`), List(Value.Lit(Tree.IntLit(uid)))), _) => 
+      case Assign(res, PureCall(Value.Ref(`returnContSymbol`), List(Value.Lit(Tree.IntLit(uid)))), _) =>
         Some(res, uid)
       case _ => None
 
   // placeholder for effectful Results, such as Call, Instantiate and anything else that could
   // return a continuation
   object ResultPlaceholder:
     private val callSymbol = freshTmp("resultPlaceholder")
-    def apply(res: Local, uid: StateId, canRet: Bool, r: Result, rest: Block) =
+    def apply(res: Local, uid: StateId, r: Result, rest: Block) =
       Assign(
         res,
-        PureCall(Value.Ref(callSymbol), List(Value.Lit(Tree.IntLit(uid)), Value.Lit(Tree.BoolLit(canRet)))),
+        PureCall(Value.Ref(callSymbol), List(Value.Lit(Tree.IntLit(uid)))),
         Assign(res, r, rest))
     def unapply(blk: Block) = blk match
       case Assign(
           res,
-          PureCall(Value.Ref(`callSymbol`), List(Value.Lit(Tree.IntLit(uid)), Value.Lit(Tree.BoolLit(canRet)))),
+          PureCall(Value.Ref(`callSymbol`), List(Value.Lit(Tree.IntLit(uid)))),
           Assign(_, c, rest)) =>
-        Some(res, uid, canRet, c, rest)
+        Some(res, uid, c, rest)
       case _ => None
 
   object StateTransition:
@@ -287,7 +281,6 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
       case TryBlock(sub, finallyDo, rest) => ??? // ignore
       case Throw(_) => PartRet(blk, Nil)
       case _: HandleBlock => lastWords("unexpected handleBlock") // already translated at this point
-      case _: HandleBlockReturn => lastWords("unexpected handleBlockReturn") // already translated at this point
 
     val result = go(blk)(using labelIds, N)
     result.states
@@ -327,26 +320,26 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
           val fun2 = applyPath(fun)
           val args2 = args.mapConserve(applyArg)
           val c2 = if (fun2 is fun) && (args2 is args) then c else Call(fun2, args2)(c.isMlsFun, c.mayRaiseEffects)
-          ResultPlaceholder(lhs, freshId(), handlerCtx.isHandlerMtd, c2, applyBlock(rest))
+          ResultPlaceholder(lhs, freshId(), c2, applyBlock(rest))
         case Assign(lhs, c @ Instantiate(cls, args), rest) =>
           val cls2 = applyPath(cls)
           val args2 = args.mapConserve(applyPath)
           val c2 = if (cls2 is cls) && (args2 is args) then c else Instantiate(cls2, args2)
-          ResultPlaceholder(lhs, freshId(), handlerCtx.isHandlerMtd, c2, applyBlock(rest))
+          ResultPlaceholder(lhs, freshId(), c2, applyBlock(rest))
         case _ => super.applyBlock(b)
       override def applyResult2(r: Result)(k: Result => Block): Block = r match
         case c @ Call(fun, args) if c.mayRaiseEffects =>
           val res = freshTmp("res")
           val fun2 = applyPath(fun)
           val args2 = args.mapConserve(applyArg)
           val c2 = if (fun2 is fun) && (args2 is args) then c else Call(fun2, args2)(c.isMlsFun, c.mayRaiseEffects)
-          ResultPlaceholder(res, freshId(), handlerCtx.isHandlerMtd, c2, k(Value.Ref(res)))
+          ResultPlaceholder(res, freshId(), c2, k(Value.Ref(res)))
         case c @ Instantiate(cls, args) =>
           val res = freshTmp("res")
           val cls2 = applyPath(cls)
           val args2 = args.mapConserve(applyPath)
           val c2 = if (cls2 is cls) && (args2 is args) then c else Instantiate(cls2, args2)
-          ResultPlaceholder(res, freshId(), handlerCtx.isHandlerMtd, c2, k(Value.Ref(res)))
+          ResultPlaceholder(res, freshId(), c2, k(Value.Ref(res)))
         case r => super.applyResult2(r)(k)
       override def applyLam(lam: Value.Lam): Value.Lam =
         Value.Lam(lam.params, translateBlock(lam.body, functionHandlerCtx(s"Cont$$lambda$$")))
@@ -389,21 +382,8 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
     val sym = BlockMemberSymbol(s"handleBlock$$", Nil)
     val lbl = freshTmp("handlerBody")
     val lblLoop = freshTmp("handlerLoop")
-    val tmp = freshTmp("retCont")
-    def prepareBody(b: Block): Block =
-      
-      val transform = new BlockTransformerShallow(SymbolSubst()):
-        override def applyBlock(b: Block): Block = 
-          b match
-          case Return(res, implct) =>
-            // In case res is effectful, it will be handled in translateBlock
-            Assign(tmp, res, Return(Instantiate(paths.retClsPath, tmp.asPath :: Nil), implct))
-          case HandleBlockReturn(res) =>
-            Return(res, false)
-          case _ => super.applyBlock(b)
-      transform.applyBlock(b)
 
-    val handlerBody = translateBlock(prepareBody(h.body), HandlerCtx(false, false, false, true,
+    val handlerBody = translateBlock(h.body, HandlerCtx(false, true,
       s"Cont$$handleBlock$$${h.lhs.nme}$$", N, state => blockBuilder
         .assignFieldN(state.res.contTrace.last, nextIdent, PureCall(state.cls, Value.Lit(Tree.IntLit(state.uid)) :: Nil))
         .ret(PureCall(paths.handleBlockImplPath, state.res :: h.lhs.asPath :: Nil))))
@@ -437,7 +417,7 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
       N, // no owner
       sym, PlainParamList(Nil) :: Nil, body)
 
-    val result = Define(defn, ResultPlaceholder(h.res, freshId(), !handlerCtx.isTopLevel, Call(sym.asPath, Nil)(true, true), h.rest))
+    val result = Define(defn, ResultPlaceholder(h.res, freshId(), Call(sym.asPath, Nil)(true, true), h.rest))
     result
 
   private def genContClass(b: Block)(using HandlerCtx): Opt[ClsLikeDefn] =
@@ -453,7 +433,7 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
       val transform = new BlockTransformerShallow(SymbolSubst()):
         override def applyBlock(b: Block): Block = b match
           case Define(_: (ClsLikeDefn | FunDefn), rst) => applyBlock(rst)
-          case ResultPlaceholder(res, uid, canRet, c, rest) =>
+          case ResultPlaceholder(res, uid, c, rest) =>
             trivial = false
             blockBuilder
               .assign(res, c)
@@ -463,12 +443,6 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
                 ReturnCont(res, uid)
               )
               .chain(ResumptionPoint(res, uid, _))
-              .staticif(canRet, _.ifthen(
-                res.asPath,
-                Case.Cls(paths.retClsSym, paths.retClsPath),
-                blockBuilder
-                  .ret(if handlerCtx.shouldUnwrapRet then res.asPath.value else res.asPath)
-              ))
               .rest(applyBlock(rest))
           case _ => super.applyBlock(b)
       transform.applyBlock(b)
@@ -558,19 +532,14 @@ class HandlerLowering(paths: HandlerPaths)(using TL, Raise, Elaborator.State, El
   private def genNormalBody(b: Block, clsSym: BlockMemberSymbol)(using HandlerCtx): Block =
     val transform = new BlockTransformerShallow(SymbolSubst()):
       override def applyBlock(b: Block): Block = b match
-        case ResultPlaceholder(res, uid, canRet, c, rest) =>
+        case ResultPlaceholder(res, uid, c, rest) =>
           blockBuilder
             .assign(res, c)
             .ifthen(
               res.asPath,
               Case.Cls(paths.effectSigSym, paths.effectSigPath),
               handlerCtx.linkAndHandle(LinkState(res.asPath, clsSym.asPath, uid))
             )
-            .staticif(canRet, _.ifthen(
-              res.asPath,
-              Case.Cls(paths.retClsSym, paths.retClsPath),
-              blockBuilder.ret(if handlerCtx.shouldUnwrapRet then res.asPath.value else res.asPath)
-            ))
             .rest(applyBlock(rest))
         case _ => super.applyBlock(b)
 

hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala

@@ -321,7 +321,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
       subTerm(rhs): par =>
         subTerms(as): asr =>
           HandleBlock(lhs, resSym, par, asr, cls, handlers,
-            term_nonTail(bod)(HandleBlockReturn(_)),
+            term_nonTail(bod)(Ret),
             k(Value.Ref(resSym)))
     case st.Blk(sts, res) => block(sts, R(res))(k)
     case Assgn(lhs, rhs) =>

hkmc2/shared/src/main/scala/hkmc2/codegen/StackSafeTransform.scala

@@ -81,7 +81,7 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
 
   def extractResTopLevel(res: Result, isTailCall: Bool, f: Result => Block, sym: Option[Symbol], curDepth: => Symbol) =
     val resSym = sym getOrElse TempSymbol(None, "res")
-    wrapStackSafe(HandleBlockReturn(res), resSym, f(resSym.asPath))
+    wrapStackSafe(Ret(res), resSym, f(resSym.asPath))
 
   // Rewrites anything that can contain a Call to increase the stack depth
   def transform(b: Block, curDepth: => Symbol, isTopLevel: Bool = false): Block =
@@ -120,7 +120,7 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
           val rst2 = applyBlock(rst)
           if isTopLevel then
             val newRes = TempSymbol(N, "res")
-            val newHandler = HandleBlock(l2, newRes, par2, args2, cls2, hdr2, bod2, HandleBlockReturn(newRes.asPath))
+            val newHandler = HandleBlock(l2, newRes, par2, args2, cls2, hdr2, bod2, Ret(newRes.asPath))
             wrapStackSafe(newHandler, res2, rst2)
           else
             HandleBlock(l2, res2, par2, args2, cls2, hdr2, bod2, rst2)

hkmc2/shared/src/main/scala/hkmc2/codegen/js/JSBuilder.scala

@@ -156,7 +156,7 @@ class JSBuilder(using TL, State, Ctx) extends CodeBuilder:
   def returningTerm(t: Block, endSemi: Bool)(using Raise, Scope): Document =
     def mkSemi = if endSemi then ";" else ""
     t match
-    case _: (HandleBlockReturn | HandleBlock) =>
+    case _: HandleBlock =>
       errStmt(msg"This code requires effect handler instrumentation but was compiled without it.")
     case Assign(l, r, rst) =>
       doc" # ${getVar(l)} = ${result(r)};${returningTerm(rst, endSemi)}"