AILang/crates/ailang-mir/src/lib.rs

//! Typed mid-level IR — the single artefact that crosses the
//! `check` → `codegen` boundary. Produced by
//! `ailang_check::lower_to_mir` from a monomorphised, typechecked
//! `Workspace`; consumed by `ailang_codegen`. Every node carries the
//! facts codegen used to re-derive: its type (`ty`), the resolved
//! callee (`Callee`, filled mir.2), parameter modes / consume counts
//! (`Mode` / `consume_count`, filled mir.3), and string representation
//! (`StrRep`, filled mir.4). MIR is never authored, never hashed,
//! never round-tripped — it is a compiler-internal derived form.

use ailang_core::ast::{Literal, Module, Pattern, Type};
use std::collections::BTreeMap;

/// A whole monomorphised program, lowered to MIR.
#[derive(Debug, Clone)]
pub struct MirWorkspace {
    pub entry: String,
    pub modules: BTreeMap<String, MirModule>,
}

#[derive(Debug, Clone)]
pub struct MirModule {
    pub name: String,
    /// The post-mono AST module. codegen reads structural data
    /// (Def::Type / Def::Const / ctor layouts / intrinsic-fn markers /
    /// imports) from here; the typed fn bodies come from `defs`. This
    /// keeps `MirWorkspace` the single artefact codegen consumes —
    /// structural reads were never a re-derivation, so carrying the AST
    /// alongside the typed bodies does not reintroduce one.
    pub ast: Module,
    pub defs: Vec<MirDef>,
    /// Typed bodies of the module's non-literal `Def::Const`s. A const
    /// is inlined at every `Var` reference site (codegen reads the
    /// body from here, keyed by name); literal consts emit a global
    /// instead and are not carried. Same role as `defs` for fns — the
    /// body is the checker-typed `MTerm` so codegen re-derives no
    /// type to lower it.
    pub consts: Vec<MirConst>,
}

/// One lowered non-literal const. `body` is the typed `MTerm` codegen
/// inlines at each reference site.
#[derive(Debug, Clone)]
pub struct MirConst {
    pub name: String,
    pub ty: Type,
    pub body: MTerm,
}

/// 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>,
}

/// One MIR node — structural counterpart of `ast::Term`, plus the
/// `Str` split. `ty` is the checker-proved type at this node.
#[derive(Debug, Clone)]
pub enum MTerm {
    Lit { lit: Literal, ty: Type },
    /// String literal — the `Str`-representation carrier. `rep` is
    /// `Static` for every literal at mir.1; mir.4 flips loop-carried
    /// seeds to `Heap`. Split out of `Lit` so the rep annotation has a
    /// home (spec §"Str split").
    Str { lit: String, rep: StrRep },
    Var { name: String, ty: Type },
    /// `callee` is `Indirect(<lowered callee>)` at mir.1; mir.2
    /// resolves static targets to `Callee::Static`.
    App { callee: Callee, args: Vec<MArg>, tail: bool, ty: Type },
    Let { name: String, mode: Mode, init: Box<MTerm>, body: Box<MTerm>, ty: Type },
    LetRec {
        name: String,
        sig: Type,
        params: Vec<String>,
        body: Box<MTerm>,
        in_term: Box<MTerm>,
        ty: Type,
    },
    If {
        cond: Box<MTerm>,
        then: Box<MTerm>,
        else_: Box<MTerm>,
        /// #63 leg 3: per-branch consume of each owned binder on the
        /// `then` / `else` branch (absolute count at the branch end).
        /// Codegen drops a param live here (`0`) but consumed on the
        /// sibling (`MirDef.consume >= 1`). Empty for const bodies.
        then_consume: BTreeMap<String, u32>,
        else_consume: BTreeMap<String, u32>,
        ty: Type,
    },
    Do { op: String, args: Vec<MArg>, tail: bool, ty: Type },
    Ctor { type_name: String, ctor: String, args: Vec<MArg>, ty: Type },
    Match { scrutinee: Box<MTerm>, arms: Vec<MArm>, ty: Type },
    Lam {
        params: Vec<String>,
        param_tys: Vec<Type>,
        ret_ty: Type,
        effects: Vec<String>,
        body: Box<MTerm>,
        ty: Type,
    },
    Seq { lhs: Box<MTerm>, rhs: Box<MTerm>, ty: Type },
    Clone { value: Box<MTerm>, ty: Type },
    ReuseAs { source: Box<MTerm>, body: Box<MTerm>, ty: Type },
    Loop { binders: Vec<MLoopBinder>, body: Box<MTerm>, ty: Type },
    Recur { args: Vec<MArg>, ty: Type },
    /// `elem` is `None` at mir.1; mir.5 carries the element type.
    New { type_name: String, elem: Option<Type>, args: Vec<MNewArg>, ty: Type },
    Intrinsic { ty: Type },
}

/// A match arm. Patterns stay structural (`ast::Pattern`) — they carry
/// no type annotation the boundary re-derives.
#[derive(Debug, Clone)]
pub struct MArm {
    pub pat: Pattern,
    pub body: MTerm,
    /// #63 leg 3: per-arm consume of each owned binder (absolute count at
    /// this arm's end). See `MTerm::If.then_consume`. Empty for const bodies.
    pub consume: BTreeMap<String, u32>,
}

#[derive(Debug, Clone)]
pub struct MLoopBinder {
    pub name: String,
    pub ty: Type,
    pub init: MTerm,
}

#[derive(Debug, Clone)]
pub enum MNewArg {
    Type(Type),
    Value(MTerm),
}

/// A call/ctor/recur argument with its mode and consume count. Both
/// fields are mir.1 defaults (`Owned` / `1`); mir.3 fills them.
#[derive(Debug, Clone)]
pub struct MArg {
    pub term: MTerm,
    pub mode: Mode,
    pub consume_count: u32,
}

/// A resolved App callee. `Static`/`Builtin` are filled by mir.2's
/// `lower_to_mir::classify_callee`, which mirrors check's `synth`
/// `Term::Var` ladder; `Indirect` is a genuinely dynamic callee
/// (fn-typed local/param, closure value, or a name shadowed by a
/// local). `sig` is the callee's fn-type — the lossless replacement
/// for the sub-term's `ty()` that codegen's drop path reads for the
/// callee `ret_mode`.
#[derive(Debug, Clone)]
pub enum Callee {
    /// A statically-resolved user/prelude fn: codegen routes it to
    /// `emit_call(module, fn_name, …)` with `module` pre-resolved
    /// (replacing codegen's `type_home_module` ladder).
    Static { module: String, fn_name: String, sig: Type },
    /// An operator / intrinsic builtin (`+`, `not`, `str_concat`, …):
    /// codegen lowers it inline by `name` (opcode selection), never
    /// via `emit_call`.
    Builtin { name: String, sig: Type },
    /// A dynamic callee — lower the boxed term to a fn-pointer.
    Indirect(Box<MTerm>),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StrRep {
    Heap,
    Static,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Mode {
    Owned,
    Borrow,
}

impl MTerm {
    /// The checker-proved type at this node. `Str` is always `Str`.
    pub fn ty(&self) -> Type {
        match self {
            MTerm::Lit { ty, .. }
            | MTerm::Var { ty, .. }
            | MTerm::App { ty, .. }
            | MTerm::Let { ty, .. }
            | MTerm::LetRec { ty, .. }
            | MTerm::If { ty, .. }
            | MTerm::Do { ty, .. }
            | MTerm::Ctor { ty, .. }
            | MTerm::Match { ty, .. }
            | MTerm::Lam { ty, .. }
            | MTerm::Seq { ty, .. }
            | MTerm::Clone { ty, .. }
            | MTerm::ReuseAs { ty, .. }
            | MTerm::Loop { ty, .. }
            | MTerm::Recur { ty, .. }
            | MTerm::New { ty, .. }
            | MTerm::Intrinsic { ty } => ty.clone(),
            MTerm::Str { .. } => Type::str_(),
        }
    }
}