Hand back a recursive group nothing was extracted from

The Lletrec arm of deep_flatten rebuilt unconditionally: it mapped every
binding into a fresh list, split that with a fold carrying a stop flag,
and reassembled through lambda_of_groups, whether or not a binding could
be lifted out of the group. It now maps with sharing, replaces the fold
with a walk that stops at the first binding referring back into the
group, and returns the original when nothing was extracted and nothing
underneath changed.

That arm was where nearly all the remaining waste was. Over a stdlib
build, runs of deep_flatten that hand back their input rise from 253 to
324 of 447 and rebuilds producing an identical tree fall from 72 to 1,
with 122 runs changing the tree either way.

The unit test covers the sharing rather than the extraction. Removing
the sharing changes no generated output, so no snapshot can catch it,
while breaking the extraction moves output the existing suite already
compares.

Signed-off-by: Cristiano Calcagno <ccrisccris@gmail.com>

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W8g8qwBARAcvW9MyuKQq8H

Author: cristianoc

compiler/core/lam_pass_deep_flatten.ml

@@ -223,15 +223,15 @@ let deep_flatten (lam : Lambda.t) : Lambda.t =
     | Llet _ ->
       let res, groups = flatten [] lam in
       lambda_of_groups res ~rev_bindings:groups
-    | Lletrec (bind_args, body) ->
+    | Lletrec (bind_args, body) as original -> (
       (* Attention: don't mess up with internal {let rec} *)
-      let rec iter bind_args groups set =
-        match bind_args with
-        | [] -> (List.rev groups, set)
-        | (id, arg) :: rest ->
-          iter rest ((id, aux arg) :: groups) (Set_ident.add set id)
+      (* Keep the mapped list so a group from which nothing can be extracted
+         remains physically shared when neither its bindings nor body change. *)
+      let groups = Ext_list.map_snd_sharing bind_args aux in
+      let collections =
+        Ext_list.fold_left groups Set_ident.empty (fun set (id, _) ->
+            Set_ident.add set id)
       in
-      let groups, collections = iter bind_args [] Set_ident.empty in
       (* Try to extract some value definitions from recursive values as [wrap],
          it will stop whenever it find it could not move forward
          {[
@@ -241,19 +241,23 @@ let deep_flatten (lam : Lambda.t) : Lambda.t =
              ...
          ]}
       *)
-      let rev_bindings, rev_wrap, _ =
-        Ext_list.fold_left groups ([], [], false)
-          (fun (inner_recursive_bindings, wrap, stop) (id, lam) ->
-            if stop || Lam_hit.hit_variables collections lam then
-              ((id, lam) :: inner_recursive_bindings, wrap, true)
-            else
-              ( inner_recursive_bindings,
-                Lam_group.Single (Strict, id, lam) :: wrap,
-                false ))
+      let rec extract rev_wrap = function
+        | [] -> (rev_wrap, [])
+        | (_, binding) :: _ as bindings
+          when Lam_hit.hit_variables collections binding ->
+          (rev_wrap, bindings)
+        | (id, binding) :: rest ->
+          extract (Lam_group.Single (Strict, id, binding) :: rev_wrap) rest
       in
-      lambda_of_groups
-        ~rev_bindings:rev_wrap (* These bindings are extracted from [letrec] *)
-        (Lambda.letrec (List.rev rev_bindings) (aux body))
+      let rev_wrap, recursive_bindings = extract [] groups in
+      let body' = aux body in
+      match rev_wrap with
+      | [] when groups == bind_args && body' == body -> original
+      | [] -> Lambda.letrec groups body'
+      | _ ->
+        lambda_of_groups
+          ~rev_bindings:rev_wrap (* Extracted bindings from [letrec]. *)
+          (Lambda.letrec recursive_bindings body'))
     | _ -> Lambda_traverse.shallow_map_sharing aux lam
   in
   aux lam

tests/ounit_tests/ounit_deep_flatten_tests.ml

@@ -0,0 +1,18 @@
+open OUnit
+
+(* Sharing is invisible to the generated output: a pass that rebuilds a term
+   into an identical one produces the same JavaScript, so no snapshot can tell.
+   This is the only place that notices. *)
+let suites =
+  __FILE__
+  >::: [
+         ( "shares an unchanged recursive group" >:: fun _ ->
+           let recursive = Ident.create "recursive" in
+           let lam =
+             Lambda.letrec
+               [(recursive, Lambda.var recursive)]
+               (Lambda.var recursive)
+           in
+           assert_bool "the recursive group is physically unchanged"
+             (Lam_pass_deep_flatten.deep_flatten lam == lam) );
+       ]

tests/ounit_tests/ounit_tests_main.ml

@@ -20,6 +20,7 @@ let suites =
       Ounit_util_tests.suites;
       Ounit_rec_check_tests.suites;
       Ounit_lambda_constant_tests.suites;
+      Ounit_deep_flatten_tests.suites;
       Ounit_sroa_tests.suites;
       Ounit_ast_mapper0_tests.suites;
       Ounit_object_mutability_tests.suites;