Add Data Representation Flattening Pass

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

@@ -28,6 +28,7 @@ case class Config(
   target: CompilationTarget,
   rewriteWhileLoops: Bool,
   etaExpansion: Opt[EtaExpansion],
+  dataRepFlatten: Opt[DataRepFlatten],
   qqEnabled: Bool,
   funcToCls: Bool,
   commentGeneratedCode: Bool,
@@ -80,6 +81,7 @@ object Config:
     rewriteWhileLoops = false,
     stageCode = false,
     etaExpansion = S(EtaExpansion.default),
+    dataRepFlatten = N,
     qqEnabled = false,
     funcToCls = false,
     commentGeneratedCode = false,
@@ -221,6 +223,13 @@ object Config:
         logAccumulator = false,
       ))
     val default: EtaExpansion = withDebug(debug = false)
+
+  case class DataRepFlatten(debug: Bool, mono: Bool)
+  object DataRepFlatten:
+    val default = DataRepFlatten(
+      debug = false,
+      mono = false,
+    )
 
   /** `altSmallThreshold` is the alternative threshold for inlining things into @inline functions.
     * Normally, we avoid inlining into @inline functions as that could lead to unexpected code bloat. */
@@ -606,6 +615,24 @@ object ConfigParser:
     case _ =>
       expect("EtaExpansion(...)")(tree)
       N
+
+  private def parseDataRepFlatten(tree: Tree, current: Opt[Config.DataRepFlatten])(using Raise): Opt[Config.DataRepFlatten] =
+    tree match
+    case Call("DataRepFlatten", args) =>
+      val base = current.getOrElse(Config.DataRepFlatten.default)
+      var debug = base.debug
+      var mono = base.mono
+      args.foreach:
+        case NamedArg("debug", value) =>
+          setFrom(value)(parseBool)(v => debug = v)
+        case NamedArg("mono", value) =>
+          setFrom(value)(parseBool)(v => mono = v)
+        case other =>
+          unsupported("DataRepFlatten", other)
+      S(Config.DataRepFlatten(debug, mono))
+    case _ =>
+      expect("DataRepFlatten(...)")(tree)
+      N
 
   /** Parse a single field override like `tailRecOpt: false`. */
   private def parseField(name: Str, value: Tree)(using Raise): Config => Config = name match
@@ -642,6 +669,10 @@ object ConfigParser:
       optionalFieldWithCurrent(value)(_.etaExpansion)(
         (tree, current) => parseEtaExpansion(tree, current)
       )(v => _.copy(etaExpansion = v))
+    case "dataRepFlatten" =>
+      optionalFieldWithCurrent(value)(_.dataRepFlatten)(
+        (tree, current) => parseDataRepFlatten(tree, current)
+      )(v => _.copy(dataRepFlatten = v))
     case "deadParamElim" =>
       optionalFieldWithCurrent(value)(_.deadParamElim)(
         (tree, current) => parseDeadParamElim(tree, current)

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

@@ -266,15 +266,15 @@ sealed abstract class Block extends Product:
 
   private def flatten(k: End => Block): Block = this match
 
-    case Match(scrut, arms, dflt, rest) =>
+    case m @ Match(scrut, arms, dflt, rest) =>
       val newRest = rest.flatten(k)
       val newArms = arms.mapConserve: arm =>
         val newBody = arm._2.flattened
         if newBody is arm._2 then arm else (arm._1, newBody)
       val newDflt = dflt.mapConserve(_.flattened)
       if (newRest is rest) && (newArms is arms) && (newDflt is dflt)
       then this
-      else Match(scrut, newArms, newDflt, newRest)
+      else Match(scrut, newArms, newDflt, newRest)(m.annotations)
 
     case Label(label, loop, body, rest) =>
       val newBody = body.flattened
@@ -371,7 +371,9 @@ case class Match(
   arms: Ls[Case -> Block],
   dflt: Opt[Block],
   rest: Block,
-) extends Block with ProductWithTail with NonBlockTail
+)(val annotations: Ls[Annot]) extends Block with ProductWithTail with NonBlockTail:
+  def matchShapes: Opt[Annot.MatchShapes] = annotations.collectFirst:
+    case annotation: Annot.MatchShapes => annotation
 
 case class Return(res: Result) extends BlockTail
 
@@ -462,6 +464,10 @@ object Define:
     case _ => new Define(defn, rest)
 
 object Match:
+  def apply(scrut: Path, arms: Ls[Case -> Block], dflt: Opt[Block], rest: Block)(annotations: Ls[Annot]): Block =
+    if annotations.nonEmpty then new Match(scrut, arms, dflt, rest)(annotations)
+    else apply(scrut, arms, dflt, rest)
+
   def apply(scrut: Path, _arms: Ls[Case -> Block], _dflt: Opt[Block], rest: Block): Block =
     val emptyDflt = _dflt.forall(_.isEmpty)
     val dflt = if emptyDflt then N else _dflt
@@ -470,7 +476,7 @@ object Match:
     else dflt match
     case S(Unreachable(_)) if scrut.isPure && arms.sizeCompare(1) === 0 =>
       Begin(arms.head._2, rest)
-    case S(Match(`scrut`, arms2, dflt2, _: End)) => // TODO: also handle non-End rest (may require a join point)
+    case S(m @ Match(`scrut`, arms2, dflt2, _: End)) if m.annotations.isEmpty => // TODO: also handle non-End rest (may require a join point)
       // * Currently, this branch does not seem used often (or at all?),
       // * because the UCS and (especially) MergeMatchArmTransformer already do a good job at merging matches
       Match(scrut, arms ::: arms2, dflt2, rest)
@@ -480,8 +486,8 @@ object Match:
         case S(d) => S(if d.isAbortive then d else Begin(d, rest))
         case N => S(rest)
       if numNonAbortive === 0 then
-        if rest.isEmpty then new Match(scrut, arms, mapDflt, rest)
-        else new Match(scrut, arms, mapDflt, End("(Unreachable:) rest of abortive match"))
+        if rest.isEmpty then new Match(scrut, arms, mapDflt, rest)(Nil)
+        else new Match(scrut, arms, mapDflt, End("(Unreachable:) rest of abortive match"))(Nil)
       else if numNonAbortive === 1 && dflt.exists(_.isAbortive) || rest.size <= 1 then
         new Match(scrut,
           arms.map: a =>
@@ -491,10 +497,10 @@ object Match:
           // * Indeed, `L: { match scrut { C => break L }; end }` can no longer be optimized
           // * if we replace `end` with `unreachable`, since the break is no longer jumping over nothing,
           // * ie no longer in tail position of the label (trying to treat it as such is unsound).
-          End("Rest moved to non-abortive branch(es)"))
+          End("Rest moved to non-abortive branch(es)"))(Nil)
       else rest match
         case Scoped(syms, body) => Scoped(syms, Match(scrut, arms, dflt, body))
-        case _ => new Match(scrut, arms, dflt, rest)
+        case _ => new Match(scrut, arms, dflt, rest)(Nil)
 
 object Begin:
   def apply(sub: Block, rest: Block): Block =
@@ -511,7 +517,8 @@ object Begin:
               "overlapping symbols when trying to merge Scoped blocks")
           Scoped(symsSub ++ symsRest, Begin(bodySub, bodyRest))
         case _ => Scoped(symsSub, Begin(bodySub, rest))
-      case Match(scrut, arms, dflt, rst) => Match(scrut, arms, dflt, Begin(rst, rest))
+      case m @ Match(scrut, arms, dflt, rst) =>
+        Match(scrut, arms, dflt, Begin(rst, rest))(m.annotations)
       case Label(lbl, loop, body, rst) => Label(lbl, loop, body, Begin(rst, rest))
       case TryBlock(sub, fin, rst) => TryBlock(sub, fin, Begin(rst, rest))
       case Assign(lhs, rhs, rst) => Assign(lhs, rhs, Begin(rst, rest))

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

@@ -937,7 +937,7 @@ class BlockSimplifier
         makeImpossibleAfter:
           super.applyBlock(b)
 
-      case Match(scrut, arms, dflt, rest) =>
+      case m @ Match(scrut, arms, dflt, rest) =>
 
         applyPath(scrut): scrut2 =>
 
@@ -1079,7 +1079,7 @@ class BlockSimplifier
           val restRewritten = applySubBlock(rest)
 
           if (scrut2 is scrut) && (newArms is arms) && (newDflt is dflt) && (restRewritten is rest) then b
-          else Match(scrut2, newArms, newDflt, restRewritten)
+          else Match(scrut2, newArms, newDflt, restRewritten)(m.annotations)
 
       case _ =>
         super.applyBlock(b)

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

@@ -46,7 +46,7 @@ class BlockTransformer(subst: SymbolSubst):
     case Throw(exc) =>
       applyResult(exc): exc2 =>
         if exc2 is exc then b else Throw(exc2)
-    case Match(scrut, arms, dflt, rst) =>
+    case m @ Match(scrut, arms, dflt, rst) =>
       def applySub(b: Block) = if rst.isEmpty then applySubBlock(b) else applySubBlockNonTail(b)
       applyPath(scrut): scrut2 =>
         applyListOf(
@@ -62,7 +62,7 @@ class BlockTransformer(subst: SymbolSubst):
             if (scrut2 is scrut) &&
                 (arms2 is arms) &&
                 (dflt2 is dflt) && (rst2 is rst)
-              then b else Match(scrut2, arms2, dflt2, rst2)
+              then b else Match(scrut2, arms2, dflt2, rst2)(m.annotations)
     case Label(lbl, loop, bod, rst) =>
       val lbl2 = lbl.subst
       val bod2 = if loop then applyScopedBlock(bod) else applySubBlock(bod)

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

@@ -62,6 +62,7 @@ class CompilationPipeline(using Config, Raise, State, Ctx, SymbolPrinter):
       else prog
     runPass("ClassParamFlattener")(ClassParamFlattener.apply)
     runPass("ReflectionInstrumenter")(ReflectionInstrumenter(using summon).apply)
+    runPass("DataRepFlattener")(DataRepFlattener.apply)
     preOptimizeHook(result)
 
     // * We run this pass here first, before inlining so that the @tailrec/@tailcall annotations