refactor

Author: CAG2Mark

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

@@ -43,7 +43,7 @@ object HandlerLowering:
         e
       case Some(v) => v
 
-  private case class LinkState(res: Local, cls: Path, uid: StateId, doUnwind: Opt[LazyVal[Path]])
+  private case class LinkState(res: Local, cls: Path, uid: Path)
 
   type FnOrCls = Either[BlockMemberSymbol, MemberSymbol[? <: ClassLikeDef] & InnerSymbol]
 
@@ -78,6 +78,9 @@ object HandlerLowering:
     def topLevel(debugNme: Str) = DebugInfo(debugNme, Set.empty, N)
 
   type StateId = BigInt
+  
+  // TODO: move somewhere else, not sure where
+  def simpleParam(sym: VarSymbol) = Param(FldFlags.empty, sym, N, Modulefulness.none)
 
 import HandlerLowering.*
 
@@ -98,18 +101,24 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
 
   private def funcLikeHandlerCtx(ctorThis: Option[Path], isHandlerMtd: Bool, contNme: Str, debugNme: Str)(using h: HandlerCtx) =
     HandlerCtx(false, false, contNme, ctorThis, h.debugInfo.copy(debugNme), state =>
-      if state.doUnwind.isEmpty then
-        lastWords("doUnwind function does not exist for this function")
-      Return(Call(
-        state.doUnwind.get.get,
-        state.res.asPath.asArg :: Value.Lit(Tree.IntLit(state.uid)).asArg :: Nil
-      )(true, false), false)
-    )
+      blockBuilder
+        .assignFieldN(state.res.asPath.contTrace.last, nextIdent, Instantiate(
+          state.cls.selN(Tree.Ident("class")),
+          state.uid :: Nil))
+        .assignFieldN(state.res.asPath.contTrace, lastIdent, state.res.asPath.contTrace.last.next)
+        .ret(state.res.asPath))
   private def functionHandlerCtx(nme: Str, debugNme: Str)(using HandlerCtx) = funcLikeHandlerCtx(N, false, nme, debugNme)
-  private def topLevelCtx(nme: Str, debugNme: Str) = HandlerCtx(true, false, nme, N, DebugInfo.topLevel(debugNme), state => Assign(
-      state.res,
-      Call(paths.topLevelEffectPath, state.res.asPath.asArg :: Value.Lit(Tree.BoolLit(opt.debug)).asArg :: Nil)(true, false),
-      End()))
+  private def topLevelCall(state: LinkState) = Call(
+      paths.topLevelEffectPath, 
+      state.res.asPath.asArg :: Value.Lit(Tree.BoolLit(opt.debug)).asArg :: Nil
+    )(true, false)
+  private def topLevelCtx(nme: Str, debugNme: Str) = HandlerCtx(
+      true, false, nme, N, DebugInfo.topLevel(debugNme), 
+      state => Assign(
+        state.res,
+        topLevelCall(state),
+      End())
+    )
   private def ctorCtx(ctorThis: Path, nme: Str, debugNme: Str)(using HandlerCtx) = funcLikeHandlerCtx(S(ctorThis), false, nme, debugNme)
   private def handlerMtdCtx(nme: Str, debugNme: Str)(using HandlerCtx) = funcLikeHandlerCtx(N, true, nme, debugNme)
   private def handlerCtx(using HandlerCtx): HandlerCtx = summon
@@ -258,8 +267,6 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     val rewrittenEntry = rewriteState(entryState)
     val rewrittenStates = states.map(rewriteState)
 
-    println(finalDests)
-    
     (rewrittenEntry, rewrittenStates)
 
   def partitionBlock(blk: Block, inclEntryPoint: Bool, labelIds: Map[Symbol, (StateId, StateId)] = Map.empty): Ls[BlockState] =
@@ -507,17 +514,9 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
         doUnwindMap += fnOrCls -> doUnwindSym.asPath
         val pcSym = VarSymbol(Tree.Ident("pc"))
         val resSym = VarSymbol(Tree.Ident("res"))
-        val doUnwindBlk = blockBuilder
-          .assignFieldN(
-              resSym.asPath.contTrace.last, nextIdent, 
-              Instantiate(
-                cls.sym.asPath.selN(Tree.Ident("class")),
-                pcSym.asPath :: Nil
-              )
-            )
-          .assignFieldN(resSym.asPath.contTrace, lastIdent, resSym.asPath.contTrace.last.next)
-          .ret(resSym.asPath)
-        
+        val doUnwindBlk = h.linkAndHandle(
+          LinkState(resSym, cls.sym.asPath, pcSym.asPath)
+        )
         def simpleParam(sym: VarSymbol) = Param(FldFlags.empty, sym, N, Modulefulness.none)
         val doUnwindDef = FunDefn(
           N, doUnwindSym,
@@ -576,11 +575,6 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     val lbl = freshTmp("handlerBody")
     val lblLoop = freshTmp("handlerLoop")
 
-    val handlerBody = translateBlock(h.body, Set.empty, L(sym), HandlerCtx(false, true,
-      s"Cont$$handleBlock$$${symToStr(h.lhs)}$$", N, handlerCtx.debugInfo.copy(debugNme = s"‹handler body of ${h.lhs.nme}›"), state => blockBuilder
-        .assignFieldN(state.res.asPath.contTrace.last, nextIdent, PureCall(state.cls, Value.Lit(Tree.IntLit(state.uid)) :: Nil))
-        .ret(PureCall(paths.handleBlockImplPath, state.res.asPath :: h.lhs.asPath :: Nil))))
-    
     val handlerMtds = h.handlers.map: handler =>
       val sym = BlockMemberSymbol("hdlrFun", Nil, true)
       val mtdBdy = translateBlock(handler.body,
@@ -605,16 +599,30 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     // NOTE: the super call is inside the preCtor
     // during resumption we need to resume both the this.x = x bindings done in JSBuilder and the ctor
 
-    val body = blockBuilder
-      .define(clsDefn)
-      .assign(h.lhs, Instantiate(Value.Ref(clsDefn.sym), Nil))
-      .rest(handlerBody)
+    val handlerBody = translateBlock(
+      h.body, Set.empty, L(sym), 
+      HandlerCtx(
+        false, true,
+        s"Cont$$handleBlock$$${symToStr(h.lhs)}$$", N, 
+        handlerCtx.debugInfo.copy(debugNme = s"‹handler body of ${h.lhs.nme}›"), 
+        state => blockBuilder
+          .assignFieldN(state.res.asPath.contTrace.last, nextIdent, PureCall(state.cls, state.uid :: Nil))
+          .ret(PureCall(paths.handleBlockImplPath, state.res.asPath :: h.lhs.asPath :: Nil))
+      )
+    )
 
     val defn = FunDefn(
       N, // no owner
-      sym, PlainParamList(Nil) :: Nil, body)
+      sym, PlainParamList(Nil) :: Nil, handlerBody)
 
-    val result = Define(defn, ResultPlaceholder(h.res, freshId(), Call(sym.asPath, Nil)(true, true), h.rest))
+    // moved all defns outside
+    val result = blockBuilder
+      .define(defn)
+      .define(clsDefn)
+      .assign(h.lhs, Instantiate(Value.Ref(clsDefn.sym), Nil))
+      .rest(
+        ResultPlaceholder(h.res, freshId(), Call(sym.asPath, Nil)(true, true), h.rest)
+      )
     result
 
   private def genContClass(b: Block)(using h: HandlerCtx): Opt[ClsLikeDefn] =
@@ -644,7 +652,6 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
       .assignFieldN(resSym.asPath.contTrace.last, nextIdent, clsSym.asPath)
       .assignFieldN(resSym.asPath.contTrace, lastIdent, clsSym.asPath)
       .ret(resSym.asPath)
-    def simpleParam(sym: VarSymbol) = Param(FldFlags.empty, sym, N, Modulefulness.none)
     val doUnwindDef = FunDefn(
       S(clsSym), doUnwindSym,
       PlainParamList(simpleParam(resSym) :: simpleParam(newPcSym) :: Nil) :: Nil,
@@ -823,13 +830,16 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
   private def genNormalBody(b: Block, clsSym: BlockMemberSymbol, doUnwind: Opt[LazyVal[Path]])(using HandlerCtx): Block =
     val transform = new BlockTransformerShallow(SymbolSubst()):
       override def applyBlock(b: Block): Block = b match
-        case ResultPlaceholder(res, uid, c, rest) =>
+        case ResultPlaceholder(res, uid, c, rest) => 
+          val (doUnwindCall, implct) = doUnwind match
+            case None => (topLevelCall(LinkState(res, clsSym.asPath, Value.Lit(Tree.IntLit(uid)))), true)
+            case Some(doUnwind) => (PureCall(doUnwind.get, res.asPath :: Value.Lit(Tree.IntLit(uid)) :: Nil), false)
           blockBuilder
             .assign(res, c)
             .ifthen(
               res.asPath,
               Case.Cls(paths.effectSigSym, paths.effectSigPath),
-              handlerCtx.linkAndHandle(LinkState(res, clsSym.asPath, uid, doUnwind))
+              Return(doUnwindCall, implct)
             )
             .rest(applyBlock(rest))
         case _ => super.applyBlock(b)

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

@@ -225,6 +225,7 @@ let Runtime1;
     };
     this.stackLimit = 0;
     this.stackDepth = 0;
+    this.stackOffset = 0;
     this.stackHandler = null;
     this.stackResume = null;
     const StackDelayHandler$class = class StackDelayHandler {
@@ -777,24 +778,34 @@ let Runtime1;
     return tmp4
   } 
   static checkDepth() {
-    let scrut, tmp, tmp1;
-    tmp = Runtime.stackDepth >= Runtime.stackLimit;
-    tmp1 = Runtime.stackHandler !== null;
-    scrut = tmp && tmp1;
+    let scrut, tmp, tmp1, tmp2;
+    tmp = Runtime.stackDepth - Runtime.stackOffset;
+    tmp1 = tmp >= Runtime.stackLimit;
+    tmp2 = Runtime.stackHandler !== null;
+    scrut = tmp1 && tmp2;
     if (scrut === true) {
       return runtime.safeCall(Runtime.stackHandler.delay())
     } else {
       return runtime.Unit
     }
   } 
   static resetDepth(tmp, curDepth) {
+    let scrut, tmp1;
     Runtime.stackDepth = curDepth;
+    scrut = curDepth < Runtime.stackOffset;
+    if (scrut === true) {
+      Runtime.stackOffset = curDepth;
+      tmp1 = runtime.Unit;
+    } else {
+      tmp1 = runtime.Unit;
+    }
     return tmp
   } 
   static runStackSafe(limit, f1) {
     let result, scrut, saved, tmp1, tmp2, tmp3;
     Runtime.stackLimit = limit;
     Runtime.stackDepth = 1;
+    Runtime.stackOffset = 0;
     Runtime.stackHandler = Runtime.StackDelayHandler;
     tmp1 = Runtime.enterHandleBlock(Runtime.StackDelayHandler, f1);
     result = tmp1;
@@ -803,6 +814,7 @@ let Runtime1;
       if (scrut === true) {
         saved = Runtime.stackResume;
         Runtime.stackResume = null;
+        Runtime.stackOffset = Runtime.stackDepth;
         tmp2 = runtime.safeCall(saved());
         result = tmp2;
         tmp3 = runtime.Unit;
@@ -814,6 +826,7 @@ let Runtime1;
     }
     Runtime.stackLimit = 0;
     Runtime.stackDepth = 0;
+    Runtime.stackOffset = 0;
     Runtime.stackHandler = null;
     return result
   }

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

@@ -345,6 +345,7 @@ fun resumeContTrace(contTrace, value) =
 // stack safety
 mut val stackLimit = 0 // How deep the stack can go before heapifying the stack
 mut val stackDepth = 0 // Tracks the virtual + real stack depth
+mut val stackOffset = 0 // How much to offset stackDepth by to get the true stack depth (i.e. the virtual depth)
 mut val stackHandler = null
 mut val stackResume = null
 
@@ -353,29 +354,34 @@ object StackDelayHandler with
     set stackResume = k
 
 fun checkDepth() =
-  if stackDepth >= stackLimit && stackHandler !== null then
+  if stackDepth - stackOffset >= stackLimit && stackHandler !== null then
     // this is a tail call to effectful function
     stackHandler.delay()
   else
     ()
 
 fun resetDepth(tmp, curDepth) =
   set stackDepth = curDepth
+  if curDepth < stackOffset do
+    set stackOffset = curDepth
   tmp
 
 fun runStackSafe(limit, f) =
   set
     stackLimit = limit
     stackDepth = 1
+    stackOffset = 0
     stackHandler = StackDelayHandler
   let result = enterHandleBlock(StackDelayHandler, f)
   while stackResume !== null do
     let saved = stackResume
     set
       stackResume = null
+      stackOffset = stackDepth
       result = saved()
   set
     stackLimit = 0
     stackDepth = 0
+    stackOffset = 0
     stackHandler = null
   result