Better Flatten (hkust-taco#276)

Author: CrescentonC

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

@@ -135,44 +135,87 @@ sealed abstract class Block extends Product with AutoLocated:
 
     (transformer.applyBlock(this), defns)
 
-  lazy val flatten: Block = 
-    // traverses a Block like a list, flatten `Begin`s using an accumulator
-    // returns the flattend but reversed Block (with the dummy tail `End("for flatten only")`) and the actual tail of the Block
-    def getReversedFlattenAndTrueTail(b: Block, acc: Block): (Block, BlockTail) = b match
-      case Match(scrut, arms, dflt, rest) => getReversedFlattenAndTrueTail(rest, Match(scrut, arms, dflt, acc))
-      case Label(label, body, rest) => getReversedFlattenAndTrueTail(rest, Label(label, body, acc))
-      case Begin(sub, rest) =>
-        val (firstBlockRev, firstTail) = getReversedFlattenAndTrueTail(sub, acc)
-        firstTail match
-          case _: End => getReversedFlattenAndTrueTail(rest, firstBlockRev)
-          // if the tail of `sub` is not `End`, ignore the `rest` of this `Begin`
-          case _ => firstBlockRev -> firstTail
-      case TryBlock(sub, finallyDo, rest) => getReversedFlattenAndTrueTail(rest, TryBlock(sub, finallyDo, acc))
-      case Assign(lhs, rhs, rest) => getReversedFlattenAndTrueTail(rest, Assign(lhs, rhs, acc))
-      case a@AssignField(lhs, nme, rhs, rest) => getReversedFlattenAndTrueTail(rest, AssignField(lhs, nme, rhs, acc)(a.symbol))
-      case AssignDynField(lhs, fld, arrayIdx, rhs, rest) => getReversedFlattenAndTrueTail(rest, AssignDynField(lhs, fld, arrayIdx, rhs, acc))
-      case Define(defn, rest) => getReversedFlattenAndTrueTail(rest, Define(defn, acc))
-      case HandleBlock(lhs, res, par, args, cls, handlers, body, rest) => getReversedFlattenAndTrueTail(rest, HandleBlock(lhs, res, par, args, cls, handlers, body, acc))
-      case t: BlockTail => acc -> t
+  lazy val flattened: Block = this.flatten(identity)
+  
+  private def flatten(k: End => Block): Block = this match
+    case 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.map(_.flattened)
+      if (newRest is rest) && (newArms is arms) && (dflt is newDflt)
+      then this
+      else Match(scrut, newArms, newDflt, newRest)
+
+    case Label(label, body, rest) =>
+      val newBody = body.flattened
+      val newRest = rest.flatten(k)
+      if (newBody is body) && (newRest is rest)
+      then this
+      else Label(label, newBody, newRest)
+      
+    case Begin(sub, rest) =>
+      sub.flatten(_ => rest.flatten(k))
 
-    // reverse the Block returnned from the previous function,
-    // which does not contain `Begin` (except for the nested ones),
-    // and whose tail must be the dummy `End("for flatten only")`
-    def rev(b: Block, t: Block): Block = b match
-      case Match(scrut, arms, dflt, rest) => rev(rest, Match(scrut, arms, dflt, t))
-      case Label(label, body, rest) => rev(rest, Label(label, body, t))
-      case TryBlock(sub, finallyDo, rest) => rev(rest, TryBlock(sub, finallyDo, t))
-      case Assign(lhs, rhs, rest) => rev(rest, Assign(lhs, rhs, t))
-      case a@AssignField(lhs, nme, rhs, rest) => rev(rest, AssignField(lhs, nme, rhs, t)(a.symbol))
-      case AssignDynField(lhs, fld, arrayIdx, rhs, rest) => rev(rest, AssignDynField(lhs, fld, arrayIdx, rhs, t))
-      case Define(defn, rest) => rev(rest, Define(defn, t))
-      case HandleBlock(lhs, res, par, args, cls, handlers, body, rest) => rev(rest, HandleBlock(lhs, res, par, args, cls, handlers, body, t))
-      case End(msg) => t
-      case _: BlockTail => ??? // unreachable
-      case Begin(sub, rest) => ??? // unreachable
+    case TryBlock(sub, finallyDo, rest) =>
+      val newSub = sub.flattened
+      val newFinallyDo = finallyDo.flattened
+      val newRest = rest.flatten(k)
+      if (newSub is sub) && (newFinallyDo is finallyDo) && (newRest is rest)
+      then this
+      else TryBlock(newSub, newFinallyDo, newRest)
+      
+    case Assign(lhs, rhs, rest) =>
+      val newRest = rest.flatten(k)
+      if newRest is rest
+      then this
+      else Assign(lhs, rhs, newRest)
+      
+    case a@AssignField(lhs, nme, rhs, rest) =>
+      val newRest = rest.flatten(k)
+      if newRest is rest
+      then this
+      else AssignField(lhs, nme, rhs, newRest)(a.symbol)
+      
+    case AssignDynField(lhs, fld, arrayIdx, rhs, rest) =>
+      val newRest = rest.flatten(k)
+      if newRest is rest
+      then this
+      else AssignDynField(lhs, fld, arrayIdx, rhs, newRest)
 
-    val (flattenRev, actualTail) = getReversedFlattenAndTrueTail(this, End("for flatten only"))
-    rev(flattenRev, actualTail)
+    case Define(defn, rest) =>
+      val newDefn = defn match
+        case d: FunDefn =>
+          val newBody = d.body.flattened
+          if newBody is d.body
+          then d
+          else d.copy(body = newBody)
+        case v: ValDefn => v
+        case c: ClsLikeDefn =>
+          val newPreCtor = c.preCtor.flattened
+          val newCtor = c.ctor.flattened
+          if (newPreCtor is c.preCtor) && (newCtor is c.ctor)
+          then c
+          else c.copy(preCtor = newPreCtor, ctor = newCtor)
+      
+      val newRest = rest.flatten(k)
+      if (newDefn is defn) && (newRest is rest)
+      then this
+      else Define(newDefn, newRest)
+    
+    case HandleBlock(lhs, res, par, args, cls, handlers, body, rest) =>
+      val newHandlers = handlers.mapConserve: h =>
+        val newBody = h.body.flattened
+        if newBody is h.body then h else h.copy(body = newBody)
+      val newBody = body.flattened
+      val newRest = rest.flatten(k)
+      if (newHandlers is handlers) && (newBody is body) && (newRest is rest)
+      then this
+      else HandleBlock(lhs, res, par, args, cls, newHandlers, newBody, newRest)
+
+    case e: End => k(e)
+    case t: BlockTail => this
 
 end Block
 

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

@@ -588,11 +588,12 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
     val stackSafe = config.stackSafety match
       case N => res
       case S(sts) => StackSafeTransform(sts.stackLimit).transformTopLevel(res)
-    
-    MergeMatchArmTransformer.applyBlock(
-      if lowerHandlers then HandlerLowering().translateTopLevel(stackSafe)
+    val withHandlers = if lowerHandlers
+      then HandlerLowering().translateTopLevel(stackSafe)
       else stackSafe
-    )
+    val flattened = withHandlers.flattened
+    
+    MergeMatchArmTransformer.applyBlock(flattened)
 
   def program(main: st): Program =
     def go(acc: Ls[Local -> Str], trm: st): Program =
@@ -762,7 +763,7 @@ object TrivialStatementsAndMatch:
 object MergeMatchArmTransformer extends BlockTransformer(new SymbolSubst()):
   override def applyBlock(b: Block): Block = super.applyBlock(b) match
     case m@Match(scrut, arms, Some(dflt), rest) =>
-      dflt.flatten match
+      dflt match
         case TrivialStatementsAndMatch(k, Match(scrutRewritten, armsRewritten, dfltRewritten, restRewritten))
           if (scrutRewritten === scrut) && (restRewritten.size * armsRewritten.length) < 10 =>
             val newArms = restRewritten match

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

@@ -31,7 +31,7 @@ object Printer:
         .map{ case (c, b) => doc"${case_doc(c)} => #{  # ${mkDocument(b)} #} " }
         .mkDocument(sep = doc" # ")
       val docDefault = dflt.map(mkDocument).getOrElse(doc"")
-      doc"match ${mkDocument(scrut)} #{  # ${docCases} # else #{  # ${docDefault} #}  #} "
+      doc"match ${mkDocument(scrut)} #{  # ${docCases} # else #{  # ${docDefault} #}  #}  # in # ${mkDocument(rest)} "
     case Return(res, implct) => doc"return ${mkDocument(res)}"
     case Throw(exc) => doc"throw ${mkDocument(exc)}"
     case Label(label, body, rest) =>

hkmc2/shared/src/test/mlscript/codegen/BlockPrinter.mls

@@ -21,18 +21,18 @@ let x = x + 1
 //│ Pretty Lowered:
 //│  
 //│ set x2 = 1 in
-//│ begin
-//│   set scrut = ==(x2, 0) in
-//│   match scrut
-//│     true =>
-//│       set tmp = 1 in
-//│       end
-//│     else
-//│       set tmp = 0 in
-//│       end;
-//│   set x3 = tmp in
-//│   set tmp1 = +(x3, 1) in
-//│   set x4 = tmp1 in
-//│   set block$res3 = undefined in
-//│   end
+//│ set scrut = ==(x2, 0) in
+//│ match scrut
+//│   true =>
+//│     set tmp = 1 in
+//│     end
+//│   else
+//│     set tmp = 0 in
+//│     end
+//│ in
+//│ set x3 = tmp in
+//│ set tmp1 = +(x3, 1) in
+//│ set x4 = tmp1 in
+//│ set block$res3 = undefined in
+//│ end 
 //│ x = 1