Preliminary deforestation implementation for hkmc2

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

@@ -406,15 +406,14 @@ enum Case:
 
 sealed trait TrivialResult extends Result
 
-type ResultId = Uid[Result]
-object ResultUidHandler extends Uid.Handler[Result]
-object ResultUid extends ResultUidHandler.State:
-  val uidToResult = collection.mutable.Map.empty[ResultId, Result]
-  def apply(id: ResultId) = uidToResult(id)
+object Result:
+  opaque type ResultId = Int
+  given Ordering[ResultId] with
+    def compare(x: ResultId, y: ResultId): Int = x.compare(y)
 
+  private def ResultId(v: Int): ResultId = v
 
 
-
 sealed abstract class Result extends AutoLocated:
 
   protected def children: List[Located] = this match
@@ -468,10 +467,14 @@ sealed abstract class Result extends AutoLocated:
     case DynSelect(qual, fld, arrayIdx) => qual.freeVarsLLIR ++ fld.freeVarsLLIR
     case Value.Rcd(args) => args.flatMap(arg => arg.idx.fold(Set.empty)(_.freeVarsLLIR) ++ arg.value.freeVarsLLIR).toSet
 
-  lazy val uid =
-    val id = ResultUid.nextUid
-    ResultUid.uidToResult.addOne(id -> this)
-    id
+  // for deforestation
+  import Result.*
+  lazy val uidValue: ResultId = ResultId(System.identityHashCode(this))
+  def uid(using d: Deforest) =
+    d.resultIdToResult.updateWith(this.uidValue):
+      case N => S(this)
+      case S(r) => assert(this is r); S(this)
+    uidValue
 
 // type Local = LocalSymbol
 type Local = Symbol

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

@@ -4,10 +4,11 @@ package codegen
 import semantics.*
 import semantics.Elaborator.State
 import syntax.{Literal, Tree}
-import utils.{TL, tl, SymbolSubst}
+import utils.*
 import mlscript.utils.*, shorthands.*
 import scala.collection.mutable
 import scala.collection.mutable.LinkedHashMap
+import Result.ResultId
 
 type StratVar
 type StratVarId = Uid[StratVar]
@@ -25,50 +26,52 @@ class StratVarState(val uid: StratVarId, val name: Str = ""):
 
   override def toString(): String = s"${if name.isEmpty() then "var" else name}@${uid}"
 
+object StratVarUidHandler extends Uid.Handler[StratVar]
 object StratVarState:
-
-  def freshVar(nme: String = "")(using vuid: Uid.Handler[StratVar]#State) =
+  def freshVar(nme: String = "")(using vuid: StratVarUidHandler.State) =
     val newId = vuid.nextUid
     val s = StratVarState(newId, nme)
     val p = s.asProdStrat
     val c = s.asConsStrat
     p -> c
 
 
-type CtorExpr = ResultId
-
 extension (i: ResultId)
-  def getResult = ResultUid(i)
-  def handleCtorIds[A](k: (ResultId, Select | Value.Ref, ClsOrModSymbol, Ls[Arg]) => A) =
+  def getResult(using d: Deforest) = d.resultIdToResult(i)
+  def handleCtorIds[A](k: (ResultId, Select | Value.Ref, ClsOrModSymbol, Ls[Arg]) => A)(using Deforest) =
+    def handleCallLike(f: Path, args: Ls[Arg]) = f match
+      case s: Select if s.symbol.flatMap(_.asCls).isDefined =>
+        Some(k(i, s, s.symbol.get.asCls.get, args))
+      case v: Value.Ref if v.l.asCls.isDefined =>
+        Some(k(i, v, v.l.asCls.get, args))
+      case _ => None
     i.getResult match
-      case Call(fun, args) => fun match
-        case s: Select if s.symbol.flatMap(_.asCls).isDefined =>
-          Some(k(i, s, s.symbol.get.asCls.get, args))
-        case v: Value.Ref if v.l.asCls.isDefined =>
-          Some(k(i, v, v.l.asCls.get, args))
-        case _ => None
+      case Call(fun, args) => handleCallLike(fun, args)
+      case Instantiate(cls, args) => handleCallLike(cls, args.map(Arg(false, _)))
       case s: Select if s.symbol.flatMap(_.asObj).isDefined =>
         Some(k(i, s, s.symbol.get.asObj.get, Nil))
       case v: Value.Ref if v.l.asObj.isDefined =>
         Some(k(i, v, v.l.asObj.get, Nil))
       case _ => None
-  def getClsSymOfUid = i.handleCtorIds((_, _, s, _) => s).get
+  def getClsSymOfUid(using Deforest) = i.handleCtorIds((_, _, s, _) => s).get
 
-case class Ctor(ctor: ClsOrModSymbol, args: Map[TermSymbol, ProdStrat], expr: CtorExpr) extends ProdStrat
+case class Ctor(ctor: ClsOrModSymbol, args: Map[TermSymbol, ProdStrat], expr: ResultId) extends ProdStrat
 case class ProdFun(l: Ls[ConsStrat], r: ProdStrat) extends ProdStrat
 case class ProdVar(s: StratVarState) extends ProdStrat with StratVarTrait(s)
 case object NoProd extends ProdStrat
 
 
 
-case class Dtor(scrut: ResultId)(val expr: Match, val outterMatch: Option[ResultId])(using d: Deforest) extends ConsStrat:
-  assert(scrut === expr.scrut.uid)
-  d.matchScrutToMatchBlock.updateWith(scrut):
+class Dtor(val expr: Match, val outterMatch: Option[ResultId])(using d: Deforest) extends ConsStrat:
+  d.matchScrutToMatchBlock.updateWith(expr.scrut.uid):
     case None => Some(expr)
-    case Some(exist) => ??? // should only update once
-  d.matchScrutToParentMatchScrut.updateWith(scrut):
+    case Some(_) => lastWords(s"should only update once (uid: ${expr.scrut.uid})")
+  d.matchScrutToParentMatchScrut.updateWith(expr.scrut.uid):
     case None => Some(outterMatch)
-    case Some(_) => ??? // should only update once
+    case Some(_) => lastWords(s"should only update once (uid: ${expr.scrut.uid})")
+object Dtor:
+  def unapply(d: Dtor)(using Deforest): Opt[ResultId] = S(d.expr.scrut.uid)
+
 
 case class FieldSel(field: Tree.Ident, consVar: ConsVar)(val expr: ResultId, val inMatching: LinkedHashMap[ResultId, ClsOrModSymbol]) extends ConsStrat with FieldSelTrait
 case class ConsFun(l: Ls[ProdStrat], r: ConsStrat) extends ConsStrat
@@ -177,7 +180,8 @@ class DeforestationFreeVarTraverserForMatch(
   selsReplacementByCurrentMatch: Map[ResultId, Symbol],
   currentMatchScrut: Symbol,
   dt: DeforestTransformer
-) extends FreeVarTraverser(alwaysDefined):  
+) extends FreeVarTraverser(alwaysDefined):
+  given Deforest = dt.d
   override def applyBlock(b: Block): Unit = b match
     // a nested match
     case m@Match(scrut, arms, dflt, rest) =>
@@ -235,8 +239,7 @@ class WillBeNonEndTailBlockTraverser(using d: Deforest) extends BlockTraverserSh
     applyBlock(b)
     flag
 
-class ReplaceLocalSymTransformer(freeVarsAndTheirNewSyms: Map[Symbol, Symbol]) extends
-  BlockTransformer(new SymbolSubst()):
+class ReplaceLocalSymTransformer(freeVarsAndTheirNewSyms: Map[Symbol, Symbol]) extends BlockTransformer(new SymbolSubst()):
   override def applyValue(v: Value): Value = v match
     case Value.Ref(l) => Value.Ref(freeVarsAndTheirNewSyms.getOrElse(l, l))
     case _ => super.applyValue(v)
@@ -267,10 +270,13 @@ extension (b: Block)
 
 class Deforest(using TL, Raise, Elaborator.State):
 
-  object StratVarUidHandler extends Uid.Handler[StratVar]()
-  given Uid.Handler[StratVar]#State = StratVarUidHandler.State()
+  given StratVarUidHandler.State = StratVarUidHandler.State()
+  given Deforest = this
   import StratVarState.freshVar
 
+
+  val resultIdToResult = mutable.Map.empty[ResultId, Result]
+
   def apply(p: Program): Opt[Program] -> String -> Int =
     val mainBlk = p.main
 
@@ -310,7 +316,7 @@ class Deforest(using TL, Raise, Elaborator.State):
 
     val fusionStat = filteredCtorDests.map:
       case (ctorUid, CtorFinalDest.Sel(s)) =>
-        "\t" + ctorUid.getClsSymOfUid.nme + " --sel--> " + s"`.${ResultUid(s).asInstanceOf[Select].name}`"
+        "\t" + ctorUid.getClsSymOfUid.nme + " --sel--> " + s"`.${resultIdToResult(s).asInstanceOf[Select].name}`"
       case (ctorUid, CtorFinalDest.Match(scrut, expr, _, _)) =>
         "\t" + ctorUid.getClsSymOfUid.nme + " --match--> " + s"`if ${expr.scrut.asInstanceOf[Value.Ref].l.nme} is ...`"
 
@@ -360,11 +366,18 @@ class Deforest(using TL, Raise, Elaborator.State):
         constrain(NoProd, store(funSymsWithoutDefn).asConsStrat)
 
 
-    // TODO: ctor as a function?
     def getStratOfSym(s: Symbol) =
       s match
         case _: BuiltinSymbol => NoProd
         case _: TopLevelSymbol => NoProd
+        // TODO: cannot fuse intermediate values created by
+        // calling data constructors passed around like functions,
+        // like `fun app(ctor) = ctor(1); if app(AA) is AA(x) then x`;
+        // immediate data constructor calls are handled directly,
+        // so if this method is called on a ClsLike symbol,
+        // it means that this constructor is passed around like a function,
+        // which we can't fuse for now
+        case _ if s.asCls.isDefined => NoProd
         case _: BlockMemberSymbol => store(s)
         case _: LocalSymbol => store(s)
     def +=(e: Symbol -> ProdVar) = store += e
@@ -383,7 +396,7 @@ class Deforest(using TL, Raise, Elaborator.State):
   ): ProdStrat = b match
     case m@Match(scrut, arms, dflt, rest) =>
       val scrutStrat = processResult(scrut)
-      constrain(scrutStrat, Dtor(scrut.uid)(m, matching.lastOption.map(_._1))(using this))
+      constrain(scrutStrat, Dtor(m, matching.lastOption.map(_._1)))
       val armsRes = if arms.forall{ case (cse, _) => cse.isInstanceOf[Case.Cls] } then
         arms.map:
           case (Case.Cls(s, _), body) => 
@@ -415,19 +428,19 @@ class Deforest(using TL, Raise, Elaborator.State):
         case FunDefn(_, sym, params, body) =>
           val funSymStratVar = symToStrat(sym)
           val param = params.head match
-            case ParamList(flags, params, restParam) => params
-          val funStrat = constrFun(param, body) // TODO: handle mutiple param list
+            case ParamList(flags, params, N) => params // TODO: handle mutiple param list
+          val funStrat = constrFun(param, body)
           constrain(funStrat, funSymStratVar.asConsStrat)
           funSymStratVar
         case v: ValDefn => throw NotDeforestableException("No support for `ValDefn` yet")
-        // only handle code inside module for now to show the
-        // todo case of if scrut being not the same as what the user writes
         case c: ClsLikeDefn => throw NotDeforestableException("No support for `ClsLikeDefn` yet")
       processBlock(rest)
     case End(msg) => NoProd
     // make it a type var instead of `NoProd` so that things like `throw match error` in
     // default else branches do not block fusion...
-    case Throw(exc) => freshVar("throw")._1
+    case Throw(exc) =>
+      processResult(exc)
+      freshVar("throw")._1
 
   def constrFun(params: Ls[Param], body: Block)(using
     inArm: LinkedHashMap[ProdVar, ClsOrModSymbol],
@@ -444,11 +457,9 @@ class Deforest(using TL, Raise, Elaborator.State):
   def processResult(r: Result)(using
     inArm: LinkedHashMap[ProdVar, ClsOrModSymbol],
     matching: LinkedHashMap[ResultId, ClsOrModSymbol]
-  ): ProdStrat = r match
-    case c@Call(f, args) =>
-      val argsTpe = args.map:
-        case Arg(false, value) => processResult(value)
-
+  ): ProdStrat =
+    def handleCallLike(f: Path, args: Ls[Path], c: Result) =
+      val argsTpe = args.map(processResult)
       f match
         case s@Select(p, nme) =>
           s.symbol.map(_.asCls) match
@@ -485,8 +496,10 @@ class Deforest(using TL, Raise, Elaborator.State):
         case Value.This(sym) => throw NotDeforestableException("No support for `this` as a callee yet")
         case Value.Lit(lit) => ???
         case Value.Arr(elems) => ???
+    r match
+    case c@Call(f, args) => handleCallLike(f, args.map {case Arg(false, value) => value}, c)
 
-    case Instantiate(cls, args) => throw NotDeforestableException("No support for `instantiate` yet")
+    case i@Instantiate(cls, args) => handleCallLike(cls, args, i)
 
     case sel@Select(p, nme) => sel.symbol match
       case Some(s) if s.asObj.isDefined =>
@@ -520,26 +533,26 @@ class Deforest(using TL, Raise, Elaborator.State):
   val lowerBounds = mutable.Map.empty[StratVarId, Ls[ProdStrat]].withDefaultValue(Nil)
 
   case class CtorDest(matches: Map[ResultId, Match], sels: Set[FieldSel], noCons: Bool)
-  case class DtorSource(ctors: Set[CtorExpr], noProd: Bool)
+  case class DtorSource(ctors: Set[ResultId], noProd: Bool)
   object ctorDests:
     val ctorDests = mutable.Map.empty[ResultId, CtorDest].withDefaultValue(CtorDest(Map.empty, Set.empty, false))
-    def update(ctor: CtorExpr, m: Match) = ctorDests.updateWith(ctor):
+    def update(ctor: ResultId, m: Match) = ctorDests.updateWith(ctor):
       case Some(CtorDest(matches, sels, noCons)) => Some(CtorDest(matches + (m.scrut.uid -> m), sels, noCons))
       case None => Some(CtorDest(Map(m.scrut.uid -> m), Set.empty, false))
-    def update(ctor: CtorExpr, s: FieldSel) = ctorDests.updateWith(ctor):
+    def update(ctor: ResultId, s: FieldSel) = ctorDests.updateWith(ctor):
       case Some(CtorDest(matches, sels, noCons)) => Some(CtorDest(matches, sels + s, noCons))
       case None => Some(CtorDest(Map.empty, Set(s), false))
-    def update(ctor: CtorExpr, n: NoCons.type) = ctorDests.updateWith(ctor):
+    def update(ctor: ResultId, n: NoCons.type) = ctorDests.updateWith(ctor):
       case Some(CtorDest(matches, sels, noCons)) => Some(CtorDest(matches, sels, true))
       case None => Some(CtorDest(Map.empty, Set.empty, true))
-    def get(ctor: CtorExpr) = ctorDests.get(ctor)
+    def get(ctor: ResultId) = ctorDests.get(ctor)
 
   object dtorSources:
     val dtorSources = mutable.Map.empty[DtorExpr, DtorSource].withDefaultValue(DtorSource(Set.empty, false))
     private def getDtorExprOfResultId(i: ResultId) = i.getResult match
       case s: Select => DtorExpr.Sel(i)
       case r: Value.Ref => DtorExpr.Match(i)
-      case _ => ??? // unreachable
+      case r => lastWords(s"try to get dtor expr from ResultId, but get $r")
     def update(dtor: ResultId, ctor: ResultId) =
       val dtorExpr = getDtorExprOfResultId(dtor)
       dtorSources.updateWith(dtorExpr):
@@ -597,7 +610,7 @@ class Deforest(using TL, Raise, Elaborator.State):
                   handle(prod -> u)
                 else
                   ()
-              case (_: ProdVar, _) => ??? // unreachable, should be handled above
+              case (_: ProdVar, _) => die
               case _ => handle(prod -> u)
         case (Ctor(ctor, args, expr), NoCons) =>
           ctorDests.update(expr, NoCons)
@@ -623,10 +636,10 @@ class Deforest(using TL, Raise, Elaborator.State):
   // ======== after resolving constraints ======
 
   lazy val resolveClashes =
-    type CtorToDtor = Map[CtorExpr, CtorDest]
+    type CtorToDtor = Map[ResultId, CtorDest]
     type DtorToCtor = Map[DtorExpr, DtorSource]
 
-    def removeCtor(ctorDests: CtorToDtor, dtorSources: DtorToCtor, rm: Set[CtorExpr]): CtorToDtor -> DtorToCtor =
+    def removeCtor(ctorDests: CtorToDtor, dtorSources: DtorToCtor, rm: Set[ResultId]): CtorToDtor -> DtorToCtor =
       if rm.isEmpty then
         ctorDests -> dtorSources
       else
@@ -666,7 +679,7 @@ class Deforest(using TL, Raise, Elaborator.State):
 
 
     val removeCycle = {
-      def getCtorInArm(ctor: CtorExpr, dtor: Match): Set[CtorExpr] =
+      def getCtorInArm(ctor: ResultId, dtor: Match): Set[ResultId] =
         val ctorSym = getClsSymOfUid(ctor)
         val arm = dtor.arms.find{ case (Case.Cls(c1, _) -> body) => c1 === ctorSym }.map(_._2).orElse(dtor.dflt).get
 
@@ -688,9 +701,9 @@ class Deforest(using TL, Raise, Elaborator.State):
         GetCtorsTraverser.applyBlock(arm)
         GetCtorsTraverser.ctors.toSet
 
-      def findCycle(ctor: CtorExpr, dtor: Match): Set[CtorExpr] =
+      def findCycle(ctor: ResultId, dtor: Match): Set[ResultId] =
         val cache = mutable.Set(ctor)
-        def go(ctorAndMatches: Set[CtorExpr -> Match]): Set[CtorExpr] =
+        def go(ctorAndMatches: Set[ResultId -> Match]): Set[ResultId] =
           val newCtorsAndNewMatches =
             ctorAndMatches.flatMap((c, m) => getCtorInArm(c, m)).flatMap: c =>
               removeClashes._1.get(c).flatMap:
@@ -717,8 +730,8 @@ class Deforest(using TL, Raise, Elaborator.State):
 
 
 
-  lazy val filteredCtorDests: Map[CtorExpr, CtorFinalDest] =
-    val res = mutable.Map.empty[CtorExpr, CtorFinalDest]
+  lazy val filteredCtorDests: Map[ResultId, CtorFinalDest] =
+    val res = mutable.Map.empty[ResultId, CtorFinalDest]
 
     // we need only one CtorFinalDest per arm for each pat mat expr
     val handledMatches = mutable.Map.empty[ResultId -> ClsOrModSymbol, Opt[CtorFinalDest]]
@@ -773,7 +786,7 @@ class Deforest(using TL, Raise, Elaborator.State):
               throw Error("more than one consumer")
               None
           )
-        else ???
+        else die
       }
       res.updateWith(ctor){_ => filteredDtor}
     }
@@ -789,7 +802,7 @@ class Deforest(using TL, Raise, Elaborator.State):
   }.toSet
 
   def rewrite(p: Block) =
-    val deforestTransformer = DeforestTransformer(using this)
+    val deforestTransformer = DeforestTransformer()
     val rest = deforestTransformer.applyBlock(p)
     val newDefsRest = deforestTransformer.matchRest.getAllFunDefs
     val newDefsArms = deforestTransformer.matchArms.getAllFunDefs
@@ -1033,13 +1046,13 @@ class DeforestTransformer(using val d: Deforest, elabState: Elaborator.State) ex
       super.applyResult(r)
     case _ => super.applyResult(r)
 
-  override def applyResult2(r: Result)(k: Result => Block): Block = r match
-    case call@Call(f, args) if d.filteredCtorDests.isDefinedAt(call.uid) =>
+  override def applyResult2(r: Result)(k: Result => Block): Block = 
+    def handleCallLike(f: Path, args: Ls[Path], uid: ResultId) =
       val c = f match
         case s: Select => s.symbol.get.asCls.get
         case Value.Ref(l) => l.asCls.get
         case _ => ???
-      d.filteredCtorDests.get(call.uid).get match
+      d.filteredCtorDests.get(uid).get match
         case CtorFinalDest.Match(scrut, expr, sels, selsMap) =>
           // use pre-determined symbols, create temp symbols for un-used fields
           val usedFieldIdentToSymbolsToBeReplaced = selsMap._1
@@ -1066,16 +1079,21 @@ class DeforestTransformer(using val d: Deforest, elabState: Elaborator.State) ex
             selsMap._1 -> selsMap._2)
 
           args.zip(assignedTempSyms.map(_._2)).foldRight[Block](k(bodyAndRestInLam)):
-            case ((a, tmp), rest) => applyResult2(a.value) { r => Assign(tmp, r, rest) }
+            case ((a, tmp), rest) => applyResult2(a) { r => Assign(tmp, r, rest) }
 
         case CtorFinalDest.Sel(s) =>
           val selFieldName = s.getResult match { case Select(p, nme) => nme }
           val idx = d.getClsFields(c).indexWhere(s => s.id === selFieldName)
-          k(args(idx).value)
+          k(args(idx))
 
+    r match
+    case call@Call(f, args) if d.filteredCtorDests.isDefinedAt(call.uid) =>
+      handleCallLike(f, args.map { case Arg(false, value) => value }, call.uid)
+    case ins@Instantiate(cls, args) if d.filteredCtorDests.isDefinedAt(ins.uid) =>
+      handleCallLike(cls, args, ins.uid)
     case _ => super.applyResult2(r)(k)
 
-  def handleObjFusing(objCallExprUid: CtorExpr, objClsSym: ModuleSymbol) =
+  def handleObjFusing(objCallExprUid: ResultId, objClsSym: ModuleSymbol) =
     // must be a pat mat on objects; no support for selection on objects yet
     val CtorFinalDest.Match(scrut, expr, sels, selsMap) = d.filteredCtorDests(objCallExprUid): @unchecked
     val body = expr.arms.find{ case (Case.Cls(m, _) -> body) => m === objClsSym }.map(_._2).orElse(expr.dflt).get