AILang/docs/plans/0117-mir.3a-consume-count-relocation.md

mir.3a — relocate consume_count into MIR, delete codegen's second uniqueness run — Implementation Plan

Parent spec: docs/specs/0060-typed-mir.md (Iteration decomposition table, mir.3 row — first of two iterations; see "Split" below)

For agentic workers: REQUIRED SUB-SKILL: use the implement skill to run this plan. Steps use - [ ] checkboxes for tracking.

Goal: Move the per-binder consume_count (the only thing codegen reads from its second infer_module_with_cross run) into MIR, so lower_to_mir computes it once on the post-mono module and codegen's three scope-close drop sites read it off MirDef.consume — then delete codegen's infer_module_with_cross call.

Architecture: lower_to_mir runs ailang_check::uniqueness:: infer_module_with_cross (the same pass, same crate) on the post-mono module it already walks, and attaches each def's binder→consume_count map to a new MirDef.consume field. Codegen builds its existing (def, binder) → consume_count lookup from those maps instead of re-running the pass; the three drop gates (own-param, let-binder, match-arm) are otherwise unchanged. This is a pure relocation of an identical computation — infer_module_with_cross ran on the post-mono module in codegen and now runs on the same post-mono module in lower_to_mir; same function, same input, same output, so drop placement is byte-identical.

Tech Stack: ailang-mir (MirDef.consume carrier), ailang-check::lower_to_mir + elaborate_workspace (producer), ailang-codegen (consumer: the three drop sites + the deleted run).

---

Split — mir.3 ships in two iterations

The spec's mir.3 row bundles "Mode + consume_count" as MIR additions and names "delete the second infer_module_with_cross run" as the deletion. Investigation during planning established that the deletion depends ONLY on consume_count: all three codegen drop sites read consume_count from self.uniqueness, while the modes they also test (param_modes for the own-param gate, scrutinee_is_owned for the match gate) come from the type, not from infer_module_with_cross. UniquenessInfo carries no mode.

So mir.3 is delivered as:

• mir.3a (this plan): relocate consume_count; delete the second infer_module_with_cross run. This is the spec's named deletion.
• mir.3b (next plan): fill MArg.mode / MTerm::Let.mode from the type's param/binder modes and switch codegen's per-arg mode reads (emit_call's anon-temp borrow-slot gate) onto MArg.mode. Additive; no further re-deriver deletion.

The split halves the change at the codebase's highest-risk surface (RC/drop correctness) and makes each half independently verifiable against the leak pins. mir.3a carries the deletion; mir.3b is pure annotation + read-relocation.

Files this plan creates or modifies:

• Modify: crates/ailang-mir/src/lib.rs:52-58 — add consume to MirDef.
• Modify: crates/ailang-check/src/lower_to_mir.rs:411 (lower_module signature + body) — run the uniqueness pass, fill MirDef.consume.
• Modify: crates/ailang-check/src/lib.rs:1317-1352 (elaborate_workspace) — build cross_module_types, thread it into lower_module.
• Modify: crates/ailang-codegen/src/lib.rs:42,848,1029,1057 — delete the import / field-type / run / init; build the lookup from MIR.
• Modify: crates/ailang-codegen/src/lib.rs:1511, :1878, crates/ailang-codegen/src/match_lower.rs:797 — the three consume_count reads.
• Test: crates/ailang-check/tests/lower_to_mir_ty.rs — producer pin (consume populated).
• Modify (docs): docs/specs/0060-typed-mir.md — record the split + the consume relocation.

---

Task 1: Add the consume carrier to MirDef

Files:

• Modify: crates/ailang-mir/src/lib.rs:52-58

• Step 1: Add the field

Replace MirDef at crates/ailang-mir/src/lib.rs:50-58:

/// One lowered top-level fn. `sig` is the declared signature; `body`
/// is the lowered fn body with `ty` on every node.
#[derive(Debug, Clone)]
pub struct MirDef {
    pub name: String,
    pub sig: Type,
    pub params: Vec<String>,
    pub body: MTerm,
}

with:

/// One lowered top-level fn. `sig` is the declared signature; `body`
/// is the lowered fn body with `ty` on every node.
#[derive(Debug, Clone)]
pub struct MirDef {
    pub name: String,
    pub sig: Type,
    pub params: Vec<String>,
    pub body: MTerm,
    /// Per-binder consume count for this def's body, keyed by binder
    /// name (let-binders, params, pattern-binders). `0` means the
    /// binder is exclusively borrowed/unused in its scope and codegen
    /// must emit the scope-close `dec` under `--alloc=rc`. Filled by
    /// `lower_to_mir` from `ailang_check::uniqueness::infer_module_with_cross`
    /// (mir.3a) — the single uniqueness engine, run once on the
    /// post-mono body. Codegen reads it instead of re-running the pass.
    pub consume: BTreeMap<String, u32>,
}

Implementer note: BTreeMap is already imported at crates/ailang-mir/src/lib.rs:12 (use std::collections::BTreeMap;).

• Step 2: Build the leaf crate (partial gate — ailang-check red until Task 2)

Run: cargo build -p ailang-mir Expected: PASS.

NOTE: ailang-check will NOT build until Task 2 — lower_module constructs MirDef without the new consume field. That is expected and closed in Task 2.

---

Task 2: Producer — lower_to_mir fills MirDef.consume

Files:

• Modify: crates/ailang-check/src/lib.rs:1317-1352 (elaborate_workspace)

• Modify: crates/ailang-check/src/lower_to_mir.rs:411 (lower_module)

• Step 1: elaborate_workspace builds cross_module_types and threads it in

In crates/ailang-check/src/lib.rs, inside elaborate_workspace, replace the post-mono lowering loop (the block at :1344-1352 beginning let env = build_check_env(&mono);):

    let env = build_check_env(&mono);
    let mut modules = std::collections::BTreeMap::new();
    for (name, m) in &mono.modules {
        let mir = lower_to_mir::lower_module(m, &env)
            .map_err(|e| vec![internal_diag(format!("lower_to_mir in `{name}`: {e}"))])?;
        modules.insert(name.clone(), mir);
    }
    Ok(ailang_mir::MirWorkspace { entry: mono.entry.clone(), modules })

with:

    let env = build_check_env(&mono);
    // Workspace-flat `module -> def -> Type` table — the shape
    // `uniqueness::infer_module_with_cross` consumes to resolve a
    // cross-module callee's `param_modes` (so a Borrow-mode arg is
    // walked as a borrow, not a consume). Identical in content to
    // codegen's `module_def_ail_types`; built here from the post-mono
    // workspace so the uniqueness pass runs in the producer (mir.3a).
    let mut cross_module_types: std::collections::BTreeMap<
        String,
        std::collections::BTreeMap<String, ailang_core::ast::Type>,
    > = std::collections::BTreeMap::new();
    for (mname, m) in &mono.modules {
        let mut def_tys = std::collections::BTreeMap::new();
        for def in &m.defs {
            if let ailang_core::ast::Def::Fn(f) = def {
                def_tys.insert(f.name.clone(), f.ty.clone());
            }
        }
        cross_module_types.insert(mname.clone(), def_tys);
    }
    let mut modules = std::collections::BTreeMap::new();
    for (name, m) in &mono.modules {
        let mir = lower_to_mir::lower_module(m, &env, &cross_module_types)
            .map_err(|e| vec![internal_diag(format!("lower_to_mir in `{name}`: {e}"))])?;
        modules.insert(name.clone(), mir);
    }
    Ok(ailang_mir::MirWorkspace { entry: mono.entry.clone(), modules })

Implementer note: confirm Def::Fn(f) carries name: String and ty: Type (it is the same FnDef codegen reads at lib.rs:334 to build module_def_ail_types). The ailang_core::ast::Type / Def::Fn paths may already be imported in lib.rs under shorter names — use whatever the file already uses (e.g. a bare Type / Def); the fully-qualified forms above are correct regardless.

• Step 2: lower_module runs the uniqueness pass and fills each MirDef.consume

In crates/ailang-check/src/lower_to_mir.rs, change the lower_module signature at :411:

pub fn lower_module(module: &Module, env: &Env) -> Result<MirModule> {

to:

pub fn lower_module(
    module: &Module,
    env: &Env,
    cross_module_types: &std::collections::BTreeMap<String, std::collections::BTreeMap<String, Type>>,
) -> Result<MirModule> {

Immediately after the permissive-import seeding loop (the for m in all_modules { env.imports.entry(...) ... } block ending at :460) and before let mut defs = Vec::new(); (:461), insert:

    // mir.3a: run the single uniqueness engine ONCE on this post-mono
    // module. The resulting per-(def, binder) consume_count is attached
    // to each `MirDef.consume` below; codegen reads it instead of
    // re-running `infer_module_with_cross`. Same pass, same post-mono
    // input as codegen used — a pure relocation, so drop placement is
    // unchanged.
    let uniqueness =
        crate::uniqueness::infer_module_with_cross(module, cross_module_types);

Then, at the MirDef construction site inside the for def in &module.defs loop (the Def::Fn arm — read the region from :461 onward to find the defs.push(ailang_mir::MirDef { ... }) literal), add the consume field to the constructed MirDef. The value is this def's slice of the uniqueness table, keyed by binder name:

            consume: uniqueness
                .iter()
                .filter(|((d, _), _)| d == &f.name)
                .map(|((_, b), info)| (b.clone(), info.consume_count))
                .collect(),

Implementer note: bind the field to whatever the Def::Fn arm calls the fn def (the plan uses f per the Def::Fn(f) convention; match the arm's actual binding — it may be fd or def). uniqueness is a BTreeMap<(String, String), UniquenessInfo>; UniquenessInfo.consume_count is a u32 (crates/ailang-check/src/uniqueness.rs:87,99). The .filter(...).map(...).collect() yields the BTreeMap<String, u32> the field wants. If the file does not already use the uniqueness module, the crate::uniqueness:: path above is fully qualified and needs no import.

• Step 3: Build the workspace (full gate — codegen still uses its own run, so green)

Run: cargo build --workspace Expected: PASS. After this task the producer fills MirDef.consume but codegen still computes its own self.uniqueness; nothing reads MirDef.consume yet, so behaviour is unchanged.

• Step 4: Confirm no behaviour change

Run: cargo test --workspace Expected: PASS — the full suite green, identical to pre-Task-2 (the new field is filled but unread).

---

Task 3: Consumer — codegen reads MirDef.consume; delete the second run

Files:

• Modify: crates/ailang-codegen/src/lib.rs:42 (import), :848 (field), :1029 (the run), :1057 (init)

• Modify: crates/ailang-codegen/src/lib.rs:1511, :1878, crates/ailang-codegen/src/match_lower.rs:797 (the three reads)

• Step 1: Replace the Emitter uniqueness field type

In crates/ailang-codegen/src/lib.rs, the Emitter field at :848:

    uniqueness: UniquenessTable,

Replace with a plain (def, binder) → consume_count map (the only thing the three drop sites read off it):

    // Per-(def, binder) consume_count, built from `MirDef.consume`
    // (mir.3a) — codegen no longer re-runs `infer_module_with_cross`.
    consume: std::collections::BTreeMap<(String, String), u32>,

• Step 2: Build consume from MIR instead of running the pass

At crates/ailang-codegen/src/lib.rs:1029, replace:

    let uniqueness = infer_module_with_cross(module, module_def_ail_types);

with a fold over the module's MirDef.consume maps:

    // mir.3a: read the per-binder consume_count off MIR (filled by
    // lower_to_mir from the single uniqueness pass) instead of
    // re-running `infer_module_with_cross`.
    let consume: std::collections::BTreeMap<(String, String), u32> = mir_module
        .defs
        .iter()
        .flat_map(|d| {
            let dname = d.name.clone();
            d.consume
                .iter()
                .map(move |(b, c)| ((dname.clone(), b.clone()), *c))
        })
        .collect();

Implementer note: mir_module is the name this plan uses for the &ailang_mir::MirModule the Emitter setup holds at this point (mir.1b wired codegen to consume MIR). Read the region around :1029 to find the actual binding — it is the MIR module whose module (the post-mono AST) was being passed to the deleted call; the MirModule carries both .ast and .defs. Use its real name.

• Step 3: Update the Emitter init

At crates/ailang-codegen/src/lib.rs:1057, replace the field init:

            uniqueness,

with:

            consume,

• Step 4: Delete the now-unused import

At crates/ailang-codegen/src/lib.rs:42:

use ailang_check::uniqueness::{infer_module_with_cross, UniquenessTable};

infer_module_with_cross and UniquenessTable are no longer referenced in codegen after Steps 1-3. Remove the import. If ParamMode or another item was glob-adjacent, keep only what survives (this line imports only those two names, so delete the whole line).

• Step 5: Rewire the three consume_count reads

Each of the three drop sites reads self.uniqueness.get(&(def, binder)).map(|info| info.consume_count).unwrap_or(u32::MAX). The value is now a bare u32, so the read becomes self.consume.get(&(def, binder)).copied().unwrap_or(u32::MAX).

1. crates/ailang-codegen/src/lib.rs:1511 (own-param scope-close dec):

                    let consume_count = self
                        .uniqueness
                        .get(&(self.current_def.clone(), pname.clone()))
                        .map(|info| info.consume_count)
                        .unwrap_or(u32::MAX);

→

                    let consume_count = self
                        .consume
                        .get(&(self.current_def.clone(), pname.clone()))
                        .copied()
                        .unwrap_or(u32::MAX);

2. crates/ailang-codegen/src/lib.rs:1878 (let-binder scope-close dec):

                    let consume_count = self
                        .uniqueness
                        .get(&(self.current_def.clone(), name.clone()))
                        .map(|info| info.consume_count)
                        .unwrap_or(u32::MAX); // Defensive: skip if missing.

→

                    let consume_count = self
                        .consume
                        .get(&(self.current_def.clone(), name.clone()))
                        .copied()
                        .unwrap_or(u32::MAX); // Defensive: skip if missing.

3. crates/ailang-codegen/src/match_lower.rs:797 (match arm-close dec):

                    let consume_count = self
                        .uniqueness
                        .get(&(self.current_def.clone(), bname.clone()))
                        .map(|info| info.consume_count)
                        .unwrap_or(u32::MAX);

→

                    let consume_count = self
                        .consume
                        .get(&(self.current_def.clone(), bname.clone()))
                        .copied()
                        .unwrap_or(u32::MAX);

• Step 6: Build the workspace (the consuming gate)

Run: cargo build --workspace Expected: PASS. Zero errors, zero unused import / unused variable warnings for module_def_ail_types (it survives — still used by emit_call's anon-temp param_modes gate; if the compiler now flags it unused, a read was missed and the leak pins in Task 4 will catch the behavioural regression).

---

Task 4: Acceptance — leak pins green, producer pin, grep-clean

Files:

• Test: crates/ailang-check/tests/lower_to_mir_ty.rs (producer pin)

• Step 1: Producer pin — MirDef.consume is populated

Add a test to crates/ailang-check/tests/lower_to_mir_ty.rs that elaborates an existing fixture with a let-binder and asserts its MirDef.consume carries that binder. Reuse the file's elaborate_fixture helper and the new_counter_user_adt fixture (its main body is (let c (new Counter 42) ...) — binder c):

#[test]
fn mirdef_consume_is_populated_for_let_binder() {
    // mir.3a: lower_to_mir runs the uniqueness pass and attaches the
    // per-binder consume_count to MirDef.consume. `main`'s `(let c …)`
    // binder must appear in the map (the exact count is validated
    // end-to-end by the RC-stats leak pins; this pin guards that the
    // producer fills the field at all, so codegen has it to read).
    let ws = elaborate_fixture("new_counter_user_adt");
    let m = ws.modules.get("new_counter_user_adt").expect("module");
    let main = m.defs.iter().find(|d| d.name == "main").expect("main");
    assert!(
        main.consume.contains_key("c"),
        "MirDef.consume must carry the `c` let-binder, got {:?}",
        main.consume,
    );
}

Implementer note: match elaborate_fixture's real name/signature (the sibling tests in this file use it). ws.modules is a BTreeMap<String, MirModule>; MirDef.consume is the BTreeMap<String, u32> added in Task 1.

• Step 2: Run the producer pin

Run: cargo test -p ailang-check --test lower_to_mir_ty Expected: PASS — including mirdef_consume_is_populated_for_let_binder.

• Step 3: Full suite — the drop-correctness acceptance

Run: cargo test --workspace Expected: PASS. The RC-stats leak pins are the value-correctness gate for the relocation (a wrong consume_count → leak or double-free): alloc_rc_loop_valued_rawbuf_binder_drops_at_scope_close (raw_buf_loop_no_leak_pin.rs, #43/#47), alloc_rc_heap_loop_binder_replaced_by_recur_does_not_leak (loop_recur_heap_binder_no_leak_pin.rs), alloc_rc_print_int_does_not_leak_show_result_str (print_no_leak_pin.rs). All live == 0. loop_recur_str_binder_no_leak_pin stays #[ignore] (#49 lifts at mir.4 — do not touch). #51/#53 build guards stay green.

Run the FULL workspace suite — drop placement is perturbed by this change (even though it should be byte-identical), so the broad RC-stats suite, not just a targeted filter, is the real gate.

• Step 4: grep-clean — the second run is gone from codegen

Run: rg -n 'infer_module_with_cross|UniquenessTable' crates/ailang-codegen/ Expected: no matches. (The pass and its table type are no longer referenced anywhere in codegen; they remain ailang-check-internal.)

---

Task 5: Record the consume relocation + the mir.3 split in the spec

Files:

• Modify: docs/specs/0060-typed-mir.md

• Step 1: Note the split and the consume relocation

In docs/specs/0060-typed-mir.md, under the Iteration-decomposition table (after the mir.2 refinement note added in the previous iteration), append:

> mir.3 split (discovered in planning, `docs/plans/0117-mir.3a-consume-count-relocation.md`):
> mir.3 ships in two iterations. **mir.3a** relocates `consume_count`
> only — the deletion of the second `infer_module_with_cross` run
> depends solely on it (all three codegen drop sites read `consume_count`
> from `self.uniqueness`; the *modes* they also test come from the
> type, not the uniqueness pass, which carries no mode). The per-binder
> `consume_count` lands on a new `MirDef.consume: BTreeMap<String, u32>`
> (binder-keyed, matching how the drop sites and the `UniquenessTable`
> key it — the spec's `MArg.consume_count` field is per-arg-position
> and does not reach the binder-keyed drop sites, so it stays a mir.1
> default). lower_to_mir runs the single uniqueness pass on the
> post-mono body — a pure relocation of the identical computation
> codegen ran, so drop placement is byte-identical. **mir.3b** fills
> `MArg.mode` / `MTerm::Let.mode` from the type modes and switches
> `emit_call`'s anon-temp borrow-slot gate onto `MArg.mode` (additive;
> no further re-deriver deletion).

• Step 2: No-op re-gate

Run: cargo test --workspace Expected: PASS (doc-only edit; fast re-confirm).

---

Self-review (planner, inline)

1. Spec coverage. mir.3 row's named deletion ("second infer_module_with_cross run") + the consume_count MIR addition are covered (Tasks 1-3); the Mode half is explicitly deferred to mir.3b and recorded in the spec (Task 5). The split is justified by the consumer analysis, not by effort.
2. Placeholder scan. No "TBD/TODO/similar to". Every code step carries exact before→after bytes. The two read-the-region points (the Def::Fn MirDef construction site in lower_module; the mir_module binding name at codegen :1029) are framed as read-then-apply with the exact transformation, not open decisions.
3. Type/name consistency. MirDef.consume, cross_module_types, infer_module_with_cross, UniquenessInfo.consume_count, self.consume, current_def, the three site paths — consistent across tasks.
4. Step granularity. Each step is one edit or one command.
5. No commit steps. None present.
6. Pin/replacement substring contiguity. The grep gate (Task 4 Step 4) asserts on infer_module_with_cross / UniquenessTable — single tokens, contiguous. The producer pin asserts consume.contains_key("c") — a string key, not a wrapped substring.
7. Compile-gate vs deferred-caller ordering. Task 1's gate is explicitly partial (-p ailang-mir); Task 2 threads every MirDef construction site (there is one fn arm) and the single lower_module caller (elaborate_workspace) in the same task, closing with a workspace build. Task 2 changes lower_module's signature and updates its only caller in the same task — no deferred caller past a build gate. Task 3's first gate is the workspace build after all three reads + the field/import deletions are threaded together.
8. Verification-command filter strings resolve. cargo test -p ailang-check --test lower_to_mir_ty targets a real file; mirdef_consume_is_populated_for_let_binder is the exact new test name; the leak pins are named by their real test fns (recon-confirmed: alloc_rc_loop_valued_rawbuf_binder_drops_at_scope_close, alloc_rc_heap_loop_binder_replaced_by_recur_does_not_leak, alloc_rc_print_int_does_not_leak_show_result_str); the workspace gates are unfiltered; the rg gate asserts on real present symbols expected to reach zero in codegen.
9. Parse-the-bytes gate. mir.3a inlines NO new surface-language .ail fixture (the producer pin reuses the existing committed examples/new_counter_user_adt.ail). The spec_validation gate has no surface body to parse — documented no-op for this plan.

---

Handoff

Plan sits unstaged at docs/plans/0117-mir.3a-consume-count-relocation.md. Hand to implement (standard mode, all five tasks). The orchestrator commits the plan when handing it forward and the iteration diff at iter close.