Skip to content

Uncurry functions via a worker/wrapper split#202

Merged
Unisay merged 6 commits into
mainfrom
issue-24/uncurry-worker-wrapper
Jul 7, 2026
Merged

Uncurry functions via a worker/wrapper split#202
Unisay merged 6 commits into
mainfrom
issue-24/uncurry-worker-wrapper

Conversation

@Unisay

@Unisay Unisay commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Implements the worker/wrapper design from #24: curried functions stop paying a Lua call plus a closure allocation per argument at their saturated call sites.

What changed

  • The IR gains the definition-side counterpart of feat: represent application as an n-ary AppN call node (#179) #199's AppN: Abs becomes the singleton pattern synonym of an n-ary AbsN, so every unary rewrite rule compiles unchanged. The Lua backend emits one n-ary function for a multi-parameter AbsN, dropping the trailing run of unused parameters and eliding the trailing run of Prim.undefined arguments on the call side (a non-trailing one lowers to nil).
  • A new pass (Language.PureScript.Backend.IR.Uncurry) splits every binding with manifest arity n ≥ 2 and at least one saturated call site into a worker f$w = AbsN [p₁…pₙ] body plus a curried wrapper f = λx₁.…λxₙ. f$w(x₁, …, xₙ), and rewrites the saturated sites to direct worker calls: add(x)(y) becomes add$w(x, y). Partial applications and higher-order uses keep going through the wrapper, so a mixed-use function wins at its saturated sites instead of being disqualified. Unreferenced wrappers are removed by the new post-uncurry optimize+dce fixpoint. Top-level candidates are collected module-wide; local (Let-bound) candidates per top-level site, since GUC is per-site. Recursive functions split too: mutual recursion becomes direct worker-to-worker calls, and a local loop like go acc n typically loses its wrapper entirely, which is the shape Lower self-recursive tail calls of uncurried workers to while loops (loopification) #181 needs.
  • One deliberate deviation from the issue design: the wrapper is not marked inline-always. The call-site inliner that plan relies on is Budgeted call-site inlining of dictionary methods to collapse non-Effect/ST monadic chains #180, which does not exist yet; today's inliner substitutes at every occurrence, which would paste lambda literals into argument positions (a closure allocation per evaluation at map f xs sites) and drop the shared top-level wrapper. The pass rewriting its own saturated sites gets the same result on exactly the sites that benefit. A second addition to the plan is the saturated-site precondition: without it a split with no sites would leave a permanent wrapper→worker indirection, because the use-once revert path does not reach recursive-group members or local sibling references.
  • betaReduce handles the exact-arity n-ary redex (subsuming betaReduceUnusedParams), which is what collapses a single-use worker pasted into its one call site. It declines on repeated parameter names, where left-to-right substitution would steal occurrences from the parameter Lua actually binds.
  • FlattenDeepBinds recognises the n-ary shapes: a chain step through w(action, \x -> rest) and spine depth through any argument of an n-ary call. Without this, Golden.LongCallbackChain and Golden.LongApplyChain would fail the Lua 5.1 nesting check the moment their sites get uncurried; the step rebuild preserves the call shape, since w(a, k) and w(a)(k) denote different programs.
  • The WellApplied invariant generalizes to exact-arity application of literal lambdas plus the trailing-ParamUnused condition, and now rides along with GUC at every pass boundary, so the checked runner verifies it on every golden. DCE blanks only a dead suffix of an AbsN parameter list to keep the trailing-run condition.
  • Minted names ($w, $p<i>, $u<i>) are deterministic and bypass the shared supply, so downstream $kont/$tmp numbering stays put.

Out of scope

Lifting runFnN/mkFnN stays with #198 (it needs the #178 lifter; AbsN unblocks the mkFnN half). Effect-shaped functions mostly stay curried for now: thunk-forcing sites only appear after magicDo, which runs later, so the pass never sees them saturated. The same applies to the $kont helpers minted by the flattening pass. Both gaps are tracked in #200 (a late, post-flattening run of the same pass). Constructors, which the codegen still emits as curried closure chains, are untouched and tracked in #201.

Verification

  • cabal test all: 599 examples, 0 failures. The pass has 11 unit tests (split, partial, over-application, recursive groups, local Lets, per-site scoping, @inline never, unused parameters, annotations) and 3 property tests over generated modules (the output lints clean for GUC and WellApplied, a second run is a no-op, and WasRewritten is signalled exactly when the module changed). Four more tests run the whole checked pipeline to fix the split's interaction with inlining and DCE: the wrapper is dropped once every site calls the worker, a used-once function leaves no residue, an exported function whose only saturated consumer is dead comes out unsplit, and an @inline always binding is claimed by the inliner before the split. Alongside are the n-ary beta and flattening cases and the mirrored AbsN tests for alphaEq, countFreeRefs, the linter, and codegen. The new tests were run red first: 9/11 pass unit tests fail against an identity pass, 5 of the beta/flattening tests fail against the pre-change implementations, and the property and pipeline tests fail against a pass mutated to mint duplicate wrapper binders and to over-report Rewritten.
  • The new Golden.Uncurry golden runs saturated, partially applied, higher-order, mutually recursive, locally recursive, over-applied and unused-parameter shapes through lua against a hand-written eval oracle. Every pre-existing eval/golden.txt is byte-identical: the semantic net for the whole transformation held.
  • Bench counters: Bench.BindChain drops one per-call FNEW, one trace-abort site and one blacklist entry with an unchanged workload result. A new Bench.CurriedStep macro pins the canonical target, a hot saturated two-argument loop. Compiling it with the pre-split pslua (main at d9d2d58) and with this branch, at n=100000 with an identical result=5000050000, the split removes the per-iteration closure the recursive go allocated: function-body FNEW drops from 6 to 5, distinct trace-abort sites from 4 to 3, and blacklisted bytecodes from 5 to 4. Median wall-clock falls from 22.0ms to 18.1ms on PUC Lua 5.1, from 22.6ms to 17.3ms on LuaJIT, and from 22.6ms to 17.9ms on LuaJIT with the JIT off. The three trace aborts that remain are the curried FFI closures (eqInt, intAdd, intSub), which is Lift the *.Uncurried wrappers to direct AppN calls #198's half.
  • Structural golden churn is large but mechanical: every golden.ir moves with the AbsN Show shape, and 17 golden.lua files change where splits fire.

Closes #24.

Unisay added 6 commits July 7, 2026 18:00
The definition-side counterpart of AppN (#199): Abs becomes the
singleton pattern synonym of AbsN, so every unary rule compiles
unchanged. Total traversals (getAnn, setAnn, subexpressions,
countFreeRefs, alphaEq, freshenBinders), the scope walkers (Uniquify,
Linker.qualifyTopRefs, DCE, Linter, Query.collectBoundNames) and the
Lua codegen handle the n-ary node; codegen drops the trailing run of
unused parameters and elides the trailing run of Prim.undefined
arguments (a non-trailing one lowers to nil).

The WellApplied invariant generalizes to the two well-formedness
conditions of Note [n-ary abstraction] — exact-arity application of
literal lambdas, ParamUnused only as a trailing run — and now rides
along with GUC at every pass boundary. DCE blanks only a dead suffix
of an AbsN parameter list, preserving the trailing-run condition.

Nothing produces a multi-parameter AbsN yet, so generated Lua is
unchanged: golden.ir churns mechanically with the AbsN Show shape,
golden.lua and eval oracles stay byte-identical.
betaReduce generalizes from the unary redex to
(λx₁…xₙ. M) N₁…Nₙ at exact arity — any other argument count on a
literal lambda head is ill-formed under WellApplied. Each
parameter/argument pair reduces by the shared inlining guard
(substitute when trivial or used at most once, let-bind otherwise),
and an argument at a ParamUnused position is dropped with its
evaluation, which subsumes betaReduceUnusedParams. On non-GUC input a
lambda with repeated parameter names is left alone: substituting left
to right would steal the occurrences belonging to the last same-named
parameter, the one Lua binds.

Nothing builds multi-parameter redexes yet, so goldens are untouched;
the rule is the reduction half the uncurrying pass relies on.
FlattenDeepBinds generalizes both strategies to the n-ary AppN shapes
an uncurrying pass emits. Strategy A recognises a chain step through
either call shape — the curried (f action) (\x -> rest) and the n-ary
f(action, \x -> rest) of an uncurried worker — and rebuilds each step
through a shape-preserving closure, since w(a, k) and w(a)(k) denote
different programs (Note [n-ary application]). Strategy B's spineDepth,
decompose and Frame follow the deepest operand of a call across the
callee and every argument position, so a deep spine running through
n-ary calls is measured and sealed into $tmp segments the same way as
a curried one.

Without this, a saturated two-argument CPS chain (Golden.LongCallbackChain)
or a deep applicative spine (Golden.LongApplyChain) would stop being
recognised the moment uncurrying rewrites their sites to n-ary calls,
and the chunk would fail the Lua 5.1 parser-nesting check. The new spec
cases fail against the previous unary-only implementation and pin both
the recognition and the call-shape preservation.
Every binding — top-level or Let-bound — whose right-hand side is a
manifest chain of n ≥ 2 unary lambdas, and which has at least one
saturated call site, splits into an n-ary worker f$w holding the
original body and a curried wrapper f delegating to it. Saturated
sites (nested unary spines of exactly n arguments) become direct
worker calls; partial applications and higher-order uses keep going
through the wrapper, so semantics is unchanged by construction — the
worker body still runs only on full saturation.

The saturated-site precondition matters: without it a split with no
sites would leave a permanent wrapper→worker indirection, because the
use-once revert path does not reach recursive-group members or a local
worker referenced from its sibling wrapper. Top-level candidates are
collected module-wide (references are Imported after qualifyTopRefs);
local candidates are collected, counted and rewritten per top-level
site, since GUC is only per-site. Minted names ($w, $p<i>, $u<i>)
are deterministic and deliberately bypass the shared supply so
downstream $kont/$tmp numbering stays put.

The pass runs between the post-merge and a new post-uncurry
optimize+dce fixpoint (which drops unreferenced wrappers and reduces
worker pastes), under the GUC + WellApplied contracts.

The Golden.Uncurry oracle runs saturated, partial, higher-order,
mutually recursive, locally recursive, over-applied and unused-param
shapes through lua; every pre-existing eval oracle is byte-identical.
Structural goldens: mutual recursion now calls worker-to-worker, local
loops lose their wrapper entirely (DCE), LongCallbackChain's CPS chain
is rebuilt from n-ary steps within the nesting cap.
The BindChain counters record the improvement (one fewer per-call FNEW,
one fewer trace-abort site and blacklist entry, result unchanged); Fib
and ArrayFoldl are untouched, as their currying lives in a unary
function and in FFI wrappers respectively. The new CurriedStep macro
pins the canonical uncurrying target — a hot two-argument loop, always
saturated — whose remaining trace aborts all come from curried FFI
(intAdd/intSub), the territory of the *.Uncurried wrapper lifting.

Also documents the pass in the pipeline overview and updates
Note [n-ary application]: the uncurrying split is the pass that
introduces multi-argument calls today.
@Unisay Unisay self-assigned this Jul 7, 2026
@Unisay Unisay requested a review from Copilot July 7, 2026 17:49
@Unisay Unisay marked this pull request as ready for review July 7, 2026 17:49

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR implements issue #24’s worker/wrapper uncurrying transformation for the IR→Lua pipeline, aimed at removing per-argument Lua call + closure-allocation overhead at saturated call sites by rewriting them to direct n-ary worker calls.

Changes:

  • Introduces definition-side n-ary lambdas (AbsN, with Abs as a singleton pattern synonym) and threads this representation through core IR utilities (alphaEq, free-ref counting, uniquify, DCE, linter invariants).
  • Adds an IR.Uncurry optimizer pass that performs the worker/wrapper split and rewrites saturated call sites; follows it with a post-uncurry optimize+DCE fixpoint to eliminate dead wrappers and reduce resulting redexes.
  • Updates Lua codegen + flattening logic to preserve semantics for n-ary calls/abstractions (including special handling of trailing Prim.undefined and unused-parameter suffixes), and adds new golden/bench coverage.

Reviewed changes

Copilot reviewed 85 out of 98 changed files in this pull request and generated no comments.

Show a summary per file
File Description
lib/Language/PureScript/Backend/IR/Types.hs Introduces AbsN/Abs singleton synonym and updates IR utilities (alphaEq, countFreeRefs, binder freshening) for n-ary abstraction.
lib/Language/PureScript/Backend/IR/Optimizer.hs Wires the new uncurrying pass into the pipeline, adds post-uncurry fixpoint, and extends beta reduction to exact-arity n-ary redexes.
lib/Language/PureScript/Backend/Lua.hs Emits n-ary Lua functions for AbsN, drops trailing unused parameters, and updates AppN lowering to elide trailing Prim.undefined and lower non-trailing ones to nil.
lib/Language/PureScript/Backend/IR/FlattenDeepBinds.hs Recognizes and preserves n-ary call shapes for chain steps/spines so flattening remains correct post-uncurry.
lib/Language/PureScript/Backend/IR/DCE.hs Blanks only a dead suffix of AbsN parameters to preserve the “trailing unused” invariant.
lib/Language/PureScript/Backend/IR/Linter.hs Generalizes WellApplied to exact-arity lambda application + trailing-unused-params discipline; updates violations accordingly.
lib/Language/PureScript/Backend/IR/Uniquify.hs / IR/Query.hs / IR/Linker.hs / IR/Pass.hs Updates binder scoping/collection/qualification and pass invariant documentation for AbsN.
lib/Language/PureScript/Backend/IR/Uncurry.hs Adds the worker/wrapper split pass (module introduced + cabal wiring).
test/Language/PureScript/Backend/IR/* and test/Language/PureScript/Backend/Lua/Spec.hs Adds/updates unit + property + pipeline tests for AbsN, n-ary beta, flattening behavior, and Lua emission rules.
test/ps/src/Golden/Uncurry/Test.purs + test/ps/output/Golden.Uncurry.Test/* Adds a new end-to-end golden (including eval oracle) covering saturated/partial/HO/recursive/unused-param shapes.
test/ps/src/Bench/CurriedStep.purs + bench/*curried_step* Adds a new macrobench and trace/fnew goldens targeting the canonical “hot saturated 2-arg loop” case.
test/ps/output/**/golden.{ir,lua} Mechanical golden churn reflecting AbsN/uncurry/codegen shape changes.
CLAUDE.md / pslua.cabal / test/Main.hs Documentation and build/test-suite wiring for the new pass and specs.

@Unisay Unisay merged commit 374b41c into main Jul 7, 2026
3 checks passed
@Unisay Unisay deleted the issue-24/uncurry-worker-wrapper branch July 7, 2026 19:03
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.

Program transformation to uncurry functions for which all applications are fully saturated.

2 participants