AILang/crates/ailang-check/tests/lower_to_mir_ty.rs

//! mir.1a: `lower_to_mir` fills `ty` on every node from the canonical
//! `synth`, and routes string literals to `MTerm::Str { rep: Static }`.
//! Loads the boundary witnesses as workspace fixtures (prelude +
//! kernel injected by `load_workspace`) and elaborates the whole
//! workspace, so the walk is exercised over the full prelude. These
//! pins protect the producer; codegen consumption arrives in mir.1b.

use ailang_check::elaborate_workspace;
use ailang_mir::{Callee, MTerm, MirWorkspace, Mode, StrRep};
use ailang_surface::load_workspace;
use ailang_test_support::examples_dir;

/// Load a witness fixture (prelude injected) and elaborate it to MIR.
fn elaborate_fixture(module: &str) -> MirWorkspace {
    let entry = examples_dir().join(format!("{module}.ail"));
    let ws = load_workspace(&entry).expect("fixture loads (prelude injected)");
    elaborate_workspace(&ws).expect("witness elaborates to MIR")
}

/// A def's lowered body in the elaborated workspace.
fn body<'a>(mir: &'a MirWorkspace, module: &str, def: &str) -> &'a MTerm {
    &mir.modules[module]
        .defs
        .iter()
        .find(|d| d.name == def)
        .expect("def present")
        .body
}

#[test]
fn loop_seed_str_literal_lowers_to_heap_str_node() {
    // #49 witness: the loop seed "x" is a Str literal → MTerm::Str.
    // mir.4: a loop-carried seed is promoted to rep = Heap so codegen
    // emits the str_clone heap promotion and the recur superseded-value
    // dec is sound (the prior alloca value is always an owned slab).
    let m = "loop_recur_str_binder_no_leak_pin";
    let mir = elaborate_fixture(m);
    assert_eq!(
        loop_seed_rep(body(&mir, m, "main")),
        Some(StrRep::Heap),
        "loop seed \"x\" must lower to MTerm::Str {{ rep: Heap }}"
    );
}

#[test]
fn loop_recur_literal_arg_lowers_to_heap_str_node() {
    // mir.4: a literal recur arg at a Str-binder position is promoted
    // to rep = Heap, mirroring the seed promotion — `(recur "reset" …)`.
    let m = "loop_str_recur_literal_no_leak_pin";
    let mir = elaborate_fixture(m);
    assert!(
        find_recur_str_arg_with_rep(body(&mir, m, "main"), StrRep::Heap),
        "the recur arg literal \"reset\" must lower to MTerm::Str {{ rep: Heap }}"
    );
}

#[test]
fn non_loop_str_literal_stays_static() {
    // Control: a Str literal NOT in a loop-carried position keeps
    // rep = Static (no spurious heap promotion). `hello` prints the
    // literal "Hello, AILang." via io/print_str.
    let m = "hello";
    let mir = elaborate_fixture(m);
    assert!(
        find_str_node_with_rep(body(&mir, m, "main"), StrRep::Static),
        "a non-loop Str literal must stay MTerm::Str {{ rep: Static }}"
    );
}

#[test]
fn loop_exit_arm_str_literal_lowers_to_heap_str_node() {
    // mir.4 loop-result leg: a Str literal in a loop exit (tail) arm is
    // promoted to rep = Heap, so the loop returns an owned heap slab the
    // caller's scope-close can safely free. `loop_str_static_exit` exits
    // with the literal "result"; it is the only Str literal in the body.
    let m = "loop_str_static_exit_no_leak_pin";
    let mir = elaborate_fixture(m);
    assert!(
        find_str_node_with_rep(body(&mir, m, "main"), StrRep::Heap),
        "the loop exit-arm literal \"result\" must lower to MTerm::Str {{ rep: Heap }}"
    );
    assert!(
        !find_str_node_with_rep(body(&mir, m, "main"), StrRep::Static),
        "no Str literal in the static-exit fixture should remain Static"
    );
}

#[test]
fn new_over_user_adt_carries_node_types() {
    // #53 witness: every node has a filled `ty`; the lowered
    // `(new Counter 42)` node's type is the user ADT `Counter`.
    //
    // A *monomorphic* `new` (no leading type-arg) is desugared to a
    // type-scoped call `(app Counter.new 42)` before lower_to_mir runs
    // (desugar.rs:1083-1098) — so the let-init lowers to `MTerm::App`,
    // not `MTerm::New`. `MTerm::New` survives only for a *polymorphic*
    // `new` carrying a written `NewArg::Type` (the #51 RawBuf case).
    // What this pin protects is the ty-fill: the lowered init node
    // carries the checker-proved result type `Counter`.
    let m = "new_counter_user_adt";
    let mir = elaborate_fixture(m);
    let MTerm::Let { init, .. } = body(&mir, m, "main") else {
        panic!("main body is a let");
    };
    let ty_str = ailang_core::pretty::type_to_string(&init.ty());
    assert!(
        ty_str.contains("Counter"),
        "lowered (new Counter 42) init node ty is Counter, got {ty_str}"
    );

    // mir.2: the `(new Counter 42)` desugars to `App{callee:
    // Var{"Counter.new"}}`; lower_to_mir must resolve it to a
    // `Callee::Static` whose module is `Counter`'s home (its own
    // module, spelled bare per the own-module-types-stay-bare rule),
    // NOT leave it `Indirect` for codegen's deleted ladder to resolve.
    let MTerm::App { callee, .. } = init.as_ref() else {
        panic!("let init is the (new Counter 42) App node");
    };
    match callee {
        ailang_mir::Callee::Static { module, fn_name, .. } => {
            assert_eq!(module, "new_counter_user_adt", "Counter.new home module");
            assert_eq!(fn_name, "new", "Counter.new fn_name");
        }
        other => panic!("expected Callee::Static for Counter.new, got {other:?}"),
    }
}

#[test]
fn rawbuf_size_only_elaborates() {
    // #51 witness: a RawBuf read only for size — must elaborate clean
    // (the element type carried only by the author's annotation does
    // not block lowering).
    let mir = elaborate_fixture("new_rawbuf_size_only");
    assert!(mir.modules.contains_key("new_rawbuf_size_only"));
}

#[test]
fn poly_free_fn_accumulating_into_own_adt_elaborates() {
    // Producer pin (mono ↔ lower_to_mir boundary): a polymorphic free
    // fn whose accumulator type-arg instantiates to the *defining
    // module's own* ADT must elaborate. `std_list.fold_left` is
    // `forall a b. ((b, a) -> b, b, List<a>) -> b`; `List_reverse`
    // calls it with `b := List<Int>` — the accumulator is std_list's
    // OWN `List`. mono normalises the observed type-arg to the
    // registry-canonical qualified form (`std_list.List`) for dedup +
    // symbol naming, but that qualified form must be localised back to
    // the module's bare convention before it is substituted into the
    // (own-bare) polymorphic signature and body. Without the
    // localisation the synthesised specialisation is self-inconsistent
    // — a qualified-`std_list.List` accumulator parameter against a
    // bare-`List` list argument and Lam body — and the post-mono
    // `synth` re-entry in `lower_to_mir` rejects it with
    // `expected std_list.List<Int>, got List<Int>`.
    let mir = elaborate_fixture("std_list_demo");
    assert!(mir.modules.contains_key("std_list"));
}

#[test]
fn callee_classification_builtin_and_static() {
    // mir.2: `lower_to_mir::classify_callee` resolves each App callee
    // against check's own synth ladder, so codegen reads the resolved
    // identity off MIR. This pins two of the three legs from the
    // `classify_pin` witness: an arithmetic op (`+`) is a `Builtin`
    // (no owning module), a bare same-module user fn (`bump`) is a
    // `Static` carrying its own module name. The `Indirect` leg (a
    // fn-typed local) is covered by the existing closure e2e tests.
    let m = "classify_pin";
    let mir = elaborate_fixture(m);

    // `bump`'s body is `(app + n 1)` → the `+` callee is a Builtin.
    match body(&mir, m, "bump") {
        MTerm::App { callee: Callee::Builtin { name, .. }, .. } => {
            assert_eq!(name, "+", "arithmetic op classified as Builtin");
        }
        other => panic!("expected Builtin `+`, got {other:?}"),
    }

    // `main`'s body is `(app bump 41)` → `bump` is a same-module user
    // fn → Static{module: "classify_pin", fn_name: "bump"}.
    match body(&mir, m, "main") {
        MTerm::App { callee: Callee::Static { module, fn_name, .. }, .. } => {
            assert_eq!(module, "classify_pin", "own-module callee module");
            assert_eq!(fn_name, "bump", "own-module callee fn_name");
        }
        other => panic!("expected Static `bump`, got {other:?}"),
    }
}

#[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,
    );
}

#[test]
fn app_arg_carries_callee_borrow_mode() {
    // mir.3b: lower_to_mir fills MArg.mode from the callee's param_modes.
    // In `(app RawBuf.get <set-result> 0)`, RawBuf.get's param 0 is a
    // borrow receiver, so the outer App's arg 0 must be Mode::Borrow —
    // the value codegen's anon-temp drop gate now reads off MIR.
    let mir = elaborate_fixture("raw_buf_drop_min");
    // main is `(app print (let buf (new RawBuf …)
    //   (app RawBuf.get (app RawBuf.set …) 0)))`, so the RawBuf.get App
    // is nested under the print App's arg and the let body. Mono mangles
    // `RawBuf.get` to `RawBuf_get__Int`.
    let outer = find_app_with_fn(body(&mir, "raw_buf_drop_min", "main"), "RawBuf_get")
        .expect("RawBuf.get App node");
    let MTerm::App { args, .. } = outer else {
        unreachable!("find_app_with_fn returns an App");
    };
    assert!(
        matches!(args[0].mode, Mode::Borrow),
        "RawBuf.get arg 0 (borrow receiver) must be Mode::Borrow, got {:?}",
        args[0].mode,
    );
}

/// First `MTerm::App` whose resolved callee (`Static`/`Builtin`) names a
/// fn equal to (or a monomorphised specialisation of) `fn_name`,
/// searching App args and Let init/body. The fixture's `RawBuf.get`
/// callee may be mangled by mono (e.g. `get__Int`), so the match is on
/// the fn-name stem.
fn find_app_with_fn<'a>(t: &'a MTerm, fn_name: &str) -> Option<&'a MTerm> {
    let names = match t {
        MTerm::App { callee: Callee::Static { fn_name: n, .. }, .. } => Some(n.as_str()),
        MTerm::App { callee: Callee::Builtin { name: n, .. }, .. } => Some(n.as_str()),
        _ => None,
    };
    if let Some(n) = names {
        if n == fn_name || n.starts_with(&format!("{fn_name}__")) {
            return Some(t);
        }
    }
    match t {
        MTerm::App { args, .. } => args.iter().find_map(|a| find_app_with_fn(&a.term, fn_name)),
        MTerm::Let { init, body, .. } => {
            find_app_with_fn(init, fn_name).or_else(|| find_app_with_fn(body, fn_name))
        }
        _ => None,
    }
}

/// The rep of the first loop binder seed that is an `MTerm::Str`,
/// searching `Let`-init / `Let`-body / the loop's own binders.
fn loop_seed_rep(t: &MTerm) -> Option<StrRep> {
    match t {
        MTerm::Loop { binders, body, .. } => binders
            .iter()
            .find_map(|b| match &b.init {
                MTerm::Str { rep, .. } => Some(*rep),
                _ => None,
            })
            .or_else(|| loop_seed_rep(body)),
        MTerm::Let { init, body, .. } => {
            loop_seed_rep(init).or_else(|| loop_seed_rep(body))
        }
        _ => None,
    }
}

/// True if any `Recur` arg in `t` is an `MTerm::Str` with the given rep.
fn find_recur_str_arg_with_rep(t: &MTerm, want: StrRep) -> bool {
    match t {
        MTerm::Recur { args, .. } => args.iter().any(|a| matches!(
            &a.term, MTerm::Str { rep, .. } if *rep == want
        )),
        MTerm::Loop { binders, body, .. } => {
            binders.iter().any(|b| find_recur_str_arg_with_rep(&b.init, want))
                || find_recur_str_arg_with_rep(body, want)
        }
        MTerm::Let { init, body, .. } => {
            find_recur_str_arg_with_rep(init, want)
                || find_recur_str_arg_with_rep(body, want)
        }
        MTerm::If { cond, then, else_, .. } => {
            find_recur_str_arg_with_rep(cond, want)
                || find_recur_str_arg_with_rep(then, want)
                || find_recur_str_arg_with_rep(else_, want)
        }
        _ => false,
    }
}

/// True if any `MTerm::Str` reachable in `t` carries the given rep
/// (used by the non-loop control: a plain literal stays Static).
fn find_str_node_with_rep(t: &MTerm, want: StrRep) -> bool {
    match t {
        MTerm::Str { rep, .. } => *rep == want,
        MTerm::Let { init, body, .. } => {
            find_str_node_with_rep(init, want) || find_str_node_with_rep(body, want)
        }
        MTerm::App { args, .. } => {
            args.iter().any(|a| find_str_node_with_rep(&a.term, want))
        }
        MTerm::Do { args, .. } => {
            args.iter().any(|a| find_str_node_with_rep(&a.term, want))
        }
        MTerm::If { cond, then, else_, .. } => {
            find_str_node_with_rep(cond, want)
                || find_str_node_with_rep(then, want)
                || find_str_node_with_rep(else_, want)
        }
        MTerm::Loop { binders, body, .. } => {
            binders.iter().any(|b| find_str_node_with_rep(&b.init, want))
                || find_str_node_with_rep(body, want)
        }
        _ => false,
    }
}

/// #63 leg 3 lock-step: the per-arm consume map attaches to the matching
/// MArm in traversal order. In the `match_arm_consume_split_no_leak_pin`
/// fixture's `run`, `b` is consumed on the `On` arm (forwarded into
/// `sink`) and live on the `Off` arm — so traversal-order attachment must
/// land MArm.consume[b] == 0 on `Off` (source-order arms[0]) and >= 1 on
/// `On` (arms[1]). A cursor desync would mis-attach and flip these.
#[test]
fn branch_consume_maps_attach_to_matching_arms() {
    let m = "match_arm_consume_split_no_leak_pin";
    let mir = elaborate_fixture(m);
    let arms = match body(&mir, m, "run") {
        MTerm::Match { arms, .. } => arms,
        other => panic!("run body is not a Match: {other:?}"),
    };
    assert_eq!(
        arms[0].consume.get("b").copied().unwrap_or(99),
        0,
        "Off arm (arms[0]) must leave `b` live (consume 0); got {:?}",
        arms[0].consume.get("b")
    );
    assert!(
        arms[1].consume.get("b").copied().unwrap_or(0) >= 1,
        "On arm (arms[1]) must consume `b` (>= 1); got {:?}",
        arms[1].consume.get("b")
    );
}