Class Lifter (#266)

Co-authored-by: Anson Yeung <[email protected]>
Co-authored-by: Lionel Parreaux <[email protected]>

Author: 3 people

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

@@ -3,6 +3,7 @@ package mlscript.utils
 import scala.annotation.tailrec
 import scala.collection.immutable.SortedMap
 
+
 object algorithms {
   final class CyclicGraphError(message: String) extends Exception(message)
 
@@ -29,4 +30,145 @@ object algorithms {
     }
     sort(toPred, Seq())
   }
+
+  /**
+    * 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. (Overlap is fine)
+    * @return A list of strongly connected components of the graph.
+    */
+  def partitionScc[A](edges: Iterable[(A, A)], nodes: Iterable[A]): List[List[A]] = {
+    case class SccNode[A](
+      val node: A,
+      val id: Int,
+      var num: Int = -1, 
+      var lowlink: Int = -1,
+      var visited: Boolean = false,
+      var onStack: Boolean = false
+    )
+
+    // 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.id).toMap
+    val nodesIdx = nodesN.map { case node => node.id -> node }.toMap
+
+    val neighbours = edges
+      .map { case (a, b) => (nodeToIdx(a), nodesIdx(nodeToIdx(b))) }
+      .groupBy(_._1)
+      .map { case (a, b) => a -> b.map(_._2) }
+      .withDefault(_ => Nil)
+
+    // Tarjan's algorithm
+
+    var stack: List[SccNode[A]] = List.empty
+    var sccs: List[List[A]] = List.empty
+    var i = 0
+
+    def dfs(node: SccNode[A], depth: Int = 0): Unit = {
+      def printlnsp(s: String) = {
+        println(s)
+      }
+      
+      node.num = i
+      node.lowlink = node.num
+      node.visited = true
+      stack = node :: stack
+      i += 1
+      for (n <- neighbours(node.id)) {
+        if (!n.visited) {
+          dfs(n, depth + 1)
+          node.lowlink = n.lowlink.min(node.lowlink)
+        } else if (!n.onStack) {
+          node.lowlink = n.num.min(node.lowlink)
+        } 
+      }
+      if (node.lowlink == node.num) {
+        var scc: List[A] = List.empty
+        var cur = stack.head
+        stack = stack.tail
+        cur.onStack = true
+        while (cur.id != node.id) {
+          scc = cur.node :: scc
+          cur = stack.head
+          stack = stack.tail
+          cur.onStack = true
+        }
+        scc = cur.node :: scc
+        sccs = scc :: sccs
+      }
+    }
+
+    for (n <- nodesN) {
+      if (!n.visited) dfs(n)
+    }
+    sccs
+  }
+  
+
+  /**
+    * Info about a graph partitioned into its strongly-connected sets. The input type must be able to
+    * be hashed efficiently as it will be used as a key.
+    *
+    * @param sccs The strongly connected sets.
+    * @param edges The edges of the strongly-connected sets. Together with `sccs`, this forms an acyclic graph.
+    * @param inDegs The in-degrees of the above described graph.
+    * @param outDegs The out-degrees of the above described graph.
+    */
+  case class SccsInfo[A](
+    sccs: Map[Int, List[A]],
+    edges: Map[Int, Iterable[Int]],
+    inDegs: Map[Int, Int],
+    outDegs: Map[Int, Int],
+  )
+
+  /**
+    * Partitions a graph into its strongly connected components and returns additional information
+    * about the partition. 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. (Overlap is fine)
+    * @return The partitioned graph and info about it.
+    */
+  def sccsWithInfo[A](edges: Iterable[(A, A)], nodes: Iterable[A]): SccsInfo[A] = {
+    val sccs = partitionScc(edges, nodes)
+    val withIdx = sccs.zipWithIndex.map(_.swap).toMap
+    val lookup = (
+      for {
+        (id, scc) <- withIdx
+        node <- scc
+      } yield node -> id
+    ).toMap
+
+    val notInSccEdges = edges.map {
+      case (a, b) => (lookup(a), lookup(b))
+    }.filter {
+      case (a, b) => a != b
+    }
+
+    val outs = notInSccEdges.groupBy {
+      case (a, b) => a
+    }
+
+    val sccEdges = withIdx.map {
+      case (a, _) => a -> Nil // add default case
+    } ++ outs.map {
+      case (a, edges) => a -> edges.map(_._2)
+    }.toMap
+    
+    val inDegs = notInSccEdges.groupBy {
+      case (a, b) => b
+    }.map {
+      case (b, edges) => b -> edges.size
+    }
+
+    val outDegs = outs.map {
+      case (a, edges) => a -> edges.size
+    }
+
+    SccsInfo(withIdx, sccEdges, inDegs, outDegs)
+  }
 }

hkmc2/shared/src/main/scala/hkmc2/Config.scala

@@ -11,6 +11,7 @@ def config(using Config): Config = summon
 case class Config(
   sanityChecks: Opt[SanityChecks],
   effectHandlers: Opt[EffectHandlers],
+  liftDefns: Opt[LiftDefns],
 ):
 
   def stackSafety: Opt[StackSafety] = effectHandlers.flatMap(_.stackSafety)
@@ -24,6 +25,7 @@ object Config:
     sanityChecks = N, // TODO make the default S
     // sanityChecks = S(SanityChecks(light = true)),
     effectHandlers = N,
+    liftDefns = N,
   )
 
   case class SanityChecks(light: Bool)
@@ -35,6 +37,8 @@ object Config:
     val default: StackSafety = StackSafety(
       stackLimit = 500,
     )
+
+  case class LiftDefns() // there may be other settings in the future, having it as a case class now
 
 end Config
 

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

@@ -91,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 AssignDynField(lhs, fld, arrayIdx, rhs, rest) => lhs.freeVars ++ fld.freeVars ++ rhs.freeVars ++ rest.freeVars
     case Define(defn, rest) => defn.freeVars ++ rest.freeVars
     case HandleBlock(lhs, res, par, args, 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, args, 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
@@ -122,15 +147,19 @@ sealed abstract class Block extends Product with AutoLocated:
   // Note that this returns the definitions in reverse order, with the bottommost definiton appearing
   // last. This is so that using defns.foldLeft later to add the definitions to the front of a block, 
   // we don't need to reverse the list again to preserve the order of the definitions.
-  def floatOutDefns =
+  def floatOutDefns(
+      ignore: Defn => Bool = _ => false, 
+      preserve: Defn => Bool = _ => false
+    ) =
     var defns: List[Defn] = Nil
     val transformer = new BlockTransformerShallow(SymbolSubst()):
       override def applyBlock(b: Block): Block = b match
-        case Define(defn, rest) => defn match
+        case Define(defn, rest) if !ignore(defn) => defn match
           case v: ValDefn => super.applyBlock(b)
           case _ =>
             defns ::= defn
-            applyBlock(rest)
+            if preserve(defn) then super.applyBlock(b)
+            else applyBlock(rest)
         case _ => super.applyBlock(b)
 
     (transformer.applyBlock(this), defns)
@@ -281,10 +310,20 @@ sealed abstract class Defn:
   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],
@@ -304,10 +343,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],
@@ -325,6 +365,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)
@@ -341,6 +382,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
@@ -353,14 +399,26 @@ sealed abstract class Result:
     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
@@ -375,6 +433,7 @@ case class Instantiate(cls: Path, args: Ls[Path]) extends Result
 
 sealed abstract class Path extends Result:
   def selN(id: Tree.Ident): Path = Select(this, id)(N)
+  def sel(id: Tree.Ident, sym: FieldSymbol): Path = Select(this, id)(S(sym))
   def selSN(id: Str): Path = selN(new Tree.Ident(id))
   def asArg = Arg(false, this)