Extend loopification to mutual recursion and let-bound join points - #300

Problem

Since #181 the code generator turns a self-recursive tail call into a while true do loop: instead of calling itself, the function reassigns its parameters and repeats. Two recursion shapes were left out and still compile to a real Lua call per step. Closes #234.

The first is mutual recursion — functions whose bodies end by calling each other. This PureScript pair

ticktock n = tick 0 n
  where
  tick acc k = if k == 0 then acc else tock (acc + 1) (k - 1)
  tock acc k = if k == 0 then acc else tick (acc + 3) (k - 1)

compiled to two closures making a real cross-call per transition ($w, rendered _S_w, marks an uncurried worker — the n-ary function the worker/wrapper split produces):

local Golden_MutualLoopification_Test_ticktock = function(n)
  local tock_S_w
  local tick_S_w
  tock_S_w = function(acc, k)
    if k == 0 then return acc else return tick_S_w(acc + 3, k - 1) end
  end
  tick_S_w = function(acc0, k0)
    if k0 == 0 then return acc0 else return tock_S_w(acc0 + 1, k0 - 1) end
  end
  return tick_S_w(0, n)
end

Each transition pays Lua's call/return machinery, and LuaJIT's trace compiler never finds a loop to anchor a trace on, so the hot path stays interpreted.

The second is join points — a local helper whose every use is a tail call from the enclosing body. Even when the helper's own recursion is already loopified, it survives as a function object plus an entry call:

local Golden_JoinPoints_Test_sumTriangles = function(m)
  local go_S_w
  go_S_w = function(acc, n)
    while true do
      if n == 0 then return acc else acc, n = acc + n * (n + 1), n - 1 end
    end
  end
  return go_S_w(0, m)
end

The transforms

Both live in the code generator's loopification module, beside #181's self-recursion rewrite; the IR keeps needing no loop constructor.

Mutual recursion → one dispatcher. Group members that tail-call each other — precisely, a cycle in the "whose body ends by calling whom" graph — merge into one function that loops over a numeric branch selector plus shared argument slots, one slot per parameter of the widest member. Every tail call between members becomes a single simultaneous assignment of the selector and the slots. The original bindings survive as one-line entry wrappers, which is what makes the cycle test purely a profitability question: a call from outside, a use as a value, or the curried wrappers that uncurrying leaves in the same recursive group all keep referring to the original names with no use analysis needed. ticktock now compiles to:

local Golden_MutualLoopification_Test_ticktock = function(n)
  local tock_S_w_S_loop
  local tock_S_w
  local acc0, k0
  tock_S_w_S_loop = function(_S_sel4, _S_a5, _S_a6)
    while true do
      if _S_sel4 == 1 then
        local acc, k = _S_a5, _S_a6
        if k == 0 then
          return acc
        else
          _S_sel4, _S_a5, _S_a6 = 2, acc + 3, k - 1
        end
      else
        local acc0, k0 = _S_a5, _S_a6
        if k0 == 0 then
          return acc0
        else
          _S_sel4, _S_a5, _S_a6 = 1, acc0 + 1, k0 - 1
        end
      end
    end
  end
  tock_S_w = function(acc, k) return tock_S_w_S_loop(1, acc, k) end
  acc0, k0 = 0, n
  return tock_S_w_S_loop(2, acc0, k0)
end

($sel/$a/$loop are freshly minted names user code cannot collide with; the argument-balancing rules from #181 carry over, so a transition into a narrower member pads its unused slots with nil.) Each branch rebinds its member's parameters from the slots as loop-body locals, re-created every iteration, so a closure created inside a branch captures that iteration's values exactly as a fresh call activation would — the capture veto #181 needs does not apply here. The entry call at the bottom has no wrapper hop left because the join-point transform below fused it.

Join points → fall-through. A helper only ever entered by tail calls from the block's exit positions loses its function shell: its parameters become plain locals, each entry call becomes an assignment that falls through into the helper's body, now placed at the end of the block. sumTriangles becomes entry assignment plus bare loop:

local Golden_JoinPoints_Test_sumTriangles = function(m)
  local acc, n
  acc, n = 0, m
  while true do
    if n == 0 then return acc else acc, n = acc + n * (n + 1), n - 1 end
  end
end

Non-recursive helpers fuse the same way — a continuation shared by two branches, classify n = let finish r = (r * 10 + r) * 2 - r in if n > 0 then finish (n + 1) else finish (0 - n), loses both calls:

local Golden_JoinPoints_Test_classify = function(n)
  local r
  if not(n < 0) and n ~= 0 then r = n + 1 else r = 0 - n end
  return (r * 10 + r) * 2 - r
end

Fusion iterates until nothing fuses: a helper entered only from another fused helper's body fuses on the next round, so chains flatten completely.

Safety conditions

The join transform commits only when it can account for every use of the helper: after rewriting the entry calls, no read of the helper's name may remain anywhere in the block — a reference from inside a nested closure, an argument-position call, or a call from a non-exit position each leaves the function untouched. Both transforms also require that control provably cannot fall off the end of a body they relocate (the value-less native-loop blocks of #233 can), since falling through a dispatcher branch would iterate instead of returning nil. Finally, the join transform refuses to bury a block-final return f(x) mid-block, because the enclosing binding's own loopification (or a mutual cycle it belongs to) may still claim that tail call.

Verification

Red-first: the first commit pins the pre-change shapes of two new runnable golden modules with hand-written output oracles, and the implementation commit flips their generated-Lua goldens while every execution oracle stays byte-identical — the diffs above are quoted from exactly that flip. Coverage: the mutual pair, differing arities, a three-state machine, mixed self+sibling calls, a non-tail pair that must stay calls, dispatch composing with join fusion, two entries into one loop, an effect-land continuation, and an escaping helper that must keep its shell. One pinned gap doubles as a canary: a Boolean-returning pair like isEven/isOdd folds to and/or chains in the IR — return n == 0 or Golden_MutualLoopification_Test_isOdd(n - 1) — which hides the sibling call inside an operand, so the pair stays call-shaped; if a later pass learns to distribute that fold back into branches, the golden flips and proves it. Existing golden churn (Golden.Uncurry, Golden.Loopification, Golden.RecursiveBindings, Golden.SpecConstr) was reviewed by hand; IR-level goldens are untouched (the change is codegen-only), 15 new unit specs pin the shapes and the vetoes, 1090 examples green.

Measurements

New mutual_step macrobenchmark — the mutual twin of curried_step: the same hot two-argument accumulator loop, split across two workers tail-calling each other. Medians at n=2e6 against the pre-change compiler:

Under LuaJIT the dispatcher trace-compiles (JLOOP, zero blacklists) to parity with the hand-written loop, ~21x over the cross-calling shape, which never traced. Under PUC the selector test and the per-branch parameter rebinds cost ~19% on this worst-case pure-transition loop against tail calls PUC already runs cheaply; substituting slot names directly for parameters instead of rebinding (at the price of a capture veto) would remove that cost and is the natural follow-up. Join-point fusion also registers on the existing benchmarks: curried_step's per-iteration closure allocations (FNEW) drop to zero and state_step loses an allocation site and a trace abort, with all compiled loop traces retained.