Skip to content

Unbox non-escaping Ref/STRef cells to mutable Lua locals - #315

Merged
Unisay merged 4 commits into
mainfrom
issue-239/unbox-ref-locals
Jul 27, 2026
Merged

Unbox non-escaping Ref/STRef cells to mutable Lua locals#315
Unisay merged 4 commits into
mainfrom
issue-239/unbox-ref-locals

Conversation

@Unisay

@Unisay Unisay commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Closes #239.

A Ref/STRef cell compiles to a one-field heap table: the foreign new allocates {value = v}, and every read/write/modify pays a table field access. This PR adds a codegen-time pass (Language.PureScript.Backend.Lua.RefUnbox) that lowers a non-escaping cell to a plain mutable Lua local: new v becomes local r = v, a read run becomes r itself, and write/modify runs become assignments. The load-bearing observation is that a Lua local captured by inner closures is an upvalue — a shared, mutable heap slot with exactly the aliasing behaviour of the {value = …} table — so occurrences under lambdas (loop bodies, nested magic-do chunks) need no special treatment, and the only thing the table provides that a local cannot is first-class identity. The analysis therefore asks one question per cell: is every occurrence the reference argument of a recognised operation? Any other occurrence (stored, returned, passed to an unknown function) keeps the cell boxed, so a miss is a missed optimisation, never a miscompile.

What lowers. Recognition is by qualified name — the identity that survives linking, the same scheme magic-do and the native-loop matcher use — over the cell primitives of Effect.Ref and Control.Monad.ST.Internal (_new/new, read, write, modifyImpl) plus the ST functor's foreign map_ as a wrapper, because void (STRef.modify f r) (the ubiquitous modify_ shape) inlines to map_ (\_ → unit) (modifyImpl f r). Only run positions lower (the application carries magic-do's EffectRunArg marker, or is the body of a literal loop lambda that NativeLoop.runStatements runs per iteration); an operation thunk used as a first-class value keeps the cell boxed. A literal modify function whose body ends in a manifest {state, value} record is beta-reduced at emission — parameter bound to the local, Let spine emitted as statements, state field assigned, value field passed on — so the per-iteration record is never allocated either.

Before/after, the ST accumulator loop from Golden.NativeLoopsST.Test (quoted verbatim from the committed goldens):

local acc = Control_Monad_ST_Internal_new(0)()
for i = 0, n + 1 - 1 do
  Control_Monad_ST_Internal_modifyImpl(function(s_S_0)
    local sPrime_S_0 = s_S_0 + i
    return { state = sPrime_S_0, value = sPrime_S_0 }
  end)(acc)()
end
return Control_Monad_ST_Internal_read(acc)()
local acc = 0
for i = 0, n + 1 - 1 do
  local s_S_0 = acc
  local sPrime_S_0 = s_S_0 + i
  acc = sPrime_S_0
end
return acc

Soundness guard. The new Golden.RefUnbox.Test pins a cell stored inside another cell (outer <- STRef.new inner): the stored inner keeps its {value = …} table while outer still unboxes, and boxed operations on an aliased cell keep going through the foreign implementation. Golden.NativeLoopsST.Test's countDown pins the per-cell granularity of the value-position restriction: its steps cell (used only in run positions) unboxes while its value cell (read under map_ inside the while condition — a first-class thunk) stays boxed. All eval/golden.txt oracles — which are never auto-accepted — are byte-identical, so runtime behaviour is unchanged across the whole golden corpus, including the write-returns-written-value ST semantics and Effect write's nil result.

Measurement (issue's verification asks for the allocation-counter drop). A new macro bench Bench.RefLoop — a hot ST loop accumulating through a local STRef — with committed counter oracles: the static table-allocation census (tnew_census) drops from 2 steady-state sites to 1 (the survivor is the dead foreign new closure body, never called), and the trace report shows the hot loop trace-compiled (JFORL) instead of aborted on per-iteration closures. Wall-clock at n=100000 on this machine: PUC Lua 5.1 0.0347s → 0.0009s, LuaJIT 0.024s → 0.0001s (medians of bench/tools/run_macro.lua). All pre-existing counter oracles are byte-identical, i.e. no other benchmark's shape moved.

Mechanics. fromIR carries a new set of unboxed cell names: the Let case decides unboxability per allocation binding (scanning the remaining bindings and body), lowers operation-run statements in place (scoped in do … end blocks whenever the lowering declares locals, keeping magic-do's ~150-statement chunks under Lua's 200-active-locals cap), and threads the extended set through everything downstream; an expression-position run compiles to a chunk ending in a return of the run's value. Two NativeLoop adjustments: runStatements now statementizes chunk-lowered runs (an unboxed-cell run returns its value in a tail return, which a run-for-effect must rewrite into evaluation statements — previously only return-free loop chunks reached that path), and dropValue additionally drops a discarded field read off the module-scope table (built by codegen as a plain table, so no metamethod can fire), which the void wrapper's unit value otherwise leaves behind as a dead local per loop iteration.

The two pre-existing -Wname-shadowing warnings in Lua.hs (both \group lambdas shadowing Relude's group) surfaced during the rebuild and are fixed by renaming to dispatchGroup, per the repo's no-tolerated-warnings rule.

Unisay added 4 commits July 27, 2026 12:06
A new Golden.RefUnbox.Test module covering the cell shapes issue #239
lowers: a loop accumulator (sumTo), modify' with distinct state/value
components (splitModify), ST write returning the written value
(writeBack), a cell stored inside another cell (nested), and an Effect
Ref with direct read/write (main). The goldens pin the current boxed
form — every cell a {value = ...} table — so the unboxing change shows
as a reviewable diff against them; the hand-written eval oracle pins
the runtime outputs the rewrite must preserve.
A Ref/STRef cell compiles to a one-field heap table ({value = ...})
with every read/write/modify going through its field. A Lua local
captured by inner closures is itself a shared mutable slot (an
upvalue), so when the cell never flows anywhere as a whole value the
table buys nothing: a Let-bound run of new whose every use is a
recognised operation now lowers to a plain local, reads to the local
itself, writes and modifies to assignments. A literal modify function
beta-reduces at emission, so its {state, value} record is never
allocated either.

Recognition is by qualified name over the Effect.Ref and
Control.Monad.ST.Internal primitives plus the ST functor's foreign
map_ (the void (modify f r) shape), in the two codegen-time head forms
the loop matcher already uses. Only run positions lower; a cell used
as a first-class value anywhere keeps its boxed form, so a miss is a
missed optimisation, never a miscompile.

NativeLoop.runStatements statementizes chunk-lowered runs (an
unboxed-cell run returns its value in a tail return, which a
run-for-effect must discard), and dropValue drops discarded reads of
module-scope table fields, which the void wrapper's unit value
otherwise leaves as a dead local per loop iteration.

Closes #239
A hot ST loop accumulating through a non-escaping local STRef, with
the hand-written Lua for-loop as its ideal. The committed counter
oracles pin the unboxed shape: one steady-state TNEW/TDUP site in the
whole artifact (the dead foreign new closure) and the hot loop
trace-compiled (JFORL) instead of aborted on the per-iteration
closures. Wall-clock on this machine: PUC Lua 5.1 0.0347s -> 0.0009s,
LuaJIT 0.024s -> 0.0001s at n=100000.
@Unisay Unisay self-assigned this Jul 27, 2026
@Unisay
Unisay marked this pull request as ready for review July 27, 2026 10:59
@Unisay
Unisay merged commit 9b4cc37 into main Jul 27, 2026
2 checks passed
@Unisay
Unisay deleted the issue-239/unbox-ref-locals branch July 27, 2026 11:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Unbox non-escaping Ref/STRef cells to mutable Lua locals

1 participant