Extract loops from toplevel function into helper functions

compiler/lib-wasm/generate.ml

@@ -2012,9 +2012,28 @@ module Generate (Target : Target_sig.S) = struct
     | (`Block _ | `Catch | `Skip) as b -> b :: context
     | `Return -> `Skip :: context
 
-  let needed_handlers (p : program) pc =
+  (* Walk the dominator subtree of [pc] (the structural region of the
+     loop body, try body, or function body that the wrap covers),
+     skipping any nested try body since it carries its own wrap. *)
+  let needed_handlers (p : program) ~dom pc =
+    let fold : 'c. _ -> _ -> (Addr.t -> 'c -> 'c) -> 'c -> 'c =
+     fun _blocks pc' f accu ->
+      let block = Addr.Map.find pc' p.blocks in
+      let try_body =
+        match block.branch with
+        | Pushtrap ((pc'', _), _, _) -> Some pc''
+        | _ -> None
+      in
+      Addr.Set.fold
+        (fun child acc ->
+          match try_body with
+          | Some pc'' when pc'' = child -> acc
+          | _ -> f child acc)
+        (Structure.get_edges dom pc')
+        accu
+    in
     Code.traverse
-      { fold = fold_children_skip_try_body }
+      { fold }
       (fun pc n ->
         let block = Addr.Map.find pc p.blocks in
         List.fold_left
@@ -2084,8 +2103,8 @@ module Generate (Target : Target_sig.S) = struct
         instr W.Unreachable
     else body ~result_typ ~fall_through ~context
 
-  let wrap_with_handlers p pc ~result_typ ~fall_through ~context body =
-    let need_zero_divide_handler, need_bound_error_handler = needed_handlers p pc in
+  let wrap_with_handlers p ~dom pc ~result_typ ~fall_through ~context body =
+    let need_zero_divide_handler, need_bound_error_handler = needed_handlers p ~dom pc in
     wrap_with_handler
       need_bound_error_handler
       bound_error_pc
@@ -2132,7 +2151,7 @@ module Generate (Target : Target_sig.S) = struct
         | Cond (_, (pc1, _), (pc2, _)) when pc' = pc1 && pc' = pc2 -> true
         | _ -> Structure.is_merge_node g pc'
       in
-      let code ~context =
+      let code ~result_typ ~fall_through ~context =
         let block = Addr.Map.find pc ctx.blocks in
         let* () = translate_instrs ctx context block.body in
         translate_node_within
@@ -2149,8 +2168,39 @@ module Generate (Target : Target_sig.S) = struct
       in
       if Structure.is_loop_header g pc
       then
-        loop { params = []; result = result_typ } (code ~context:(`Block pc :: context))
-      else code ~context
+        let outermost_toplevel_loop =
+          Option.is_none name_opt
+          && not
+               (List.exists
+                  ~f:(function
+                    | `Block pc' when pc' >= 0 -> Structure.is_loop_header g pc'
+                    | _ -> false)
+                  context)
+        in
+        loop
+          { params = []; result = result_typ }
+          (if outermost_toplevel_loop
+           then
+             (* The outermost loops of the toplevel function are later
+                hoisted into helper functions by [Hoist_loops], which
+                requires them to be self-contained (no [Br] escaping
+                the loop body). Bounds and zero-divide checks normally
+                branch to handler blocks installed at the top of the
+                function; we install handlers around the loop body
+                itself so those branches stay inside the loop. Nested
+                loops do not need their own wrap: only the outermost
+                one is extracted, and its body is copied verbatim into
+                the helper along with the surrounding handler blocks. *)
+             wrap_with_handlers
+               p
+               ~dom
+               pc
+               ~result_typ
+               ~fall_through
+               ~context:(`Block pc :: context)
+               code
+           else code ~result_typ ~fall_through ~context:(`Block pc :: context))
+      else code ~result_typ ~fall_through ~context
     and translate_node_within ~result_typ ~fall_through ~pc ~l ~context =
       match l with
       | pc' :: rem ->
@@ -2229,6 +2279,7 @@ module Generate (Target : Target_sig.S) = struct
                 ~context:(extend_context fall_through context)
                 (wrap_with_handlers
                    p
+                   ~dom
                    (fst cont)
                    (fun ~result_typ ~fall_through ~context ->
                      translate_branch result_typ fall_through pc cont context))
@@ -2300,6 +2351,7 @@ module Generate (Target : Target_sig.S) = struct
            let* () =
              wrap_with_handlers
                p
+               ~dom
                pc
                ~result_typ:[ Option.value ~default:Type.value (unboxed_type return_type) ]
                ~fall_through:`Return
@@ -2446,6 +2498,7 @@ module Generate (Target : Target_sig.S) = struct
         functions
     in
     global_context.init_code <- [];
+    let functions = Hoist_loops.f ~toplevel:toplevel_name functions in
     global_context.other_fields <- List.rev_append functions global_context.other_fields;
     let js_code = StringMap.bindings global_context.fragments in
     global_context.fragments <- StringMap.empty;