Class Lifter

core/shared/main/scala/utils/algorithms.scala

@@ -2,6 +2,8 @@ package mlscript.utils
 
 import scala.annotation.tailrec
 import scala.collection.immutable.SortedMap
+import scala.collection.mutable.ArrayBuffer
+
 
 object algorithms {
   final class CyclicGraphError(message: String) extends Exception(message)
@@ -29,4 +31,75 @@ object algorithms {
     }
     sort(toPred, Seq())
   }
+
+  private case class SccNode[A](
+    val node: A,
+    val num: Int, 
+    var lowlink: Int = -1,
+    var visited: Boolean = false,
+    var processed: Boolean = false
+  ) 
+
+  /**
+    * Partitions a graph into its strongly connected components. The input type must be able to
+    * be hashed efficiently as it will be used as a key.
+    *
+    * @param edges The edges of the graph.
+    * @param nodes Any additional nodes that are not necessarily in the edges list.
+    * @return A list of strongly connected components of the graph.
+    */
+  def partitionScc[A](edges: Iterable[(A, A)], nodes: Iterable[A]): List[List[A]] = {
+    // pre-process: assign each node an id
+    val edgesSet = edges.toSet
+    val nodesUniq = (edgesSet.flatMap { case (a, b) => Set(a, b) } ++ nodes.toSet).toList
+    val nodesN = nodesUniq.zipWithIndex.map { case (node, idx) => SccNode(node, idx) }
+    val nodeToIdx = nodesN.map(node => node.node -> node.num).toMap
+    val nodesIdx = nodeToIdx.map { case (node, idx) => idx -> SccNode(node, idx) }
+
+    val neighbours = edges
+      .map { case (a, b) => (nodeToIdx(a), nodesIdx(nodeToIdx(b))) }
+      .groupBy(_._1)
+      .map { case (a, b) => a -> b.map(_._2) }
+
+    var stack: List[SccNode[A]] = List.empty
+    var sccs: List[List[A]] = List.empty
+    var i = 0
+
+    def dfs(node: SccNode[A]): Unit = {
+      node.lowlink = node.num
+      node.visited = true
+      stack = node :: stack      
+      i += 1
+      for (n <- neighbours(node.num)) {
+        if (!n.visited) {
+          dfs(n)
+          node.lowlink = n.lowlink.min(node.lowlink)
+        } else if (!n.processed) {
+          node.lowlink = n.num.min(node.lowlink)
+        } 
+      }
+      node.processed = true
+      if (node.lowlink == node.num) {
+        var scc: List[A] = List.empty
+        var cur = stack.head
+        stack = stack.tail
+        while (cur.num != node.num) {
+          scc = node.node :: scc
+          cur = stack.head
+          stack = stack.tail
+        }
+        sccs = scc :: sccs
+      }
+    }
+
+    for (n <- nodesN) {
+      if (!n.visited) dfs(n)
+    }
+    sccs
+  }
+
+  // TODO
+  def sscsWithInfo[A](edges: Iterable[(A, A)], nodes: Iterable[A]): List[List[A]] = {
+    ???
+  }
 }

hkmc2/jvm/src/test/scala/hkmc2/CompileTestRunner.scala

@@ -54,4 +54,3 @@ class CompileTestRunner
 
 end CompileTestRunner
 
-

hkmc2/shared/src/main/scala/hkmc2/MLsCompiler.scala

@@ -74,7 +74,7 @@ class MLsCompiler(preludeFile: os.Path):
       val parsed = mainParse.resultBlk
       val (blk, newCtx) = elab.importFrom(parsed)
       val low = ltl.givenIn:
-        codegen.Lowering(lowerHandlers = false, stackLimit = None) // TODO: properly hook up stack limit
+        codegen.Lowering(lowerHandlers = false, stackLimit = None, lift = false) // TODO: properly hook up stack limit
       val jsb = ltl.givenIn:
         codegen.js.JSBuilder()
       val le = low.program(blk)

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

@@ -55,6 +55,7 @@ sealed abstract class Block extends Product with AutoLocated:
     case TryBlock(sub, fin, rst) => 1 + sub.size + fin.size + rst.size
     case Label(_, bod, rst) => 1 + bod.size + rst.size
     case HandleBlock(lhs, res, par, cls, handlers, bdy, rst) => 1 + handlers.map(_.body.size).sum + bdy.size + rst.size
+    case AssignDynField(l, n, ai, r, rst) => 1 + rst.size
 
   // TODO conserve if no changes
   def mapTail(f: BlockTail => Block): Block = this match
@@ -90,12 +91,37 @@ sealed abstract class Block extends Product with AutoLocated:
     case TryBlock(sub, finallyDo, rest) => sub.freeVars ++ finallyDo.freeVars ++ rest.freeVars
     case Assign(lhs, rhs, rest) => Set(lhs) ++ rhs.freeVars ++ rest.freeVars
     case AssignField(lhs, nme, rhs, rest) => lhs.freeVars ++ rhs.freeVars ++ rest.freeVars
-    case Define(defn, rest) => defn.freeVars ++ rest.freeVars
+    case AssignDynField(lhs, fld, arrayIdx, rhs, rest) => lhs.freeVars ++ fld.freeVars ++ rhs.freeVars ++ rest.freeVars
+    case Define(defn, rest) => defn.freeVarsLLIR ++ rest.freeVars
     case HandleBlock(lhs, res, par, cls, hdr, bod, rst) =>
       (bod.freeVars - lhs) ++ rst.freeVars ++ hdr.flatMap(_.freeVars)
     case HandleBlockReturn(res) => res.freeVars
     case End(msg) => Set.empty
 
+  // TODO: freeVarsLLIR skips `fun` and `cls` in `Call` and `Instantiate` respectively, which is needed in some
+  // other places. However, adding them breaks some LLIR tests. Supposedly, once the IR uses the new symbol system, 
+  // this should no longer happen. This version should be removed once that is resolved.
+  lazy val freeVarsLLIR: Set[Local] = this match
+    case Match(scrut, arms, dflt, rest) =>
+      scrut.freeVarsLLIR ++ dflt.toList.flatMap(_.freeVarsLLIR) ++ rest.freeVarsLLIR
+      ++ arms.flatMap:
+        (pat, arm) => arm.freeVarsLLIR -- pat.freeVars
+    case Return(res, implct) => res.freeVarsLLIR
+    case Throw(exc) => exc.freeVarsLLIR
+    case Label(label, body, rest) => (body.freeVarsLLIR - label) ++ rest.freeVarsLLIR 
+    case Break(label) => Set(label)
+    case Continue(label) => Set(label)
+    case Begin(sub, rest) => sub.freeVarsLLIR ++ rest.freeVarsLLIR
+    case TryBlock(sub, finallyDo, rest) => sub.freeVarsLLIR ++ finallyDo.freeVarsLLIR ++ rest.freeVarsLLIR
+    case Assign(lhs, rhs, rest) => Set(lhs) ++ rhs.freeVarsLLIR ++ rest.freeVarsLLIR
+    case AssignField(lhs, nme, rhs, rest) => lhs.freeVarsLLIR ++ rhs.freeVarsLLIR ++ rest.freeVarsLLIR
+    case AssignDynField(lhs, fld, arrayIdx, rhs, rest) => lhs.freeVarsLLIR ++ fld.freeVarsLLIR ++ rhs.freeVarsLLIR ++ rest.freeVarsLLIR
+    case Define(defn, rest) => defn.freeVarsLLIR ++ rest.freeVarsLLIR
+    case HandleBlock(lhs, res, par, cls, hdr, bod, rst) =>
+      (bod.freeVarsLLIR - lhs) ++ rst.freeVarsLLIR ++ hdr.flatMap(_.freeVars)
+    case HandleBlockReturn(res) => res.freeVarsLLIR
+    case End(msg) => Set.empty
+
   lazy val subBlocks: Ls[Block] = this match
     case Match(p, arms, dflt, rest) => p.subBlocks ++ arms.map(_._2) ++ dflt.toList :+ rest
     case Begin(sub, rest) => sub :: rest :: Nil
@@ -112,7 +138,7 @@ sealed abstract class Block extends Product with AutoLocated:
     case Throw(r) => r.subBlocks
 
     case _: Return | _: Throw | _: Label | _: 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`
   // and `while` statements are moved out.
@@ -192,14 +218,24 @@ sealed abstract class Defn:
     case _: ValDefn => Nil
     case ClsLikeDefn(preCtor = preCtor, ctor = ctor, methods = mtds) =>
       preCtor :: ctor :: mtds.flatMap(_.subBlocks)
-  
+
   lazy val freeVars: Set[Local] = this match
     case FunDefn(own, sym, params, body) => body.freeVars -- params.flatMap(_.paramSyms) - sym
     case ValDefn(owner, k, sym, rhs) => rhs.freeVars
-    case ClsLikeDefn(own, isym, sym, k, paramsOpt, parentSym, methods, privateFields, publicFields, preCtor, ctor) =>
+    case ClsLikeDefn(own, isym, sym, k, paramsOpt, auxParams, parentSym, 
+      methods, privateFields, publicFields, preCtor, ctor) =>
       preCtor.freeVars
         ++ ctor.freeVars ++ methods.flatMap(_.freeVars)
-        -- privateFields -- publicFields.map(_.sym)
+        -- privateFields -- publicFields.map(_.sym) -- auxParams.flatMap(_.paramSyms)
+
+  lazy val freeVarsLLIR: Set[Local] = this match
+    case FunDefn(own, sym, params, body) => body.freeVarsLLIR -- params.flatMap(_.paramSyms) - sym
+    case ValDefn(owner, k, sym, rhs) => rhs.freeVarsLLIR
+    case ClsLikeDefn(own, isym, sym, k, paramsOpt, auxParams, parentSym, 
+      methods, privateFields, publicFields, preCtor, ctor) =>
+      preCtor.freeVarsLLIR
+        ++ ctor.freeVarsLLIR ++ methods.flatMap(_.freeVarsLLIR)
+        -- privateFields -- publicFields.map(_.sym) -- auxParams.flatMap(_.paramSyms)
 
 final case class FunDefn(
     owner: Opt[InnerSymbol],
@@ -219,10 +255,11 @@ final case class ValDefn(
 
 final case class ClsLikeDefn(
     owner: Opt[InnerSymbol],
-    isym: MemberSymbol[? <: ClassLikeDef],
+    isym: MemberSymbol[? <: ClassLikeDef] & InnerSymbol,
     sym: BlockMemberSymbol,
     k: syntax.ClsLikeKind,
     paramsOpt: Opt[ParamList],
+    auxParams: List[ParamList],
     parentPath: Opt[Path],
     methods: Ls[FunDefn],
     privateFields: Ls[TermSymbol],
@@ -238,6 +275,7 @@ final case class Handler(
     params: Ls[ParamList],
     body: Block,
 ):
+  lazy val freeVarsLLIR: Set[Local] = body.freeVarsLLIR -- params.flatMap(_.paramSyms) - sym - resumeSym
   lazy val freeVars: Set[Local] = body.freeVars -- params.flatMap(_.paramSyms) - sym - resumeSym
 
 /* Represents either unreachable code (for functions that must return a result)
@@ -254,6 +292,11 @@ enum Case:
     case Cls(_, path) => path.freeVars
     case Tup(_, _) => Set.empty
 
+  lazy val freeVarsLLIR: Set[Local] = this match
+    case Lit(_) => Set.empty
+    case Cls(_, path) => path.freeVarsLLIR
+    case Tup(_, _) => Set.empty
+
 sealed abstract class Result:
 
   // TODO rm Lam from values and thus the need for this method
@@ -264,16 +307,28 @@ sealed abstract class Result:
     case Value.Lam(params, body) => body :: Nil
     case Value.Arr(elems) => elems.flatMap(_.value.subBlocks)
     case _ => Nil
-  
+
   lazy val freeVars: Set[Local] = this match
-    case Call(fun, args) => args.flatMap(_.value.freeVars).toSet
-    case Instantiate(cls, args) => args.flatMap(_.freeVars).toSet
+    case Call(fun, args) => fun.freeVars ++ args.flatMap(_.value.freeVars).toSet
+    case Instantiate(cls, args) => cls.freeVars ++ args.flatMap(_.freeVars).toSet
     case Select(qual, name) => qual.freeVars 
     case Value.Ref(l) => Set(l)
     case Value.This(sym) => Set.empty
     case Value.Lit(lit) => Set.empty
     case Value.Lam(params, body) => body.freeVars -- params.paramSyms
     case Value.Arr(elems) => elems.flatMap(_.value.freeVars).toSet
+    case DynSelect(qual, fld, arrayIdx) => qual.freeVars ++ fld.freeVars
+
+  lazy val freeVarsLLIR: Set[Local] = this match
+    case Call(fun, args) => args.flatMap(_.value.freeVarsLLIR).toSet
+    case Instantiate(cls, args) => args.flatMap(_.freeVarsLLIR).toSet
+    case Select(qual, name) => qual.freeVarsLLIR 
+    case Value.Ref(l) => Set(l)
+    case Value.This(sym) => Set.empty
+    case Value.Lit(lit) => Set.empty
+    case Value.Lam(params, body) => body.freeVarsLLIR -- params.paramSyms
+    case Value.Arr(elems) => elems.flatMap(_.value.freeVarsLLIR).toSet
+    case DynSelect(qual, fld, arrayIdx) => qual.freeVarsLLIR ++ fld.freeVarsLLIR
 
 // type Local = LocalSymbol
 type Local = Symbol

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

@@ -83,6 +83,13 @@ class BlockTransformer(subst: SymbolSubst):
       if (l2 is l) && (res2 is res) && (par2 is par) && (cls2 is cls) &&
           (hdr2 is hdr) && (bod2 is bod) && (rst2 is rst)
         then b else HandleBlock(l2, res2, par2, cls2, hdr2, bod2, rst2)
+    case AssignDynField(l, fld, arrayIdx, r, rst) =>
+      applyResult2(r): r2 =>
+        val l2 = applyPath(l)
+        val fld2 = applyPath(fld)
+        val rst2 = applyBlock(rst)
+        if (l2 is l) && (fld2 is fld) && (r2 is r) && (rst2 is rst)
+          then b else AssignDynField(l2, fld2, arrayIdx, r2, rst2)
 
   def applyResult2(r: Result)(k: Result => Block): Block = k(applyResult(r))
 
@@ -98,6 +105,10 @@ class BlockTransformer(subst: SymbolSubst):
     case p: Path => applyPath(p)
 
   def applyPath(p: Path): Path = p match
+    case DynSelect(qual, fld, arrayIdx) =>
+      val qual2 = applyPath(qual)
+      val fld2 = applyPath(fld)
+      if (qual2 is qual) && (fld2 is fld) then p else DynSelect(qual2, fld2, arrayIdx) 
     case p @ Select(qual, name) =>
       val qual2 = applyPath(qual)
       val sym2 = p.symbol.mapConserve(_.subst)
@@ -135,11 +146,13 @@ class BlockTransformer(subst: SymbolSubst):
       val rhs2 = applyPath(rhs)
       if (owner2 is owner) && (sym2 is sym) && (rhs2 is rhs)
         then defn else ValDefn(owner2, k, sym2, rhs2)
-    case ClsLikeDefn(own, isym, sym, k, paramsOpt, parentPath, methods, privateFields, publicFields, preCtor, ctor) =>
+    case ClsLikeDefn(own, isym, sym, k, paramsOpt, auxParams, parentPath, methods, 
+      privateFields, publicFields, preCtor, ctor) =>
       val own2 = own.mapConserve(_.subst)
       val isym2 = isym.subst
       val sym2 = sym.subst
       val paramsOpt2 = paramsOpt.mapConserve(applyParamList)
+      val auxParams2 = auxParams.mapConserve(applyParamList)
       val parentPath2 = parentPath.mapConserve(applyPath)
       val methods2 = methods.mapConserve(applyFunDefn)
       val privateFields2 = privateFields.mapConserve(_.subst)
@@ -148,12 +161,14 @@ class BlockTransformer(subst: SymbolSubst):
       val ctor2 = applyBlock(ctor)
       if (own2 is own) && (isym2 is isym) && (sym2 is sym) &&
           (paramsOpt2 is paramsOpt) &&
+          (auxParams2 is auxParams) &&
           (parentPath2 is parentPath) &&
           (methods2 is methods) &&
           (privateFields2 is privateFields) &&
           (publicFields2 is publicFields) &&
           (preCtor2 is preCtor) && (ctor2 is ctor)
-        then defn else ClsLikeDefn(own, isym2, sym2, k, paramsOpt, parentPath2, methods2, privateFields2, publicFields2, preCtor2, ctor2)
+        then defn else ClsLikeDefn(own2, isym2, sym2, k, paramsOpt2, 
+          auxParams2, parentPath2, methods2, privateFields2, publicFields2, preCtor2, ctor2)
 
   def applyArg(arg: Arg): Arg =
     val val2 = applyPath(arg.value)
@@ -221,3 +236,25 @@ class BlockTransformerShallow(subst: SymbolSubst) extends BlockTransformer(subst
         then b else HandleBlock(l2, res2, par2, cls2, hdr2, bod, rst2)
     case _ => super.applyBlock(b)
 
+// does not traverse into any other block
+class BlockTransformerNoRec(subst: SymbolSubst) extends BlockTransformerShallow(subst):
+  override def applyBlock(b: Block): Block = b match
+    case Match(scrut, arms, dflt, rest) => 
+      val scrut2 = applyPath(scrut)
+      if (scrut is scrut2) then b else Match(scrut2, arms, dflt, rest)
+    case Assign(lhs, rhs, rest) => 
+      applyResult2(rhs): rhs2 =>
+        val lhs2 = lhs.subst
+        if (lhs is lhs2) && (rhs is rhs2) then b else Assign(lhs2, rhs2, rest)
+    case AssignField(lhs, ident, rhs, rest) => 
+      applyResult2(rhs): rhs2 =>
+        val lhs2 = applyPath(lhs)
+        if rhs is rhs2 then b else AssignField(lhs2, ident, rhs2, rest)(N)
+    case AssignDynField(l, fld, arrayIdx, r, rst) =>
+      applyResult2(r): r2 =>
+        val l2 = applyPath(l)
+        val fld2 = applyPath(fld)
+        if (l2 is l) && (fld2 is fld) && (r2 is r)
+          then b else AssignDynField(l2, fld2, arrayIdx, r2, rst)
+    case _: Label | _: Begin | _: TryBlock | _: Define | _: HandleBlock | _: End => b
+    case _: Return | _: Break | _: Continue | _: HandleBlockReturn | _: Throw => super.applyBlock(b)

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

@@ -53,9 +53,8 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
     HandlerCtx(true, false, ctorThis, state =>
       val tmp = freshTmp()
       blockBuilder
-        .assignFieldN(state.res.tail, nextIdent, Instantiate(
-          state.cls.selN(Tree.Ident("class")),
-          Value.Lit(Tree.IntLit(state.uid)) :: Nil))
+        .assignFieldN(state.res.tail, nextIdent, Call(
+          state.cls, Value.Lit(Tree.IntLit(state.uid)).asArg :: Nil)(true))
         .assignFieldN(state.res, tailIdent, state.res.tail.next)
         .ret(state.res))
   private val functionHandlerCtx = funcLikeHandlerCtx(N)
@@ -267,6 +266,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, rest), parts)
       case Return(_, _) => PartRet(blk, Nil)
       // ignored cases
       case TryBlock(sub, finallyDo, rest) => ??? // ignore
@@ -370,7 +372,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       transform.applyBlock(b)
 
     val handlerBody = translateBlock(prepareBody(h.body), HandlerCtx(false, false, N, state => blockBuilder
-      .assignFieldN(state.res.tail, nextIdent, Instantiate(state.cls, Value.Lit(Tree.IntLit(state.uid)) :: Nil))
+      .assignFieldN(state.res.tail, nextIdent, Call(state.cls, Value.Lit(Tree.IntLit(state.uid)).asArg :: Nil)(true))
       .ret(SimpleCall(handleBlockImplPath, state.res :: h.lhs.asPath :: Nil))))
 
     val handlers = h.handlers.map: handler =>
@@ -380,21 +382,21 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       val tmp = freshTmp()
       FunDefn(
         S(h.cls),
-        handler.sym, handler.params, Return(SimpleCall(mkEffectPath, h.lhs.asPath :: lam :: Nil), false))
+        handler.sym, handler.params, Return(SimpleCall(mkEffectPath, h.cls.asPath :: lam :: Nil), false))
 
     // TODO: it seems that our current syntax didn't know how to call super, calling it with empty param list now
     val clsDefn = ClsLikeDefn(
       N, // no owner
       h.cls,
       BlockMemberSymbol(h.cls.id.name, Nil),
       syntax.Cls,
-      N,
+      S(PlainParamList(Nil)), Nil,
       S(h.par), handlers, Nil, Nil,
       Assign(freshTmp(), SimpleCall(Value.Ref(State.builtinOpsMap("super")), Nil), End()), End())
 
     val body = blockBuilder
       .define(clsDefn)
-      .assign(h.lhs, Instantiate(Value.Ref(clsDefn.sym), Nil))
+      .assign(h.lhs, Call(clsDefn.sym.asPath, Nil)(true))
       .rest(handlerBody)
 
     val defn = FunDefn(
@@ -506,6 +508,7 @@ class HandlerLowering(using TL, Raise, Elaborator.State, Elaborator.Ctx):
       BlockMemberSymbol(clsSym.nme, Nil),
       syntax.Cls,
       S(PlainParamList(Param(FldFlags.empty, pcVar, N) :: Nil)),
+      Nil,
       S(contClsPath),
       resumeFnDef :: Nil,
       Nil,