Fix stack safety

Author: CAG2Mark

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

@@ -35,12 +35,19 @@ class StackSafeTransform(depthLimit: Int)(using State):
         .ret(res)
     else
       val tmp = TempSymbol(None, "tmp")
+      val offsetGtDepth = TempSymbol(None, "offsetGtDepth")
       val prevDepth = TempSymbol(None, "prevDepth")
       blockBuilder
         .assign(prevDepth, stackDepthPath)
         .assignFieldN(predefPath, STACK_DEPTH_IDENT, op("+", stackDepthPath, intLit(1)))
         .assign(tmp, res)
         .assignFieldN(predefPath, STACK_DEPTH_IDENT, prevDepth.asPath)
+        .assign(offsetGtDepth, op("<", prevDepth.asPath, stackOffsetPath))
+        .ifthen(
+          offsetGtDepth.asPath, 
+          Case.Lit(Tree.BoolLit(true)), 
+          blockBuilder.assignFieldN(predefPath, STACK_OFFSET_IDENT, prevDepth.asPath).end
+          )
         .rest(f(tmp.asPath))
 
   def extractResTopLevel(res: Result, isTailCall: Bool, f: Result => Block) =
@@ -66,14 +73,14 @@ class StackSafeTransform(depthLimit: Int)(using State):
             let curOffset = stackOffset
             stackOffset = stackDepth
             let ret = resume()
-            stackOffset = curOffset
+            // stackOffset = curOffset // this line is wrong
             ret
         */
         blockBuilder
           .assign(curOffsetSym, stackOffsetPath)
           .assignFieldN(predefPath, STACK_OFFSET_IDENT, stackDepthPath)
           .assign(handlerRes, Call(Value.Ref(resumeSym), Nil)(true))
-          .assignFieldN(predefPath, STACK_OFFSET_IDENT, curOffsetSym.asPath)
+          // .assignFieldN(predefPath, STACK_OFFSET_IDENT, curOffsetSym.asPath) // this line is wrong
           .ret(handlerRes.asPath)
       ) :: Nil,
       blockBuilder

hkmc2/shared/src/test/mlscript/handlers/StackSafety.mls

@@ -104,7 +104,6 @@ hi(0)
 //│             contLoop: while (true) {
 //│               if (this.pc === 4) {
 //│                 res2 = res3;
-//│                 globalThis.Predef.__stackOffset = curOffset;
 //│                 return res2;
 //│               }
 //│               break;
@@ -121,7 +120,6 @@ hi(0)
 //│           return res3;
 //│         }
 //│         res2 = res3;
-//│         globalThis.Predef.__stackOffset = curOffset;
 //│         return res2;
 //│       });
 //│     }
@@ -179,7 +177,7 @@ sum(10000)
 //│ JS (unsanitized):
 //│ let sum3, res1, handleBlock$3;
 //│ sum3 = function sum(n) {
-//│   let scrut, tmp, tmp1, tmp2, prevDepth, diff, diffGeqLimit, handlerExists, scrut1, dummy, res2, res3, Cont$;
+//│   let scrut, tmp, tmp1, tmp2, offsetGtDepth, prevDepth, diff, diffGeqLimit, handlerExists, scrut1, dummy, res2, res3, Cont$;
 //│   Cont$ = function Cont$(pc1) { return new Cont$.class(pc1); };
 //│   Cont$.class = class Cont$ extends globalThis.Predef.__Cont.class {
 //│     constructor(pc) {
@@ -194,7 +192,7 @@ sum(10000)
 //│         res2 = value$;
 //│       }
 //│       contLoop: while (true) {
-//│         if (this.pc === 3) {
+//│         if (this.pc === 4) {
 //│           scrut = n == 0;
 //│           if (scrut === true) {
 //│             return 0;
@@ -217,11 +215,20 @@ sum(10000)
 //│         } else if (this.pc === 1) {
 //│           tmp2 = res3;
 //│           globalThis.Predef.__stackDepth = prevDepth;
+//│           offsetGtDepth = prevDepth < globalThis.Predef.__stackOffset;
+//│           if (offsetGtDepth === true) {
+//│             globalThis.Predef.__stackOffset = prevDepth;
+//│             this.pc = 3;
+//│             continue contLoop;
+//│           }
+//│           this.pc = 3;
+//│           continue contLoop;
+//│         } else if (this.pc === 3) {
 //│           tmp1 = tmp2;
 //│           return n + tmp1;
 //│         } else if (this.pc === 0) {
 //│           dummy = res2;
-//│           this.pc = 3;
+//│           this.pc = 4;
 //│           continue contLoop;
 //│         }
 //│         break;
@@ -257,6 +264,10 @@ sum(10000)
 //│     }
 //│     tmp2 = res3;
 //│     globalThis.Predef.__stackDepth = prevDepth;
+//│     offsetGtDepth = prevDepth < globalThis.Predef.__stackOffset;
+//│     if (offsetGtDepth === true) {
+//│       globalThis.Predef.__stackOffset = prevDepth;
+//│     }
 //│     tmp1 = tmp2;
 //│     return n + tmp1;
 //│   }
@@ -279,13 +290,12 @@ sum(10000)
 //│             this.pc = pc;
 //│           }
 //│           resume(value$) {
-//│             if (this.pc === 5) {
+//│             if (this.pc === 6) {
 //│               res4 = value$;
 //│             }
 //│             contLoop: while (true) {
-//│               if (this.pc === 5) {
+//│               if (this.pc === 6) {
 //│                 res3 = res4;
-//│                 globalThis.Predef.__stackOffset = curOffset;
 //│                 return res3;
 //│               }
 //│               break;
@@ -297,12 +307,11 @@ sum(10000)
 //│         globalThis.Predef.__stackOffset = globalThis.Predef.__stackDepth;
 //│         res4 = resume();
 //│         if (res4 instanceof globalThis.Predef.__EffectSig.class) {
-//│           res4.tail.next = new Cont$1.class(5);
+//│           res4.tail.next = new Cont$1.class(6);
 //│           res4.tail = res4.tail.next;
 //│           return res4;
 //│         }
 //│         res3 = res4;
-//│         globalThis.Predef.__stackOffset = curOffset;
 //│         return res3;
 //│       });
 //│     }
@@ -317,11 +326,11 @@ sum(10000)
 //│       this.pc = pc;
 //│     }
 //│     resume(value$) {
-//│       if (this.pc === 4) {
+//│       if (this.pc === 5) {
 //│         res2 = value$;
 //│       }
 //│       contLoop: while (true) {
-//│         if (this.pc === 4) {
+//│         if (this.pc === 5) {
 //│           return res2;
 //│         }
 //│         break;
@@ -334,7 +343,7 @@ sum(10000)
 //│   globalThis.Predef.__stackHandler = stackHandler;
 //│   res2 = sum3(10000);
 //│   if (res2 instanceof globalThis.Predef.__EffectSig.class) {
-//│     res2.tail.next = new Cont$(4);
+//│     res2.tail.next = new Cont$(5);
 //│     return globalThis.Predef.__handleBlockImpl(res2, stackHandler);
 //│   }
 //│   return res2;
@@ -475,7 +484,6 @@ hi(0)
 //│ ═══[RUNTIME ERROR] Error: This code requires effect handler instrumentation but was compiled without it.
 
 
-:fixme
 :stackSafe 1000
 :handler
 :expect 100010000
@@ -485,4 +493,4 @@ fun sum(n) =
     n + sum(n - 1)
 fun bad() = sum(10000) + sum(10000)
 bad()
-//│ ═══[RUNTIME ERROR] RangeError: Maximum call stack size exceeded
+//│ = 100010000

hkmc2/shared/src/test/mlscript/handlers/ZCombinator.mls

@@ -23,7 +23,7 @@ fact(10)
 fact(10000)
 //│ ═══[RUNTIME ERROR] RangeError: Maximum call stack size exceeded
 
-
+:stackSafe 1000
 fun mkrec(g) =
   selfApp of self =>
     g of y => selfApp(self)(y)
@@ -91,16 +91,16 @@ let fact =
 :expect 3628800
 :stackSafe 5
 fact(10)
-//│ > 2 0
-//│ > 4 0
-//│ > 6 5
-//│ > 8 5
-//│ > 10 10
-//│ > 12 10
-//│ > 14 10
-//│ > 16 15
-//│ > 18 15
-//│ > 20 20
+//│ > 2 993
+//│ > 4 1
+//│ > 6 4
+//│ > 8 4
+//│ > 10 8
+//│ > 12 8
+//│ > 14 12
+//│ > 16 12
+//│ > 18 16
+//│ > 20 16
 //│ > resumed: 1
 //│ > resumed: 2
 //│ > resumed: 3