Stack safety overhaul and effect handlers debloating MK1 (hkust-taco#307)

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

Author: 3 people

…src/main/scala/mlscript/utils/Lazy.scala core/shared/main/scala/utils/Lazy.scalashared/src/main/scala/mlscript/utils/Lazy.scala renamed to core/shared/main/scala/utils/Lazy.scala shared/src/main/scala/mlscript/utils/Lazy.scala renamed to core/shared/main/scala/utils/Lazy.scala

@@ -16,6 +16,7 @@ class Eager[+A](val value: A) extends Box[A] {
 
 class Lazy[A](thunk: => A) extends Box[A] {
   def isComputing = _isComputing
+  def isEmpty: Bool = _value.isEmpty
   private var _isComputing = false
   private var _value: Opt[A] = N
   def get = if (_isComputing) N else S(get_!)

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

@@ -151,6 +151,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     * @param localPaths The path to access a particular local (possibly belonging to a previous function) in the current scope
     * @param iSymPaths The path to access a particular `innerSymbol` (possibly belonging to a previous class) in the current scope
     * @param replacedDefns Ignored (unlifted) definitions that have been rewritten and need to be replaced at the definition site.
+    * @param firstClsFns Nested functions which are used as first-class functions.
     * @param companionMap Map from companion object symbols to the corresponding regular class symbol.
     */
   case class LifterCtx private (
@@ -173,6 +174,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     val localPaths: Map[Local, LocalPath] = Map.empty,
     val isymPaths: Map[InnerSymbol, LocalPath] = Map.empty,
     val replacedDefns: Map[BlockMemberSymbol, Defn] = Map.empty,
+    val firstClsFns: Set[BlockMemberSymbol] = Set.empty,
     val companionMap: Map[InnerSymbol, InnerSymbol] = Map.empty
   ):
     // gets the function to which a local belongs
@@ -197,6 +199,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     def withNestedDefns(mp: Map[BlockMemberSymbol, List[Defn]]) = copy(nestedDefns = mp)
     def withAccesses(mp: Map[BlockMemberSymbol, AccessInfo]) = copy(accessInfo = mp)
     def withInScopes(mp: Map[BlockMemberSymbol, Set[BlockMemberSymbol]]) = copy(inScopeDefns = mp)
+    def withFirstClsFns(fns: Set[BlockMemberSymbol]) = copy(firstClsFns = fns)
     def withCompanionMap(mp: Map[InnerSymbol, InnerSymbol]) = copy(companionMap = mp)
     def addFnLocals(f: FreeVars) = copy(prevFnLocals = prevFnLocals ++ f)
     def addClsDefn(c: ClsLikeDefn) = copy(prevClsDefns = c :: prevClsDefns)
@@ -267,7 +270,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
 
     val sortedVars = cap.toArray.sortBy(_.uid).map: sym =>
       val id = fresh.make
-      val nme = sym.nme + id + "$"
+      val nme = sym.nme + "$capture$" + id
 
       val ident = new Tree.Ident(nme)
       val varSym = VarSymbol(ident)
@@ -367,10 +370,18 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     val extraDefns: List[Defn],
   )
 
+  private case class LifterMetadata(
+    unliftable: Set[BlockMemberSymbol],
+    modules: List[ClsLikeDefn],
+    objects: List[ClsLikeDefn],
+    firstClsFns: Set[BlockMemberSymbol]
+  )
+  
   // d is a top-level definition
   // returns (ignored classes, modules, objects)
-  def createMetadata(d: Defn, ctx: LifterCtx): (Set[BlockMemberSymbol], List[ClsLikeDefn], List[ClsLikeDefn]) =
+  private def createMetadata(d: Defn, ctx: LifterCtx): LifterMetadata =
     var ignored: Set[BlockMemberSymbol] = Set.empty
+    var firstClsFns: Set[BlockMemberSymbol] = Set.empty
     var unliftable: Set[BlockMemberSymbol] = Set.empty
     var clsSymToBms: Map[Local, BlockMemberSymbol] = Map.empty
     var modules: List[ClsLikeDefn] = Nil
@@ -483,6 +494,10 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
           applyBlock(ctor)
           mod.foreach(applyClsLikeBody)
 
+      def isFun(d: Defn) = d match
+        case FunDefn(owner, sym, params, body) => true
+        case _ => false
+      
       override def applyValue(v: Value): Unit = v match
         case RefOfBms(l) if clsSyms.contains(l) && !modOrObj(ctx.defns(l)) =>
           raise(WarningReport(
@@ -491,6 +506,9 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
           ))
           ignored += l
           unliftable += l
+        case RefOfBms(l) if ctx.defns.contains(l) && isFun(ctx.defns(l)) =>
+          // naked reference to a function definition
+          firstClsFns += l
         case _ => super.applyValue(v)
 
     // analyze the extends graph
@@ -512,8 +530,8 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
         dfs(b)
     for s <- ignored do
       dfs(s)
-  
-    (ignored ++ newUnliftable, modules.toList, objects.toList)
+    
+    LifterMetadata(ignored ++ newUnliftable, modules.toList, objects.toList, firstClsFns)
 
   extension (b: Block)
     private def floatOut(ctx: LifterCtx) =
@@ -902,7 +920,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
 
       d match
         case f: FunDefn =>
-          val createSym = (nme: String) =>
+          def createSym(nme: String) =
             val vsym = VarSymbol(Tree.Ident(nme))
             (vsym, LocalPath.Sym(vsym))
           val (extraParams, newCtx, _) = createSymbolsUpdateCtx(createSym)
@@ -932,18 +950,17 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
           val mainDefn = FunDefn(f.owner, f.sym, PlainParamList(extraParamsCpy) :: headPlistCopy :: Nil, bdy)
           val auxDefn = FunDefn(N, singleCallBms, flatPlist, lifted.body)
 
-
-          Lifted(mainDefn, auxDefn :: extras)
+          if ctx.firstClsFns.contains(f.sym) then
+            Lifted(mainDefn, auxDefn :: extras)
+          else
+            Lifted(auxDefn, extras) // we can include just the flattened defn
         case c: ClsLikeDefn =>
-          val createSym: String => (VarSymbol, LocalPath.PubField) = 
-            // due to the possibility of capturing a TempSymbol in HandlerLowering, it is necessary to generate a discriminator
-            val fresh = FreshInt()
-            (nme: String) =>
-              val id = fresh.make
-              (
-                VarSymbol(Tree.Ident(nme + "$" + id)),
-                LocalPath.PubField(c.isym, BlockMemberSymbol(nme + "$" + id, Nil, true))
-              )
+          val fresh = FreshInt()
+          def createSym(nme: String): (VarSymbol, LocalPath.PubField) = 
+            (
+              VarSymbol(Tree.Ident(nme)),
+              LocalPath.PubField(c.isym, BlockMemberSymbol(nme, Nil, true))
+            )
           val (extraParams, newCtx, flds) = createSymbolsUpdateCtx(createSym)
 
           // add aux params, private fields, update preCtor
@@ -1263,7 +1280,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
     val walker1 = new BlockTransformerShallow(SymbolSubst()):
       override def applyBlock(b: Block): Block = b match
         case Define(d, rest) =>
-          val (unliftable, modules, objects) = createMetadata(d, ctx)
+          val LifterMetadata(unliftable, modules, objects, firstClsFns) = createMetadata(d, ctx)
 
           val modObjLocals = (modules ++ objects).map: c =>
               analyzer.nestedIn.get(c.sym) match
@@ -1280,6 +1297,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
           val ctxx = ctx
             .addIgnored(unliftable)
             .withModObjLocals(modObjLocals)
+            .withFirstClsFns(firstClsFns)
 
           val Lifted(lifted, extra) = d match
             case f: FunDefn => 

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

@@ -1029,13 +1029,15 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
     val desug = LambdaRewriter.desugar(blk)
 
     val handlerPaths = new HandlerPaths
+    
+    val (withHandlers, doUnwindPaths) = config.effectHandlers.fold((desug, Map.empty)): opt =>
+      HandlerLowering(handlerPaths, opt).translateTopLevel(desug)
+      
     val stackSafe = config.stackSafety match
-      case N => desug
-      case S(sts) => StackSafeTransform(sts.stackLimit, handlerPaths).transformTopLevel(desug)
-    val withHandlers = config.effectHandlers.fold(stackSafe): opt =>
-      HandlerLowering(handlerPaths, opt).translateTopLevel(stackSafe)
+      case N => withHandlers
+      case S(sts) => StackSafeTransform(sts.stackLimit, handlerPaths, doUnwindPaths).transformTopLevel(withHandlers)
 
-    val flattened = withHandlers.flattened
+    val flattened = stackSafe.flattened
 
     val lifted = 
       if lift then Lifter(S(handlerPaths)).transform(flattened)

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

@@ -7,13 +7,18 @@ import hkmc2.codegen.*
 import hkmc2.semantics.Elaborator.State
 import hkmc2.semantics.*
 import hkmc2.syntax.Tree
+import hkmc2.codegen.HandlerLowering.FnOrCls
 
-class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
+class StackSafeTransform(depthLimit: Int, paths: HandlerPaths, doUnwindMap: Map[FnOrCls, Path])(using State):
   private val STACK_DEPTH_IDENT: Tree.Ident = Tree.Ident("stackDepth")
+  
+  val doUnwindFns = doUnwindMap.values.collect:
+      case s: Select if s.symbol.isDefined => s.symbol.get
+      case Value.Ref(sym) => sym
+    .toSet
 
   private val runtimePath: Path = State.runtimeSymbol.asPath
   private val checkDepthPath: Path = runtimePath.selN(Tree.Ident("checkDepth"))
-  private val resetDepthPath: Path = runtimePath.selN(Tree.Ident("resetDepth"))
   private val runStackSafePath: Path = runtimePath.selN(Tree.Ident("runStackSafe"))
   private val stackDepthPath: Path = runtimePath.selN(STACK_DEPTH_IDENT)
 
@@ -24,18 +29,13 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
 
   // Increases the stack depth, assigns the call to a value, then decreases the stack depth
   // then binds that value to a desired block
-  def extractRes(res: Result, isTailCall: Bool, f: Result => Block, sym: Option[Symbol], curDepth: => Symbol) =
-    if isTailCall then
-      blockBuilder
-        .assignFieldN(runtimePath, STACK_DEPTH_IDENT, op("+", stackDepthPath, intLit(1)))
-        .ret(res)
+  def extractRes(res: Result, isTailCall: Bool, f: Result => Block, sym: Option[Symbol], curDepth: => Symbol): Block =
+    if isTailCall then Return(res, false)
     else
       val tmp = sym getOrElse TempSymbol(None, "tmp")
-      val offsetGtDepth = TempSymbol(None, "offsetGtDepth")
       blockBuilder
-        .assignFieldN(runtimePath, STACK_DEPTH_IDENT, op("+", stackDepthPath, intLit(1)))
         .assign(tmp, res)
-        .assign(tmp, Call(resetDepthPath, tmp.asPath.asArg :: curDepth.asPath.asArg :: Nil)(true, false))
+        .assignFieldN(runtimePath, STACK_DEPTH_IDENT, curDepth.asPath)
         .rest(f(tmp.asPath))
 
   def wrapStackSafe(body: Block, resSym: Local, rest: Block) =
@@ -51,6 +51,7 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
   def transform(b: Block, curDepth: => Symbol, isTopLevel: Bool = false): Block =
     def usesStack(r: Result) = r match
       case Call(Value.Ref(_: BuiltinSymbol), _) => false
+      case c: Call if !c.mayRaiseEffects => false // a call can only trigger a stack delay if it can raise effects
       case _: Call | _: Instantiate => true
       case _ => false
 
@@ -75,39 +76,20 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
               extract(r, false, _ => applyBlock(rest), S(lhs), curDepth)
           else
             super.applyBlock(b)
-        case HandleBlock(l, res, par, args, cls, hdr, bod, rst) =>
-          val l2 = applyLocal(l)
-          val res2 = applyLocal(res)
-          applyPath(par): par2 =>
-            applyListOf(args, applyPath(_)(_)): args2 =>
-              val cls2 = cls.subst
-              val hdr2 = hdr.mapConserve(applyHandler)
-              val bod2 = rewriteBlk(bod)
-              val rst2 = applyBlock(rst)
-              if isTopLevel then
-                val newRes = TempSymbol(N, "res")
-                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)
+        
+        case HandleBlock(l, res, par, args, cls, hdr, bod, rst) => lastWords("HandleBlock in stack safe transformation")
 
         case _ => super.applyBlock(b)
 
-        override def applyHandler(hdr: Handler): Handler =
-          val sym2 = hdr.sym.subst
-          val resumeSym2 = hdr.resumeSym.subst
-          val params2 = hdr.params.mapConserve(applyParamList)
-          val body2 = rewriteBlk(hdr.body)
-          Handler(sym2, resumeSym2, params2, body2)
+        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
           extract(r, false, k, N, curDepth)
         else
           super.applyResult(r)(k)
 
-      override def applyLam(lam: Lambda): Lambda =
-        Lambda(lam.params, rewriteBlk(lam.body))
+      override def applyLam(lam: Lambda): Lambda = lastWords("Lambda in stack safe transformation")
 
     transform.applyBlock(b)
 
@@ -121,40 +103,75 @@ class StackSafeTransform(depthLimit: Int, paths: HandlerPaths)(using State):
         case _ => ()
     trivial
 
-  def rewriteCls(defn: ClsLikeDefn, isTopLevel: Bool): ClsLikeDefn =
-    val ClsLikeDefn(owner, isym, sym, k, paramsOpt, auxParams,
-      parentPath, methods, privateFields, publicFields, preCtor, ctor, mod) = defn
-    ClsLikeDefn(
-      owner, isym, sym, k, paramsOpt, auxParams, parentPath,
-      methods.map(rewriteFn),
-      privateFields,
-      publicFields, rewriteBlk(preCtor),
-      rewriteBlk(ctor),
-      mod.map(rewriteObjBody(_, isTopLevel)),
-    )
+  def rewriteCls(defn: ClsLikeDefn, isTopLevel: Bool): ClsLikeDefn = defn.parentPath match
+    case Some(value) if value eq paths.contClsPath => defn
+    case _ =>
+      val ClsLikeDefn(owner, isym, sym, k, paramsOpt, auxParams,
+        parentPath, methods, privateFields, publicFields, preCtor, ctor, mod) = defn
+      ClsLikeDefn(
+        owner, isym, sym, k, paramsOpt, auxParams, parentPath,
+        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),
+        mod.map(rewriteObjBody(_, isTopLevel)),
+      )
 
   def rewriteObjBody(defn: ClsLikeBody, isTopLevel: Bool): ClsLikeBody =
     ClsLikeBody(
       defn.isym,
       defn.methods.map(rewriteFn),
       defn.privateFields,
       defn.publicFields,
-      if isTopLevel then transformTopLevel(defn.ctor) else rewriteBlk(defn.ctor),
+      if isTopLevel then transformTopLevel(defn.ctor) else rewriteBlk(defn.ctor, R(defn.isym), 1),
     )
 
-  def rewriteBlk(blk: Block) =
-    val curDepth =
+  // fnOrCls points us to the doUnwind function
+  def rewriteBlk(blk: Block, fnOrCls: FnOrCls, increment: Int) =
+    var usedDepth = false
+    lazy val curDepth =
+      usedDepth = true
       TempSymbol(None, "curDepth")
+      
+    val doUnwindPath = doUnwindMap.get(fnOrCls)
     val newBody = transform(blk, curDepth)
+    
     if isTrivial(blk) then
       newBody
-    else
+    else if doUnwindPath.isEmpty then
       val resSym = TempSymbol(None, "stackDelayRes")
       blockBuilder
-        .assign(curDepth, stackDepthPath)
+        .staticif(usedDepth, _.assign(curDepth, stackDepthPath))
+        .rest(newBody)
+    else
+      val resSym = TempSymbol(None, "stackDelayRes")
+      val rewritten = blockBuilder
+        .staticif(usedDepth, _.assign(curDepth, stackDepthPath))
+        .assignFieldN(runtimePath, STACK_DEPTH_IDENT, op("+", stackDepthPath, intLit(increment)))
         .assign(resSym, Call(checkDepthPath, Nil)(true, true))
+        .ifthen(
+          resSym.asPath,
+          Case.Cls(paths.effectSigSym, paths.effectSigPath),
+          Return(
+            Call(doUnwindPath.get, resSym.asPath.asArg :: intLit(0).asArg :: Nil)(true, false),
+            false
+          )
+        )
         .rest(newBody)
-  
-  def rewriteFn(defn: FunDefn) = FunDefn(defn.owner, defn.sym, defn.params, rewriteBlk(defn.body))
+      // Float out defns, including the doUnwind function, so that they appear at the top of the block
+      // This is because the doUnwind function must appear before the checks inserted by the stack
+      // safety pass.
+      // However, due to how tightly coupled the stack safety and handler lowering are, it might be
+      // better to simply merge the two passes in the future.
+      val (blk, defns) = doUnwindPath.get match
+        case Value.Ref(sym) => rewritten.floatOutDefns()
+        case _ => (rewritten, Nil)
+      defns.foldLeft(blk)((acc, defn) => Define(defn, acc))
+
+     
+  def rewriteFn(defn: FunDefn) = 
+    if doUnwindFns.contains(defn.sym) then defn
+    else FunDefn(defn.owner, defn.sym, defn.params, rewriteBlk(defn.body, L(defn.sym), 1))
 
   def transformTopLevel(b: Block) = transform(b, TempSymbol(N), true)

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

@@ -173,7 +173,7 @@ class UsedVarAnalyzer(b: Block, handlerPaths: Opt[HandlerPaths])(using State):
         blkAccessesShallow(f.body).withoutLocals(fVars)
       case c: ClsLikeDefn =>
         val methodSyms = c.methods.map(_.sym).toSet
-        c.methods.foldLeft(blkAccessesShallow(c.preCtor) ++ blkAccessesShallow(c.ctor)):
+        val ret = c.methods.foldLeft(blkAccessesShallow(c.preCtor) ++ blkAccessesShallow(c.ctor)):
           case (acc, fDefn) =>
             // class methods do not need to be lifted, so we don't count calls to their methods.
             // a previous reference to this class's block member symbol is enough to assume any
@@ -182,6 +182,11 @@ class UsedVarAnalyzer(b: Block, handlerPaths: Opt[HandlerPaths])(using State):
             // however, we must keep references to the class itself!
             val defnAccess = findAccessesShallow(fDefn)
             acc ++ defnAccess.withoutBms(methodSyms)
+        if c.parentPath.isDefined && isHandlerClsPath(c.parentPath.get) then
+          // for continuation classes, treat them like they only read variables
+          AccessInfo(ret.accessed ++ ret.mutated, Set.empty, ret.refdDefns)
+        else
+          ret
       case _: ValDefn => AccessInfo.empty
 
     accessedCache.getOrElseUpdate(defn.sym, create)
@@ -333,16 +338,7 @@ class UsedVarAnalyzer(b: Block, handlerPaths: Opt[HandlerPaths])(using State):
 
         def handleCalledBms(called: BlockMemberSymbol): Unit = defnSyms.get(called) match
           case None => ()
-          case Some(defn) =>
-            // special case continuation classes
-            defn match
-              case c: ClsLikeDefn => c.parentPath match
-                case S(path) if isHandlerClsPath(path) => return
-                    // treat the continuation class as if it does not exist
-                case _ => ()
-              case _ => ()
-            
-
+          case Some(defn) => 
             val AccessInfo(accessed, muted, refd) = accessMap(defn.sym)
             val muts = muted.intersect(thisVars)
             val reads = defn.freeVars.intersect(thisVars) -- muts