Stack safe recursion

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

@@ -288,24 +288,24 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
         case b: HandleBlock =>
           val rest = applyBlock(b.rest)
           translateHandleBlock(b.copy(rest = rest))
-        case Return(c: Call, implct) if handlerCtx.isHandleFree =>
-          val fun2 = applyPath(c.fun)
-          val args2 = c.args.map(applyArg)
-          val c2 = if (fun2 is c.fun) && (args2 zip c.args).forall(_ is _) then c else Call(fun2, args2)(c.isMlsFun)
+        case Return(c @ Call(fun, args), implct) if handlerCtx.isHandleFree && !handlerCtx.isTopLevel =>
+          val fun2 = applyPath(fun)
+          val args2 = args.map(applyArg)
+          val c2 = if (fun2 is fun) && (args2 zip args).forall(_ is _) then c else Call(fun2, args2)(c.isMlsFun)
           if c2 is c then b else Return(c2, implct)
         case _ => super.applyBlock(b)
       override def applyResult2(r: Result)(k: Result => Block): Block = r match
         case r @ Call(Value.Ref(_: BuiltinSymbol), _) => super.applyResult2(r)(k)
         case c @ Call(fun, args) =>
           val res = freshTmp("res")
           val fun2 = applyPath(fun)
-          val args2 = c.args.map(applyArg)
+          val args2 = args.map(applyArg)
           val c2 = if (fun2 is fun) && (args2 zip args).forall(_ is _) then c else Call(fun2, args2)(c.isMlsFun)
           ResultPlaceholder(res, freshId(), false, c2, k(Value.Ref(res)))
         case c @ Instantiate(cls, args) =>
           val res = freshTmp("res")
           val cls2 = applyPath(cls)
-          val args2 = c.args.map(applyPath)
+          val args2 = args.map(applyPath)
           val c2 = if (cls2 is cls) && (args2 zip args).forall(_ is _) then c else Instantiate(cls2, args2)
           ResultPlaceholder(res, freshId(), false, c2, k(Value.Ref(res)))
         case r => super.applyResult2(r)(k)
@@ -374,7 +374,6 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
         case b1 -> b2 => b1.asCls match
           case Some(value) => 
             val newSym = ClassSymbol(value.tree, Tree.Ident(b2.nme))
-            newSym.defn = value.defn
             S(value -> newSym)
           case None => None
       .collect:
@@ -386,7 +385,6 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
         case b1 -> b2 => b1.asMod match
           case Some(value) => 
             val newSym = ModuleSymbol(value.tree, Tree.Ident(b2.nme))
-            newSym.defn = value.defn
             S(value -> newSym)
           case None => None
       .collect:
@@ -455,7 +453,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       h.cls,
       BlockMemberSymbol(h.cls.id.name, Nil),
       syntax.Cls,
-      h.cls.defn.get.paramsOpt,
+      N,
       S(h.par), handlers, Nil, Nil,
       Assign(freshTmp(), SimpleCall(Value.Ref(State.builtinOpsMap("super")), Nil), End()), End())
 
@@ -478,15 +476,6 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
     )
 
     val pcVar = VarSymbol(Tree.Ident("pc"))
-    clsSym.defn = S(ClassDef(
-      N,
-      syntax.Cls,
-      clsSym,
-      BlockMemberSymbol(clsSym.nme, Nil),
-      Nil,
-      S(PlainParamList(Param(FldFlags.empty, pcVar, N) :: Nil)),
-      ObjBody(Term.Blk(Nil, Term.Lit(Tree.UnitLit(true)))),
-      List()))
 
     var trivial = true
     def prepareBlock(b: Block): Block =
@@ -582,7 +571,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       clsSym,
       BlockMemberSymbol(clsSym.nme, Nil),
       syntax.Cls,
-      clsSym.defn.get.paramsOpt,
+      S(PlainParamList(Param(FldFlags.empty, pcVar, N) :: Nil)),
       S(contClsPath),
       resumeFnDef :: Nil,
       Nil,

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

@@ -12,11 +12,11 @@ class StackSafeTransform(depthLimit: Int)(using State):
   private val STACK_OFFSET_IDENT: Tree.Ident = Tree.Ident("__stackOffset")
   private val STACK_HANDLER_IDENT: Tree.Ident = Tree.Ident("__stackHandler")
 
-  private val stackDelayClsPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef")).selN(Tree.Ident("__StackDelay")).selN(Tree.Ident("class"))
-  private val stackDepthPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef")).selN(STACK_DEPTH_IDENT)
-  private val stackOffsetPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef")).selN(STACK_OFFSET_IDENT)
-  private val stackHandlerPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef")).selN(STACK_HANDLER_IDENT)
   private val predefPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef"))
+  private val stackDelayClsPath: Path = predefPath.selN(Tree.Ident("__StackDelay")).selN(Tree.Ident("class"))
+  private val stackDepthPath: Path = predefPath.selN(STACK_DEPTH_IDENT)
+  private val stackOffsetPath: Path = predefPath.selN(STACK_OFFSET_IDENT)
+  private val stackHandlerPath: Path = predefPath.selN(STACK_HANDLER_IDENT)
 
   private def intLit(n: BigInt) = Value.Lit(Tree.IntLit(n))
 
@@ -53,7 +53,6 @@ class StackSafeTransform(depthLimit: Int)(using State):
       Tree.TypeDef(syntax.Cls, Tree.Error(), N, N),
       Tree.Ident("StackDelay$")
     )
-    clsSym.defn = S(ClassDef(N, syntax.Cls, clsSym, BlockMemberSymbol(clsSym.nme, Nil), Nil, N, ObjBody(Term.Blk(Nil, Term.Lit(Tree.UnitLit(true)))), Nil))
 
     // the global stack handler is created here
     HandleBlock(
@@ -96,33 +95,25 @@ class StackSafeTransform(depthLimit: Int)(using State):
 
     val extract = if isTopLevel then extractResTopLevel else extractRes
 
-    val transform = new BlockTransformerShallow(SymbolSubst()):
-      override def applyDefn(defn: Defn): Defn = rewriteDefn(defn)
+    val transform = new BlockTransformer(SymbolSubst()):
+
+      override def applyFunDefn(fun: FunDefn): FunDefn = rewriteFn(fun)
+      override def applyDefn(defn: Defn): Defn = defn match
+        case defn: ClsLikeDefn => rewriteCls(defn)
+        case _: FunDefn | _: ValDefn => super.applyDefn(defn)
 
       override def applyBlock(b: Block): Block = b match
-        case Return(res, implct) if usesStack(res) => 
-          applyResult2(res): res =>
-            extract(res, true, Return(_, implct))
-        case Assign(lhs, rhs, rest) if usesStack(rhs) => 
-          applyResult2(rhs): res =>
-            extract(res, false, Assign(lhs, _, applyBlock(rest)))
-        case b @ AssignField(lhs, nme, rhs, rest) if usesStack(rhs) => 
-          applyResult2(rhs): res =>
-            extract(res, false, AssignField(lhs, nme, _, applyBlock(rest))(b.symbol))
-        case Define(defn, rest) => 
-          Define(rewriteDefn(defn), applyBlock(rest))
-        case HandleBlock(lhs, res, par, cls, handlers, body, rest) =>
-          HandleBlock(
-            lhs, res, par, cls, handlers.map(h => Handler(h.sym, h.resumeSym, h.params, applyBlock(h.body))),
-            applyBlock(body), applyBlock(rest)
-          )
-        case HandleBlockReturn(res) if usesStack(res) => 
-          applyResult2(res): res =>
-            extract(res, true, HandleBlockReturn(_))
+        case Return(res, implct) if usesStack(res) =>
+          extract(applyResult(res), true, Return(_, implct))
         case _ => super.applyBlock(b)
-
-      override def applyValue(v: Value): Value = 
-        v match
+
+      override def applyResult2(r: Result)(k: Result => Block): Block =
+        if usesStack(r) then
+          extract(r, false, k)
+        else
+          super.applyResult2(r)(k)
+
+      override def applyValue(v: Value): Value = v match
         case Value.Lam(params, body) => 
           Value.Lam(params, rewriteBlk(body))
         case _ => super.applyValue(v)
@@ -153,16 +144,13 @@ class StackSafeTransform(depthLimit: Int)(using State):
       case HandleBlockReturn(res) => true
       case End(msg) => true
 
-  def rewriteDefn(defn: Defn) = 
-    defn match
-    case d: FunDefn => rewriteFn(d)
-    case _: ValDefn => defn
-    case ClsLikeDefn(owner, isym, sym, k, paramsOpt, 
-      parentPath, methods, privateFields, publicFields, preCtor, ctor) =>
-      ClsLikeDefn(
-        owner, isym, sym, k, paramsOpt, parentPath, methods.map(rewriteFn), privateFields,
-        publicFields, rewriteBlk(preCtor), rewriteBlk(ctor)
-      )
+  def rewriteCls(defn: ClsLikeDefn): ClsLikeDefn = 
+    val ClsLikeDefn(owner, isym, sym, k, paramsOpt, 
+      parentPath, methods, privateFields, publicFields, preCtor, ctor) = defn
+    ClsLikeDefn(
+      owner, isym, sym, k, paramsOpt, parentPath, methods.map(rewriteFn), privateFields,
+      publicFields, rewriteBlk(preCtor), rewriteBlk(ctor)
+    )
 
   def rewriteBlk(blk: Block) =
     val newBody = transform(blk)
@@ -192,4 +180,4 @@ class StackSafeTransform(depthLimit: Int)(using State):
 
   def rewriteFn(defn: FunDefn) = FunDefn(defn.owner, defn.sym, defn.params, rewriteBlk(defn.body))
 
-  def transformTopLevel(b: Block) = transform(b, true)
+  def transformTopLevel(b: Block) = transform(b, true)

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

@@ -419,7 +419,7 @@ const Predef$class = class Predef {
   } 
   __resume(cur2, tail) {
     return (value) => {
-      let scrut, cont, scrut1, scrut2, tmp, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6;
+      let scrut, cont, scrut1, scrut2, scrut3, tmp, tmp1, tmp2, tmp3, tmp4, tmp5, tmp6, tmp7;
       scrut = cur2.resumed;
       if (scrut === true) {
         throw globalThis.Error("Multiple resumption");
@@ -428,43 +428,50 @@ const Predef$class = class Predef {
       }
       cur2.resumed = true;
       cont = cur2.next;
-      tmp7: while (true) {
+      tmp8: while (true) {
         if (cont instanceof this.__Cont.class) {
           tmp1 = cont.resume(value) ?? null;
           value = tmp1;
           if (value instanceof this.__EffectSig.class) {
-            value.tail = tail;
-            scrut1 = cur2.handleBlockList.next !== null;
+            scrut1 = value.tail.next === null;
             if (scrut1 === true) {
-              value.handleBlockList.tail.next = cur2.handleBlockList.next;
-              value.handleBlockList.tail = cur2.handleBlockList.tail;
+              value.tail.next = cont.next;
               tmp2 = null;
             } else {
               tmp2 = null;
             }
+            value.tail = tail;
+            scrut2 = cur2.handleBlockList.next !== null;
+            if (scrut2 === true) {
+              value.handleBlockList.tail.next = cur2.handleBlockList.next;
+              value.handleBlockList.tail = cur2.handleBlockList.tail;
+              tmp3 = null;
+            } else {
+              tmp3 = null;
+            }
             return value;
           } else {
             cont = cont.next;
-            tmp3 = null;
+            tmp4 = null;
           }
-          tmp4 = tmp3;
-          continue tmp7;
+          tmp5 = tmp4;
+          continue tmp8;
         } else {
-          tmp4 = null;
+          tmp5 = null;
         }
         break;
       }
-      scrut2 = cur2.handleBlockList.next === null;
-      if (scrut2 === true) {
+      scrut3 = cur2.handleBlockList.next === null;
+      if (scrut3 === true) {
         return value;
       } else {
-        tmp5 = this.__resumeHandleBlocks(cur2.handleBlockList.next, cur2.handleBlockList.tail, value);
-        cur2 = tmp5;
+        tmp6 = this.__resumeHandleBlocks(cur2.handleBlockList.next, cur2.handleBlockList.tail, value);
+        cur2 = tmp6;
         if (cur2 instanceof this.__EffectSig.class) {
           cur2.tail = tail;
-          tmp6 = null;
+          tmp7 = null;
         } else {
-          tmp6 = null;
+          tmp7 = null;
         }
         return cur2;
       }
@@ -478,7 +485,7 @@ const Predef$class = class Predef {
         tmp = handleBlock.contHead.next.resume(value) ?? null;
         value = tmp;
         if (value instanceof this.__EffectSig.class) {
-          scrut2 = handleBlock.contHead.next !== value.tail.next;
+          scrut2 = value.tail.next === null;
           if (scrut2 === true) {
             handleBlock.contHead.next = handleBlock.contHead.next.next;
             tmp1 = null;

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

@@ -207,6 +207,10 @@ fun __resume(cur, tail)(value) =
     cont is __Cont do
       set value = cont.resume(value)
       if value is __EffectSig then
+        // in tail call optimization, the continuation is not appended, we append it here
+        if value.tail.next === null do
+          // since it is a tail call, it is already completed so we append the continuation after it
+          set value.tail.next = cont.next
         // we are returning to handler, which performs unwinding so tail is needed
         set value.tail = tail
         if cur.handleBlockList.next !== null do
@@ -234,7 +238,8 @@ fun __resumeHandleBlocks(handleBlock, tailHandleBlock, value) =
       set value = handleBlock.contHead.next.resume(value)
       if value is __EffectSig then
         // this checks when continuation resume results in tail call to effectful func
-        if handleBlock.contHead.next !== value.tail.next do
+        // when a tail call happens, the continuation will not be appended so this will be null
+        if value.tail.next === null do
           // if this is a tail call that results in effect, the continuation is already completed
           // and should be removed
           set handleBlock.contHead.next = handleBlock.contHead.next.next
@@ -255,4 +260,4 @@ val __stackDepth = 0 // Tracks the virtual + real stack depth
 val __stackOffset = 0 // How much to offset __stackDepth by to get the true stack depth (i.e. the virtual depth)
 val __stackHandler = null
 abstract class __StackDelay() with
-  fun perform()// Private definitions for algebraic effects
+  fun perform()