Deforestation with nofibs in the Benchmark Sub-project

build.sbt

@@ -72,11 +72,16 @@ lazy val hkmc2DiffTests = project.in(file("hkmc2DiffTests"))
     Test/run/fork := true, // so that CTRL+C actually terminates the watcher
   )
 
-lazy val hkmc2AllTests = project.in(file("hkmc2AllTests"))
+lazy val hkmc2MainTests = project.in(file("hkmc2MainTests"))
   .settings(
     Test / test := ((hkmc2DiffTests / Test / test) dependsOn (hkmc2JVM / Test / test)).value
   )
 
+lazy val hkmc2AllTests = project.in(file("hkmc2AllTests"))
+  .settings(
+    Test / test := ((hkmc2Benchmarks / Test / test) dependsOn hkmc2MainTests / Test / test).value
+  )
+
 lazy val core = crossProject(JSPlatform, JVMPlatform).in(file("core"))
   .settings(
     sourceDirectory := baseDirectory.value.getParentFile()/"shared",

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(
+              seethroughModules = Set(workingDir/"hkmc2Benchmarks"/"src"/"test"/"bench"/"mlscript-compile"/"nofib-deforest"/"NofibPrelude.mls"),
+              seeThroughLazySymbolsNames = Set("lazy"),
+              seeThroughForceSymbolsNames = Set("force")
+            )))
+          else
+            Config.default
 
         val compiler = MLsCompiler(
           preludePath,

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

@@ -20,6 +20,7 @@ case class Config(
   sanityChecks: Opt[SanityChecks],
   effectHandlers: Opt[EffectHandlers],
   liftDefns: Opt[LiftDefns],
+  deforest: Opt[Deforestation],
   stageCode: Bool,
   target: CompilationTarget,
 ):
@@ -36,6 +37,7 @@ object Config:
     // sanityChecks = S(SanityChecks(light = true)),
     effectHandlers = N,
     liftDefns = N,
+    deforest = N,
     stageCode = false,
     target = CompilationTarget.JS
   )
@@ -52,6 +54,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 seethroughModules: Set[os.Path],
+    val seeThroughLazySymbolsNames: Set[String],
+    val seeThroughForceSymbolsNames: Set[String]
+  )
+
 end Config
 
 

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

@@ -8,8 +8,11 @@ import utils.*
 import hkmc2.semantics.MemberSymbol
 import hkmc2.semantics.Elaborator
 import hkmc2.semantics.Resolver
+import hkmc2.semantics.BlockMemberSymbol
 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 +40,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"
@@ -96,16 +99,38 @@ class MLsCompiler(preludeFile: os.Path, mkOutput: ((Str => Unit) => Unit) => Uni
       val jsb = ltl.givenIn:
         codegen.js.JSBuilder()
       val le = low.program(blk)
+
+      val nme = file.baseName
+      val originalExportedSymbol = parsed.definedSymbols.find(_._1 === nme).map(_._2)
+      val (lowered -> exportedSymbol) =
+        if cfg.deforest.isEmpty then
+          le -> originalExportedSymbol
+        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)
+          deforestResult match
+            case Right(msg) => le -> originalExportedSymbol
+            case Left(prog -> renewSym) =>
+              prog ->
+              originalExportedSymbol.map: s =>
+                renewSym.applyLocal(s).asInstanceOf[BlockMemberSymbol]
+
       val baseScp: utils.Scope =
         utils.Scope.empty
       // * This line serves for `import.meta.url`, which retrieves directory and file names of mjs files.
       // * Having `module id"import" with ...` in `prelude.mls` will generate `globalThis.import` that is undefined.
       baseScp.addToBindings(Elaborator.State.importSymbol, "import", shadow = false)
       val nestedScp = baseScp.nest
-      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,8 @@ object Uid:
         curUid
       def reset = curUid = -1
   object Symbol extends Handler[semantics.Symbol]
+  object StratVar extends Handler[codegen.deforest.StratVarState]
+  object Result extends Handler[codegen.Result]
 
 extension [T] (x: Uid[T])
   def <=(rhs: Uid[T]) = x <= rhs

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

@@ -228,9 +228,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)
@@ -540,6 +549,16 @@ sealed abstract class Result extends AutoLocated:
     case Value.Lit(lit) => Set.empty
     case DynSelect(qual, fld, arrayIdx) => qual.freeVarsLLIR ++ fld.freeVarsLLIR
 
+
+  def uid(using ds: deforest.Deforest.State): Uid[Result] =
+    ds.resultToResultId.getOrElseUpdate(
+      this,
+      locally:
+        val id = ds.ResultUidState.nextUid
+        ds.resultIdToResult.addOne(id -> this)
+        id)
+
+
 // type Local = LocalSymbol
 type Local = Symbol
 

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

@@ -13,7 +13,9 @@ object LambdaRewriter:
 
     def rewriteOneBlk(b: Block) = b match
       case Assign(lhs, Lambda(params, body), rest) if !lhs.isInstanceOf[TempSymbol] =>
-        val newSym = BlockMemberSymbol(lhs.nme, Nil,
+        // use dummy TermDef of kind syntax.Fun here so that we know this symbol refers to a function
+        // for deforestation
+        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 +28,10 @@ object LambdaRewriter:
         val lambdaRewriter = new BlockDataTransformer(SymbolSubst()):
           override def applyResult(r: Result)(k: Result => Block): Block = r match
             case lam: Lambda => 
-              val sym = BlockMemberSymbol("lambda", Nil, nameIsMeaningful = false)
+              val sym = BlockMemberSymbol(
+                "lambda",
+                syntax.Tree.DummyTermDef(syntax.Fun) :: Nil,
+                nameIsMeaningful = false)
               lambdasList ::= (sym -> super.applyLam(lam))
               k(Value.Ref(sym))
             case _ => super.applyResult(r)(k)

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

@@ -573,7 +573,7 @@ class Lifter(handlerPaths: Opt[HandlerPaths])(using State, Raise):
         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,

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

@@ -447,7 +447,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
           val isAnd = sym is State.andSymbol
           val isOr = sym is State.orSymbol
           if isAnd || isOr then
-            val lamSym = BlockMemberSymbol("lambda", Nil, false)
+            val lamSym = BlockMemberSymbol("lambda", syntax.Tree.DummyTermDef(syntax.Fun) :: Nil, false)
             val lamDef = FunDefn(N, lamSym, PlainParamList(Nil) :: Nil, returnedTerm(arg2))
             Define(
               lamDef,
@@ -566,7 +566,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
       if k.isInstanceOf[TailOp] || bodyBlock.size <= 5
       then k(Lambda(paramLists.head, bodyBlock))
       else
-        val lamSym = new BlockMemberSymbol("lambda", Nil, false)
+        val lamSym = new BlockMemberSymbol("lambda", syntax.Tree.DummyTermDef(syntax.Fun) :: Nil, false)
         val lamDef = FunDefn(N, lamSym, paramLists, bodyBlock)
         Define(
           lamDef,
@@ -902,7 +902,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
       case v: Value => k(v)
       case p: Path => k(p)
       case Lambda(params, body) =>
-        val lamSym = BlockMemberSymbol("lambda", Nil, false)
+        val lamSym = BlockMemberSymbol("lambda", syntax.Tree.DummyTermDef(syntax.Fun) :: Nil, false)
         val lamDef = FunDefn(N, lamSym, params :: Nil, body)
         Define(lamDef, k(lamSym |> Value.Ref.apply))
       case r =>