AILang/docs/specs/0068-branch-consume-split-drop.md

Branch-Consume-Split Drop Soundness — Design Spec

Date: 2026-06-02 Status: Draft — awaiting user spec review Authors: orchestrator + Claude

Goal

Close the last open leg (leg 3) of bug #63: an (own ...) heap parameter that is consumed on one branch of a match/if but stays live on a sibling branch is never dropped on the live branch, leaking one slab per call. Closing this lets the series milestone #61 close — its headline series_sma example currently leaks live=2 through this path.

The leak is not Series-specific and not match-specific. The root is general: the per-fn aggregate consume_count (the worst-case max over all branches, computed by uniqueness::merge_states) is the only consume signal codegen has, and every owned-param drop site gates on it. When a param is consumed on some branch, the aggregate is >= 1, so every drop site skips the param — including the branch where it stays live. if and match share this root exactly: both collapse per-branch consume via the same max merge, and if is a first-class MTerm::If (it does not desugar to match), so it needs the same fix in its own node. This cycle fixes the whole branch-consume-split class.

Architecture

The fix threads per-branch consume liveness from the check pass to codegen, so each branch's drop sites gate on that branch's consume count instead of the per-fn aggregate.

The single load-bearing constraint: the check pass stays the sole consume authority. This is the mir.3a direction ("codegen reads the consume count instead of re-running infer_module"). The per-branch consume is therefore the same snapshot the uniqueness walker already computes per branch and currently discards at the max merge — never a re-derivation in codegen and never a second walk with re-seeded context.

Three pieces:

1. Capture (check): uniqueness already walks each branch into a fresh self.binders before merging (if at uniqueness.rs:297-305 via the after_then snapshot; match at :305-321 via the per-arm self.binders before merge_states). Today those per-branch snapshots are folded to max and dropped. The pass instead retains them, emitting one per-branch consume map alongside the aggregate.

2. Carry (MIR): the per-branch map rides additively on the MIR branch nodes — MArm.consume per arm, MTerm::If.then_consume / .else_consume per branch. The aggregate MirDef.consume is unchanged (it still drives binder classification and the single-path fn-return drop). lower_to_mir attaches the per-branch maps as it lowers each arm/branch; because neither the uniqueness walk nor lower_to_mir carries a node id (Term::Match/If/Arm are purely structural — verified: no span/id field), the correspondence is structural traversal order (both walk arms left-to-right over the same post-mono AST). That lock-step ordering is the load-bearing invariant the planner encodes, with a test that the per-branch maps land on the matching arms.

3. Gate (codegen): the owned-param drop decision is split by the param's aggregate consume, with the per-branch map only ever adding drops for the leak class:

   • pre-tail-call Own-param dec (match_lower.rs:823): its gate source switches from the per-fn aggregate self.consume to the per-arm arm.consume (drop iff arm.consume[p] == 0). This is identical to the aggregate gate for legs 1/2 (a param consumed on no arm has arm.consume == 0 on every arm) and additionally drops a param that is live on this tail-call arm but consumed on a sibling. A tail-call arm terminates its block, so the fn-return dec never also fires for it.
   • a new fall-through Own-param drop, emitted at each non-terminated match-arm / if-branch close before the br to the join, firing only for the leak class: branch_consume[p] == 0 (live on this branch) AND self.consume[p] >= 1 (consumed on a sibling branch). It excludes the scrutinee and any param that is the branch's own tail value, and honours moved_slots.
   • the fn-return Own-param dec (lib.rs:1484) and the arm-close pattern-binder dec (match_lower.rs:910) are UNCHANGED. The fn-return dec keeps firing for self.consume[p] == 0 params (live on every path). The pattern-binder dec already reads a per-fn aggregate that, for an arm-local pattern binder, equals its per-arm value, so there is nothing to switch.

The drop invariant (double-free + use-after-free safety)

An owned heap param is dropped exactly once per control-flow path. The drop sites partition by aggregate consume, so no param is reachable by two of them:

• self.consume[p] == 0 (live on every path): dropped by the existing machinery — the fn-return dec, or the pre-tail-call dec on a tail-call arm. Unchanged from legs 1/2.
• self.consume[p] >= 1 (consumed on some sibling branch): dropped per-branch, only on the branches where it stays live (branch_consume[p] == 0), before that branch reaches the join. Never at the join (a merge point cannot tell the consumed path from the live path).

Two independent guarantees make this safe:

• Double-free safety — disjointness by aggregate. The new per-branch drop requires self.consume[p] >= 1; every existing drop site that could collide (fn-return) requires self.consume[p] == 0. The two predicates are mutually exclusive, so no param is dropped by both — no fn-return disable and no tail-position analysis are needed.
• Use-after-free safety — the consume checker forecloses the hazard. A per-branch drop of a live param would be unsafe only if that param were used after the branch construct. But the leak class requires self.consume[p] >= 1 (consumed on a sibling), and AILang's checker rejects any use of an owned value after it was consumed on any branch (error: [use-after-consume], pinned green by ailang-check/tests/workspace.rs::use_after_consume_on_own_param_is_reported and the match-then-double-consume case harden_ownership_heap_double_consume_still_errors at :701/:749). So an aggregate >= 1 param is provably never live past the construct; the per-branch drop is path-terminal regardless of whether the construct sits in the fn's tail position. The aggregate >= 1 guard therefore does double duty: it secures both disjointness and the no-after-use property.

Type eligibility — the per-branch drop fires only for heap-RC ADTs

The new fall-through drop is the first drop site that can reach an aggregate >= 1 owned param — and an aggregate >= 1 closure param is routinely a static closure-pair (a top-level fn ref like inc lowered to a .rodata constant with no rc-header), not a heap RC slab. The existing drop sites never hit this because they all require aggregate == 0, and a called closure has aggregate >= 1; their soundness for Type::Fn / static-Str params rests on a codegen-level invariant (move-tracking + the aggregate == 0 gate) that keeps static constants out of the bare ailang_rc_dec route. The new site would violate that invariant, so it carries an explicit type precondition: it drops a param only when field_drop_call resolves to a real per-type heap-ADT drop fn (drop_<owner>_<T>), i.e. not the bare ailang_rc_dec fallback that Type::Fn, static-Str, and Type::Var route to. This is sound (no underflow on a static constant) and leak-correct for the real cases: a static closure / static Str allocates nothing, so there is nothing to drop. (A heap-capturing closure param in the exact leak-class position would leak rather than be freed — but the existing machinery cannot free it either, so this is not a regression; soundness is preferred over completeness here.) Witnessed by the std_either_demo / std_list_demo / poly_rec_capture_demo e2e tests, which pass closure refs through a match that consumes them on one arm and leaves them live on another.

Why this is behaviour-preserving for legs 1/2

Legs 1 and 2 (already shipped) drop an owned param whose aggregate consume is 0 — consumed on no branch. The only existing site this cycle touches is the pre-tail-call dec (:823), whose gate source changes from self.consume to arm.consume. For an aggregate == 0 param the per-arm count is 0 on every arm (a max of 0 is 0), so arm.consume[p] == 0 yields the identical decision the aggregate gate gave — the leg-1/2 cases are unchanged. The new fall-through drop never fires for them (it requires aggregate >= 1), and the fn-return dec is untouched. Legs 1/2 are the aggregate == 0 special case of one model; their green pins stay green.

Concrete code shapes

User-facing programs (the evidence)

(1) match-arm split — existing RED (examples/match_arm_consume_split_no_leak_pin.ail, pin crates/ail/tests/match_arm_consume_split_no_leak_pin.rs, currently #[ignore]d). b is consumed on the On arm, live on the Off arm; run b Off runs Off, so b stays live. Today: allocs=2 frees=1 live=1. Target: live=0.

(fn run
  (type (fn-type (params (own (con Box)) (own (con Flag))) (ret (own (con Unit))) (effects IO)))
  (params b f)
  (body
    (match f
      (case (pat-ctor Off) (do io/print_str ""))
      (case (pat-ctor On) (app sink b)))))

(2) if-branch split — new RED this cycle (examples/if_branch_consume_split_no_leak_pin.ail, pin crates/ail/tests/if_branch_consume_split_no_leak_pin.rs). b is consumed on the else branch, live on the then branch; run_if b 0 takes then (ge 0 0 is true), so b stays live. Verified today: allocs=1 frees=0 live=1. Target: live=0.

(module if_branch_consume_split_no_leak_pin
  (data Box (ctor B (con Int)))
  (fn sink
    (type (fn-type (params (own (con Box))) (ret (own (con Unit))) (effects IO)))
    (params b)
    (body (match b (case (pat-ctor B n) (do io/print_str "x")))))
  (fn run_if
    (type (fn-type (params (own (con Box)) (own (con Int))) (ret (own (con Unit))) (effects IO)))
    (params b k)
    (body
      (if (app ge k 0)
          (do io/print_str "")
          (app sink b))))
  (fn main
    (type (fn-type (params) (ret (own (con Unit))) (effects IO)))
    (params)
    (body
      (let b (term-ctor Box B 7)
        (seq
          (app run_if b 0)
          (do io/print_str "done\n"))))))

Implementation shape (before → after, secondary)

// ailang-mir
pub struct MArm { pub pat: Pattern, pub body: MTerm }                 // before
pub struct MArm { pub pat: Pattern, pub body: MTerm,                  // after
                  pub consume: BTreeMap<String, u32> }

MTerm::If { cond, then, else_, ty }                                   // before
MTerm::If { cond, then, else_, ty,                                    // after
            then_consume: BTreeMap<String,u32>,
            else_consume: BTreeMap<String,u32> }
// MirDef.consume (the per-fn aggregate) is UNCHANGED.

// (a) pre-tail-call dec (match_lower.rs:823): gate source aggregate -> per-arm
let cc = self.consume.get(&(self.current_def, p))   // before: per-fn aggregate
let cc = arm.consume.get(p)                          // after:  per-arm
if cc != 0 { continue }                              // shape unchanged

// (b) NEW fall-through drop (match arms + if branches), before `br join`:
//     fires only for the leak class — live here AND consumed on a sibling.
let branch_cc = branch_consume.get(p).copied().unwrap_or(u32::MAX);
let agg       = self.consume.get(&(self.current_def, p)).copied().unwrap_or(0);
if branch_cc == 0 && agg >= 1 && p_ssa != tail_val && Some(p) != scrutinee {
    /* emit the moved_slots-aware drop (same machinery as fn-return) */
}

// (c) fn-return dec (lib.rs:1484): UNCHANGED — still `self.consume[p] == 0`.

Components

• crates/ailang-check/src/uniqueness.rs — retain + emit per-branch consume snapshots (the values currently discarded at merge_states). merge_states itself stays (it still computes the aggregate for classification).
• crates/ailang-mir/src/lib.rs — additive fields on MArm and MTerm::If.
• crates/ailang-check/src/lower_to_mir.rs — attach the per-branch maps to the MIR branch nodes in traversal-order lock-step.
• crates/ailang-codegen/src/match_lower.rs + src/lib.rs — switch the pre-tail-call dec's gate source to arm.consume; add the new fall-through leak-class drop (match arms + if branches), guarded by branch_consume == 0 && aggregate >= 1. The fn-return dec and the arm-close pattern-binder dec are left unchanged.

Data flow

post-mono AST → uniqueness (one walk: per-branch snapshots and aggregate) → lower_to_mir (aggregate → MirDef.consume; per-branch → MArm.consume / MTerm::If.{then,else}_consume, lock-step) → codegen (pre-tail-call dec reads per-arm; new fall-through drop reads per-branch gated by aggregate >= 1; fn-return dec still reads the aggregate for the == 0 class — the two classes are disjoint by aggregate).

Error handling

No new diagnostics. This is a codegen drop-placement correctness fix; the type/uniqueness analysis (what ail check reports) is unchanged — check already exits 0 on both repros today. The observable change is the runtime RC ledger (live=0) under --alloc=rc.

Lockstep pairs (per CLAUDE.md): neither lockstep pair is touched — no Pattern::Lit reject path and no INTERCEPTS/(intrinsic) change. The planner verifies this stays true (the change is confined to consume threading + drop siting).

Testing strategy

• RED (match): un-#[ignore] the committed match_arm_consume_split_no_leak_pin.rs; it asserts live==0 / allocs==frees. Goes green with the fix.
• RED (if): new pin if_branch_consume_split_no_leak_pin.rs + fixture examples/if_branch_consume_split_no_leak_pin.ail (shape above), same live==0 assertion. Authored RED-first; goes green with the fix.
• Regression — legs 1/2: the existing tail_recur_owned_param_no_leak_pin.rs and tail_recur_let_wrapped_no_leak_pin.rs must stay green (the behaviour-preserving argument above predicts this).
• Series close: series_sma (the #61 headline) must reach live=0 under AILANG_RC_STATS — the residual live=2 is this leg.
• Lock-step correspondence: a unit test that per-branch consume maps attach to the correct arms (guards the traversal-order invariant).
• UAF-safety precondition stays green: the existing use_after_consume_on_own_param_is_reported (and the match-then-consume rejection cases) must remain green — they are the foundation of the no-after-use guarantee the per-branch drop relies on. No new test needed; they are an existing invariant this cycle leans on.
• Full workspace suite green (the typed-MIR re-synth strictness family means a MIR-node change can ripple — run all crates, not just the new pins).

Acceptance criteria

Feature-acceptance criterion (CLAUDE.md): measurably improves correctness. This removes a drop-soundness leak that the project's own headline streaming pattern hits — a direct correctness gain, not an ergonomic one. The two repros are the empirical evidence: both report live=1 today and must report live=0 after.

Done when:

1. Both RED pins (match un-ignored, if new) are green.
2. Legs 1/2 pins and the full workspace suite stay green.
3. series_sma reports live=0, clearing the #61 leak tail.
4. No lockstep-pair drift; no new diagnostics; ail check output unchanged on all fixtures.