Runtime rewrite for handlers

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

@@ -68,13 +68,16 @@ object Config:
     // Whether we check `Instantiate` nodes for effects. Currently, effects cannot be raised in constructors.
     checkInstantiateEffect: Bool = false,
     // A debug option that allows codegen to continue even if an unlifted definition is encountered.
-    softLifterError: Bool = false
+    softLifterError: Bool = false,
+    // Skips instrumenting module constructors, this can be used when the file is statically known to not
+    // raise any effect and cannot use the Runtime.mls module during module construction due to cyclic dependency.
+    doNotInstrumentTopLevelModCtor: Bool = false,
   )
 
   case class StackSafety(stackLimit: Int)
   object StackSafety:
     val default: StackSafety = StackSafety(
-      stackLimit = 500,
+      stackLimit = 1000,
     )
 
   case class LiftDefns() // there may be other settings in the future, having it as a case class now

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

@@ -83,12 +83,18 @@ object HandlerLowering:
   private enum HandlerCtx:
     case FunctionLike(ctx: FunctionCtx)
     case Ctor
-    case ModCtor
+    case ModCtor(trulyNested: Bool)
     case TopLevel
 
-    def inCtor = this === Ctor || this === ModCtor
+    def inCtor = this === Ctor || this.isInstanceOf[ModCtor]
     def inTopLevel = this === TopLevel
-    def allowDefn = inTopLevel || this === ModCtor
+    def allowDefn = inTopLevel || this.isInstanceOf[ModCtor]
+    def innerDefIsTrulyNested = this match
+      case FunctionLike(_) => true
+      case Ctor => true
+      case ModCtor(trulyNested) => trulyNested
+      case TopLevel => false
+
 
   // currentFun: path to the current function for resumption
   // thisPath: path to `this` binding if the function is a method, `this` will be rebinded on resumption
@@ -125,8 +131,6 @@ import HandlerLowering.*
 
 class HandlerPaths(using Elaborator.State):
   val runtimePath: Path = State.runtimeSymbol.asPath
-  val effectSigPath: Path = runtimePath.selSN("EffectSig").selSN("class")
-  val effectSigSym: ClassSymbol = State.effectSigSymbol
   val contClsPath: Path = runtimePath.selSN("FunctionContFrame").selSN("class")
   val mkEffectPath: Path = runtimePath.selSN("mkEffect")
   val handleBlockImplPath: Path = runtimePath.selSN("handleBlockImpl")
@@ -138,14 +142,14 @@ class HandlerPaths(using Elaborator.State):
   val stackDepthPath: Path = runtimePath.selN(stackDepthIdent)
   val fnLocalsPath: Path = runtimePath.selSN("FnLocalsInfo").selSN("class")
   val localVarInfoPath: Path = runtimePath.selSN("LocalVarInfo").selSN("class")
-  val unwindPath: Path = runtimePath.selSN("unwind")
   val curEffect: Path = runtimePath.selSN("curEffect")
+  val unwindPath: Path = runtimePath.selSN("unwind")
   val resumePc: Path = runtimePath.selSN("resumePc")
   val resumeIdx: Path = runtimePath.selSN("resumeIdx")
   val resumeValueIdent = new Tree.Ident("resumeValue")
   val resumeValue: Path = runtimePath.selN(resumeValueIdent)
 
-type StackSafetyMap = collection.Map[FnOrCls, Block]
+type StackSafetyMap = collection.Map[FnOrCls, (Int, Block)]
 
 class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise, Elaborator.State, Elaborator.Ctx):
 
@@ -474,7 +478,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       .toList
       .map(locals(_))
 
-  val stackSafetyMap: mutable.Map[FnOrCls, Block] = mutable.HashMap.empty
+  val stackSafetyMap: mutable.Map[FnOrCls, (Int, Block)] = mutable.HashMap.empty
 
   private def lifterReport(using Line, FileName)(msgs: Ls[Message -> Opt[Loc]])(using Name) =
     if opt.softLifterError then
@@ -524,13 +528,13 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
             raise(lifterReport(msg"Unexpected nested class: lambdas may not function correctly." -> isym.toLoc :: Nil))
           val debugInfos = mutable.ArrayBuffer.empty[(Local, List[Arg])]
           val newMtds = methods.map: f =>
-            val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, Value.Ref(isym).sel(new Tree.Ident(f.sym.nme), f.sym.asTrm.get),
+            val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, Value.Ref(isym).sel(new Tree.Ident(f.sym.nme), f.dSym),
               S(Value.Ref(isym)), s"${sym.nme}#${f.sym.nme}")
             debugInfos += debugInfoSym -> debugInfo
             fun2
           val companion2 = companion.map: bod =>
             val newMtds = bod.methods.map: f =>
-              val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, Value.Ref(bod.isym).sel(new Tree.Ident(f.sym.nme), f.sym.asTrm.get),
+              val (debugInfoSym, debugInfo, fun2) = translateFunLike(f, Value.Ref(bod.isym).sel(new Tree.Ident(f.sym.nme), f.dSym),
                 S(Value.Ref(bod.isym)), s"${sym.nme}.${f.sym.nme}")
               debugInfos += debugInfoSym -> debugInfo
               fun2
@@ -539,7 +543,8 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
             // TODO: Companion's ctor is more well behaved so it is possible to handle it
             // However, JSBuilder inserts extra statements between preCtor and ctor and it's not possible to replicate the exact behavior
             // without many special handling.
-            val newCtor = translateCtorLike(bod.ctor, bod.isym.asPath, true)
+            val newCtor = if opt.doNotInstrumentTopLevelModCtor && !h.innerDefIsTrulyNested then bod.ctor else
+              translateCtorLike(bod.ctor, bod.isym.asPath, true)
             tl.log(s"companion name: ${bod.isym.nme}")
             ClsLikeBody(bod.isym, newMtds, bod.privateFields, bod.publicFields, newCtor)
           val c2 = ClsLikeDefn(owner, isym, sym, ctorSym, kind, paramsOpt, auxParams, parentPath, newMtds, privateFields, publicFields,
@@ -561,6 +566,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     val parts = partitionBlock(b)
     stackSafetyMap += ctx.resumeInfo.currentStackSafetySym ->
       (
+        1,
         ctx.doUnwind(ctx.resumeInfo.currentStackSafetySym.fold(_.toLoc, _.toLoc).fold(unit)(locToStr(_)), -1, Nil)(using paths)
       )
     if parts.states.size <= 1 then
@@ -616,7 +622,7 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
         mainLoop))
 
   private def translateCtorLike(b: Block, thisPath: Path, isModCtor: Bool)(using h: HandlerCtx): Block =
-    translateBlock(b, if isModCtor then HandlerCtx.ModCtor else HandlerCtx.Ctor, Set.empty)
+    translateBlock(b, if isModCtor then HandlerCtx.ModCtor(h.innerDefIsTrulyNested) else HandlerCtx.Ctor, Set.empty)
 
   private def translateIllegalEffectCtx(b: Block, onEffect: Call)(using HandlerCtx): Block =
     def effectCheck(l: Local, r: Result, rst: Block): Block =

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

@@ -326,7 +326,9 @@ class Lifter(topLevelBlk: Block)(using State, Raise, Config):
                   def join2: Block =
                     resolveDefnRef(l, d, r) match
                       case Some(value) => k(c.copy(fun = value, args = newArgs)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall).withLoc(c.toLoc))
-                      case None => super.applyResult(c)(k)
+                      case None => super.applyPath(c.fun): fun2 =>
+                        if (fun2 is c.fun) && (args is newArgs) then k(c)
+                        else k(c.copy(fun = fun2, args = newArgs)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall).withLoc(c.toLoc))
                   r match
                     // function call
                     case f: LiftedFunc => k(f.rewriteCall(c, newArgs))
@@ -978,7 +980,9 @@ class Lifter(topLevelBlk: Block)(using State, Raise, Config):
     private val aux = Lazy[Defn](mkAuxDefn)
 
     def rewriteCall(c: Call, args: List[Arg])(using ctx: LifterCtxNew): Call =
-      if isTrivial then c
+      if isTrivial then
+        if args is c.args then c
+        else c.copy(args = args)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall).withLocOf(c)
       else
         Call(
           Value.Ref(mainSym, S(mainDsym)),
@@ -1112,7 +1116,7 @@ class Lifter(topLevelBlk: Block)(using State, Raise, Config):
       if isTrivial then
         val path = Value.Ref(cls.sym, S(cls.isym))
         if (inst.cls === path) && (inst.args is args) then inst
-        else inst.copy(cls = path, args = args)
+        else inst.copy(cls = path, args = args).withLocOf(inst)
       else
         flat.force // force computation
         Call(
@@ -1124,7 +1128,7 @@ class Lifter(topLevelBlk: Block)(using State, Raise, Config):
       if obj.isObj then lastWords("tried to rewrite instantiate for an object")
       if isTrivial then
         if c.args is args then c
-        else c.copy(args = args)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall)
+        else c.copy(args = args)(c.isMlsFun, c.mayRaiseEffects, c.explicitTailCall).withLocOf(c)
       else
         flat.force // force computation
         Call(

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

@@ -535,7 +535,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
               k(Call(
                 Value.Ref(State.runtimeSymbol).selN(Tree.Ident(if isAnd then "short_and" else "short_or")),
                 Arg(N, ar1) :: Arg(N, lamDef.asPath) :: Nil
-              )(true, false, false)))
+              )(true, true, false)))
           else
             subTerm_nonTail(arg2): ar2 =>
               val target = wasmIntrinsicPath(sym, unary = false)
@@ -1063,7 +1063,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
 
     val stackSafe = config.stackSafety match
       case N => withHandlers2
-      case S(sts) => StackSafeTransform(sts.stackLimit, handlerPaths, stackSafetyInfo).transformTopLevel(withHandlers2)
+      case S(sts) => StackSafeTransform(sts.stackLimit, config.effectHandlers.exists(_.doNotInstrumentTopLevelModCtor), handlerPaths, stackSafetyInfo).transformTopLevel(withHandlers2)
 
     val flattened = stackSafe.flattened
 

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

@@ -107,7 +107,7 @@ object ScopeData:
           clsBody.privateFields.toSet + clsBody.isym
         case _: ClassCtor => Set.empty
         case Func(fun, _) => fun.params.flatMap: p =>
-            p.params.map(_.sym)
+            p.restParam.map(_.sym) ++ p.params.map(_.sym)
           .toSet
         case ScopedBlock(_, block) => block.syms.toSet
         case _: Loop => Set.empty

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

@@ -9,7 +9,7 @@ import hkmc2.semantics.*
 import hkmc2.syntax.Tree
 import hkmc2.codegen.HandlerLowering.FnOrCls
 
-class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, stackSafetyMap: StackSafetyMap)(using State):
+class StackSafeTransform(depthLimit: Int, doNotInstrumentTopLevelModCtor: Bool, paths: HandlerPaths, stackSafetyMap: StackSafetyMap)(using State):
   private val STACK_DEPTH_IDENT: Tree.Ident = Tree.Ident("stackDepth")
 
   private val runtimePath: Path = State.runtimeSymbol.asPath
@@ -78,7 +78,7 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, stackSafetyMap: S
 
         case _ => super.applyBlock(b)
 
-        override def applyHandler(hdr: Handler): Handler = lastWords("HandleBlock in stack safe transformation")
+      override def applyHandler(hdr: Handler): Handler = lastWords("HandleBlock in stack safe transformation")
 
       override def applyResult(r: Result)(k: Result => Block): Block =
         if usesStack(r) then
@@ -111,8 +111,8 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, stackSafetyMap: S
         methods.map(rewriteFn),
         privateFields,
         publicFields, 
-        rewriteBlk(preCtor, L(BlockMemberSymbol("TODO", Nil)), 1), // TODO: preCtor is not translated in handler lowering
-        if isTopLevel && (defn.k is syntax.Mod) then transformTopLevel(ctor) else rewriteBlk(ctor, R(isym), 1),
+        preCtor,
+        if isTopLevel && (defn.k is syntax.Mod) then transformTopLevel(ctor) else ctor,
         mod.map(rewriteObjBody(_, isTopLevel)),
         bufferable,
       )
@@ -123,13 +123,15 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, stackSafetyMap: S
       defn.methods.map(rewriteFn),
       defn.privateFields,
       defn.publicFields,
-      if isTopLevel then transformTopLevel(defn.ctor) else rewriteBlk(defn.ctor, R(defn.isym), 1),
+      if isTopLevel then
+        if doNotInstrumentTopLevelModCtor then defn.ctor else transformTopLevel(defn.ctor)
+      else rewriteBlk(defn.ctor, R(defn.isym)),
     )
 
   // fnOrCls points us to the doUnwind function
-  def rewriteBlk(blk: Block, fnOrCls: FnOrCls, increment: Int) =
+  def rewriteBlk(blk: Block, fnOrCls: FnOrCls) =
     (stackSafetyMap.get(fnOrCls), isTrivial(blk)) match
-    case (S(doUnwindBlk), false) =>
+    case (S((increment, doUnwindBlk)), false) =>
       var usedDepth = false
       lazy val curDepth =
         usedDepth = true
@@ -153,6 +155,6 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, stackSafetyMap: S
 
 
   def rewriteFn(defn: FunDefn) = 
-    FunDefn(defn.owner, defn.sym, defn.dSym, defn.params, rewriteBlk(defn.body, L(defn.sym), 1))(defn.forceTailRec)
+    FunDefn(defn.owner, defn.sym, defn.dSym, defn.params, rewriteBlk(defn.body, L(defn.sym)))(defn.forceTailRec)
 
   def transformTopLevel(b: Block) = transform(b, TempSymbol(N), true)

hkmc2/shared/src/main/scala/hkmc2/semantics/Elaborator.scala

@@ -258,7 +258,6 @@ object Elaborator:
     val blockSymbol = TempSymbol(N, "Block")
     val shapeSymbol = TempSymbol(N, "Shape")
     val wasmSymbol = TempSymbol(N, "wasm")
-    val effectSigSymbol = ClassSymbol(DummyTypeDef(syntax.Cls), Ident("EffectSig"))
     val nonLocalRetHandlerTrm =
       val id = new Ident("NonLocalReturn")
       val sym = ClassSymbol(DummyTypeDef(syntax.Cls), id)

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

@@ -277,4 +277,5 @@ fun fromStack(stack) = Iterable of () =>
 
 fun toStack(xs) = xs rightFolded of Nil, Cons
 
-fun isArrayLike = _ is Array | IterableBase
+// Important: This function is used by runtime directly. Do not call any function or make any lambda here.
+fun isArrayLike(xs) = xs is Array | IterableBase