Skip to content

Replace local mutable blocks in one pass over the function #9258

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cknitt wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Replace local mutable blocks in one pass over the function #9258

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
52c1aa6
Replace local mutable blocks in one pass over the function
cknitt Nov 7, 2026
ba4c395
Analyze blocks without allocating a closure per node
cknitt Nov 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/linear-sroa.polish.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
- Compile functions with many local `ref`s or mutable records in linear time.
287 changes: 207 additions & 80 deletions compiler/core/lam_pass_sroa.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,56 +21,47 @@

type field_use = {mutable read: bool; mutable written: bool}

let valid_field uses index = index >= 0 && index < Array.length uses
(* A local mutable block, with how its fields are used so far. [order]
numbers blocks in the order of their declarations. [fields] is set when
the block's scope has been analyzed and every occurrence was a direct,
in-range field read or write. *)
type candidate = {
order: int;
uses: field_use array;
mutable escapes: bool;
mutable fields: Ident.t array option;
mutable slots: slot list;
}

(* 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_traverse.shallow_exists
(fun child -> not (analyze block uses child))
lam)
(* An initializer of a field of [owner], or a value written to it, without
side occurrences. If the field is never read, the rewrite drops the value, and
the occurrences of blocks declared before [owner] in it don't count:
[occurrences] holds them until [owner] is decided. [parent] is the slot that
contains this one. *)
and slot = {
owner: candidate;
field: int;
parent: slot option;
mutable occurrences: occurrence list;
}

and occurrence = Escape of candidate | Read of candidate * int

let valid_field uses index = index >= 0 && index < Array.length uses

let discard_value value body =
if Lam_analysis.no_side_effects value then body else Lambda.seq value body
let is_candidate kind (init : Lambda.t) =
match (kind, init) with
| ( (Lambda.Strict | StrictOpt),
Lprim {primitive = Pmakeblock info; args = _ :: _} ) ->
not (Lambda.is_immutable_block info)
| _ -> false

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_traverse.shallow_map_sharing (rewrite block fields uses) lam
(* Drop [value], a value of a field that is never read, keeping it only for
its side effects, before [body]. [analyze] doesn't count the occurrences in
the values dropped here, so they are not rewritten. *)
let discard_value rewrite value body =
if Lam_analysis.no_side_effects value then body
else Lambda.seq (rewrite value) body

let fields_for_block block info field_count =
let fallback () =
Expand All @@ -97,18 +88,140 @@ let fields_for_block block info field_count =
| 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})
(* Find the blocks whose every occurrence is a direct, in-range field read or
write, recording how each field is used, in one traversal of the term.
[rewrite] below handles exactly the occurrences accepted here, and asserts
on the rest; keep their cases in step.

The result is the one of analyzing each block's scope after rewriting the
blocks declared in it, as if blocks were replaced one at a time from the
innermost out, so pending [slot]s keep the occurrences that such a rewrite
would drop. A block is decided when its scope has been traversed, so inner
blocks are decided (and their fields named) before outer ones. *)
let analyze ~considered (lam : Lambda.t) =
let candidates = Hash_ident.create 16 in
let count = ref 0 in
let apply = function
| Escape c -> c.escapes <- true
| Read (c, index) -> c.uses.(index).read <- true
in
let rec record slot occurrence =
let target =
match occurrence with
| Escape c | Read (c, _) -> c
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
match slot with
| Some slot when target.order < slot.owner.order ->
slot.occurrences <- occurrence :: slot.occurrences
| _ -> apply occurrence
and decide c =
if not c.escapes then (
List.iter
(fun slot ->
if c.uses.(slot.field).read then
List.iter (record slot.parent) slot.occurrences)
c.slots;
true)
else (
List.iter
(fun slot -> List.iter (record slot.parent) slot.occurrences)
c.slots;
false)
in
(* The innermost pending slot containing the term being traversed. A
reference rather than an argument of [go], so that traversing allocates
no closures. *)
let slot = ref None in
let rec go (lam : Lambda.t) =
match lam with
| Lvar id -> (
match Hash_ident.find_opt candidates id with
| Some c -> record !slot (Escape c)
| None -> ())
| Lassign (id, value) ->
(match Hash_ident.find_opt candidates id with
| Some c -> record !slot (Escape c)
| None -> ());
go value
| Lprim {primitive = Pfield (index, _); args = [Lvar id]}
when Hash_ident.mem candidates id ->
let c = Hash_ident.find candidates id in
if valid_field c.uses index then record !slot (Read (c, index))
else record !slot (Escape c)
| Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]}
when Hash_ident.mem candidates id ->
let c = Hash_ident.find candidates id in
if valid_field c.uses index then (
c.uses.(index).written <- true;
go_value c index value)
else (
record !slot (Escape c);
go value)
| Llet (kind, block, init, body)
when is_candidate kind init && considered block -> (
match init with
| Lprim {primitive = Pmakeblock info; args = initializers} ->
let field_count = List.length initializers in
let c =
{
order = !count;
uses =
Array.init field_count (fun _ -> {read = false; written = false});
escapes = false;
fields = None;
slots = [];
}
in
incr count;
List.iteri (fun index init -> go_value c index init) initializers;
Hash_ident.add candidates block c;
go body;
if decide c then
c.fields <- Some (fields_for_block block info field_count)
| _ -> assert false)
| _ -> ignore (Lambda_traverse.shallow_exists go_child lam)
and go_child lam =
go lam;
false
(* Traverse a value of field [index] of [c], in a slot of its own if the
rewrite drops the value when the field is never read. *)
and go_value c index value =
if Lam_analysis.no_side_effects value then (
let parent = !slot in
let s = {owner = c; field = index; parent; occurrences = []} in
c.slots <- s :: c.slots;
slot := Some s;
go value;
slot := parent)
else go value
in
go lam;
candidates

let rewrite candidates (lam : Lambda.t) =
let replaced id =
match Hash_ident.find_opt candidates id with
| Some {fields = Some fields; uses} -> Some (fields, uses)
| _ -> None
in
let rec rewrite (lam : Lambda.t) =
match lam with
| Lprim {primitive = Pfield (index, _); args = [Lvar id]} -> (
match replaced id with
| Some (fields, _) -> Lambda.var fields.(index)
| None -> lam)
| Lprim {primitive = Psetfield (index, _); args = [Lvar id; value]} -> (
match replaced id with
| Some (fields, uses) ->
if not uses.(index).read then
discard_value rewrite value Lambda.lambda_unit
else Lambda.assign fields.(index) (rewrite value)
| None -> Lambda_traverse.shallow_map_sharing rewrite lam)
| Llet
(_, block, Lprim {primitive = Pmakeblock _; args = initializers}, body)
when replaced block <> None ->
let fields, uses = Option.get (replaced block) in
let body = rewrite body in
let rec bind_fields index initializers body =
match initializers with
| [] -> body
Expand All @@ -119,28 +232,42 @@ let replace ~block ~info ~initializers body =
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
if not use.read then discard_value rewrite 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
Lam_util.refine_let ~kind:Strict fields.(index) (rewrite init) body
else Lambda.let_ Variable fields.(index) (rewrite init) body
in
Some (bind_fields 0 initializers body)
bind_fields 0 initializers body
(* Unreachable: [analyze] rejected the block for both of these. 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 replaced id <> None -> assert false
| Lassign (id, _) when replaced id <> None -> assert false
| _ -> Lambda_traverse.shallow_map_sharing rewrite lam
in
rewrite lam

(* Analyzing and rewriting each block separately would traverse its whole
scope once per block, which is quadratic for a function with many local
references. *)
let replace ~block ~info ~initializers body =
let init =
Lambda.prim ~primitive:(Pmakeblock info) ~args:initializers Location.none
in
if not (is_candidate Strict init) then None
else
let lam = Lambda.let_ Strict block init body in
let candidates = analyze ~considered:(Ident.same block) lam in
match Hash_ident.find_opt candidates block with
| Some {fields = Some _} -> Some (rewrite candidates lam)
| _ -> None

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_traverse.shallow_map_sharing simplify lam
let simplify (lam : Lambda.t) =
let candidates = analyze ~considered:(fun _ -> true) lam in
if
Hash_ident.fold
(fun _ c found -> found || c.fields <> None)
candidates false
then rewrite candidates lam
else lam
Loading
Loading