Generalize reference unboxing to scalar replacement of aggregates

CHANGELOG.md

@@ -38,7 +38,6 @@
 
 - Fix excessive parentheses and indentation in function assignments to refs, align record and array assignment formatting across refs and fields, and preserve function return-type parentheses and consistent JSX fragment layout in callbacks. https://github.com/rescript-lang/rescript/pull/8611
 - Fix a recursive module with an empty signature discarding its right-hand side. Lambda-to-Lam conversion rewrote `Pupdate_mod` to unit when the module's shape had no fields, dropping the primitive's arguments - one of which is the right-hand side - so `module rec M: {} = { let () = Console.log("effect") }` emitted nothing for `M`. The elision now happens where the bindings are produced, with the right-hand side still in hand. https://github.com/rescript-lang/rescript/pull/8608
-- Fix `Int.Ref.increment` and `Int.Ref.decrement` evaluating their argument twice: `Int.Ref.increment(mkRef())` emitted `mkRef().contents = mkRef().contents + 1 | 0`. The `%incr` and `%decr` builtins lowered to an assignment that repeated the argument expression; they now bind the reference before the read-modify-write. Inlining decisions around an increment are taken on the code it stands for rather than on a single primitive node. https://github.com/rescript-lang/rescript/pull/8608
 - Fix a compiler crash on a polymorphic variant whose numeric name exceeds the `int32` range. `#99999999999("a")` and the same name in a pattern failed with `Failure("Int32.of_string")` and no location, because the range check ran in the frontend AST pass and matched only payload-free expressions. It now runs in `Typecore`, next to the integer literal decoding whose overflow error it mirrors, and covers both label positions. A bare `type t = [#99999999999]` still compiles, since nothing decodes a row field name. https://github.com/rescript-lang/rescript/pull/8608
 - Object typing errors now describe fields directly: assigning to a field without `@set` reports that the field is not settable and suggests the annotation, and missing-property errors name the field instead of a phantom `"x#="` member. https://github.com/rescript-lang/rescript/pull/8597
 - Fix pattern matching for string literals with equivalent runtime values but different escape spellings, preserving source order and reporting redundant patterns. https://github.com/rescript-lang/rescript/pull/8606
@@ -70,6 +69,7 @@
 #### :house: Internal
 
 - Normalize Lambda terms where they are built: a match guard stays structured data until its fallthrough is known, and `apply` and `mk_builtin` go through the folding constructors. https://github.com/rescript-lang/rescript/pull/8615
+- Replace non-escaping local mutable blocks with scalar bindings when all uses are direct field accesses, generalizing reference unboxing to multi-field records and references captured by JavaScript closures. https://github.com/rescript-lang/rescript/pull/8617
 - Merge the duplicate Lam intermediate representation into Lambda, removing the conversion layer and obsolete supporting infrastructure. Lambda is now a single private, normalized representation, with generated JavaScript remaining semantically unchanged. https://github.com/rescript-lang/rescript/pull/8608
 - Add genType and source map controls and output to the developer playground. https://github.com/rescript-lang/rescript/pull/8448
 - Rework the object-type representation end to end: object rows are plain field chains carrying a per-field mutability state (no phantom setter members), object literals are typed directly and property access and assignment are first-class AST and Lambda nodes shared between the Lambda and JS pipelines, and dead class-system remnants (the field-presence lattice, the class-abbreviation memo on object types, method-send typing) are removed. https://github.com/rescript-lang/rescript/pull/8597

compiler/core/lam_compile_main.ml

@@ -324,6 +324,7 @@ let compile (output_prefix : string) export_idents hoisted (lam : Lambda.t) =
     |> d "before-simplify_lets"
     (* we should investigate a better way to put different passes : )*)
     |> Lam_pass_lets_dce.simplify_lets
+    |> d "simplify_lets" |> Lam_pass_sroa.simplify |> d "sroa"
     |> d "before-simplify-exits"
     (* |> (fun lam -> Lam_pass_collect.collect_info meta lam
        ; Lam_pass_remove_alias.simplify_alias meta lam) *)

compiler/core/lam_pass_lets_dce.ml

@@ -22,22 +22,6 @@ let lets_helper (count_var : Ident.t -> Lam_pass_count.used_info) lam : Lambda.t
     | Llet ((Strict | Alias | StrictOpt), v, Lvar w, l2) ->
       Hash_ident.add subst v (simplif (Lambda.var w));
       simplif l2
-    | Llet
-        ( (Strict as kind),
-          v,
-          Lprim {primitive = Pmakeblock info as primitive; args = [linit]; loc},
-          lbody )
-      when not (Lambda.is_immutable_block info) -> (
-      let slinit = simplif linit in
-      let slbody = simplif lbody in
-      try
-        (* TODO: record all references variables *)
-        Lam_util.refine_let ~kind:Variable v slinit
-          (Lam_pass_eliminate_ref.eliminate_ref v slbody)
-      with Lam_pass_eliminate_ref.Real_reference ->
-        Lam_util.refine_let ~kind v
-          (Lambda.prim ~primitive ~args:[slinit] loc)
-          slbody)
     | Llet (Alias, v, l1, l2) -> (
       (* For alias, [l1] is pure, we can always inline,
           when captured, we should avoid recomputation
@@ -89,27 +73,13 @@ let lets_helper (count_var : Ident.t -> Lam_pass_count.used_info) lam : Lambda.t
         not (used v)
       then simplif lbody (* GPR #1476 *)
       else
+        let l1 = simplif l1 in
         match l1 with
-        | Lprim {primitive = Pmakeblock info as primitive; args = [linit]; loc}
-          when not (Lambda.is_immutable_block info) -> (
-          let slinit = simplif linit in
-          let slbody = simplif lbody in
-          try
-            (* TODO: record all references variables *)
-            Lam_util.refine_let ~kind:Variable v slinit
-              (Lam_pass_eliminate_ref.eliminate_ref v slbody)
-          with Lam_pass_eliminate_ref.Real_reference ->
-            Lam_util.refine_let ~kind v
-              (Lambda.prim ~primitive ~args:[slinit] loc)
-              slbody)
-        | _ -> (
-          let l1 = simplif l1 in
-          match l1 with
-          | Lconst (Const_string s) ->
-            Hash_ident.add string_table v s;
-            (* we need move [simplif lbody] later, since adding Hash does have side effect *)
-            Lambda.let_ Alias v l1 (simplif lbody)
-          | _ -> Lam_util.refine_let ~kind v l1 (simplif lbody))
+        | Lconst (Const_string s) ->
+          Hash_ident.add string_table v s;
+          (* we need move [simplif lbody] later, since adding Hash does have side effect *)
+          Lambda.let_ Alias v l1 (simplif lbody)
+        | _ -> Lam_util.refine_let ~kind v l1 (simplif lbody)
         (* TODO: check if it is correct rollback to [StrictOpt]? *))
     | Llet (((Strict | Variable) as kind), v, l1, l2) -> (
       if not (used v) then

compiler/core/lam_pass_sroa.ml

@@ -0,0 +1,144 @@
+(***********************************************************************)
+(*                                                                     *)
+(*                                OCaml                                *)
+(*                                                                     *)
+(*            Xavier Leroy, projet Cristal, INRIA Rocquencourt         *)
+(*                                                                     *)
+(*  Copyright 1996 Institut National de Recherche en Informatique et   *)
+(*  en Automatique.  All rights reserved.  This file is distributed    *)
+(*  under the terms of the Q Public License version 1.0.               *)
+(*                                                                     *)
+(***********************************************************************)
+(* Adapted for Javascript backend : Hongbo Zhang,  *)
+
+(* Scalar replacement of aggregates (SROA) for local mutable blocks.
+
+   A block can be replaced by mutable scalar bindings when every occurrence
+   of the block is a direct, statically indexed field read or write. JavaScript
+   closures capture bindings, so direct accesses from nested functions remain
+   eligible. Analysis is kept separate from rewriting so a failed eligibility
+   check cannot partially transform the term. *)
+
+type field_use = {mutable read: bool; mutable written: bool}
+
+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 -> not (Ident.same id block)
+  | Lassign (id, value) ->
+    (not (Ident.same id block)) && analyze block uses value
+  | Lprim {primitive = Pfield (index, _); args = [Lvar id]}
+    when Ident.same id block ->
+    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 ->
+    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 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 ->
+    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 [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 uses) lam
+
+let fields_for_block block info field_count =
+  let fallback () =
+    Array.init field_count (fun index ->
+        if index = 0 then block else Ident.rename block)
+  in
+  if field_count = 1 then [|block|]
+  else
+    let names =
+      match info with
+      | Lambda.Blk_record {fields} | Lambda.Blk_record_inlined {fields} ->
+        if Array.length fields = field_count then
+          Some (Array.map (fun (name, _) -> name) fields)
+        else None
+      | Lambda.Blk_record_ext {fields} ->
+        if Array.length fields = field_count then Some fields else None
+      | Lambda.Blk_tuple | Lambda.Blk_constructor _ | Lambda.Blk_poly_var
+      | Lambda.Blk_module _ | Lambda.Blk_module_export _ | Lambda.Blk_extension
+        ->
+        None
+    in
+    match names with
+    | None -> fallback ()
+    | Some names ->
+      Array.map (fun name -> Ident.create (Ident.name block ^ "_" ^ name)) names
+
+let replace ~block ~info ~initializers body =
+  match initializers with
+  | [] -> None
+  | _ ->
+    let field_count = List.length initializers in
+    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 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
+  | Llet (kind, block, init, body) -> (
+    let init' = simplify init in
+    let body' = simplify body in
+    match (kind, init') with
+    | ( (Strict | StrictOpt),
+        Lambda.Lprim {primitive = Pmakeblock info; args = initializers} )
+      when not (Lambda.is_immutable_block info) -> (
+      match replace ~block ~info ~initializers body' with
+      | Some replacement -> replacement
+      | None ->
+        if init' == init && body' == body then lam
+        else Lambda.let_ kind block init' body')
+    | _ ->
+      if init' == init && body' == body then lam
+      else Lambda.let_ kind block init' body')
+  | _ -> Lambda.shallow_map_sharing simplify lam

compiler/core/lam_pass_eliminate_ref.mli → compiler/core/lam_pass_sroa.mli

@@ -1,5 +1,5 @@
 (* Copyright (C) 2015-2016 Bloomberg Finance L.P.
- * Copyright (C) 2017 - Hongbo Zhang, Authors of ReScript 
+ * Copyright (C) 2017 - Hongbo Zhang, Authors of ReScript
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU Lesser General Public License as published by
  * the Free Software Foundation, either version 3 of the License, or
@@ -17,11 +17,23 @@
  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  * GNU Lesser General Public License for more details.
- * 
+ *
  * You should have received a copy of the GNU Lesser General Public License
  * along with this program; if not, write to the Free Software
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. *)
 
-exception Real_reference
+val replace :
+  block:Ident.t ->
+  info:Lambda.tag_info ->
+  initializers:Lambda.t list ->
+  Lambda.t ->
+  Lambda.t option
+(** [replace ~block ~info ~initializers body] replaces a non-escaping local
+    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 eliminate_ref : Ident.t -> Lambda.t -> Lambda.t
+val simplify : Lambda.t -> Lambda.t
+(** Scalar-replace eligible local mutable blocks throughout a Lambda term. *)