Non-tail match statements

[arthurpaulino]

Goal

Enable let x = match y { 0 => return (0, 0), 1 => z, }; use(x) — match
expressions in non-tail position, including branches with early returns.

Key constraint: avoid selector blowup

A pure Term-level hoist (pushing continuations into branches) turns independent
matches into a nested match tree: two sequential 3-armed matches produce 9 leaf
selectors instead of 6. With k sequential n-armed matches this is n^k instead
of k·n. Unacceptable for a circuit DSL.

The fix: a bytecode-level non-tail match that allocates its own selectors,
merges branch results via conditional select, and continues with a shared
continuation. Each match contributes its selectors independently.

Circuit semantics of non-tail match

For let x = match tag { 0 => a, 1 => b }; use(x):

1. Evaluate tag
2. Allocate 2 selectors (sel₀, sel₁) for the match
3. Evaluate a under sel₀, evaluate b under sel₁
4. Merge: x = sel₀ · a + sel₁ · b (conditional select)
5. Continue with use(x) in the outer selector context

Each branch's constraints are gated by its selector. The merge produces values
usable by the continuation without duplication. Operations with global effects
(store/load, I/O, assertions) are already selector-gated.

Early returns: if a branch is return v, it exits the function directly (gets
its own leaf selector via Ctrl::Return). Non-returning branches yield values
for the merge. The continuation only runs when a yielding branch is active.

Implementation plan

Step 1: Bytecode (Ix/Aiur/Bytecode.lean + src/aiur/bytecode.rs)

Add two new Ctrl variants:

-- Lean
| yield : SelIdx → Array ValIdx → Ctrl
| matchContinue : ValIdx → Array (G × Block) → Option Block
    → (outputSize : Nat) → Block → Ctrl

// Rust
Yield(SelIdx, Vec<ValIdx>),
MatchContinue(ValIdx, FxIndexMap<G, Block>, Option<Box<Block>>, usize, Box<Block>),

yield sel vals: terminates a non-returning branch of a matchContinue. Like
return it assigns a selector, but instead of exiting the function it provides
values for the merge.

matchContinue tag cases default outputSize continuation: like match but
non-returning branches end in yield instead of return. After the match, the
yielded values are merged (conditional select) and the continuation block
executes with the merged values available as the next outputSize value indices.

Also update Block.for_each_op and Toplevel.needs_circuit (Rust) to recurse
into matchContinue branches and continuation.

Step 2: Execution (src/aiur/execute.rs)

Ctrl::Yield during execution is straightforward — it means "this branch
produced its values, now the continuation runs." The execution model already
dispatches exactly one branch. For the taken branch:

• If it ends in Yield: the branch values are the match result. Continue
  executing the continuation block.
• If it ends in Return: exit the function (same as today).

Implementation: when processing MatchContinue, dispatch to the matched branch
(same logic as Match). Then:

• If the branch's ctrl is Yield: push the continuation block's entries onto
  the execution stack, with the yielded values in the map.
• If the branch's ctrl is Return: handle as a normal function return.

The push_block_exec_entries! macro already handles pushing a block. For
MatchContinue, after the branch block finishes (hitting a Yield), push the
continuation block.

Step 3: Constraints (src/aiur/constraints.rs)

Ctrl::Yield constraint generation:

• Same as Return for selector assignment: slice.selectors[sel] = 1.
• But NO function return lookup (the function hasn't returned yet).
• Just record that this selector gates this branch's constraints.

Ctrl::MatchContinue constraint generation:

Ctrl::MatchContinue(var, cases, def, output_size, continuation) => {
  let (var, _) = state.map[*var].clone();
  let init = state.save();
  
  // Process branches exactly like Ctrl::Match:
  // - For each case: branch_sel * (var - value) = 0
  // - For default: branch_sel * ((var - value) * inverse - 1) = 0
  // - Collect branch constraints with branch_sel gating
  
  // After all branches, emit merge constraints:
  // For each output slot i (0..output_size):
  //   merged_i = Σⱼ (yield_sel_j · yield_val_j_i)
  // This is implemented as auxiliary values with constraints:
  //   sel * (merged_i - Σⱼ (yield_sel_j · yield_val_j_i)) = 0
  
  // Push merged values to state.map
  // Collect constraints from continuation block
  continuation.collect_constraints(sel, state);
}

The branch_sel gating works exactly like Ctrl::Match. The new part is the
merge: for each yielding branch j and output slot i, emit
yield_sel_j * yield_val_j_i and sum them. The sum is the merged value for
slot i. The continuation's constraints use these merged values.

Key detail: the state.save()/restore() pattern used for branches means each
branch starts from the same auxiliary column index. The merge happens after
restoring to the post-branch state, using fresh auxiliaries for the merged
values.

Step 4: Trace (src/aiur/trace.rs)

Ctrl::Yield trace generation:

• Set slice.selectors[sel] = G::ONE (same as Return).
• Do NOT emit a function return lookup (unlike Return).

Ctrl::MatchContinue trace generation:

Ctrl::MatchContinue(var, cases, def, output_size, continuation) => {
  let val = map[*var].0;
  let branch = /* dispatch same as Ctrl::Match, with default inverse witnesses */;
  
  // Populate branch row (sets selector, computes auxiliaries)
  branch.populate_row(map, index, slice, context, io_buffer);
  
  // After branch: the branch ended with Yield, so map now has the yielded values.
  // The merge in the trace is trivial: only one branch ran, so the merged values
  // ARE the yielded values (the conditional select is identity for the active branch).
  // Push output_size auxiliaries for the merged values:
  for i in 0..*output_size {
    let merged = /* yielded value i from the active branch */;
    slice.push_auxiliary(index, merged);
  }
  
  // Continue with continuation block
  continuation.populate_row(map, index, slice, context, io_buffer);
}

The trace is simpler than constraints because only one branch executes. The
merge auxiliary columns take the actual yielded values (no conditional select
needed at trace time — the constraint system ensures the algebra is correct).

Step 5: Lowering (Ix/Aiur/Compiler/Lower.lean)

In toIndex: replace the | .match .. => throw "Non-tail ..." error with:

1. Evaluate scrutinee via toIndex
2. Extract accumulated ops
3. For each branch:
   • If term.escapes: compile as a block ending in Ctrl.return (early exit)
   • Else: evaluate branch body via toIndex, create block ending in
     Ctrl.yield sel vals
4. Emit merge ops after the match:
   • For each output slot i: merged_i = Σⱼ (sel_j * val_j_i)
   • This is outputSize values, each a sum of products
5. Return the merged Array ValIdx

In TypedTerm.compile: handle .let pat val bod where val.inner is
.match:

1. Compile the match as above (yields merged values)
2. Bind merged values to pat in the bindings map
3. Compile bod as the continuation → produces a Block
4. Package as Ctrl.matchContinue tag cases default outputSize continuation

Step 6: Layout (Ix/Aiur/Compiler/Layout.lean)

Add cases in blockLayout:

• Ctrl.yield: bump selectors (+1), bump lookups (+1), same as return.
• Ctrl.matchContinue: process branch blocks (save/restore + maximals, same as
  match), then account for merge auxiliaries (outputSize extra), then
  process the continuation block.

Step 7: Dedup (Ix/Aiur/Compiler/Dedup.lean)

• skeletonBlock/skeletonCtrl: recurse into matchContinue branches and
  continuation; handle yield like return.
• collectCalleesBlock/collectCalleesCtrl: recurse into continuation.
• rewriteBlock/rewriteCtrl: rewrite FunIdx in continuation and branches.

Step 8: Lean FFI / Bytecode serialization

The Lean-to-Rust bridge must serialize the new Ctrl variants. Check how
Ctrl.match and Ctrl.return are currently serialized and add cases for
yield and matchContinue.

Step 9: Validation test (Tests/Aiur/Aiur.lean)

• Basic non-tail match: let x = match y { 0 => a, 1 => b }; use(x)
• Early return: let x = match y { 0 => return v, 1 => z }; use(x)
• Sequential non-tail matches (verifies selector count is additive not
  multiplicative)
• Nested non-tail matches
• Non-tail match with store/load in branches (verifies lookup gating)

Success criteria: lake test -- --ignored aiur passes.

[gabriel-barrett]

I believe the design is correct for bytecode, except that I would use Rc instead of Box, as otherwise you will have to copy the continue block for each yield, leading to an exponential blowup.

However, when it comes to selectors, I'm not quite sure. Here's the overall design I believe should work:

1. Assign one selector per Block, which should be constrained as to be a boolean. At this point, it doesn't mean the selector will be a distinguished column.
2. Whenever a Block contains a match statement as its ctrl field, we know that it should be selected if and only if one of the immediate children blocks (not including the continue block) is selected. We should not worry whether more than one children are selected, since match statements are designed in a way that they have mutually exclusive constraints
3. Whenever a Block is a continue of some MatchContinue, there should be a constraint that says that it should be selected if and only if one of the MatchContinue blocks has a yield that points to it (note that the yield might come from a descendant block). In particular, if there are no blocks, then the selector should be constrained to 0.

Now, regarding which selectors are columns, which are expressions. Perhaps the algorithm should initially pretend the selectors are all columns and collect the constraints. Then, you go through a simplify phase: whenever you see a constraint of form sel = A, you can immediately remove that constraint and simply identify sel with A itself, updating the constraint set. This should be done recursively so that you find the minimum number of "real" selectors. Perhaps there's a more intelligent algorithm I'm not aware of, we should investigate this

Lastly, the constraint "A is selected iff B or C or D is selected" can be described by a simple sum, since every selector should be a boolean. In the current scheme, without non-tail match statements, we only constrain the real selectors and the topmost selector (which is a sum of all real selectors) to be booleans. With non-tail matches, we have to be a little bit more careful. The safe option is to simply constrain all block selectors to be booleans. This should not impact performance in a significant way, as the constraints comprise only a small portion of the prover time