Compilation Pipeline

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

@@ -27,17 +27,13 @@ case class Config(
   stageCode: Bool,
   target: CompilationTarget,
   rewriteWhileLoops: Bool,
-  tailRecOpt: Bool,
-  deforest: Opt[Deforest],
   etaExpansion: Opt[EtaExpansion],
-  inlining: Opt[Inliner],
-  deadBranchRemoval: Bool,
   qqEnabled: Bool,
   funcToCls: Bool,
   commentGeneratedCode: Bool,
   noFreeze: Bool,
   noModuleCheck: Bool,
-  deadParamElim: Opt[DeadParamElim],
+  optimizer: Optimizer,
 ):
 
   def stackSafety: Opt[StackSafety] = effectHandlers.flatMap(_.stackSafety)
@@ -55,6 +51,18 @@ case class Config(
   def shouldRewriteWhile: Bool =
     rewriteWhileLoops
 
+  def tailRecOpt: Bool = optimizer.tailRecOpt
+
+  def deforest: Opt[Deforest] = optimizer.deforest
+
+  def inlining: Opt[Inliner] = optimizer.inlining
+
+  def deadBranchRemoval: Bool = optimizer.deadBranchRemoval
+
+  def deadParamElim: Opt[DeadParamElim] = optimizer.deadParamElim
+
+  def mapOptimizer(f: Optimizer => Optimizer): Config = copy(optimizer = f(optimizer))
+
 end Config
 
 
@@ -71,17 +79,13 @@ object Config:
     target = CompilationTarget.JS,
     rewriteWhileLoops = false,
     stageCode = false,
-    tailRecOpt = true,
-    deforest = N,
     etaExpansion = S(EtaExpansion.default),
-    inlining = S(Inliner(default.inlineThreshold)),
-    deadBranchRemoval = default.deadBranchRemoval,
     qqEnabled = false,
     funcToCls = false,
     commentGeneratedCode = false,
     noFreeze = false,
     noModuleCheck = false,
-    deadParamElim = S(DeadParamElim.default)
+    optimizer = Optimizer.FastOpt,
   )
   object default:
     val patMatConsequentSharingThreshold = S(15)
@@ -229,6 +233,37 @@ object Config:
     .foldLeft(identity[Config]): (acc, modify) =>
       cfg => modify(acc(cfg))
     configModify(config)
+
+  case class Optimizer(
+    deforest: Opt[Deforest],
+    tailRecOpt: Bool,
+    inlining: Opt[Inliner],
+    deadBranchRemoval: Bool,
+    deadCodeElim: Bool,
+    dataFlowAnalysis: Bool,
+    deadParamElim: Opt[DeadParamElim],
+  )
+
+  object Optimizer:
+    val NoOpt = Optimizer(
+      N,
+      false,
+      N,
+      false,
+      false,
+      false,
+      N
+    )
+
+    val FastOpt = Optimizer(
+      deforest = N,
+      tailRecOpt = true,
+      inlining = S(Inliner(default.inlineThreshold)),
+      deadBranchRemoval = default.deadBranchRemoval,
+      deadCodeElim = true,
+      dataFlowAnalysis = true,
+      deadParamElim = S(DeadParamElim.default),
+    )
 
 end Config
 
@@ -575,8 +610,10 @@ object ConfigParser:
   /** Parse a single field override like `tailRecOpt: false`. */
   private def parseField(name: Str, value: Tree)(using Raise): Config => Config = name match
     case "language" => parseLanguageOverride(value)
+    case "noOpt" =>
+      parsedField(value)(parseBool)(v => _.mapOptimizer(_ => Optimizer.NoOpt))
     case "tailRecOpt" =>
-      parsedField(value)(parseBool)(v => _.copy(tailRecOpt = v))
+      parsedField(value)(parseBool)(v => _.mapOptimizer(_.copy(tailRecOpt = v)))
     case "noFreeze" =>
       parsedField(value)(parseBool)(v => _.copy(noFreeze = v))
     case "noModuleCheck" =>
@@ -600,23 +637,23 @@ object ConfigParser:
     case "deforest" =>
       optionalFieldWithCurrent(value)(_.deforest)(
         (tree, current) => parseDeforest(tree, current)
-      )(v => _.copy(deforest = v))
+      )(v => _.mapOptimizer(_.copy(deforest = v)))
     case "etaExpansion" =>
       optionalFieldWithCurrent(value)(_.etaExpansion)(
         (tree, current) => parseEtaExpansion(tree, current)
       )(v => _.copy(etaExpansion = v))
     case "deadParamElim" =>
       optionalFieldWithCurrent(value)(_.deadParamElim)(
         (tree, current) => parseDeadParamElim(tree, current)
-      )(v => _.copy(deadParamElim = v))
+      )(v => _.mapOptimizer(_.copy(deadParamElim = v)))
     case "sanityChecks" =>
       optionalField(value)(_ => S(Config.SanityChecks(light = true, checkUnreachable = true)))(v => _.copy(sanityChecks = v))
     case "patMatConsequentSharingThreshold" =>
       parsedField(value)(parseInt)(v => _.copy(patMatConsequentSharingThreshold = S(v)))
     case "inlining" =>
-      optionalField(value)(parseInt)(v => _.copy(inlining = v.map(Inliner(_))))
+      optionalField(value)(parseInt)(v => _.mapOptimizer(_.copy(inlining = v.map(Inliner(_)))))
     case "deadBranchRemoval" =>
-      parsedField(value)(parseBool)(v => _.copy(deadBranchRemoval = v))
+      parsedField(value)(parseBool)(v => _.mapOptimizer(_.copy(deadBranchRemoval = v)))
     case _ =>
       raise(ErrorReport(
         msg"Unknown config field '${name}'" -> value.toLoc :: Nil,

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/BlockSimplifier.scala

@@ -47,25 +47,27 @@ class BlockSimplifier
 
       log(s"⬤ Simplif. iter. $iteration")
 
-      // * Running DCE once sometimes produces more DCE opportunities;
-      // * it is important to apply all of them so that later passes, such as COC,
-      // * are not impeded by things like unused labels from inlining.
-      var dceIteration = 0
-      while
-        val dce = new DeadCodeElim()
-        res = dce.apply(res)
-        changed ||= dce.changed
-        if dce.changed then
-          log("▶ DCE:\n" + printRes)
-          dceIteration += 1
-          dceIteration < MaxDCEIterationsPerIter
-        else false
-      do ()
+      if summon[Config].optimizer.deadCodeElim then
+        // * Running DCE once sometimes produces more DCE opportunities;
+        // * it is important to apply all of them so that later passes, such as COC,
+        // * are not impeded by things like unused labels from inlining.
+        var dceIteration = 0
+        while
+          val dce = new DeadCodeElim()
+          res = dce.apply(res)
+          changed ||= dce.changed
+          if dce.changed then
+            log("▶ DCE:\n" + printRes)
+            dceIteration += 1
+            dceIteration < MaxDCEIterationsPerIter
+          else false
+        do ()
 
-      val vp = new DataFlowAnalysis(LocalVars.analyze(res.main))
-      res = vp.apply(res)
-      changed ||= vp.changed
-      if vp.changed then log("▶ VP:\n" + printRes)
+      if summon[Config].optimizer.dataFlowAnalysis then
+        val vp = new DataFlowAnalysis(LocalVars.analyze(res.main))
+        res = vp.apply(res)
+        changed ||= vp.changed
+        if vp.changed then log("▶ VP:\n" + printRes)
 
       summon[Config].inlining.foreach: cfg =>
 

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

@@ -34,6 +34,7 @@ class ClassParamFlattener(using State) extends BlockTransformer(SymbolSubst.Id):
 
   /** Normalize class params so that `paramsOpt = N` and `auxParams` has exactly one element. */
   private def flattenClsParams(cls: ClsLikeDefn): ClsLikeDefn =
+    if cls.paramsOpt.isEmpty && cls.auxParams.sizeIs == 1 then return cls
     val paramss = cls.paramsOpt.toList ::: cls.auxParams
     val flatAux = paramss match
       case Nil => PlainParamList(Nil)
@@ -105,6 +106,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,64 @@
+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) = ()
+
+  def passHook(passName: Str, before: Program, after: Program) = ()
+
+  private inline def blockPass(inline pass: Block => Block)(prog: Program): Program =
+    val blk = pass(prog.main)
+    if blk is prog.main then prog else Program(prog.imports, blk)
+
+  def run(prog: Program, printer: Program => Str, symbolsToPreserve: Set[BoundSymbol], otl: TL)(using TL): Program =
+    var result = prog
+    inline def runPass(passName: Str)(inline transform: Program => Program) =
+      val before = result
+      result = transform(before)
+      passHook(passName, before, result)
+
+    runPass("LambdaRewriter")(LambdaRewriter.desugar)
+    runPass("Deforest"): prog =>
+      val outterTl = tl
+      config.deforest match
+        case None => prog
+        case Some(dCfg) =>
+          flowAnalysis.FlowAnalysis.mkTraceLogger(dCfg.config, "deforest > ", outterTl).givenIn:
+            deforest.Deforest(prog)
+    runPass("EtaExpansion")(EtaExpansion.apply)
+    runPass("Lifter"): prog =>
+      if config.liftDefns.isDefined then
+        blockPass(Lifter(_).transform)(prog)
+      else prog
+    runPass("HandlerLowering"): prog =>
+      config.effectHandlers.fold(prog): opt =>
+        HandlerLowering(new HandlerPaths, opt).translateProgram(prog)
+    runPass("Flattening")(blockPass(_.flattened))
+    runPass("BufferableTransform")(BufferableTransform().transform)
+    runPass("MergeMatchArmTransformer")(MergeMatchArmTransformer.applyProgram)
+    runPass("FirstClassFunctionTransformer"): prog =>
+      if config.funcToCls then
+        blockPass(FirstClassFunctionTransformer().transform(_))(prog)
+      else prog
+    runPass("Lifter"): prog =>
+      if config.funcToCls then
+        blockPass(Lifter(_).transform)(prog)
+      else prog
+    runPass("ClassParamFlattener")(ClassParamFlattener.apply)
+    runPass("ReflectionInstrumenter")(ReflectionInstrumenter(using summon).apply)
+    runPass("TailRecOpt")(TailRecOpt().transform)
+    preOptimizeHook(result)
+    runPass("WorkerWrapper")(WorkerWrapper(symbolsToPreserve, otl, printer))
+    runPass("BlockSimplifier")(BlockSimplifier(symbolsToPreserve, otl, printer).apply)
+    runPass("DeadParamElim")(otl.givenIn(DeadParamElim.apply))
+
+    result

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

@@ -89,7 +89,10 @@ class FirstClassFunctionTransformer
 
   override def applyResult(r: Result)(k: Result => Block): Block = r match
     case c @ Call(fun, argss) => applyListOf(argss, (args, k2) => applyArgs(args)(k2)): argss2 =>
-      def call(f: Path) = Call(f, argss2.ne_!)(c.metadata)
+      def call(f: Path) =
+        if (f is fun) && (argss is argss2)
+        then c
+        else Call(f, argss2.ne_!)(c.metadata)
       fun match
         case ref @ Value.SimpleRef(sym) => sym match
           case _: VarSymbol |  _: TempSymbol => k(call(ref.selSN("call")))

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)