Deforestation with nofibs in the Benchmark Sub-project

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

@@ -48,7 +48,15 @@ class CompileTestRunner
 
         val preludePath = mainTestDir/"mlscript"/"decls"/"Prelude.mls"
 
-        given Config = Config.default
+        given Config =
+          if file.segments.contains("nofib-deforest") then
+            Config.default.copy(deforest = S(Config.Deforestation(
+              importedPublicModNames = Set("NofibPrelude"),
+              seeThroughLazySymbolsNames = Set("lazy"),
+              seeThroughForceSymbolsNames = Set("force")
+            )))
+          else
+            Config.default
 
         val compiler = MLsCompiler(
           preludePath,

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

@@ -14,6 +14,7 @@ case class Config(
   sanityChecks: Opt[SanityChecks],
   effectHandlers: Opt[EffectHandlers],
   liftDefns: Opt[LiftDefns],
+  deforest: Opt[Deforestation]
 ):
 
   def stackSafety: Opt[StackSafety] = effectHandlers.flatMap(_.stackSafety)
@@ -28,6 +29,7 @@ object Config:
     // sanityChecks = S(SanityChecks(light = true)),
     effectHandlers = N,
     liftDefns = N,
+    deforest = N
   )
 
   case class SanityChecks(light: Bool)
@@ -42,6 +44,12 @@ object Config:
 
   case class LiftDefns() // there may be other settings in the future, having it as a case class now
 
+  case class Deforestation(
+    val importedPublicModNames: Set[String],
+    val seeThroughLazySymbolsNames: Set[String],
+    val seeThroughForceSymbolsNames: Set[String]
+  )
+
 end Config
 
 

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

@@ -10,6 +10,8 @@ import hkmc2.semantics.Elaborator
 import hkmc2.semantics.Resolver
 import hkmc2.syntax.Keyword.`override`
 import semantics.Elaborator.{Ctx, State}
+import hkmc2.codegen.deforest.Deforest
+import hkmc2.Config.LiftDefns
 
 
 class ParserSetup(file: os.Path, dbgParsing: Bool)(using Elaborator.State, Raise):
@@ -37,7 +39,7 @@ class ParserSetup(file: os.Path, dbgParsing: Bool)(using Elaborator.State, Raise
 
 
 // * The weird type of `mkOutput` is to allow wrapping the reporting of diagnostics in synchronized blocks
-class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Unit)(using Config):
+class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Unit)(using cfg: Config):
 
   val runtimeFile: os.Path = preludeFile/os.up/os.up/os.up/"mlscript-compile"/"Runtime.mjs"
   val termFile: os.Path = preludeFile/os.up/os.up/os.up/"mlscript-compile"/"Term.mjs"
@@ -94,6 +96,32 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
       val jsb = ltl.givenIn:
         codegen.js.JSBuilder()
       val le = low.program(blk)
+
+      val lowered =
+        if cfg.deforest.isEmpty then
+          le
+        else
+          val deforestLow = ltl.givenIn:
+            cfg.copy(liftDefns = S(LiftDefns())).givenIn:
+              new codegen.Lowering()
+          val deforestResult = Deforest.apply(deforestLow.program(blk), wd)(using
+            cfg,
+            ltl,
+            raise,
+            newCtx,
+            new Deforest.State(),
+            State,
+            preludeFile)
+          deforestResult match
+            case Right(msg) =>
+              println(msg) // TODO: no println
+              le
+            case Left(prog -> summary -> detail) =>
+              // if summary.nonEmpty then
+                // println("-----summary-----")
+                // println(summary.mapLines(l => s"\t$l"))
+              prog
+
       val baseScp: utils.Scope =
         utils.Scope.empty
       // * This line serves for `import.meta.url`, which retrieves directory and file names of mjs files.
@@ -103,7 +131,7 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
       val nme = file.baseName
       val exportedSymbol = parsed.definedSymbols.find(_._1 === nme).map(_._2)
       val je = nestedScp.givenIn:
-        jsb.program(le, exportedSymbol, wd)
+        jsb.program(lowered, exportedSymbol, wd)
       val jsStr = je.stripBreaks.mkString(100)
       val out = file / os.up / (file.baseName + ".mjs")
       os.write.over(out, jsStr)

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

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

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

@@ -227,9 +227,18 @@ sealed abstract class Block extends Product:
             case f@FunDefn(owner, sym, params, body) =>
               val newBody = body.flattened
               if newBody is body then f else f.copy(body = newBody)
-          if (newPreCtor is c.preCtor) && (newCtor is c.ctor) && (newMethods is c.methods)
+          val newCompanion = c.companion.mapConserve:
+            case b@ClsLikeBody(isym, methods, privateFields, publicFields, ctor) =>
+              val newMethods = methods.mapConserve:
+                case f@FunDefn(owner, sym, params, body) =>
+                  val newBody = body.flattened
+                  if newBody is body then f else f.copy(body = newBody)
+              val newCtor = ctor.flattened
+              if (newMethods is methods) && (newCtor is ctor) then b
+              else b.copy(methods = newMethods, ctor = newCtor)
+          if (newPreCtor is c.preCtor) && (newCtor is c.ctor) && (newMethods is c.methods) && (newCompanion is c.companion)
           then c
-          else c.copy(preCtor = newPreCtor, ctor = newCtor, methods = newMethods)
+          else c.copy(preCtor = newPreCtor, ctor = newCtor, methods = newMethods, companion = newCompanion)
 
       val newRest = rest.flatten(k)
       if (newDefn is defn) && (newRest is rest)
@@ -478,6 +487,11 @@ enum Case:
 
 sealed trait TrivialResult extends Result
 
+object Result:
+  opaque type ResultId = Int
+  private def ResultId(v: Int): ResultId = v
+
+
 sealed abstract class Result extends AutoLocated:
 // // * Used for debugging locations:
 // sealed abstract class Result extends AutoLocated with ProductWithExtraInfo:
@@ -538,6 +552,13 @@ sealed abstract class Result extends AutoLocated:
       args.flatMap(arg => arg.idx.fold(Set.empty)(_.freeVarsLLIR) ++ arg.value.freeVarsLLIR).toSet
     case DynSelect(qual, fld, arrayIdx) => qual.freeVarsLLIR ++ fld.freeVarsLLIR
 
+  def uid =
+    import Result.*
+    val uidValue = ResultId(System.identityHashCode(this))
+    uidValue
+
+
+
 // type Local = LocalSymbol
 type Local = Symbol
 

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

@@ -13,7 +13,7 @@ object LambdaRewriter:
 
     def rewriteOneBlk(b: Block) = b match
       case Assign(lhs, Value.Lam(params, body), rest) if !lhs.isInstanceOf[TempSymbol] =>
-        val newSym = BlockMemberSymbol(lhs.nme, Nil,
+        val newSym = BlockMemberSymbol(lhs.nme, syntax.Tree.DummyTermDef(syntax.Fun) :: Nil,
           nameIsMeaningful = true // TODO: lhs.nme is not always meaningful
         )
         val blk = blockBuilder
@@ -26,7 +26,10 @@ object LambdaRewriter:
         val lambdaRewriter = new BlockDataTransformer(SymbolSubst()):
           override def applyValue(v: Value): Value = v match
             case lam: Value.Lam => 
-              val sym = BlockMemberSymbol("lambda", Nil, nameIsMeaningful = false)
+              val sym = BlockMemberSymbol(
+                "lambda",
+                syntax.Tree.DummyTermDef(syntax.Fun) :: Nil,
+                nameIsMeaningful = false)
               lambdasList ::= (sym -> super.applyLam(lam))
               Value.Ref(sym)
             case _ => super.applyValue(v)

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

@@ -86,7 +86,7 @@ object Lifter:
   def getVars(d: Defn)(using state: State): Set[Local] = d match
     case f: FunDefn =>
       (f.body.definedVars ++ f.params.flatMap(_.paramSyms)).collect:
-        case s: FlowSymbol if !(s is state.runtimeSymbol) => s
+        case s: FlowSymbol if !(s is state.runtimeSymbol) => s // FIXME: doesn't this test always fail?
     case c: ClsLikeDefn =>      
       val companionVars = c.companion.fold(Set.empty)(_.ctor.definedVars)
       (companionVars ++ c.preCtor.definedVars ++ c.ctor.definedVars).collect:
@@ -551,9 +551,9 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
       val fakeCtorBms = d match
         case c: ClsLikeDefn if !modLocal.isDefined => S(BlockMemberSymbol(d.sym.nme + "$ctor", Nil))
         case _ => N
-      
-      val singleCallBms = BlockMemberSymbol(d.sym.nme + "$", Nil)
-      
+
+      val singleCallBms = BlockMemberSymbol(d.sym.nme + "$", Tree.DummyTermDef(syntax.Fun) :: Nil)
+
       val info = LiftedInfo(
         includedCaptures, includedLocals, clsCaptures,
         refBms, fakeCtorBms, singleCallBms

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

@@ -311,7 +311,9 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
     case st.Asc(lhs, rhs) =>
       term(lhs, inStmtPos = inStmtPos)(k)
     case st.Tup(fs) =>
-      args(fs)(args => k(Value.Arr(mut = false, args)))
+      args(fs): args =>
+        val tmpSym = TempSymbol(N, "arr")
+        Assign(tmpSym, Value.Arr(mut = false, args), k(Value.Ref(tmpSym)))
     case Mut(st.Tup(fs)) =>
       args(fs)(args => k(Value.Arr(mut = true, args)))
     case st.CtxTup(fs) =>