Uncurry functions via a worker/wrapper split#202
Merged
Conversation
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.
Contributor
There was a problem hiding this comment.
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, withAbsas a singleton pattern synonym) and threads this representation through core IR utilities (alphaEq, free-ref counting, uniquify, DCE, linter invariants). - Adds an
IR.Uncurryoptimizer 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.undefinedand 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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
AppN:Absbecomes the singleton pattern synonym of an n-aryAbsN, so every unary rewrite rule compiles unchanged. The Lua backend emits one n-ary function for a multi-parameterAbsN, dropping the trailing run of unused parameters and eliding the trailing run ofPrim.undefinedarguments on the call side (a non-trailing one lowers tonil).Language.PureScript.Backend.IR.Uncurry) splits every binding with manifest arity n ≥ 2 and at least one saturated call site into a workerf$w = AbsN [p₁…pₙ] bodyplus a curried wrapperf = λx₁.…λxₙ. f$w(x₁, …, xₙ), and rewrites the saturated sites to direct worker calls:add(x)(y)becomesadd$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 likego acc ntypically loses its wrapper entirely, which is the shape Lower self-recursive tail calls of uncurried workers to while loops (loopification) #181 needs.map f xssites) 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.betaReducehandles the exact-arity n-ary redex (subsumingbetaReduceUnusedParams), 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.FlattenDeepBindsrecognises the n-ary shapes: a chain step throughw(action, \x -> rest)and spine depth through any argument of an n-ary call. Without this,Golden.LongCallbackChainandGolden.LongApplyChainwould fail the Lua 5.1 nesting check the moment their sites get uncurried; the step rebuild preserves the call shape, sincew(a, k)andw(a)(k)denote different programs.WellAppliedinvariant generalizes to exact-arity application of literal lambdas plus the trailing-ParamUnusedcondition, 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 anAbsNparameter list to keep the trailing-run condition.$w,$p<i>,$u<i>) are deterministic and bypass the shared supply, so downstream$kont/$tmpnumbering stays put.Out of scope
Lifting
runFnN/mkFnNstays with #198 (it needs the #178 lifter;AbsNunblocks themkFnNhalf). 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$konthelpers 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, localLets, 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, andWasRewrittenis 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 alwaysbinding is claimed by the inliner before the split. Alongside are the n-ary beta and flattening cases and the mirroredAbsNtests 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-reportRewritten.Golden.Uncurrygolden runs saturated, partially applied, higher-order, mutually recursive, locally recursive, over-applied and unused-parameter shapes throughluaagainst a hand-written eval oracle. Every pre-existingeval/golden.txtis byte-identical: the semantic net for the whole transformation held.Bench.BindChaindrops one per-callFNEW, one trace-abort site and one blacklist entry with an unchanged workload result. A newBench.CurriedStepmacro pins the canonical target, a hot saturated two-argument loop. Compiling it with the pre-splitpslua(main atd9d2d58) and with this branch, atn=100000with an identicalresult=5000050000, the split removes the per-iteration closure the recursivegoallocated: function-bodyFNEWdrops 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.golden.irmoves with theAbsNShow shape, and 17golden.luafiles change where splits fire.Closes #24.