Jb/control fold

compiler/stz-el.stanza

@@ -9,6 +9,7 @@ defpackage stz/el :
   import stz/el-var-table
   import stz/el-utils
   import stz/el-infer
+  import stz/el-basic-blocks
   import stz/dyn-tree
   import stz/conversion-utils
   import stz/set-utils
@@ -92,13 +93,15 @@ defn lower (epackage:EPackage, optimize?:True|False) -> EPackage :
     run-pass("Beta Reduce", beta-reduce, "beta-reduce1", false)
     run-pass("Box Unbox", box-unbox-fold, "box-unbox1", false)
     run-pass("Eliminate Dead Code", eliminate-dead-code, "dead-code1", false)
+    run-pass("Control Fold", control-fold, "control-fold1", false)
     ;Phase 2
     run-pass("Simple Inline", simple-inline, "inlined2", false)
     run-pass("Within Package Inline", within-package-inline{_, false}, "wp-inlined2", false)
     run-pass("Cleanup Labels", cleanup-labels, "cleanup-labels2", false)
     run-pass("Beta Reduce", beta-reduce, "beta-reduce2", false)
     run-pass("Box Unbox", box-unbox-fold, "box-unbox2", false)
     run-pass("Eliminate Dead Code", eliminate-dead-code, "dead-code2", false)
+    run-pass("Control Fold", control-fold, "control-fold2", false)
     ;Finish by removing checks 
     run-pass("Force Remove Checks", force-remove-checks, "removed-checks", false)
   run-pass("Resolve Matches", resolve-matches, "resolved-matches", false)
@@ -2331,6 +2334,73 @@ defn eliminate-dead-code (epackage:EPackage) -> EPackage :
     else : eliminate-in-item(e, elim-vars) as ETExp
   sub-exps(epackage, texps*)
 
+;============================================================
+;==================== CONTROL FOLD ==========================
+;============================================================
+;Inline single control flow instruction blocks into goto's
+;setting up box-unbox of match/if's to fold away overhead
+;Also prune unreachable blocks by walking basic-blocks
+defn control-fold (epackage:EPackage) -> EPackage :
+  defn fold-body (body:EBody) -> EBody :
+    val block-table = analyze-basic-blocks(ins(body))
+
+    ;Find blocks with a single control flow instruction
+    val inlineable-blocks = IntTable<List<EIns>>()
+    for block in block-table do :
+      val insts = to-tuple $ for i in instructions(block) filter : i is-not ELive
+      if length(insts) == 2 :
+        match(insts[0],insts[1]) :
+          (i0:ELabel,i1:EGoto|EIf|EMatch) :
+            inlineable-blocks[n(i0)] = tail(to-list $ instructions(block))
+          (i0,i1) : false
+
+    val buffer = BodyBuffer(body)
+
+    ;Record all defined variables
+    val defs = IntSet()
+    for block in block-table do :
+      for i in instructions(block) do :
+        for v in varlocs(i) do : add(defs, n(v))
+
+    ;Keep all the locals that are not eliminated.
+    for l in locals(body) do :
+      emit(buffer,l) when defs[n(l)]
+
+    ;Record all the used objects and fns during this sweep.
+    val used-objs-and-fns = IntSet()
+
+    ;Inline single control flow blocks into goto instructions
+    for block in block-table do :
+      for i in instructions(block) do :
+        match(i:EGoto) :
+          for i in get?(inlineable-blocks, n(i), List(i)) do :
+            emit(buffer, i)
+        else :
+          match(i) :
+            (i:ENew) : add(used-objs-and-fns, n(i))
+            (i:EInitClosures) : do(add{used-objs-and-fns, n(_)}, xs(i))
+            (i) : false
+          emit(buffer, i)
+
+    ;Keep only the fns and objects that are not eliminated.
+    for f in localfns(body) do :
+      emit(buffer, f) when used-objs-and-fns[n(f)]
+    for obj in localobjs(body) do :
+      emit(buffer, obj) when used-objs-and-fns[n(obj)]
+    to-body(buffer, false, false, false)
+
+  defn fold-exp (e:ETExp) -> ETExp :    
+    defn fold-in-bodies (e:ELBigItem) :
+      match(map(fold-in-bodies, e)) :
+        (e:EBody) : fold-body(e)
+        (e) : e
+    fold-in-bodies(e) as ETExp
+
+  val texps* = for e in exps(epackage) map :
+    fold-exp(e) as ETExp
+
+  sub-exps(epackage, texps*)
+
 ;============================================================
 ;=================== RESOLVE MATCHES ========================
 ;============================================================