Deforestation mk4

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

@@ -50,8 +50,12 @@ object algorithms {
     )
 
     // 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 nodesUniq = {
+      val seen = collection.mutable.LinkedHashSet.empty[A]
+      for ((a, b) <- edges) { seen.add(a); seen.add(b) }
+      for (n <- nodes) seen.add(n)
+      seen.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

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

@@ -26,6 +26,7 @@ case class Config(
   target: CompilationTarget,
   rewriteWhileLoops: Bool,
   tailRecOpt: Bool,
+  deforest: Opt[Deforest],
   qqEnabled: Bool,
   funcToCls: Bool
 ):
@@ -60,6 +61,7 @@ object Config:
     rewriteWhileLoops = false,
     stageCode = false,
     tailRecOpt = true,
+    deforest = N,
     qqEnabled = false,
     funcToCls = false,
   )
@@ -90,6 +92,11 @@ object Config:
 
   case class LiftDefns() // there may be other settings in the future, having it as a case class now
 
+  case class Deforest(val debug: Boolean)
+
+  object Deforest:
+    val default = Deforest(true)
+
 end Config
 
 

hkmc2/shared/src/main/scala/hkmc2/Uid.scala

@@ -13,6 +13,8 @@ object Uid:
         curUid
       def reset = curUid = -1
   object Symbol extends Handler[semantics.Symbol]
+  object Result extends Handler[codegen.Result]
+  object StratVar extends Handler[codegen.deforest.StratVar]
 
 extension [T] (x: Uid[T])
   def <=(rhs: Uid[T]) = x <= rhs

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

@@ -479,7 +479,9 @@ final case class FunDefn(
   val asPath = Value.Ref(sym, S(dSym))
 object FunDefn:
   def withFreshSymbol(owner: Opt[InnerSymbol], sym: BlockMemberSymbol, params: Ls[ParamList], body: Block)(forceTailRec: Bool)(using State) =
-    FunDefn(owner, sym, TermSymbol(syntax.Fun, owner, Tree.Ident(sym.nme)), params, body)(forceTailRec)
+    val tSym = TermSymbol(syntax.Fun, owner, Tree.Ident(sym.nme))
+    sym.tsym = S(tSym)
+    FunDefn(owner, sym, tSym, params, body)(forceTailRec)
 
 final case class ValDefn(
     tsym: TermSymbol,

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

@@ -13,13 +13,11 @@ object LambdaRewriter:
   def desugar(b: Block)(using State) =
     val transformer = new BlockTransformer(SymbolSubst()):
       override def applyResult(r: Result)(k: Result => Block): Block = r match
-        case lam: Lambda => 
+        case lam: Lambda =>
           val sym = BlockMemberSymbol("lambda", Nil, nameIsMeaningful = false)
-          val tSym = TermSymbol.fromFunBms(sym, N)
-          val lamDefn =
-            val Lambda(params, body) = super.applyLam(lam)
-            FunDefn(N, sym, tSym, params :: Nil, body)(false)
-          Scoped(Set(sym), Define(lamDefn, k(Value.Ref(sym, S(tSym)))))
+          val Lambda(params, body) = super.applyLam(lam)
+          val lamDefn = FunDefn.withFreshSymbol(N, sym, params :: Nil, body)(false)
+          Scoped(Set.single(sym), Define(lamDefn, k(lamDefn.asPath)))
         case _ => super.applyResult(r)(k)
 
       override def applyBlock(b: Block): Block = b match
@@ -32,7 +30,7 @@ object LambdaRewriter:
             .define(defn)
             .assign(lhs, defn.asPath)
             .rest(applyBlock(rest))
-          Scoped(Set(newSym), blk)
+          Scoped(Set.single(newSym), blk)
         case _ => super.applyBlock(b)
 
     transformer.applyBlock(b)

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

@@ -1056,9 +1056,27 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
 
     val desug = LambdaRewriter.desugar(blk)
 
+    val deforested =
+      val outterTl = tl
+      config.deforest match
+        case None => desug
+        case Some(dCfg) =>
+          /*
+          // * For some weird reason (Scala bug?),
+          // * the version below leads to a stack overflows during its initialization
+          given TraceLogger with
+            override def doTrace: Bool = dCfg.debug
+            override def emitDbg(str: Str): Unit = outterTl.emitDbg(s"deforest > $str")
+          */
+          (new TraceLogger:
+            override def doTrace: Bool = dCfg.debug
+            override def emitDbg(str: Str): Unit = outterTl.emitDbg(s"deforest > $str")
+          ).givenIn:
+            deforest.Deforest(Program(imps.map(imp => imp.sym -> imp.str), desug)).main
+
     val handlerPaths = new HandlerPaths
 
-    val withHandlers1 = config.effectHandlers.fold(desug): opt =>
+    val withHandlers1 = config.effectHandlers.fold(deforested): opt =>
       HandlerLowering(handlerPaths, opt).translateHandleBlocks(desug)
 
     val shouldFlattenScopes = config.effectHandlers.isDefined