AILang/crates/ailang-check/src/lower_to_mir.rs

//! Post-monomorphisation lowering from the typechecked `ast::Term`
//! to typed `MTerm`. Runs after `monomorphise_workspace`. Carries no
//! second type engine: each node's type comes from a fresh-scaffolding
//! call to the canonical `crate::synth` (lib.rs:3223). The only state
//! threaded across nodes is the lexical `locals` / `loop_stack`,
//! maintained exactly as `synth` maintains it (e.g. the `Term::Let`
//! push/pop at lib.rs:3703-3715).

use ailang_core::ast::{Literal, Module, Term, Type};
use ailang_mir::{
    Callee, MArg, MArm, MLoopBinder, MNewArg, MTerm, MirDef, MirModule, Mode, StrRep,
};
use indexmap::IndexMap;
use std::collections::BTreeSet;

use crate::{Env, Result};

/// Per-walk scaffolding for the canonical `synth` re-entry. `env` and
/// `in_def` are fixed per fn; `locals` / `loop_stack` are mutated as
/// the walk descends and restored on the way up, mirroring `synth`.
struct Ctx<'a> {
    env: &'a Env,
    in_def: &'a str,
    locals: IndexMap<String, Type>,
    loop_stack: Vec<Vec<(String, Type)>>,
}

impl<'a> Ctx<'a> {
    /// The canonical type of `t` in the current lexical scope, with no
    /// side effects on real check state. Fresh inference scaffolding
    /// per call; the result has its substitution applied so downstream
    /// reads see a resolved type.
    fn synth_pure(&mut self, t: &Term) -> Result<Type> {
        let mut effects: BTreeSet<String> = BTreeSet::new();
        let mut subst = crate::Subst::default();
        let mut counter: u32 = 0;
        let mut residuals: Vec<crate::ResidualConstraint> = Vec::new();
        let mut free_fn_calls: Vec<crate::FreeFnCall> = Vec::new();
        let mut warnings: Vec<crate::diagnostic::Diagnostic> = Vec::new();
        let ty = crate::synth(
            t,
            self.env,
            &mut self.locals,
            &mut self.loop_stack,
            &mut effects,
            self.in_def,
            &mut subst,
            &mut counter,
            &mut residuals,
            &mut free_fn_calls,
            &mut warnings,
        )?;
        Ok(wildcard_residual_metavars(&subst.apply(&ty)))
    }
}

/// Rewrite every residual `$m`-prefixed inference metavar to the
/// `$u` wildcard codegen's monomorphisation unifier expects.
///
/// `synth_pure` synthesises each node's type in isolation, so a
/// type-arg that the full program would pin only through a *sibling*
/// node stays an unbound metavar here — e.g. a bare `Nil`'s element
/// type in `(Cons 3 Nil)`: synthesised alone, `Nil : List<$m0>`,
/// because nothing in the `Nil` node itself constrains the element.
/// Post-monomorphisation any such residual is, by construction, an
/// unconstrained *phantom* position (a nullary ctor carries no value
/// of that type, so no runtime representation depends on it) — every
/// rep-bearing type was already pinned by its own value or by mono.
/// Codegen already has a wildcard for exactly this: `$u`, which
/// `subst::unify_for_subst` accepts on either side without binding, so
/// a sibling arg pins the type var instead (`Cons`'s declared
/// `List<a>` field unifies against `List<$u>` without the head's
/// `a := Int` binding then clashing on `$u`). The canonical checker
/// emits `$m` metavars, never `$u` — `$u` was the now-deleted
/// codegen-side synth's own spelling — so the typed-MIR boundary
/// normalises to it here, once, rather than teaching every node-type
/// consumer to also tolerate a raw `$m`.
fn wildcard_residual_metavars(t: &Type) -> Type {
    match t {
        Type::Var { name } if name.starts_with("$m") => Type::Var {
            name: name.replacen("$m", "$u", 1),
        },
        Type::Var { .. } => t.clone(),
        Type::Con { name, args } => Type::Con {
            name: name.clone(),
            args: args.iter().map(wildcard_residual_metavars).collect(),
        },
        Type::Fn { params, param_modes, ret, ret_mode, effects } => Type::Fn {
            params: params.iter().map(wildcard_residual_metavars).collect(),
            param_modes: param_modes.clone(),
            ret: Box::new(wildcard_residual_metavars(ret)),
            ret_mode: ret_mode.clone(),
            effects: effects.clone(),
        },
        Type::Forall { vars, constraints, body } => Type::Forall {
            vars: vars.clone(),
            constraints: constraints.clone(),
            body: Box::new(wildcard_residual_metavars(body)),
        },
    }
}

/// Default arg wrapper — mode/consume_count are mir.1 placeholders
/// (mir.3 fills them).
fn arg(term: MTerm) -> MArg {
    MArg { term, mode: Mode::Owned, consume_count: 1 }
}

/// The MIR `Mode` of a callee's parameter at position `i`, read off
/// the callee's `Type::Fn.param_modes`. `Borrow` maps to
/// `Mode::Borrow`; `Own` and `Implicit` (the `Implicit ≡ Own`
/// contract, `ast.rs` ParamMode) both map to `Mode::Owned`. Out of
/// range / non-fn → `Owned` (the consume default). Used by the App arm
/// to fill `MArg.mode` so codegen reads the slot mode off MIR.
fn param_mode_at(sig: &Type, i: usize) -> Mode {
    match sig {
        Type::Fn { param_modes, .. } => match param_modes.get(i) {
            Some(ailang_core::ast::ParamMode::Borrow) => Mode::Borrow,
            _ => Mode::Owned,
        },
        _ => Mode::Owned,
    }
}

/// The resolved identity of an App callee, mirroring synth's
/// `Term::Var` resolution ladder (`lib.rs:3409-3582`). `lower_term`
/// turns this into the matching `Callee` variant. Resolution uses
/// check's own env (the single engine) — never a copy of codegen's
/// `is_static_callee` allowlist.
enum CalleeClass {
    Builtin { name: String },
    Static { module: String, fn_name: String },
    Indirect,
}

fn classify_callee(ctx: &Ctx, callee: &Term) -> CalleeClass {
    // A non-`Var` callee (lambda, applied expression, …) is dynamic.
    let name = match callee {
        Term::Var { name } => name.clone(),
        _ => return CalleeClass::Indirect,
    };
    // rung 1: shadowed by a local binder → dynamic (synth `lib.rs:3409`).
    if ctx.locals.contains_key(&name) {
        return CalleeClass::Indirect;
    }
    // rung 5: dotted `T.fn` / `Mod.fn` — the TypeDef-first ladder
    // (synth `lib.rs:3557-3582`). `T.fn` resolves to T's home module;
    // `Mod.fn` to the imported module. This replaces codegen's
    // `type_home_module`.
    if name.matches('.').count() == 1 {
        let (prefix, suffix) = name.split_once('.').expect("checked");
        let type_home = ctx
            .env
            .module_types
            .iter()
            .find_map(|(m, types)| types.contains_key(prefix).then(|| m.clone()));
        let target = type_home.or_else(|| ctx.env.imports.get(prefix).cloned());
        return match target {
            Some(module) => CalleeClass::Static { module, fn_name: suffix.to_string() },
            // Unreachable for typechecked input (check would have
            // raised TypeScopedReceiverNotAType); defensive dynamic.
            None => CalleeClass::Indirect,
        };
    }
    // rung 2: same-module global. It is a user fn IFF this module's
    // declared globals carry the name; otherwise it is a builtin —
    // both live in `env.globals` (synth `lib.rs:3416-3425`).
    if ctx.env.globals.contains_key(&name) {
        let is_user_fn = ctx
            .env
            .module_globals
            .get(&ctx.env.current_module)
            .is_some_and(|m| m.contains_key(&name));
        return if is_user_fn {
            CalleeClass::Static { module: ctx.env.current_module.clone(), fn_name: name }
        } else {
            CalleeClass::Builtin { name }
        };
    }
    // rung 3: implicit-import (prelude-fallback) fn (synth `lib.rs:3428-3435`).
    if let Some(module) = ctx.env.imports.values().find_map(|m| {
        ctx.env
            .module_globals
            .get(m)
            .filter(|g| g.contains_key(&name))
            .map(|_| m.clone())
    }) {
        return CalleeClass::Static { module, fn_name: name };
    }
    // Typechecked input always resolves above; defensive dynamic.
    CalleeClass::Indirect
}

/// Lower one term to `MTerm`, filling `ty` from `synth_pure`.
fn lower_term(ctx: &mut Ctx, t: &Term) -> Result<MTerm> {
    let ty = ctx.synth_pure(t)?;
    Ok(match t {
        // ---- String literal: the rep carrier (Static at mir.1) ----
        Term::Lit { lit: Literal::Str { value } } => {
            MTerm::Str { lit: value.clone(), rep: StrRep::Static }
        }
        Term::Lit { lit } => MTerm::Lit { lit: lit.clone(), ty },

        Term::Var { name } => MTerm::Var { name: name.clone(), ty },

        // mir.2: classify the callee against check's own resolution
        // ladder and emit the resolved `Callee`. `sig` is the callee's
        // fn-type (mode-preserving via `synth_pure`), carried so
        // codegen's drop path reads the callee `ret_mode`.
        Term::App { callee, args, tail } => {
            let class = classify_callee(ctx, callee);
            let sig = ctx.synth_pure(callee)?;
            // mir.3b: each App arg carries the callee's slot mode, read
            // off the resolved callee fn-type. codegen's anon-temp drop
            // gate reads this instead of re-looking-up the callee's
            // param_modes from a sig table. Built before `m_callee`
            // consumes `sig` into the resolved `Callee` variant.
            let m_args = args
                .iter()
                .enumerate()
                .map(|(i, a)| {
                    Ok(MArg {
                        term: lower_term(ctx, a)?,
                        mode: param_mode_at(&sig, i),
                        consume_count: 1,
                    })
                })
                .collect::<Result<Vec<_>>>()?;
            let m_callee = match class {
                CalleeClass::Builtin { name } => Callee::Builtin { name, sig },
                CalleeClass::Static { module, fn_name } => {
                    Callee::Static { module, fn_name, sig }
                }
                CalleeClass::Indirect => {
                    Callee::Indirect(Box::new(lower_term(ctx, callee)?))
                }
            };
            MTerm::App { callee: m_callee, args: m_args, tail: *tail, ty }
        }

        // scope-affecting: insert binder type, lower body, restore —
        // exactly synth's Term::Let arm (lib.rs:3703-3715).
        Term::Let { name, value, body } => {
            let m_init = lower_term(ctx, value)?;
            let v_ty = ctx.synth_pure(value)?;
            let prev = ctx.locals.insert(name.clone(), v_ty);
            let m_body = lower_term(ctx, body)?;
            match prev {
                Some(p) => {
                    ctx.locals.insert(name.clone(), p);
                }
                None => {
                    ctx.locals.shift_remove(name);
                }
            }
            MTerm::Let {
                name: name.clone(),
                mode: Mode::Owned,
                init: Box::new(m_init),
                body: Box::new(m_body),
                ty,
            }
        }

        Term::If { cond, then, else_ } => MTerm::If {
            cond: Box::new(lower_term(ctx, cond)?),
            then: Box::new(lower_term(ctx, then)?),
            else_: Box::new(lower_term(ctx, else_)?),
            ty,
        },

        Term::Do { op, args, tail } => {
            let m_args = args
                .iter()
                .map(|a| Ok(arg(lower_term(ctx, a)?)))
                .collect::<Result<Vec<_>>>()?;
            MTerm::Do { op: op.clone(), args: m_args, tail: *tail, ty }
        }

        Term::Ctor { type_name, ctor, args } => {
            let m_args = args
                .iter()
                .map(|a| Ok(arg(lower_term(ctx, a)?)))
                .collect::<Result<Vec<_>>>()?;
            MTerm::Ctor {
                type_name: type_name.clone(),
                ctor: ctor.clone(),
                args: m_args,
                ty,
            }
        }

        // scope-affecting: each arm binds its pattern vars. Mirror
        // synth's Match arm exactly — including its own helper
        // `type_check_pattern`, which resolves the binder types from
        // the scrutinee's ADT. Push, lower the arm body, restore
        // (innermost-first), per arm.
        Term::Match { scrutinee, arms } => {
            let m_scrut = Box::new(lower_term(ctx, scrutinee)?);
            let s_ty = ctx.synth_pure(scrutinee)?;
            let mut m_arms = Vec::with_capacity(arms.len());
            for a in arms {
                let bindings = crate::type_check_pattern(&a.pat, &s_ty, ctx.env)?;
                let mut pushed = Vec::new();
                for (n, t) in &bindings {
                    let prev = ctx.locals.insert(n.clone(), t.clone());
                    pushed.push((n.clone(), prev));
                }
                let m_body = lower_term(ctx, &a.body)?;
                for (n, prev) in pushed.into_iter().rev() {
                    match prev {
                        Some(p) => {
                            ctx.locals.insert(n, p);
                        }
                        None => {
                            ctx.locals.shift_remove(&n);
                        }
                    }
                }
                m_arms.push(MArm { pat: a.pat.clone(), body: m_body });
            }
            MTerm::Match { scrutinee: m_scrut, arms: m_arms, ty }
        }

        // scope-affecting: params enter scope for the body. Mirror
        // synth's Lam arm: insert each param:param_ty into locals for
        // the body, restore after.
        Term::Lam { params, param_tys, ret_ty, effects, body } => {
            let saved = ctx.locals.clone();
            for (p, pty) in params.iter().zip(param_tys.iter()) {
                ctx.locals.insert(p.clone(), pty.clone());
            }
            let m_body = Box::new(lower_term(ctx, body)?);
            ctx.locals = saved;
            MTerm::Lam {
                params: params.clone(),
                param_tys: param_tys.clone(),
                ret_ty: (**ret_ty).clone(),
                effects: effects.clone(),
                body: m_body,
                ty,
            }
        }

        Term::Seq { lhs, rhs } => MTerm::Seq {
            lhs: Box::new(lower_term(ctx, lhs)?),
            rhs: Box::new(lower_term(ctx, rhs)?),
            ty,
        },

        Term::Clone { value } => MTerm::Clone {
            value: Box::new(lower_term(ctx, value)?),
            ty,
        },

        Term::ReuseAs { source, body } => MTerm::ReuseAs {
            source: Box::new(lower_term(ctx, source)?),
            body: Box::new(lower_term(ctx, body)?),
            ty,
        },

        // scope-affecting: binders enter scope and a loop frame is
        // pushed for the body so an inner Recur resolves. Mirror
        // synth's Loop arm: push the binder frame onto loop_stack and
        // the binder types onto locals before lowering the body,
        // pop/restore after.
        Term::Loop { binders, body } => {
            let mut m_binders = Vec::with_capacity(binders.len());
            for b in binders {
                m_binders.push(MLoopBinder {
                    name: b.name.clone(),
                    ty: b.ty.clone(),
                    init: lower_term(ctx, &b.init)?,
                });
            }
            let saved_locals = ctx.locals.clone();
            let frame: Vec<(String, Type)> =
                binders.iter().map(|b| (b.name.clone(), b.ty.clone())).collect();
            for b in binders {
                ctx.locals.insert(b.name.clone(), b.ty.clone());
            }
            ctx.loop_stack.push(frame);
            let m_body = Box::new(lower_term(ctx, body)?);
            ctx.loop_stack.pop();
            ctx.locals = saved_locals;
            MTerm::Loop { binders: m_binders, body: m_body, ty }
        }

        Term::Recur { args } => {
            let m_args = args
                .iter()
                .map(|a| Ok(arg(lower_term(ctx, a)?)))
                .collect::<Result<Vec<_>>>()?;
            MTerm::Recur { args: m_args, ty }
        }

        // elem is None at mir.1 (mir.5 carries the element type).
        Term::New { type_name, args } => {
            let m_args = args
                .iter()
                .map(|a| {
                    Ok(match a {
                        ailang_core::ast::NewArg::Type(t) => MNewArg::Type(t.clone()),
                        ailang_core::ast::NewArg::Value(v) => {
                            MNewArg::Value(lower_term(ctx, v)?)
                        }
                    })
                })
                .collect::<Result<Vec<_>>>()?;
            MTerm::New { type_name: type_name.clone(), elem: None, args: m_args, ty }
        }

        // scope-affecting: synth's LetRec arm binds `name: ty` for both
        // `body` and `in_term`. Mirror it.
        Term::LetRec { name, ty: rec_ty, params, body, in_term } => {
            let saved = ctx.locals.clone();
            ctx.locals.insert(name.clone(), rec_ty.clone());
            let m_body = Box::new(lower_term(ctx, body)?);
            let m_in = Box::new(lower_term(ctx, in_term)?);
            ctx.locals = saved;
            MTerm::LetRec {
                name: name.clone(),
                sig: rec_ty.clone(),
                params: params.clone(),
                body: m_body,
                in_term: m_in,
                ty,
            }
        }

        Term::Intrinsic => MTerm::Intrinsic { ty },
    })
}

/// Lower one post-mono module to MIR. `env` is the workspace check
/// env (built by `build_check_env`); this fn sets `env.imports` for
/// the module and seeds per-fn `locals` from the declared params,
/// mirroring mono.rs:771-816.
pub fn lower_module(
    module: &Module,
    env: &Env,
    cross_module_types: &std::collections::BTreeMap<String, std::collections::BTreeMap<String, Type>>,
) -> Result<MirModule> {
    let mut env = env.clone();
    env.current_module = module.name.clone();
    // Seed `env.globals` from the current module's fns so `synth`'s
    // `Term::Var` lookup resolves bare same-module references —
    // including the monomorphic specialisations mono appended (e.g.
    // `compare__Int`). Mirrors `mono::collect_mono_targets`
    // (mono.rs:769-781); per-module because top-level fn names are
    // only per-module-unique.
    if let Some(g) = env.module_globals.get(&module.name).cloned() {
        for (n, t) in g {
            env.globals.insert(n, t);
        }
    }
    // Seed `env.imports` from the current module's import list so
    // `synth`'s qualified-var path resolves `Mod.fn` references.
    // Mirrors mono.rs:782-787.
    if let Some(im) = env.module_imports.get(&module.name).cloned() {
        env.imports = im;
    }
    // Post-mono permissive seeding. Monomorphisation synthesises
    // `<module>.<def>` references that do NOT obey user-source import
    // discipline — in particular *downward* class-method dispatch: a
    // method mono'd into the class's home module (e.g. prelude's
    // `print__IntBox`) references the *instance's* module
    // (`show_user_adt`) that the home module never imports. The code has
    // already passed `check_workspace`; this re-synth only re-derives
    // types, it does not re-check import discipline, so resolution may
    // reach any workspace module. Seed every module name as an identity
    // import — without overwriting a real author alias — so the canonical
    // `synth`'s qualified-var path (lib.rs ~3556, which otherwise resolves
    // only via TypeDef-home or `env.imports`) resolves these synthesised
    // cross-module references. The canonical `synth` itself stays strict:
    // only this post-mono producer is permissive, mirroring the
    // module-name fallback the now-deleted codegen `synth_arg_type` carried
    // for exactly this case.
    // Skip the current module: a module does not import itself, and
    // seeding a self-import would route its own type-cons through the
    // cross-module qualification path (`show_user_adt.IntBox`),
    // violating the own-module-types-stay-bare invariant the canonical
    // `synth` upholds (`IntBox` stays bare in its home module).
    let all_modules: Vec<String> = env
        .module_globals
        .keys()
        .filter(|m| *m != &module.name)
        .cloned()
        .collect();
    for m in all_modules {
        env.imports.entry(m.clone()).or_insert(m);
    }
    // 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);
    let mut defs = Vec::new();
    let mut consts = Vec::new();
    for def in &module.defs {
        // Non-literal consts are inlined by codegen at each reference
        // site; lower their bodies to typed MTerm so codegen re-derives
        // no type. Literal consts emit a global and need no MIR body.
        if let ailang_core::ast::Def::Const(c) = def {
            if !matches!(c.value, Term::Lit { .. }) {
                let mut ctx = Ctx {
                    env: &env,
                    in_def: &c.name,
                    locals: IndexMap::new(),
                    loop_stack: Vec::new(),
                };
                let body = lower_term(&mut ctx, &c.value)?;
                consts.push(ailang_mir::MirConst {
                    name: c.name.clone(),
                    ty: c.ty.clone(),
                    body,
                });
            }
            continue;
        }
        let ailang_core::ast::Def::Fn(f) = def else { continue };
        // Polymorphic (`Type::Forall`) defs are stale source defs kept
        // for round-trip / `ail diff`; the mono pass synthesised their
        // monomorphic counterparts as separate defs. Codegen skips them
        // (`emit_module`'s `Type::Forall` guard), so producing a MirDef
        // would be dead work — and worse, their bodies may carry
        // mono-rewritten call names (e.g. `fold_left__Unit__Int`) whose
        // specialisations were never synthesised because the poly def
        // itself is never instantiated at those types. The synth
        // re-entry on such a body would hit `UnknownIdentifier`. Skip,
        // mirroring `emit_module`.
        if matches!(&f.ty, Type::Forall { .. }) {
            continue;
        }
        // An `(intrinsic)` body is signature-only — codegen supplies it
        // via the intercept registry, never walking a body. Skip the
        // lower (which would hit synth's `Term::Intrinsic` unreachable
        // guard), exactly as `mono::collect_mono_targets` skips its
        // synth re-entry. No `MirDef` is produced for such a fn.
        if crate::is_intrinsic_body(f) {
            continue;
        }
        // param types from the fn signature (same source mono uses).
        let param_tys = crate::fn_param_types(&f.ty);
        let mut locals: IndexMap<String, Type> = IndexMap::new();
        for (n, t) in f.params.iter().zip(param_tys.iter()) {
            locals.insert(n.clone(), t.clone());
        }
        let mut ctx = Ctx {
            env: &env,
            in_def: &f.name,
            locals,
            loop_stack: Vec::new(),
        };
        let body = lower_term(&mut ctx, &f.body)?;
        defs.push(MirDef {
            name: f.name.clone(),
            sig: f.ty.clone(),
            params: f.params.clone(),
            body,
            consume: uniqueness
                .iter()
                .filter(|((d, _), _)| d == &f.name)
                .map(|((_, b), info)| (b.clone(), info.consume_count))
                .collect(),
        });
    }
    Ok(MirModule { name: module.name.clone(), ast: module.clone(), defs, consts })
}