feat(ir): absorb the magic-do thunk into early-split effect workers - #358

Closes #265.

The problem

pslua compiles PureScript's Effect and ST the way every PureScript backend does: an action is a nullary thunk, and running it is calling that thunk. A do block is a chain of bind applications with lexically nested continuations, and the magic-do pass (Language.PureScript.Backend.IR.MagicDo) flattens such a chain into one thunk holding a flat statement sequence — function() local x = m1(); …; return last() end.

An earlier pass in the pipeline, the uncurrying worker/wrapper split (Language.PureScript.Backend.IR.Uncurry), turns a curried binding into a worker — one n-ary Lua function holding the original body — plus a wrapper that keeps the original name and delegates to it, and rewrites every saturated call site to a direct worker call. That split runs before magic-do. An effect function of two or more real arguments is therefore already saturated at its real arity when the split measures it, so the split fires there and magic-do afterwards rewrites the worker's body into the thunk. Every fully applied statement site ends up as two Lua calls with a closure allocation in between:

local Golden_EffectWorkerThunk_Test_report_S_w = function(tag, n)
  return function()
    local _ = Effect_Console_log(tag)()
    local _ = Effect_Console_log(Data_Show_showIntImpl(n))()
    return Effect_Console_log("-")()
  end
end
local _ = Golden_EffectWorkerThunk_Test_report_S_w("a", 1)()

The worker call allocates the thunk over tag/n; the trailing () immediately forces it. Nothing else ever sees that closure.

The second, late run of the uncurry pass (added for #200) cannot repair this. It splits manifest chains of unary lambdas, and here the thunk sits inside a worker that is already n-ary, so there is no chain left to split. That run only absorbs the thunk parameter of the actions the early run left alone — the unary ones, below its arity floor of two.

The change

A new pass, Language.PureScript.Backend.IR.AbsorbEffectThunk, widens such a worker in place: the thunk's parameter joins the worker's parameter list and the thunk's body becomes the worker's body.

w = AbsN [p₁…pₙ] (λ_. body)   ↦   w = AbsN [p₁…pₙ, _] body

Each forced site loses its outer call — w(a₁…aₙ)(run) becomes w(a₁…aₙ, run). That is still recognised as an effect run by its trailing marker (isEffectRun, so dead-code elimination keeps a statement whose binder is unreferenced) and it is still one Lua call, because the backend erases the marker from an n-ary argument list. The site above becomes:

local Golden_EffectWorkerThunk_Test_report_S_w = function(tag, n)
  local _ = Effect_Console_log(tag)()
  local _ = Effect_Console_log(Data_Show_showIntImpl(n))()
  return Effect_Console_log("-")()
end
local _ = Golden_EffectWorkerThunk_Test_report_S_w("a", 1)

Each wrapper — the curried delegate the split left under the original name — grows one parameter, which the delegate call passes on, so a partial application still evaluates to a closure (the action) with the new innermost lambda playing the thunk's role:

M.Golden_EffectWorkerThunk_Test_tally = function(tally_S_p1)
  return function(tally_S_p2)
    return function(tally_S_p2_S_t)
      return Golden_EffectWorkerThunk_Test_tally_S_w(tally_S_p1, tally_S_p2, tally_S_p2_S_t)
    end
  end
end

Both rewrites are local, and the pass keys on the shape rather than on the $w naming scheme, so any producer of an n-ary worker whose body is a thunk is covered — including the $kont helpers the deep-bind flattening mints and the copies call-pattern specialization makes.

The precondition

The extension fires only for a worker with at least one forced site and whose every reference is either a forced site or a wrapper's delegate call. Any other reference vetoes the whole binding, because the Lua backend drops a trailing unused parameter run: the widened worker still compiles to function(p₁…pₙ), so a call left under-applied by the wider arity is a saturated Lua call and would run the effect at construction time instead of returning the action. The shape that most often vetoes a binding is an action bound to a name and run later:

local held_S_0 = Golden_EffectWorkerThunk_Test_deferred_S_w("d", 4)
local _ = Effect_Console_log("before held")()
local _ = held_S_0()
return held_S_0()

Golden.EffectWorkerThunk pins that case unchanged next to the three that do fire.

The pass runs last in the IR pipeline, after the late uncurry run and the dead-code pass that follows it. Nothing later moves a call, so the reference census the precondition needs is final; and the wrappers whose sites all went to their workers are already gone, so no doomed wrapper is grown.

A tail call falls out

Taking the run marker into the argument list also turns a recursive driver's self-call into a genuine tail call, which the native-loop lowering then turns into a Lua while. Bench.EffectStep's driver — a two-argument effect action, so exactly this pass's case — went from a forward-declared recursive closure

local Bench_EffectStep_go_S_w
Bench_EffectStep_go_S_w = function(i, ref)
  return function()
    local _ = Bench_EffectStep_step_S_w(ref)
    if i >= 1 and i ~= 1 then
      return Bench_EffectStep_go_S_w(i - 1, ref)()
    else
      return Control_Monad_ST_Internal_read(ref)()
    end
  end
end

to a flat loop:

local Bench_EffectStep_go_S_w = function(i, ref)
  while true do
    local _ = Bench_EffectStep_step_S_w(ref)
    if i >= 1 and i ~= 1 then
      i, ref = i - 1, ref
    else
      return Control_Monad_ST_Internal_read(ref)()
    end
  end
end

Measurements

Bench.EffectStep2 is new here: a hot ST loop whose per-iteration step is a two-argument effect action, the sibling of the existing unary Bench.EffectStep. Medians from bench/tools/run_macro.lua, n=100000, identical result=1500000 on both sides:

The other ten macrobenchmarks are unchanged within noise, and every one returns the same value as before.

The committed LuaJIT counter oracles record the JIT-level reason. effect_step2 loses the two closure-allocation trace aborts and its driver ends compiled rather than interpreted:

 aborts (distinct site -- reason):
   Bench.EffectStep2.lua:11 -- NYI: bytecode FNEW
-  Bench.EffectStep2.lua:29 -- NYI: bytecode FNEW
-  Bench.EffectStep2.lua:40 -- NYI: bytecode FNEW
-  Bench.EffectStep2.lua:56 -- NYI: bytecode FNEW
+  Bench.EffectStep2.lua:59 -- NYI: bytecode FNEW
 …
-  Bench.EffectStep2.lua:32 IFUNCF
-  Bench.EffectStep2.lua:33 IFUNCF
+  Bench.EffectStep2.lua:30 ILOOP
-counts: aborts=6 compiled=3 blacklisted=7
+counts: aborts=4 compiled=3 blacklisted=5

effect_step moves the same way (aborts=5 → 4, blacklisted=6 → 5, IFUNCF → ILOOP). The TNEW/TDUP table-allocation censuses are byte-identical everywhere: the pass touches closures, not tables.

Tests

Test-first, in three commits: the golden pinning the two-call shape, the benchmark with its oracles at that state, then the pass — so the feature commit carries the whole codegen delta as a reviewable diff.

• Golden.EffectWorkerThunk — four shapes in one module: an action whose wrapper is dead (the plain win), one also passed as a value so its wrapper survives and must grow, a let-bound local worker, and an action bound to a name and run later (vetoed). eval/golden.txt is the hand-verified runtime oracle and is unchanged by the pass in either direction.
• Golden.NativeLoopsGuard is the one pre-existing golden carrying this shape; its whileE$w loses the closure.
• IR AbsorbEffectThunk — six focused cases over the pass itself: worker extension, wrapper growth, local extension, and the three veto shapes (no forced site, an unforced saturated call, a bare reference). Against an identity-stubbed pass the three positive cases fail and the three vetoes pass, which is what makes the suite evidence rather than assertion; both structural goldens fail too, while the eval oracle stays green.
• Full suite green (1246 examples). The IR Optimizer Hedgehog group and the uncurry/pass/DCE groups were each re-run twelve times with fresh seeds. The pipeline's per-pass linter checks the new pass's output for well-scopedness, globally unique binders, and the n-ary application/abstraction invariants on every golden.

Sibling audit

Three passes genuinely mint multi-parameter AbsN nodes: the uncurry split, the CPR result split, and call-pattern specialization. A CPR worker's every return path builds a constructor lowered to Lua multiple values, never a lambda, so it cannot carry this shape; specialization copies a candidate's body, so a specialized copy of a thunk-returning worker is covered by the same structural recogniser as the original. The other AbsN sites (linker, DCE, CSE, renumber, uniquify) rebuild an existing node and cannot introduce one. No unhandled instance of the shape remains — which is why the recogniser keys on structure rather than on the $w suffix.

The audit did surface the dual shape, out of scope here: a literal thunk immediately forced, (function() … end)(), 141 occurrences across the golden corpus. Codegen already collapses it when the thunk's body is a single application; the remaining cases have a Let body and are deliberately left thunked, because splicing their locals into the parent function threatens Lua 5.1's ~200-local ceiling that magic-do's chunking exists to respect. A budgeted version — splice when the parent's local count allows — needs a new local-count analysis, so it is follow-up work rather than part of this diff.