Miscellaneous Fixes (hkust-taco#280)

Co-authored-by: Anson Yeung <nanolify@gmail.com>

Author: CAG2Mark

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

@@ -106,13 +106,14 @@ sealed abstract class Block extends Product with AutoLocated:
     case AssignDynField(_, _, _, rhs, rest) => rhs.subBlocks ::: rest :: Nil
     case Define(d, rest) => d.subBlocks ::: rest :: Nil
     case HandleBlock(_, _, par, args, _, handlers, body, rest) => par.subBlocks ++ args.flatMap(_.subBlocks) ++ handlers.map(_.body) :+ body :+ rest
+    case Label(_, body, rest) => body :: rest :: Nil
 
     // TODO rm Lam from values and thus the need for these cases
     case Return(r, _) => r.subBlocks
     case HandleBlockReturn(r) => r.subBlocks
     case Throw(r) => r.subBlocks
 
-    case _: Return | _: Throw | _: Label | _: Break | _: Continue | _: End | _: HandleBlockReturn => Nil
+    case _: Return | _: Throw | _: Break | _: Continue | _: End | _: HandleBlockReturn => Nil
 
   // Moves definitions in a block to the top. Only scans the top-level definitions of the block;
   // i.e, definitions inside other definitions are not moved out. Definitions inside `match`/`if`

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

@@ -85,13 +85,13 @@ class BlockTransformer(subst: SymbolSubst):
           (cls2 is cls) && (hdr2 is hdr) && (bod2 is bod) && (rst2 is rst)
         then b else HandleBlock(l2, res2, par2, args2, cls2, hdr2, bod2, rst2)
     case AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
-      val lhs2 = applyPath(lhs)
-      val fld2 = applyPath(fld)
-      val rhs2 = applyResult(rhs)
-      val rest2 = applyBlock(rest)
-      if (lhs2 is lhs) && (fld2 is fld) && (rhs2 is rhs) && (rest2 is rest)
-      then b
-      else AssignDynField(lhs2, fld2, arrayIdx, rhs2, rest2)
+      applyResult2(rhs): rhs2 =>
+        val lhs2 = applyPath(lhs)
+        val fld2 = applyPath(fld)
+        val rest2 = applyBlock(rest)
+        if (lhs2 is lhs) && (fld2 is fld) && (rhs2 is rhs) && (rest2 is rest)
+        then b
+        else AssignDynField(lhs2, fld2, arrayIdx, rhs2, rest2)
 
 
   def applyResult2(r: Result)(k: Result => Block): Block = k(applyResult(r))

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

@@ -59,6 +59,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
         .assignFieldN(state.res, tailIdent, state.res.tail.next)
         .ret(state.res))
   private val functionHandlerCtx = funcLikeHandlerCtx(N)
+  private val topLevelCtx = HandlerCtx(true, true, N, _ => rtThrowMsg("Unhandled effects"))
   private def ctorCtx(ctorThis: Path) = funcLikeHandlerCtx(S(ctorThis))
   private def handlerCtx(using HandlerCtx): HandlerCtx = summon
   private val predefPath: Path = State.globalThisSymbol.asPath.selN(Tree.Ident("Predef"))
@@ -267,6 +268,9 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       case blk @ AssignField(lhs, nme, rhs, rest) =>
         val PartRet(head, parts) = go(rest)
         PartRet(AssignField(lhs, nme, rhs, head)(blk.symbol), parts)
+      case AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
+        val PartRet(head, parts) = go(rest)
+        PartRet(AssignDynField(lhs, fld, arrayIdx, rhs, head), parts)
       case Return(_, _) => PartRet(blk, Nil)
       // ignored cases
       case TryBlock(sub, finallyDo, rest) => ??? // ignore
@@ -336,7 +340,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       override def applyLam(lam: Value.Lam): Value.Lam = Value.Lam(lam.params, translateBlock(lam.body, functionHandlerCtx))
       override def applyDefn(defn: Defn): Defn = defn match
         case f: FunDefn => translateFun(f)
-        case c: ClsLikeDefn => translateCls(c)
+        case c: ClsLikeDefn => translateCls(c, handlerCtx.isTopLevel)
         case _: ValDefn => super.applyDefn(defn)
     transformer.applyBlock(b)
 
@@ -356,9 +360,11 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
   private def translateFun(f: FunDefn): FunDefn =
     FunDefn(f.owner, f.sym, f.params, translateBlock(f.body, functionHandlerCtx))
 
-  private def translateCls(cls: ClsLikeDefn): ClsLikeDefn =
+  private def translateCls(cls: ClsLikeDefn, isTopLevel: Bool): ClsLikeDefn =
+    val curCtorCtx = if isTopLevel && (cls.k is syntax.Mod) then topLevelCtx else
+      ctorCtx(cls.sym.asClsLike.getOrElse(wat("asClsLike", cls.sym)).asPath)
     cls.copy(methods = cls.methods.map(translateFun),
-      ctor = translateBlock(cls.ctor, ctorCtx(cls.sym.asClsLike.getOrElse(wat("asClsLike", cls.sym)).asPath)))
+      ctor = translateBlock(cls.ctor, curCtorCtx))
 
   // Handle block becomes a FunDefn and CallPlaceholder
   private def translateHandleBlock(h: HandleBlock): Block =
@@ -548,5 +554,5 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
     transform.applyBlock(b)
 
   def translateTopLevel(b: Block): Block =
-    translateBlock(b, HandlerCtx(true, true, N, _ => rtThrowMsg("Unhandled effects")))
+    translateBlock(b, topLevelCtx)
 

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

@@ -98,7 +98,7 @@ class StackSafeTransform(depthLimit: Int)(using State):
       override def applyFunDefn(fun: FunDefn): FunDefn = rewriteFn(fun)
 
       override def applyDefn(defn: Defn): Defn = defn match
-        case defn: ClsLikeDefn => rewriteCls(defn)
+        case defn: ClsLikeDefn => rewriteCls(defn, isTopLevel)
         case _: FunDefn | _: ValDefn => super.applyDefn(defn)
 
       override def applyBlock(b: Block): Block = b match
@@ -143,12 +143,13 @@ class StackSafeTransform(depthLimit: Int)(using State):
     walker.applyBlock(b)
     trivial
 
-  def rewriteCls(defn: ClsLikeDefn): ClsLikeDefn = 
+  def rewriteCls(defn: ClsLikeDefn, isTopLevel: Bool): 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)
+      publicFields, rewriteBlk(preCtor),
+      if isTopLevel && (defn.k is syntax.Mod) then transformTopLevel(ctor) else rewriteBlk(ctor)
     )
 
   def rewriteBlk(blk: Block) =

hkmc2/shared/src/test/mlscript/codegen/While.mls

@@ -239,19 +239,23 @@ f(Cons(1, Cons(2, Cons(3, 0))))
 //│ > Cons(3, 0)
 //│ > 0
 
+while false do
+  class A
+  new A
+
 
 // ——— FIXME: ———
 
 
 () => while true then 0
-//│ = [function block$res26]
+//│ = [function block$res27]
 
 :fixme
 while print("Hello World"); false
   then 0(0)
   else 1
 //│ ╔══[PARSE ERROR] Unexpected 'then' keyword here
-//│ ║  l.251: 	  then 0(0)
+//│ ║  l.255: 	  then 0(0)
 //│ ╙──       	  ^^^^
 //│ ═══[ERROR] Unrecognized term split (false literal).
 //│ > Hello World

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

@@ -127,4 +127,17 @@ let oops =
 oops.h
 //│ = Handler$h$
 
+let obj = {}
 
+handle h = Effect with
+  fun perform(arg)(k) =
+    k()
+    print(arg)
+set obj.["s"] = h.perform("k")
+//│ > k
+
+fun f() = ()
+
+module A with
+  f()
+  () // defeats tail call optimization

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

@@ -424,6 +424,17 @@ in
   foo(h)
 //│ = 50005000
 
+:effectHandlers
+:stackSafe 100
+:expect 50005000
+module A with
+  val f = n =>
+    if n <= 0 then 0
+    else n + f(n-1)
+  val x = f(10000)
+A.x
+//│ = 50005000
+
 :re
 :effectHandlers
 fun foo(h) = h.perform(2)