Compilation Pipeline

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

@@ -36,6 +36,7 @@ case class Config(
   funcToCls: Bool,
   commentGeneratedCode: Bool,
   noFreeze: Bool,
+  noOpt: Bool,
   noModuleCheck: Bool,
   deadParamElim: Opt[DeadParamElim],
 ):
@@ -80,6 +81,7 @@ object Config:
     funcToCls = false,
     commentGeneratedCode = false,
     noFreeze = false,
+    noOpt = false,
     noModuleCheck = false,
     deadParamElim = S(DeadParamElim.default)
   )
@@ -577,6 +579,8 @@ object ConfigParser:
     case "language" => parseLanguageOverride(value)
     case "tailRecOpt" =>
       parsedField(value)(parseBool)(v => _.copy(tailRecOpt = v))
+    case "noOpt" =>
+      parsedField(value)(parseBool)(v => _.copy(noOpt = v))
     case "noFreeze" =>
       parsedField(value)(parseBool)(v => _.copy(noFreeze = v))
     case "noModuleCheck" =>

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

@@ -6,6 +6,7 @@ import hkmc2.utils.*, shorthands.*
 import hkmc2.io
 import utils.*
 
+import hkmc2.codegen.CompilationPipeline
 import hkmc2.semantics.*
 import hkmc2.syntax.Keyword.`override`
 import semantics.Elaborator.{Ctx, State}
@@ -117,16 +118,13 @@ class MLsCompiler
         new codegen.Lowering()
       val jsb = ltl.givenIn:
         codegen.js.JSBuilder()
-      val lowered = low.program(blk, symbolsToPreserve = Set.empty)
-      var optimized = lowered
+      val lowered = ltl.givenIn:
+        low.program(blk, symbolsToPreserve = Set.empty)
       val nme = file.baseName
       val exportedSymbol = parsed.definedSymbols.find(_._1 === nme).map(_._2)
-      optimized =
-        val printer = (p: codegen.Program) => p.showAsTree // TODO: proper printing like in diff-tests
-        optimized = codegen.WorkerWrapper(exportedSymbol.toSet, dtl, printer)(optimized)
-        codegen.BlockSimplifier(exportedSymbol.toSet, dtl, printer)(optimized)
-      ltl.givenIn:
-        optimized = codegen.DeadParamElim(optimized)
+      val optimized = ltl.givenIn:
+        val printer = (p: codegen.Program) => p.showAsTree
+        CompilationPipeline().run(lowered, printer, exportedSymbol.toSet, dtl)
       val baseScp: utils.Scope =
         utils.Scope.empty(utils.Scope.Cfg.default)
       // * This line serves for `import.meta.url`, which retrieves directory and file names of mjs files.

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

@@ -15,7 +15,7 @@ import hkmc2.Message.MessageContext
 import hkmc2.syntax.Tree.DummyTypeDef
 
 class BufferableTransform()(using Ctx, State, Raise):
-  def transform(blk: Block): Block =
+  def transform(prog: Program): Program =
     val transformer = new BlockTransformer(SymbolSubst.Id):
       override def applyDefn(defn: Defn)(k: Defn => Block): Block = defn match
         case cls: ClsLikeDefn if cls.k is syntax.Cls =>
@@ -113,4 +113,4 @@ class BufferableTransform()(using Ctx, State, Raise):
                 cls.bufferable,
               )(cls.configOverride, cls.annotations)
         case _ => super.applyDefn(defn)(k)
-    transformer.applyBlock(blk)
+    transformer.applyProgram(prog)

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

@@ -105,6 +105,3 @@ end ClassParamFlattener
 object ClassParamFlattener:
   def apply(program: Program)(using State): Program =
     new ClassParamFlattener().applyProgram(program)
-
-  def apply(block: Block)(using State): Block =
-    new ClassParamFlattener().applyBlock(block)

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

@@ -0,0 +1,63 @@
+package hkmc2
+package codegen
+
+import hkmc2.utils.*, shorthands.*
+import utils.*
+
+import hkmc2.Config
+import hkmc2.semantics.Elaborator.{Ctx, State}
+import hkmc2.semantics.SymbolPrinter
+import hkmc2.utils.TL
+
+class CompilationPipeline(using Config, Raise, State, Ctx, SymbolPrinter):
+
+  def preOptimizeHook(prog: Program): Program =
+    prog
+
+  def run(prog: Program, printer: Program => Str, symbolsToPreserve: Set[BoundSymbol], otl: TL)(using TL): Program =
+
+    var result = prog
+
+    result = LambdaRewriter.desugar(result)
+
+    result =
+      val outterTl = tl
+      config.deforest match
+        case None => result
+        case Some(dCfg) =>
+          flowAnalysis.FlowAnalysis.mkTraceLogger(dCfg.config, "deforest > ", outterTl).givenIn:
+            deforest.Deforest(result)
+
+    result = EtaExpansion(result)
+
+    if config.liftDefns.isDefined then
+      result = Program(result.imports, Lifter(result.main).transform)
+
+    result = config.effectHandlers.fold(result): opt =>
+      HandlerLowering(new HandlerPaths, opt).translateProgram(result)
+
+    result = Program(result.imports, result.main.flattened)
+
+    result = BufferableTransform().transform(result)
+
+    // * TODO[Anto]: Can we remove MergeMatchArmTransformer? Seems no longer necessary
+    result = Program(result.imports, MergeMatchArmTransformer.applyBlock(result.main))
+
+    if config.funcToCls then
+      result = Program(result.imports, Lifter(FirstClassFunctionTransformer().transform(result.main)).transform)
+
+    result = ClassParamFlattener(result)
+
+    result = ReflectionInstrumenter(using summon).apply(result)
+
+    if config.tailRecOpt then
+      result = TailRecOpt().transform(result)
+
+    result = preOptimizeHook(result)
+
+    if !summon[Config].noOpt then
+      result = WorkerWrapper(symbolsToPreserve, otl, printer)(result)
+      result = BlockSimplifier(symbolsToPreserve, otl, printer)(result)
+      result = otl.givenIn(DeadParamElim(result))
+
+    result

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

@@ -777,12 +777,17 @@ class HandlerLowering(paths: HandlerPaths, opt: EffectHandlers)(using TL, Raise,
     topLevelPostTransform.applyBlock(b)
 
 
-  def translateTopLevel(b: Block): Block =
+  def translateProgram(prog: Program): Program =
     val ctx = HandlerCtx.TopLevel
-    val transformed = translateBlock(b, ctx, Set.empty)
-    blockBuilder
-      .staticif(
-        !opt.doNotInstrumentTopLevelModCtor,
-        _.assign(NoSymbol, Call(paths.resetEffects, Nil ne_:: Nil)(CallMetadata.defaultMlsFun))
+    val transformed = blockBuilder
+        .staticif(
+          !opt.doNotInstrumentTopLevelModCtor,
+          _.assign(NoSymbol, Call(paths.resetEffects, Nil ne_:: Nil)(CallMetadata.defaultMlsFun))
+        )
+        .rest(translateBlock(prog.main, ctx, Set.empty))
+    if transformed is prog.main then prog
+    else
+      Program(
+        prog.imports,
+        transformed
       )
-      .rest(transformed)

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

@@ -10,7 +10,7 @@ import hkmc2.syntax.Tree
 
 object LambdaRewriter:
 
-  def desugar(b: Block)(using State) =
+  def desugar(b: Program)(using State) =
 
     val transformer = new BlockTransformer(SymbolSubst.Id):
 
@@ -38,6 +38,6 @@ object LambdaRewriter:
           Scoped(Set.single(newSym), blk)
         case _ => super.applyBlock(b)
 
-    transformer.applyBlock(b)
+    transformer.applyProgram(b)
 
 

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

@@ -93,7 +93,6 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
       case t => t
 
   val lowerHandlers: Bool = config.effectHandlers.isDefined
-  val lift: Bool = config.liftDefns.isDefined
 
   private lazy val wasmBinaryIntrinsicMap: Map[Str, Str] = Map(
     "+" -> "plus_impl",
@@ -1364,46 +1363,9 @@ class Lowering()(using Config, TL, Raise, State, Ctx, SymbolPrinter):
       inScopedBlockExcept(symbolsToPreserve)(using LoweringCtx.empty):
         block(funs ::: rest, R(main.res))(ImplctRet)
 
-    val desug = LambdaRewriter.desugar(blk)
-
-    val deforested =
-      val outterTl = tl
-      config.deforest match
-        case None => desug
-        case Some(dCfg) =>
-          flowAnalysis.FlowAnalysis.mkTraceLogger(dCfg.config, "deforest > ", outterTl).givenIn:
-            deforest.Deforest(Program(imps.map(imp => imp.sym -> imp.str), desug)).main
-
-    val etaExpanded =
-      EtaExpansion(Program(imps.map(imp => imp.sym -> imp.str), deforested)).main
-
-    val lifted =
-      if lift then Lifter(etaExpanded).transform
-      else etaExpanded
-
-    val withHandlers = config.effectHandlers.fold(lifted): opt =>
-      HandlerLowering(new HandlerPaths, opt).translateTopLevel(lifted)
-
-    val bufferable = BufferableTransform().transform(withHandlers.flattened)
-
-    // * TODO[Anto]: Can we remove MergeMatchArmTransformer? Seems no longer necessary
-    val merged = MergeMatchArmTransformer.applyBlock(bufferable)
-
-    val funcToCls =
-      if config.funcToCls then Lifter(FirstClassFunctionTransformer().transform(merged)).transform
-      else merged
-
-    val flatClassParams = ClassParamFlattener(funcToCls)
-
-    val staged = ReflectionInstrumenter(using summon).apply(flatClassParams)
-
-    val res =
-      if config.tailRecOpt then TailRecOpt().transform(staged)
-      else staged
-
     Program(
       imps.map(imp => imp.sym -> imp.str),
-      res
+      blk
     )
 
 

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

@@ -430,6 +430,6 @@ class ReflectionInstrumenter(using State, Raise, Ctx) extends BlockTransformer(n
         case _ => super.applyDefn(defn)
     transformer.applyBlock(b)
 
-  def apply(b: Block) =
-    mkDefnMap(b)
-    applyBlock(b)
+  def apply(prog: Program) =
+    mkDefnMap(prog.main)
+    applyProgram(prog)

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

@@ -597,7 +597,7 @@ class TailRecOpt(using State, TL, Raise):
         comp.copy(methods = cMtds)
       c.copy(methods = mtds, companion = companion)(c.configOverride, c.annotations)
 
-  def transform(b: Block) =
+  def transform(prog: Program) =
     /* To avoid `x` being overridden in the following when the lifter is not run:
      * 
      * let lam
@@ -608,7 +608,7 @@ class TailRecOpt(using State, TL, Raise):
      * we need to do some analysis on what nested functions use what variables. We
      * re-use the analysis from the lifter to do this.
      */
-
+    val b = prog.main
     given (ScopeData, AccessMap) = 
       // IgnoredScoes can be an empty set, since that information is only relevant for lifting
       given IgnoredScopes = IgnoredScopes(S(Set.empty))
@@ -665,4 +665,6 @@ class TailRecOpt(using State, TL, Raise):
         case _ => super.applyDefn(defn)
     .applyBlock(result)
 
-    result
+    if result is b
+    then prog
+    else Program(prog.imports, result)

hkmc2DiffTests/src/test/scala/hkmc2/JSBackendDiffMaker.scala

@@ -121,36 +121,37 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
         new codegen.Lowering()
           with codegen.LoweringTraceLog(traceJS.isSet)
 
-      var lowered = low.program(blk, symbolsToPreserve = symbolsToPreserve)
+      val lowered = ltl.givenIn:
+        low.program(blk, symbolsToPreserve = symbolsToPreserve)
 
-      var optimized = lowered
+      var preOptimize: Opt[Program] = N
 
-      if showLoweredTree.isSet then
-        outputSeparator("Lowered IR Tree")
-        output(optimized.showAsTree)
-
-      if showIR.isSet || showIRLines.isSet then
-        given ShowCfg = ShowCfg(
-          showExpansionMappings = false,
-          showFlowSymbols = true,
-          debug = debug.isSet,
-        )
-        val irStr = Printer().worksheet(optimized)(using irPrintingScp).mkString(output.ColWidth)
-        val sloc = irStr.count(_ == '\n') + 1
-        if showIRLines.isSet then output(s"Lines of IR: ${sloc}")
-        if showIR.isSet then
-          outputSeparator("Lowered IR")
-          output(irStr)
-
-      if noOptimizations.isUnset then
-        optimized = WorkerWrapper(symbolsToPreserve, dtl, print)(optimized)
-
-        optimized = BlockSimplifier(symbolsToPreserve, dtl, print)(optimized)
-        ltl.givenIn:
-          optimized = DeadParamElim(optimized)
+      val optimized = ltl.givenIn:
+        val customPipeline = new CompilationPipeline:
+          override def preOptimizeHook(prog: Program) =
+            if showLoweredTree.isSet then
+              outputSeparator("Lowered IR Tree")
+              output(prog.showAsTree)
+            if showIR.isSet || showIRLines.isSet then
+              given ShowCfg = ShowCfg(
+                showExpansionMappings = false,
+                showFlowSymbols = true,
+                debug = debug.isSet,
+              )
+              val irStr = Printer().worksheet(prog)(using irPrintingScp).mkString(output.ColWidth)
+              val sloc = irStr.count(_ == '\n') + 1
+              if showIRLines.isSet then output(s"Lines of IR: ${sloc}")
+              if showIR.isSet then
+                outputSeparator("Lowered IR")
+                output(irStr)
+            preOptimize = S(prog)
+            super.preOptimizeHook(prog)
+        customPipeline.run(lowered, print, symbolsToPreserve, dtl)
 
+      // TODO: remove this and make sure all transformer preserve object identity
+      val loweredTransformed = preOptimize.get
       // TODO: Test that transformers retain object identity when there are no changes
-      if (optimized isnt lowered) && (optimized === lowered) then
+      if (optimized isnt loweredTransformed) && (optimized === loweredTransformed) then
         output("/!\\ Warning: object identity between equal objects was not preserved by BlockSimplifier or DeadParamElim")
         def rec(lhs: Block, rhs: Block): Bool =
           (lhs is rhs) || {
@@ -162,7 +163,7 @@ abstract class JSBackendDiffMaker extends MLsDiffMaker:
               false
             else false
           }
-        rec(optimized.main, lowered.main)
+        rec(optimized.main, loweredTransformed.main)
       if checkIR.isSet then
         BlockChecker().applyProgram(optimized)
 

hkmc2DiffTests/src/test/scala/hkmc2/MLsDiffMaker.scala

@@ -188,6 +188,7 @@ abstract class MLsDiffMaker extends DiffMaker:
       funcToCls = funcToCls.isSet,
       commentGeneratedCode = debug.isSet,
       noFreeze = noFreeze.isSet,
+      noOpt = noOptimizations.isSet,
       noModuleCheck = noModuleCheck.isSet,
       deadParamElim =
         if deadParamElim.isUnset then S(DeadParamElim.default)