Read back the sharing changes for clarity

Reviewing the sharing series on clarity rather than allocation found
three places where chasing the property made the code worse, all of them
mine.

simplify_alias's string switch had become a when-guard containing a
match, which ran the same lookup twice and left an arm the guard makes
unreachable. It now finds the constant once and branches on that.

regroups_binding mirrors flatten's cases one for one, including why a
null conversion of a variable is left alone while any other one is
split. Nothing said the two have to stay in step, or that drifting costs
the flattening silently, because the binding then takes the fast path
and never reaches flatten at all.

Two passes bound `as original` for a value already in scope as `lam`,
giving one idiom two spellings across seven passes.

The traversal every pass delegates to had no test. Breaking the sharing
in its Lapply and Lswitch arms leaves every fixture in the repository
byte for byte identical and no test failing, while all seven passes
quietly lose the property. The new test checks each constructor twice:
that an identity map hands the node back, and that a replacing map does
not, since a node whose children were never visited would pass the first
by doing nothing.

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

@@ -131,7 +131,13 @@ let rec rhs_is_beta_residue (lam : Lambda.t) =
 (* [flatten] restructures a binding only when it hoists something out of the
    right hand side, splits a null conversion, or eliminates a tuple. Every
    other binding comes back as the same binding, so it can be rebuilt in place
-   and shared instead of taken apart and reassembled. *)
+   and shared instead of taken apart and reassembled.
+
+   This mirrors [flatten]'s own cases one for one, including why a null
+   conversion of a variable is left alone while any other one is split, so a
+   case added there that restructures has to be added here too. Drifting apart
+   costs the flattening, silently: the binding takes the fast path and is never
+   handed to [flatten] at all. *)
 let regroups_binding (str : Lambda.let_kind) (id : Ident.t) (arg : Lambda.t) =
   if rhs_is_beta_residue arg then false
   else
@@ -223,7 +229,7 @@ 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) as original -> (
+    | Lletrec (bind_args, body) -> (
       (* Attention: don't mess up with internal {let rec} *)
       (* Keep the mapped list so a group from which nothing can be extracted
          remains physically shared when neither its bindings nor body change. *)
@@ -252,7 +258,7 @@ let deep_flatten (lam : Lambda.t) : Lambda.t =
       let rev_wrap, recursive_bindings = extract [] groups in
       let body' = aux body in
       match rev_wrap with
-      | [] when groups == bind_args && body' == body -> original
+      | [] when groups == bind_args && body' == body -> lam
       | [] -> Lambda.letrec groups body'
       | _ ->
         lambda_of_groups

compiler/core/lam_pass_exits.ml

@@ -156,7 +156,7 @@ let subst_helper (subst : subst_tbl) (query : int -> int) (lam : Lambda.t) :
     Lambda.t =
   let rec simplif (lam : Lambda.t) =
     match lam with
-    | Lstaticcatch (l1, (i, xs), l2) as original -> (
+    | Lstaticcatch (l1, (i, xs), l2) -> (
       let i_occur = query i in
       match (i_occur, l2) with
       | 0, _ -> simplif l1
@@ -184,7 +184,7 @@ let subst_helper (subst : subst_tbl) (query : int -> int) (lam : Lambda.t) :
           simplif l1)
         else
           let l1' = simplif l1 in
-          if l1' == l1 && l2' == l2 then original
+          if l1' == l1 && l2' == l2 then lam
           else Lambda.staticcatch l1' (i, xs) l2')
     | Lstaticraise (i, []) -> (
       match Hash_int.find_opt subst i with

compiler/core/lam_pass_remove_alias.ml

@@ -236,17 +236,22 @@ let simplify_alias (meta : Lam_stats.t) (lam : Lambda.t) : Lambda.t =
     (*   *\) *)
     (*   when  Ext_list.same_length params args -> *)
     (*   simpl (Lam_beta_reduce.propogate_beta_reduce meta params body args) *)
-    | Lstringswitch (Lvar s, sw, d)
-      when match Hash_ident.find_opt meta.ident_tbl s with
-           | Some (Constant _) -> true
-           | Some _ | None -> false -> (
-      (* The scrutinee is a known constant, so switch on it directly. *)
-      match Hash_ident.find_opt meta.ident_tbl s with
-      | Some (Constant c) ->
+    | Lstringswitch (l, sw, d) -> (
+      let known_constant =
+        match l with
+        | Lvar s -> (
+          match Hash_ident.find_opt meta.ident_tbl s with
+          | Some (Constant c) -> Some c
+          | Some _ | None -> None)
+        | _ -> None
+      in
+      match known_constant with
+      | Some c ->
+        (* Switch on the constant the scrutinee is bound to. *)
         Lambda.stringswitch (Lambda.const c)
           (Ext_list.map_snd sw simpl)
           (Ext_option.map d simpl)
-      | Some _ | None -> Lambda_traverse.shallow_map_sharing simpl lam)
+      | None -> Lambda_traverse.shallow_map_sharing simpl lam)
     | _ -> Lambda_traverse.shallow_map_sharing simpl lam
   in
   simpl lam

tests/ounit_tests/ounit_lambda_traverse_tests.ml

@@ -0,0 +1,72 @@
+open OUnit
+
+let loc = Location.none
+let x = Ident.create "x"
+let y = Ident.create "y"
+let debugger = Lambda.prim ~primitive:Pdebugger ~args:[] loc
+let var = Lambda.var x
+
+(* One node per constructor that has children. A leaf shares trivially, so it
+   would pass either check below without exercising anything. *)
+let nodes : (string * Lambda.t) list =
+  [
+    ("apply", Lambda.apply var [var] {ap_loc = loc; ap_inlined = Default_inline});
+    ( "function",
+      Lambda.function_ ~loc ~attr:Lambda.default_function_attribute ~params:[x]
+        ~body:debugger );
+    ("let", Lambda.let_ Strict y debugger var);
+    ("letrec", Lambda.letrec [(y, debugger)] var);
+    ("prim", Lambda.prim ~primitive:Pdebugger ~args:[var] loc);
+    ( "switch",
+      Lambda.switch var
+        {
+          sw_consts_full = false;
+          sw_consts = [(Switch_int 0, debugger)];
+          sw_blocks_full = false;
+          sw_blocks = [];
+          sw_failaction = Some debugger;
+          sw_dispatch = Switch_direct;
+        } );
+    ("stringswitch", Lambda.stringswitch var [("a", debugger)] (Some debugger));
+    ("staticraise", Lambda.staticraise 1 [var]);
+    ( "staticcatch",
+      Lambda.staticcatch (Lambda.staticraise 1 []) (1, []) debugger );
+    ("trywith", Lambda.try_ debugger y var);
+    ("ifthenelse", Lambda.if_ var debugger debugger);
+    ("sequence", Lambda.seq debugger var);
+    ("while", Lambda.while_ var debugger);
+    ("for", Lambda.for_ y var var Upto debugger);
+    ("for_of", Lambda.for_of y var debugger);
+    ("for_await_of", Lambda.for_await_of y var debugger);
+    ("assign", Lambda.assign x debugger);
+  ]
+
+(* Every optimization pass routes its "nothing to do here" case through
+   [shallow_map_sharing], so an arm of it that stops sharing silently costs the
+   property in all of them. That is invisible to generated output: breaking the
+   [Lapply] and [Lswitch] arms leaves every fixture in the repository byte for
+   byte identical. Add a node above when adding a Lambda constructor. *)
+let suites =
+  __FILE__
+  >::: [
+         ( "an unchanged child is not rebuilt" >:: fun _ ->
+           List.iter
+             (fun (name, node) ->
+               assert_bool
+                 (name ^ " should be handed back when nothing changed")
+                 (Lambda_traverse.shallow_map_sharing (fun lam -> lam) node
+                 == node))
+             nodes );
+         ( "a changed child is rebuilt" >:: fun _ ->
+           (* Without this, a node whose children were never visited would pass
+              the check above by doing nothing at all. *)
+           List.iter
+             (fun (name, node) ->
+               assert_bool
+                 (name ^ " should be rebuilt when a child changed")
+                 (Lambda_traverse.shallow_map_sharing
+                    (fun _ -> Lambda.const Lambda.const_unit)
+                    node
+                 != node))
+             nodes );
+       ]

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_lambda_traverse_tests.suites;
       Ounit_deep_flatten_tests.suites;
       Ounit_exits_tests.suites;
       Ounit_sroa_tests.suites;