Compiler: exit loop early, even leading loops (#2077)

Author: hhugo

CHANGES.md

@@ -1,3 +1,8 @@
+# dev
+
+## Features/Changes
+* Compiler: exit-loop-early in more cases (#2077)
+
 # 6.1.1 (2025-07-07) - Lille
 
 ## Bug fixes

compiler/lib-wasm/generate.ml

@@ -1520,6 +1520,7 @@ let start () = make_context ~value_type:Gc_target.Type.value
 
 let f ~context ~unit_name p ~live_vars ~in_cps ~deadcode_sentinal ~global_flow_data =
   let state, info = global_flow_data in
+  let p = Structure.norm p in
   let types = Typing.f ~state ~info ~deadcode_sentinal p in
   let t = Timer.make () in
   let p = fix_switch_branches p in

compiler/lib/generate.ml

@@ -2399,6 +2399,7 @@ let f
     ~should_export
     ~warn_on_unhandled_effect
     ~deadcode_sentinal =
+  let p = Structure.norm p in
   let mutated_vars = Freevars.f_mutable p in
   let freevars = Freevars.f p in
   let t' = Timer.make () in

compiler/lib/structure.ml

@@ -277,3 +277,67 @@ let build_graph blocks pc =
   let g = build_graph blocks pc in
   shrink_loops blocks g;
   g
+
+(* Ensure that all loops have a predecessor block. Function
+   shrink_loops assumes this. *)
+let norm p =
+  let free_pc = ref p.free_pc in
+  let visited = BitSet.create' p.free_pc in
+  let rec mark_used ~function_start pc =
+    if not (BitSet.mem visited pc)
+    then (
+      if not function_start then BitSet.set visited pc;
+      let block = Addr.Map.find pc p.blocks in
+      List.iter
+        ~f:(fun i ->
+          match i with
+          | Let (_, Closure (_, (pc', _), _)) -> mark_used ~function_start:true pc'
+          | _ -> ())
+        block.body;
+      fold_children p.blocks pc (fun pc' () -> mark_used ~function_start:false pc') ())
+  in
+  mark_used ~function_start:true p.start;
+  let closure_need_update = function
+    | Let (_, Closure (_, (pc, _), _)) -> BitSet.mem visited pc
+    | _ -> false
+  in
+  let rewrite_cont cont blocks =
+    let npc = !free_pc in
+    incr free_pc;
+    let body =
+      let b = Addr.Map.find (fst cont) blocks in
+      match b.body with
+      | (Event _ as e) :: _ -> [ e ]
+      | _ -> []
+    in
+    let blocks = Addr.Map.add npc { body; params = []; branch = Branch cont } blocks in
+    (npc, []), blocks
+  in
+  let blocks =
+    Addr.Map.fold
+      (fun pc block blocks ->
+        if List.exists block.body ~f:closure_need_update
+        then
+          let blocks = ref blocks in
+          let body =
+            List.map block.body ~f:(function
+              | Let (x, Closure (params, cont, loc)) as i when closure_need_update i ->
+                  let cont', blocks' = rewrite_cont cont !blocks in
+                  blocks := blocks';
+                  Let (x, Closure (params, cont', loc))
+              | i -> i)
+          in
+          Addr.Map.add pc { block with body } !blocks
+        else blocks)
+      p.blocks
+      p.blocks
+  in
+  if BitSet.mem visited p.start
+  then (
+    let npc = !free_pc in
+    incr free_pc;
+    let blocks =
+      Addr.Map.add npc { body = []; params = []; branch = Branch (p.start, []) } blocks
+    in
+    { blocks; free_pc = !free_pc; start = npc })
+  else { blocks; free_pc = !free_pc; start = p.start }

compiler/lib/structure.mli

@@ -24,3 +24,7 @@ val sort_in_post_order : t -> Addr.t list -> Addr.t list
 val blocks_in_reverse_post_order : t -> Code.Addr.t list
 
 val get_nodes : t -> Addr.Set.t
+
+val norm : program -> program
+(** [norm p] normalizes a program [p] to accommodate [Structure.build_graph] logic.
+    In practice, it ensures that all loops have a predecessor block and allows to exit loops early. *)

compiler/tests-compiler/gh1559.ml

@@ -195,37 +195,36 @@ let () = my_ref := 2
         my_ref = [0, 1],
         nesting = 1;
        function handle_state(t$1){
-        var t$0 = t$1;
-        for(;;){
-         let t$1 = t$0;
-         var
-          this_will_be_undefined =
-            function(param){
-             a:
-             {if(t$1 && 1 === t$1[1]){var _c_ = 1; break a;} var _c_ = 0;}
-             return _c_ ? 1 : 2;
-            };
-         if(t$0) var i = t$0[1], match = i; else var match = - 1;
-         if(0 === match){
-          var
-           g = function(param){return 2 + this_will_be_undefined(0) | 0;},
-           _b_ = g(0);
-          return g(0) + _b_ | 0;
-         }
-         if(1 === match){
-          if(caml_call2(Stdlib_Int[8], nesting, 0)) return nesting;
+        a:
+        {
+         var t$0 = t$1;
+         for(;;){
+          let t$1 = t$0;
           var
-           g$0 =
+           this_will_be_undefined =
              function(param){
-              return 1 < caml_call1(Stdlib_Random[5], 3)
-                      ? 2 + this_will_be_undefined(0) | 0
-                      : 1;
-             },
-           _c_ = g$0(0);
-          return g$0(0) + _c_ | 0;
+              a:
+              {if(t$1 && 1 === t$1[1]){var _c_ = 1; break a;} var _c_ = 0;}
+              return _c_ ? 1 : 2;
+             };
+          if(t$0) var i = t$0[1], match = i; else var match = - 1;
+          if(0 === match) break;
+          if(1 === match) break a;
+          t$0 = t;
          }
-         t$0 = t;
+         var
+          g = function(param){return 2 + this_will_be_undefined(0) | 0;},
+          _b_ = g(0);
+         return g(0) + _b_ | 0;
+        }
+        if(caml_call2(Stdlib_Int[8], nesting, 0)) return nesting;
+        function g$0(param){
+         return 1 < caml_call1(Stdlib_Random[5], 3)
+                 ? 2 + this_will_be_undefined(0) | 0
+                 : 1;
         }
+        var _c_ = g$0(0);
+        return g$0(0) + _c_ | 0;
        }
        var _a_ = handle_state(init), _b_ = caml_call1(Stdlib_Int[12], _a_);
        caml_call1(Stdlib[46], _b_);

compiler/tests-compiler/loops.ml

@@ -85,15 +85,7 @@ let rec fun_with_loop acc = function
        return caml_call1
                (list_rev, caml_call1(list_rev, caml_call1(list_rev, acc)));
       var x = param[1];
-      if(1 === x && ! param[2]){
-       var a$0 = [0, acc], i$0 = 0;
-       for(;;){
-        a$0[1] = [0, 1, a$0[1]];
-        var _b_ = i$0 + 1 | 0;
-        if(10 === i$0) return a$0[1];
-        i$0 = _b_;
-       }
-      }
+      if(1 === x && ! param[2]) break;
       var xs = param[2], a = [0, acc], i = 0;
       for(;;){
        a[1] = [0, 1, a[1]];
@@ -105,6 +97,13 @@ let rec fun_with_loop acc = function
       acc = acc$0;
       param = xs;
      }
+     var a$0 = [0, acc], i$0 = 0;
+     for(;;){
+      a$0[1] = [0, 1, a$0[1]];
+      var _b_ = i$0 + 1 | 0;
+      if(10 === i$0) return a$0[1];
+      i$0 = _b_;
+     }
     }
     //end
     |}]