Extract loops from toplevel function into helper functions

CHANGES.md

@@ -14,6 +14,7 @@
 * Lib: add `Performance` module (#2221)
 * Put more values into global variables (#2211)
 * Runtime: intial support for quickjs-ng
+* Extract loops from toplevel function into helper functions (#2245)
 
 ## Bug fixes
 * Compiler: fix reference unboxing (#2210)

compiler/lib/code.ml

@@ -776,6 +776,15 @@ let poptraps blocks pc =
   in
   loop blocks pc Addr.Set.empty 0 Addr.Set.empty |> fst
 
+let map_branch_conts f branch =
+  match branch with
+  | Branch c -> Branch (f c)
+  | Cond (x, c1, c2) -> Cond (x, f c1, f c2)
+  | Switch (x, a) -> Switch (x, Array.map a ~f)
+  | Pushtrap (c1, x, c2) -> Pushtrap (f c1, x, f c2)
+  | Poptrap c -> Poptrap (f c)
+  | (Return _ | Raise _ | Stop) as b -> b
+
 let fold_children blocks pc f accu =
   let block = Addr.Map.find pc blocks in
   match block.branch with

compiler/lib/code.mli

@@ -290,6 +290,10 @@ val fold_closures_outermost_first :
     outermost closures first. Unlike with {!fold_closures}, only the closures
     reachable from [p.start] are considered. *)
 
+val map_branch_conts : (cont -> cont) -> last -> last
+(** Apply [f] to every continuation in [branch], leaving non-branching
+    terminators unchanged. *)
+
 val fold_children : 'c fold_blocs
 
 val fold_children_skip_try_body : 'c fold_blocs

compiler/lib/driver.ml

@@ -715,6 +715,7 @@ let optimize ~shapes ~profile ~keep_flow_data p =
       | O2 -> o2
       | O3 -> o3)
     +> specialize_js_once_after
+    +> Hoist_loops.f
     +> effects_and_exact_calls ~keep_flow_data ~deadcode_sentinel ~shapes profile
     +> map_fst5
          (match Config.target (), Config.effects () with