Classify field use in scalar replacement

The pass replaced every field of an eligible block with a mutable binding,
whether or not the field was ever read. Recording how each field is used costs
nothing - the eligibility walk already visits every occurrence - and decides
what each field needs.

A field that is never read needs no storage: its initializer and its writes are
kept only when they have effects. A field that is read but never written is
immutable, so a refined let will do. Only a field that is both read and written
needs a mutable scalar.

test_ramification shows the shape this is for. A ref written in both branches
of a match and never read afterwards disappears, and what is left folds:

    let v = ref(0)
    let y = switch x { | A(_) => v := 1; 3 | B(_) => v := 1; 4 }

    -  let v = 0;                             let y;
    -  if (x.TAG === "A") { v = 1; y = 3; }   y = x.TAG === "A" ? 3 : 4;
    -  else { v = 1; y = 4; }

escapes becomes analyze, returning eligibility rather than escape so it can
report through the same walk. The short circuit still holds: on success every
occurrence has been visited, so the use table is complete, and on failure it is
discarded with the rejection.

Read-only fields go through refine_let, which may substitute the initializer at
its use sites - but only when is_safe_to_alias admits it, which is variables,
constants and module field reads. An effectful initializer read five times in a
loop, or three times through a closure, is still evaluated once.

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_sroa.ml

@@ -19,41 +19,56 @@
    eligible. Analysis is kept separate from rewriting so a failed eligibility
    check cannot partially transform the term. *)
 
-let valid_field field_count index = index >= 0 && index < field_count
+type field_use = {mutable read: bool; mutable written: bool}
 
-(* Does the block appear anywhere other than as a direct, in-range field read
-   or write? [escapes] and [rewrite] below are a matched pair: [rewrite] handles
-   exactly the occurrences [escapes] accepts, and asserts on the rest. Extending
-   one without the other is a compiler crash rather than a type error, so keep
-   their cases in step. *)
-let rec escapes block field_count (lam : Lambda.t) =
+let valid_field uses index = index >= 0 && index < Array.length uses
+
+(* Does the block appear only as direct, in-range field reads and writes? While
+   answering, record how every field is used. [analyze] and [rewrite] below are
+   a matched pair: [rewrite] handles exactly the occurrences [analyze] accepts,
+   and asserts on the rest. Extending one without the other is a compiler crash
+   rather than a type error, so keep their cases in step. *)
+let rec analyze block uses (lam : Lambda.t) =
   match lam with
-  | Lvar id -> Ident.same id block
+  | Lvar id -> not (Ident.same id block)
   | Lassign (id, value) ->
-    Ident.same id block || escapes block field_count value
+    (not (Ident.same id block)) && analyze block uses value
   | Lprim {primitive = Pfield (index, _); args = [Lvar id]}
     when Ident.same id block ->
-    not (valid_field field_count index)
+    if valid_field uses index then (
+      uses.(index).read <- true;
+      true)
+    else false
   | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]}
     when Ident.same id block ->
-    (not (valid_field field_count index)) || escapes block field_count value
-  | _ -> Lambda.shallow_exists (escapes block field_count) lam
+    if valid_field uses index then (
+      uses.(index).written <- true;
+      analyze block uses value)
+    else false
+  | _ ->
+    not
+      (Lambda.shallow_exists (fun child -> not (analyze block uses child)) lam)
 
-let rec rewrite block fields (lam : Lambda.t) =
+let discard_value value body =
+  if Lam_analysis.no_side_effects value then body else Lambda.seq value body
+
+let rec rewrite block fields uses (lam : Lambda.t) =
   match lam with
   | Lprim {primitive = Pfield (index, _); args = [Lvar id]}
     when Ident.same id block ->
     Lambda.var fields.(index)
   | Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]}
     when Ident.same id block ->
-    Lambda.assign fields.(index) (rewrite block fields value)
-  (* Unreachable: [escapes] rejected the block for both of these, so [replace]
+    let value = rewrite block fields uses value in
+    if not uses.(index).read then discard_value value Lambda.lambda_unit
+    else Lambda.assign fields.(index) value
+  (* Unreachable: [analyze] rejected the block for both of these, so [replace]
      never reaches the rewrite. They are kept as assertions rather than dropped
-     so that a future occurrence form added to [escapes] but not here fails
+     so that a future occurrence form added to [analyze] but not here fails
      loudly instead of silently losing the write. *)
   | Lvar id when Ident.same id block -> assert false
   | Lassign (id, _) when Ident.same id block -> assert false
-  | _ -> Lambda.shallow_map_sharing (rewrite block fields) lam
+  | _ -> Lambda.shallow_map_sharing (rewrite block fields uses) lam
 
 let fields_for_block block info field_count =
   let fallback () =
@@ -85,13 +100,29 @@ let replace ~block ~info ~initializers body =
   | [] -> None
   | _ ->
     let field_count = List.length initializers in
-    if escapes block field_count body then None
+    let uses =
+      Array.init field_count (fun _ -> {read = false; written = false})
+    in
+    if not (analyze block uses body) then None
     else
       let fields = fields_for_block block info field_count in
-      let body = rewrite block fields body in
-      Some
-        (Ext_list.fold_right2 (Array.to_list fields) initializers body
-           (fun field init body -> Lambda.let_ Variable field init body))
+      let body = rewrite block fields uses body in
+      let rec bind_fields index initializers body =
+        match initializers with
+        | [] -> body
+        | init :: rest ->
+          let body = bind_fields (index + 1) rest body in
+          let use = uses.(index) in
+          (* A never-read field needs no storage; its initializer and writes are
+             retained only when they have effects. A read-only field can use a
+             normal refined let, while a field that is both read and written
+             still needs a mutable scalar binding. *)
+          if not use.read then discard_value init body
+          else if not use.written then
+            Lam_util.refine_let ~kind:Strict fields.(index) init body
+          else Lambda.let_ Variable fields.(index) init body
+      in
+      Some (bind_fields 0 initializers body)
 
 let rec simplify (lam : Lambda.t) =
   match lam with

compiler/core/lam_pass_sroa.mli

@@ -29,9 +29,11 @@ val replace :
   Lambda.t ->
   Lambda.t option
 (** [replace ~block ~info ~initializers body] replaces a non-escaping local
-    block with one mutable binding per field. The initializer order is
-    preserved. Returns [None] when the block is used other than by direct field
-    access. *)
+    block with independent scalar values. Fields that are never read need no
+    storage, but effects from their initializers and writes are preserved in
+    order. Read-only fields use immutable bindings; fields that are both read
+    and written use mutable bindings. Returns [None] when the block is used
+    other than by direct field access. *)
 
 val simplify : Lambda.t -> Lambda.t
 (** Scalar-replace eligible local mutable blocks throughout a Lambda term. *)

tests/ounit_tests/ounit_sroa_tests.ml

@@ -14,6 +14,32 @@ let write block index value =
     ~args:[Lambda.var block; value]
     loc
 
+let debugger = Lambda.prim ~primitive:Pdebugger ~args:[] loc
+
+let count_debuggers lam =
+  let count = ref 0 in
+  let rec loop (lam : Lambda.t) =
+    (match lam with
+    | Lprim {primitive = Pdebugger} -> incr count
+    | _ -> ());
+    ignore
+      (Lambda.shallow_exists
+         (fun child ->
+           loop child;
+           false)
+         lam)
+  in
+  loop lam;
+  !count
+
+let contains_storage lam =
+  let rec loop (lam : Lambda.t) =
+    match lam with
+    | Llet _ | Lassign _ -> true
+    | _ -> Lambda.shallow_exists loop lam
+  in
+  loop lam
+
 let pair_info =
   Lambda.Blk_record
     {
@@ -33,36 +59,37 @@ let suites =
              ]
            in
            let body =
-             Lambda.seq
-               (write block 1 (Lambda.const (Lambda.const_int 30)))
-               (read block 0)
+             Lambda.seq (write block 1 (read block 0)) (read block 1)
            in
            match
              Lam_pass_sroa.replace ~block ~info:pair_info ~initializers body
            with
            | Some
                (Llet
-                  ( Variable,
+                  ( Alias,
                     field0,
                     Lconst (Const_int 10l),
                     Llet
                       ( Variable,
                         field1,
                         Lconst (Const_int 20l),
-                        Lsequence (Lassign (assigned, Lconst _), Lvar returned)
+                        Lsequence (Lassign (assigned, Lvar read), Lvar returned)
                       ) )) ->
              assert_equal "pair_left" (Ident.name field0);
              assert_equal "pair_right" (Ident.name field1);
              assert_bool "field one gets its own binding"
                (Ident.same field1 assigned);
-             assert_bool "the read uses field zero" (Ident.same field0 returned)
+             assert_bool "the assignment reads field zero"
+               (Ident.same field0 read);
+             assert_bool "the result reads field one"
+               (Ident.same field1 returned)
            | _ -> assert_failure "expected two scalar bindings" );
          ( "replaces fields captured by a closure" >:: fun _ ->
            let block = Ident.create "pair" in
            let closure =
              Lambda.function_ ~loc ~attr:Lambda.default_function_attribute
                ~params:[]
-               ~body:(write block 1 (read block 0))
+               ~body:(Lambda.seq (write block 1 (read block 0)) (read block 1))
            in
            match
              Lam_pass_sroa.replace ~block ~info:pair_info
@@ -75,19 +102,24 @@ let suites =
            with
            | Some
                (Llet
-                  ( Variable,
+                  ( Alias,
                     field0,
                     _,
                     Llet
                       ( Variable,
                         field1,
                         _,
-                        Lfunction {body = Lassign (assigned, Lvar returned)} )
-                  )) ->
+                        Lfunction
+                          {
+                            body =
+                              Lsequence
+                                (Lassign (assigned, Lvar read), Lvar returned);
+                          } ) )) ->
              assert_bool "the closure writes field one"
                (Ident.same field1 assigned);
-             assert_bool "the closure reads field zero"
-               (Ident.same field0 returned)
+             assert_bool "the closure reads field zero" (Ident.same field0 read);
+             assert_bool "the closure reads field one"
+               (Ident.same field1 returned)
            | _ -> assert_failure "expected scalar closure captures" );
          ( "preserves unrelated closure subtrees" >:: fun _ ->
            let block = Ident.create "pair" in
@@ -101,13 +133,58 @@ let suites =
                ~initializers:[Lambda.const (Lambda.const_int 10)]
                body
            with
-           | Some
-               (Llet (Variable, _, _, Lsequence (Lvar _, preserved_unrelated)))
+           | Some (Llet (Alias, _, _, Lsequence (Lvar _, preserved_unrelated)))
              ->
              assert_bool "the unrelated subtree is physically shared"
                (preserved_unrelated == unrelated)
            | _ -> assert_failure "expected a scalar read followed by a closure"
          );
+         ( "keeps effectful read-only fields strict" >:: fun _ ->
+           let block = Ident.create "cell" in
+           let body = Lambda.seq (read block 0) (read block 0) in
+           match
+             Lam_pass_sroa.replace ~block ~info:Lambda.ref_tag_info
+               ~initializers:[debugger] body
+           with
+           | Some
+               (Llet
+                  ( Strict,
+                    field,
+                    Lprim {primitive = Pdebugger},
+                    Lsequence (Lvar first, Lvar second) )) ->
+             assert_bool "both reads use the strict scalar"
+               (Ident.same field first && Ident.same field second)
+           | _ -> assert_failure "expected one strict scalar binding" );
+         ( "removes write-only storage but preserves effects" >:: fun _ ->
+           let block = Ident.create "pair" in
+           let body =
+             Lambda.seq (write block 0 debugger) (write block 1 debugger)
+           in
+           match
+             Lam_pass_sroa.replace ~block ~info:pair_info
+               ~initializers:[Lambda.const (Lambda.const_int 10); debugger]
+               body
+           with
+           | Some replacement ->
+             assert_equal 3 (count_debuggers replacement);
+             assert_bool "write-only fields have no scalar storage"
+               (not (contains_storage replacement))
+           | None -> assert_failure "expected write-only fields to scalarize" );
+         ( "removes write-only storage captured by a closure" >:: fun _ ->
+           let block = Ident.create "cell" in
+           let closure =
+             Lambda.function_ ~loc ~attr:Lambda.default_function_attribute
+               ~params:[] ~body:(write block 0 debugger)
+           in
+           match
+             Lam_pass_sroa.replace ~block ~info:Lambda.ref_tag_info
+               ~initializers:[debugger] closure
+           with
+           | Some replacement ->
+             assert_equal 2 (count_debuggers replacement);
+             assert_bool "the closure captures no scalar storage"
+               (not (contains_storage replacement))
+           | None -> assert_failure "expected the closure write to scalarize" );
          ( "rejects a whole-block use" >:: fun _ ->
            let block = Ident.create "pair" in
            assert_equal None