Add Scoped blocks

compiler/shared/main/scala/mlscript/compiler/optimizer/TailRecOpt.scala

@@ -13,6 +13,8 @@ import compiler.ir.Node._
 
 DOCUMENTATION OF SEMANTICS OF @tailcall and @tailrec
 
+FIXME: This doc is a bit outdated, as we have not ported the "modulo-cons" optimization yet.
+
 @tailcall: Used to annotate specific function calls. Calls annotated with @tailcall 
 must be tail calls or tail modulo-cons calls. These calls must be optimized to not
 consume additional stack space. If such an optimization is not possible, then the

hkmc2/shared/src/main/scala/hkmc2/Config.scala

@@ -28,6 +28,16 @@ case class Config(
 
   def stackSafety: Opt[StackSafety] = effectHandlers.flatMap(_.stackSafety)
 
+  // NOTE: We force the rewriting of while loops to functions when handler lowering is on
+  // to prevent the "floating out" of definitions done by handler lowering,
+  // which currently does not respect scopes introduced by `Scoped` blocks.
+  // Currently, this is only a problem with while loops because we do not yet
+  // construct nested Scoped blocks in other places (but will in the future).
+  // see https://github.com/hkust-taco/mlscript/pull/356#discussion_r2579529893
+  // and https://github.com/hkust-taco/mlscript/pull/356#discussion_r2585183902
+  def shouldRewriteWhile: Bool =
+    rewriteWhileLoops || effectHandlers.isDefined
+
 end Config
 
 
@@ -39,7 +49,7 @@ object Config:
     effectHandlers = N,
     liftDefns = N,
     target = CompilationTarget.JS,
-    rewriteWhileLoops = true,
+    rewriteWhileLoops = false,
     stageCode = false,
     tailRecOpt = true,
   )

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

@@ -29,7 +29,38 @@ sealed abstract class Block extends Product:
     case _: End => true
     case _ => false
 
+  /** This variation of `definedVars` excludes the `syms` in `Scoped` blocks.
+    * It is used in JSBuilder now: names are allocated in JSBuilder for these `definedVarsNoScoped` symbols.
+    * This is needed now because
+    *   - there are symbols that are not collected in the `Scoped` blocks (due to some passes that have not yet been adapted to Scoped),
+    *     and they still need to be allocated a name in JSBuilder
+    *   - if we don't exlude the `Scoped` symbols, they may be allocated a name
+    *     prematurely and decalred in wrong places, e.g. symbols inside while bodies
+    *     may be declared in an outer level
+    */
+  lazy val definedVarsNoScoped: Set[Local] = this match
+    case _: Return | _: Throw => Set.empty
+    case Begin(sub, rst) => sub.definedVarsNoScoped ++ rst.definedVarsNoScoped
+    case Assign(l: TermSymbol, r, rst) => rst.definedVarsNoScoped
+    case Assign(l, r, rst) => rst.definedVarsNoScoped + l
+    case AssignField(l, n, r, rst) => rst.definedVarsNoScoped
+    case AssignDynField(l, n, ai, r, rst) => rst.definedVarsNoScoped
+    case Match(scrut, arms, dflt, rst) =>
+      arms.flatMap(_._2.definedVarsNoScoped).toSet ++ dflt.toList.flatMap(_.definedVarsNoScoped) ++ rst.definedVarsNoScoped
+    case End(_) => Set.empty
+    case Break(_) => Set.empty
+    case Continue(_) => Set.empty
+    case Define(defn, rst) =>
+      val rest = rst.definedVarsNoScoped
+      if defn.isOwned then rest else rest + defn.sym
+    // Note that the handler's LHS and body are not part of the current block, so we do not consider them here.
+    case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) => rst.definedVarsNoScoped + res
+    case TryBlock(sub, fin, rst) => sub.definedVarsNoScoped ++ fin.definedVarsNoScoped ++ rst.definedVarsNoScoped
+    case Label(lbl, _, bod, rst) => bod.definedVarsNoScoped ++ rst.definedVarsNoScoped
+    case Scoped(syms, body) => body.definedVarsNoScoped -- syms
 
+  // * Note: there is a good chance that historical users of `definedVars` do not properly respect Scoped blocks
+  // * and should adapt their logic to use `definedVarsNoScoped` instead.
   lazy val definedVars: Set[Local] = this match
     case _: Return | _: Throw => Set.empty
     case Begin(sub, rst) => sub.definedVars ++ rst.definedVars
@@ -49,6 +80,7 @@ sealed abstract class Block extends Product:
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) => rst.definedVars + res
     case TryBlock(sub, fin, rst) => sub.definedVars ++ fin.definedVars ++ rst.definedVars
     case Label(lbl, _, bod, rst) => bod.definedVars ++ rst.definedVars
+    case Scoped(syms, body) => body.definedVars
 
   lazy val size: Int = this match
     case _: Return | _: Throw | _: End | _: Break | _: Continue => 1
@@ -62,9 +94,11 @@ sealed abstract class Block extends Product:
     case TryBlock(sub, fin, rst) => 1 + sub.size + fin.size + rst.size
     case Label(_, _, bod, rst) => 1 + bod.size + rst.size
     case HandleBlock(lhs, res, par, args, cls, handlers, bdy, rst) => 1 + handlers.map(_.body.size).sum + bdy.size + rst.size
+    case Scoped(_, body) => body.size
 
   // TODO conserve if no changes
   def mapTail(f: BlockTail => Block): Block = this match
+    case Scoped(syms, body) => Scoped(syms, body.mapTail(f))
     case b: BlockTail => f(b)
     case Begin(sub, rst) => Begin(sub, rst.mapTail(f))
     case Assign(lhs, rhs, rst) => Assign(lhs, rhs, rst.mapTail(f))
@@ -101,6 +135,7 @@ sealed abstract class Block extends Product:
     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 Scoped(syms, body) => body.freeVars
     case End(msg) => Set.empty
 
   lazy val freeVarsLLIR: Set[Local] = this match
@@ -121,6 +156,7 @@ sealed abstract class Block extends Product:
     case Define(defn, rest) => defn.freeVarsLLIR ++ (rest.freeVarsLLIR - defn.sym)
     case HandleBlock(lhs, res, par, args, cls, hdr, bod, rst) =>
       (bod.freeVarsLLIR - lhs) ++ rst.freeVarsLLIR ++ hdr.flatMap(_.freeVarsLLIR)
+    case Scoped(syms, body) => body.freeVarsLLIR
     case End(msg) => Set.empty
 
   lazy val subBlocks: Ls[Block] = this match
@@ -133,6 +169,7 @@ sealed abstract class Block extends Product:
     case Define(d, rest) => d.subBlocks ::: rest :: Nil
     case HandleBlock(_, _, par, args, _, handlers, body, rest) => par.subBlocks ++ args.flatMap(_.subBlocks) ++ handlers.map(_.body) :+ body :+ rest
     case Label(_, _, body, rest) => body :: rest :: Nil
+    case Scoped(_, body) => body :: Nil
 
     // TODO rm Lam from values and thus the need for these cases
     case Return(r, _) => r.subBlocks
@@ -147,10 +184,10 @@ sealed abstract class Block extends Product:
   // Note that this returns the definitions in reverse order, with the bottommost definiton appearing
   // last. This is so that using defns.foldLeft later to add the definitions to the front of a block, 
   // we don't need to reverse the list again to preserve the order of the definitions.
-  def floatOutDefns(
-      ignore: Defn => Bool = _ => false, 
-      preserve: Defn => Bool = _ => false
-    ) =
+  def extractDefns(
+        ignore: Defn => Bool = _ => false, 
+        preserve: Defn => Bool = _ => false
+      ): (Block, List[Defn]) =
     var defns: List[Defn] = Nil
     val transformer = new BlockTransformerShallow(SymbolSubst()):
       override def applyBlock(b: Block): Block = b match
@@ -164,6 +201,12 @@ sealed abstract class Block extends Product:
 
     (transformer.applyBlock(this), defns)
 
+  def gatherDefns(
+      ignore: Defn => Bool = _ => false, 
+      preserve: Defn => Bool = _ => false
+    ): List[Defn] = extractDefns(ignore, preserve)._2 // TODO: fix this very inefficient implementation
+
+
   lazy val flattened: Block = this.flatten(identity)
 
   private def flatten(k: End => Block): Block = this match
@@ -183,6 +226,13 @@ sealed abstract class Block extends Product:
       if (newBody is body) && (newRest is rest)
       then this
       else Label(label, loop, newBody, newRest)
+
+    // * Do not omit the `Begin`s that are used for nested scopes
+    case Begin(e: End, Scoped(syms, body)) =>
+      val newBody = body.flatten(k)
+      if newBody is body
+      then this
+      else new Begin(e, new Scoped(syms, newBody)) 
 
     case Begin(sub, rest) =>
       sub.flatten(_ => rest.flatten(k))
@@ -247,6 +297,12 @@ sealed abstract class Block extends Product:
       then this
       else HandleBlock(lhs, res, par, args, cls, newHandlers, newBody, newRest)
 
+    case Scoped(syms, body) =>
+      val newBody = body.flatten(k)
+      if newBody is body
+      then this
+      else Scoped(syms, newBody)
+
     case e: End => k(e)
     case t: BlockTail => this
 
@@ -272,6 +328,9 @@ case class Label(label: Local, loop: Bool, body: Block, rest: Block) extends Blo
 case class Break(label: Local) extends BlockTail
 case class Continue(label: Local) extends BlockTail
 
+
+case class Scoped(syms: collection.Set[Local], body: Block) extends BlockTail
+
 // TODO: remove this form?
 case class Begin(sub: Block, rest: Block) extends Block with ProductWithTail
 
@@ -286,6 +345,47 @@ case class AssignDynField(lhs: Path, fld: Path, arrayIdx: Bool, rhs: Result, res
 
 case class Define(defn: Defn, rest: Block) extends Block with ProductWithTail
 
+object Match:
+  def apply(scrut: Path, arms: Ls[Case -> Block], dflt: Opt[Block], rest: Block): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, Match(scrut, arms, dflt, body))
+    case _ => new Match(scrut, arms, dflt, rest)
+object Label:
+  def apply(label: Local, loop: Bool, body: Block, rest: Block): Block = rest match
+    case Scoped(syms, rest) => Scoped(syms, Label(label, loop, body, rest))
+    case _ => new Label(label, loop, body, rest)
+object Scoped:
+  def apply(syms: collection.Set[Local], body: Block): Block = body match
+    case Scoped(syms2, body) =>
+      if syms2.isEmpty && syms.isEmpty then Scoped(Set.empty, body) else Scoped(syms ++ syms2, body)
+    case _ =>
+      if syms.isEmpty then body else new Scoped(syms, body)
+object Begin:
+  def apply(sub: Block, rest: Block): Block = (sub, rest) match
+    case (Scoped(symsSub, bodySub), Scoped(symsRest, bodyRest)) =>
+      Scoped(symsSub ++ symsRest, Begin(bodySub, bodyRest))
+    case (Scoped(symsSub, bodySub), _) => Scoped(symsSub, Begin(bodySub, rest))
+    case (_, Scoped(symsRest, bodyRest)) => Scoped(symsRest, Begin(sub, bodyRest))
+    case _ => new Begin(sub, rest)
+object TryBlock:
+  def apply(sub: Block, finallyDo: Block, rest: Block): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, TryBlock(sub, finallyDo, body))
+    case _ => new TryBlock(sub, finallyDo, rest)
+object Assign:
+  def apply(lhs: Local, rhs: Result, rest: Block): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, Assign(lhs, rhs, body))
+    case _ => new Assign(lhs, rhs, rest)
+object AssignField:
+  def apply(lhs: Path, nme: Tree.Ident, rhs: Result, rest: Block)(symbol: Opt[MemberSymbol]): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, AssignField(lhs, nme, rhs, body)(symbol))
+    case _ => new AssignField(lhs, nme, rhs, rest)(symbol)
+object AssignDynField:
+  def apply(lhs: Path, fld: Path, arrayIdx: Bool, rhs: Result, rest: Block): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, AssignDynField(lhs, fld, arrayIdx, rhs, body))
+    case _ => new AssignDynField(lhs, fld, arrayIdx, rhs, rest)
+object Define:
+  def apply(defn: Defn, rest: Block): Block = rest match
+    case Scoped(syms, body) => Scoped(syms, Define(defn, body))
+    case _ => new Define(defn, rest)
 
 case class HandleBlock(
     lhs: Local,
@@ -298,6 +398,40 @@ case class HandleBlock(
     rest: Block
 ) extends Block with ProductWithTail
 
+object HandleBlock:
+  def apply(
+    lhs: Local,
+    res: Local,
+    par: Path,
+    args: Ls[Path],
+    cls: ClassSymbol,
+    handlers: Ls[Handler],
+    body: Block,
+    rest: Block
+  ) = rest match
+      case Scoped(syms, rest) =>
+        Scoped(
+          syms,
+          new HandleBlock(
+            lhs,
+            res,
+            par,
+            args,
+            cls,
+            handlers,
+            body,
+            rest
+          ))
+      case _ => new HandleBlock(
+        lhs,
+        res,
+        par,
+        args,
+        cls,
+        handlers,
+        body,
+        rest)
+
 
 sealed abstract class Defn:
   val innerSym: Opt[MemberSymbol]

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

@@ -95,6 +95,9 @@ class BlockTransformer(subst: SymbolSubst):
             if (lhs2 is lhs) && (fld2 is fld) && (rhs2 is rhs) && (rest2 is rest)
             then b
             else AssignDynField(lhs2, fld2, arrayIdx, rhs2, rest2)
+    case Scoped(s, bd) =>
+      val nb = applySubBlock(bd)
+      if nb is bd then b else Scoped(s, nb)
 
   def applyRcdArg(rcdArg: RcdArg)(k: RcdArg => Block): Block =
     val RcdArg(idx, p) = rcdArg

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

@@ -51,6 +51,7 @@ class BlockTraverser:
       applyResult(rhs)
       applyPath(fld)
       applySubBlock(rest)
+    case Scoped(_, body) => applySubBlock(body)
 
   def applyResult(r: Result): Unit = r match
     case r @ Call(fun, args) => applyPath(fun); args.foreach(applyArg)

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

@@ -420,6 +420,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       // ignored cases
       case TryBlock(sub, finallyDo, rest) => ??? // ignore
       case Throw(_) => PartRet(blk, Nil)
+      case Scoped(_, body) => go(body)
       case _: HandleBlock => lastWords("unexpected handleBlock") // already translated at this point
 
     val PartRet(head, states) = go(blk)(using labelIds, N)
@@ -581,7 +582,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
   private def thirdPass(b: Block): Block =
     // to ensure the fun and class references in the continuation class are properly scoped,
     // we move all function defns to the top level of the handler block
-    val (blk, defns) = b.floatOutDefns()
+    val (blk, defns) = b.extractDefns()
     defns.foldLeft(blk)((acc, defn) => Define(defn, acc))
 
   private def locToStr(l: Loc): Str =

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

@@ -546,8 +546,10 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
 
   extension (b: Block)
     private def floatOut(ctx: LifterCtx) =
-      b.floatOutDefns(preserve = defn => ctx.isModOrObj(defn.sym) || ctx.ignored(defn.sym))
-
+      b.extractDefns(preserve = defn => ctx.isModOrObj(defn.sym) || ctx.ignored(defn.sym))
+    private def gather(ctx: LifterCtx) =
+      b.gatherDefns(preserve = defn => ctx.isModOrObj(defn.sym) || ctx.ignored(defn.sym))
+
 
   def createLiftInfoCont(d: Defn, parentCls: Opt[ClsLikeDefn], ctx: LifterCtx): Map[BlockMemberSymbol, LiftedInfo] =
     val AccessInfo(accessed, _, refdDefns) = ctx.getAccesses(d.sym)
@@ -600,7 +602,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     defns.flatMap(createLiftInfoCont(_, N, ctx.addFnLocals(ctx.usedLocals(f.sym)))).toMap
 
   def createLiftInfoCls(c: ClsLikeDefn, ctx: LifterCtx): Map[BlockMemberSymbol, LiftedInfo] =
-    val defns = c.preCtor.floatOut(ctx)._2 ++ c.ctor.floatOut(ctx)._2 ++ c.companion.fold(Nil)(_.ctor.floatOut(ctx)._2)
+    val defns = c.preCtor.gather(ctx) ++ c.ctor.gather(ctx) ++ c.companion.fold(Nil)(_.ctor.gather(ctx))
     val newCtx = if (c.companion.isDefined) && !ctx.ignored(c.sym) then ctx else ctx.addClsDefn(c)
     val staticMtdInfo = c.companion.fold(Map.empty):
       case value => value.methods.flatMap(f => createLiftInfoFn(f, newCtx))