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

//! Uniqueness inference over the typed AST.
//!
//! This pass classifies every binder of every fn body as either
//! [`Uniqueness::Unique`] or [`Uniqueness::Shared`]. Unlike the
//! `crate::linearity` check it is **not** a user-facing diagnostic
//! pass — it is internal codegen input. Codegen consults the side
//! table to decide where to emit `ailang_rc_inc` and `ailang_rc_dec`
//! calls under `--alloc=rc` (Iter 18c.3).
//!
//! ## Why a separate pass from `linearity`
//!
//! - Different scope. `linearity` is opt-in by signature (every param
//!   mode must be explicit) and is emitted as user diagnostics.
//!   Uniqueness inference runs unconditionally on every fn body and
//!   is never reported.
//! - Different output. `linearity` emits `Diagnostic`s; uniqueness
//!   produces a [`UniquenessTable`] keyed by `(def_name, binder_name)`.
//! - Different consumption point. `linearity` runs from
//!   [`crate::check_workspace`]; the uniqueness side table is built
//!   on demand from codegen via [`infer_module`].
//!
//! ## What "unique" means here
//!
//! A binder `x` is `Unique` if every reachable path from its
//! introduction to the end of its scope passes through **at most one**
//! consume use of `x`. Borrow-position uses (Match scrutinee,
//! `(clone)` argument, `Borrow`-mode call argument) do not count as
//! consumes. A binder is `Shared` otherwise.
//!
//! For the codegen consumer, a closely related side-output is more
//! useful: [`UniquenessInfo::consume_count`] — the maximum, across
//! all reachable paths in the binder's scope, of the number of
//! consume uses of the binder. `consume_count == 0` means the binder
//! is exclusively borrowed (or unused) inside its scope, which is
//! the case where emitting `dec` at scope close is unambiguously
//! safe under `--alloc=rc` (no callee / outer term consumed the
//! value first; the binder owns the only outstanding ref to its
//! allocation).
//!
//! ## Position model
//!
//! Identical to `crate::linearity`'s, modulo that this pass has no
//! activation gate (it runs on every fn). See that module's docs for
//! the propagation rules. `Term::Lam` captures are conservatively
//! treated as Consume — the same conservative decision the linearity
//! check makes.
//!
//! ## What this pass does NOT do
//!
//! - **Cross-fn reasoning.** The inference is whole-fn local. It does
//!   not reason about a callee's body when classifying a binder passed
//!   into the callee. Cross-fn ownership currently flows through
//!   explicit mode annotations on the callee's signature (18a).
//! - **Per-type drop fns and recursive `dec` cascades.** Those are
//!   emitted by codegen (`emit_drop_fn_for_type`), not consulted from
//!   this side table. The inference's role is "should codegen emit a
//!   dec on this binder leaving scope at all?"; the *shape* of the
//!   dec (uniform `drop_<m>_<T>` call vs inlined partial drop) is a
//!   codegen-side decision driven by `moved_slots`.

use ailang_core::ast::{Arm, Def, FnDef, Module, NewArg, ParamMode, Pattern, Term, Type};
use std::collections::BTreeMap;

/// Per-binder classification produced by the inference. See module
/// docs for semantics.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Uniqueness {
    /// Every reachable path from the binder's introduction to the end
    /// of its scope passes through at most one consume use.
    Unique,
    /// Some reachable path consumes the binder more than once.
    Shared,
}

/// Per-binder result of the inference. The codegen consumer uses
/// both fields: `uniqueness` for the high-level classification and
/// `consume_count` for the precise "is anyone else going to dec
/// this" answer at scope close.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct UniquenessInfo {
    pub uniqueness: Uniqueness,
    /// Maximum number of consume-position uses of this binder
    /// across any reachable path in its scope. `0` means the binder
    /// is only borrowed (or unused) — under `--alloc=rc` the only
    /// outstanding reference at scope close is the one the binder
    /// itself holds, so dec at scope close is safe and necessary.
    pub consume_count: u32,
}

/// Side table indexed by `(def_name, binder_name)`. Built fresh for
/// each codegen invocation; the side table is not cached.
///
/// The key is injective per fn: desugaring alpha-renames any binder
/// whose name shadows an enclosing binding to a fresh `<name>$<n>`
/// (see `ailang-core::desugar`), so within one `def_name` every
/// `binder_name` denotes exactly one binding. Codegen consumers look
/// up by current-scope binder name and always resolve the binding
/// they mean — there is no shadow collapse.
pub type UniquenessTable = BTreeMap<(String, String), UniquenessInfo>;

/// Top-level entry: walk every fn def in `m` and produce per-binder
/// classifications keyed by `(fn_name, binder_name)`.
///
/// Pre-condition: `m` is fully typechecked. Inputs that have not
/// been type-checked (unresolved `Var` lookups, mismatched ctor
/// arities) may produce a noisy / incorrect side table — the
/// codegen never runs on such inputs because `ail build` runs
/// typecheck first.
pub fn infer_module(m: &Module) -> UniquenessTable {
    infer_module_with_cross(m, &BTreeMap::new())
}

/// Same as [`infer_module`], but with access to the workspace-flat
/// cross-module signature table (`module -> def -> Type`, the shape
/// codegen builds as `module_def_ail_types`). The per-module pass
/// only registers same-module + builtin signatures in `globals`; a
/// call to a *cross-module* callee (a qualified `<mod>.<def>` name,
/// e.g. the monomorphised intrinsic `raw_buf.RawBuf_get__Int` reached
/// from a `main` that lives in a different module) misses the
/// per-module `globals` and so previously defaulted every arg to
/// `Position::Consume` — mis-counting a `Borrow`-mode arg as a
/// consume and inflating the binder's `consume_count`, which gates
/// off the scope-close `dec` and leaks the slab. Threading the
/// cross-module table lets `callee_arg_modes` resolve the qualified
/// name positionally to its registered `Type::Fn.param_modes` and
/// walk borrow args as `Position::Borrow`.
///
/// This generalises to any cross-module borrow callee, not just
/// RawBuf — it removes a latent over-conservative consume default.
/// The same-module `globals` lookup stays the first resort; the
/// cross-module table is consulted only on a miss.
pub fn infer_module_with_cross(
    m: &Module,
    cross_module_types: &BTreeMap<String, BTreeMap<String, Type>>,
) -> UniquenessTable {
    let mut globals: BTreeMap<String, Type> = BTreeMap::new();
    // register builtin signatures so the App-arg walker
    // sees their `param_modes` and walks `Borrow`-mode args as
    // `Position::Borrow` (not `Consume`). Without this, a builtin
    // like `str_clone : (Str borrow) -> Str` invoked on a heap-Str
    // let-binder would falsely classify the binder as Consumed,
    // gating off the scope-close `ailang_rc_dec` and leaking the
    // slab. Symmetric to the class-method registration the
    // linearity pass added in iter 23.4-prep.
    let mut builtins_env = crate::Env::default();
    crate::builtins::install(&mut builtins_env);
    for (name, ty) in &builtins_env.globals {
        globals.insert(name.clone(), strip_forall(ty).clone());
    }
    for def in &m.defs {
        match def {
            Def::Fn(f) => {
                let ty = strip_forall(&f.ty).clone();
                globals.insert(f.name.clone(), ty);
            }
            Def::Const(c) => {
                globals.insert(c.name.clone(), strip_forall(&c.ty).clone());
            }
            Def::Type(_) => {}
            // class/instance defs are skipped here. Class
            // method signatures contribute global names once
            // monomorphisation (22b.3) materialises them; the
            // pre-monomorphisation form has no concrete `Type` to
            // hand to the uniqueness pass.
            Def::Class(_) | Def::Instance(_) => {}
        }
    }

    let mut table: UniquenessTable = BTreeMap::new();
    for def in &m.defs {
        if let Def::Fn(f) = def {
            infer_fn(f, &globals, cross_module_types, &mut table);
        }
    }
    table
}

fn strip_forall(t: &Type) -> &Type {
    match t {
        Type::Forall { body, .. } => body,
        other => other,
    }
}

/// Per-fn driver: install the parameters as binders, walk the body,
/// snapshot the per-binder consume counts into the side table.
fn infer_fn(
    f: &FnDef,
    globals: &BTreeMap<String, Type>,
    cross_module_types: &BTreeMap<String, BTreeMap<String, Type>>,
    out: &mut UniquenessTable,
) {
    let inner = strip_forall(&f.ty);
    if !matches!(inner, Type::Fn { .. }) {
        return;
    }

    let param_states = {
        let mut walker = Walker {
            globals,
            cross_module_types,
            binders: BTreeMap::new(),
            def_name: f.name.clone(),
            out,
        };

        // Install fn parameters. Their consume counter starts at 0;
        // the walk over the body increments it.
        for name in f.params.iter() {
            walker.binders.insert(name.clone(), BinderState::default());
        }

        walker.walk(&f.body, Position::Consume);

        // Snapshot params before dropping the walker.
        f.params
            .iter()
            .filter_map(|name| walker.binders.get(name).map(|s| (name.clone(), s.clone())))
            .collect::<Vec<_>>()
    };

    // Walker dropped — `out` is borrowable again.
    for (name, s) in param_states {
        let info = UniquenessInfo {
            uniqueness: classify(s.consume_count),
            consume_count: s.consume_count,
        };
        out.insert((f.name.clone(), name), info);
    }
}

#[inline]
fn classify(consume_count: u32) -> Uniqueness {
    if consume_count <= 1 {
        Uniqueness::Unique
    } else {
        Uniqueness::Shared
    }
}

/// Position in which a [`Term::Var`] occurs. Consume uses count
/// against the binder's `consume_count`; Borrow uses do not.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum Position {
    Consume,
    Borrow,
}

/// Per-binder mutable state during the walk. Future extensions
/// (e.g. last-use markers, per-arm split) can grow more fields.
#[derive(Default, Clone, Debug)]
struct BinderState {
    consume_count: u32,
}

struct Walker<'a> {
    globals: &'a BTreeMap<String, Type>,
    /// Workspace-flat `module -> def -> Type` table. Consulted only
    /// when `globals` (same-module + builtins) misses on a qualified
    /// cross-module callee. See [`infer_module_with_cross`].
    cross_module_types: &'a BTreeMap<String, BTreeMap<String, Type>>,
    binders: BTreeMap<String, BinderState>,
    def_name: String,
    out: &'a mut UniquenessTable,
}

impl<'a> Walker<'a> {
    fn walk(&mut self, t: &Term, pos: Position) {
        match t {
            Term::Lit { .. } => {}
            Term::Var { name } => self.use_var(name, pos),
            Term::App { callee, args, .. } => {
                self.walk(callee, Position::Consume);
                let arg_modes = self.callee_arg_modes(callee, args.len());
                for (i, arg) in args.iter().enumerate() {
                    let arg_pos = match arg_modes.get(i) {
                        Some(ParamMode::Borrow) => Position::Borrow,
                        _ => Position::Consume,
                    };
                    self.walk(arg, arg_pos);
                }
            }
            Term::Let { name, value, body } => {
                self.walk(value, Position::Consume);
                let prev = self.binders.insert(name.clone(), BinderState::default());
                self.walk(body, pos);
                let final_state = self.binders.remove(name).unwrap_or_default();
                if let Some(p) = prev {
                    self.binders.insert(name.clone(), p);
                }
                self.record_binder(name, &final_state);
            }
            Term::LetRec { name, body, in_term, .. } => {
                let prev = self.binders.insert(name.clone(), BinderState::default());
                self.walk(body, Position::Consume);
                self.walk(in_term, pos);
                let final_state = self.binders.remove(name).unwrap_or_default();
                if let Some(p) = prev {
                    self.binders.insert(name.clone(), p);
                }
                self.record_binder(name, &final_state);
            }
            Term::If { cond, then, else_ } => {
                self.walk(cond, Position::Consume);
                let saved = self.binders.clone();
                self.walk(then, pos);
                let after_then = std::mem::replace(&mut self.binders, saved);
                self.walk(else_, pos);
                merge_states(&mut self.binders, &after_then);
            }
            Term::Match { scrutinee, arms } => {
                self.walk(scrutinee, Position::Borrow);
                let saved = self.binders.clone();
                let mut acc: Option<BTreeMap<String, BinderState>> = None;
                for arm in arms {
                    self.binders = saved.clone();
                    self.walk_arm(arm, pos);
                    match acc.as_mut() {
                        None => acc = Some(self.binders.clone()),
                        Some(m) => merge_states(m, &self.binders),
                    }
                }
                if let Some(merged) = acc {
                    self.binders = merged;
                } else {
                    self.binders = saved;
                }
            }
            Term::Seq { lhs, rhs } => {
                self.walk(lhs, Position::Consume);
                self.walk(rhs, pos);
            }
            Term::Ctor { args, .. } => {
                for a in args {
                    self.walk(a, Position::Consume);
                }
            }
            Term::Do { args, .. } => {
                for a in args {
                    self.walk(a, Position::Borrow);
                }
            }
            Term::Lam { params, body, .. } => {
                // Lam captures cross the linearity boundary — every
                // free var of the body is conservatively a consume.
                // Walking the body with each param as a fresh binder
                // produces the right intra-Lam classification.
                let mut saved_for_params: BTreeMap<String, Option<BinderState>> = BTreeMap::new();
                for p in params {
                    saved_for_params.insert(p.clone(), self.binders.remove(p));
                    self.binders.insert(p.clone(), BinderState::default());
                }
                self.walk(body, Position::Consume);
                for p in params {
                    let final_state = self.binders.remove(p).unwrap_or_default();
                    if let Some(prev) = saved_for_params.remove(p).flatten() {
                        self.binders.insert(p.clone(), prev);
                    }
                    self.record_binder(p, &final_state);
                }
            }
            Term::Clone { value } => {
                self.walk(value, Position::Borrow);
            }
            Term::ReuseAs { source, body } => {
                // `(reuse-as SRC BODY)` is a Consume of
                // `SRC` (the binder's slot is being taken over) and
                // a normal walk of `BODY`. So a bare-Var SRC bumps
                // its `consume_count` by one. Linearity is the
                // user-visible enforcement; this pass only feeds the
                // RC codegen its bookkeeping.
                self.walk(source, Position::Consume);
                self.walk(body, pos);
            }
            Term::Loop { binders, body } => {
                for b in binders {
                    self.walk(&b.init, Position::Consume);
                }
                self.walk(body, pos);
            }
            Term::Recur { args } => {
                for a in args {
                    self.walk(a, Position::Consume);
                }
            }
            // prep.2 (kernel-extension-mechanics): each NewArg::Value
            // is consumed by the construction call — same rule as a
            // Ctor or App arg in `Position::Consume`. NewArg::Type
            // carries no runtime term.
            Term::New { args, .. } => {
                for arg in args {
                    if let NewArg::Value(v) = arg {
                        self.walk(v, Position::Consume);
                    }
                }
            }
            Term::Intrinsic => {}
        }
    }

    fn walk_arm(&mut self, arm: &Arm, pos: Position) {
        let names = collect_pattern_binders(&arm.pat);
        let mut saved: BTreeMap<String, Option<BinderState>> = BTreeMap::new();
        for n in &names {
            saved.insert(n.clone(), self.binders.remove(n));
            self.binders.insert(n.clone(), BinderState::default());
        }
        self.walk(&arm.body, pos);
        for n in &names {
            let final_state = self.binders.remove(n).unwrap_or_default();
            if let Some(prev) = saved.remove(n).flatten() {
                self.binders.insert(n.clone(), prev);
            }
            self.record_binder(n, &final_state);
        }
    }

    fn use_var(&mut self, name: &str, pos: Position) {
        if let Some(s) = self.binders.get_mut(name) {
            if matches!(pos, Position::Consume) {
                s.consume_count = s.consume_count.saturating_add(1);
            }
        }
    }

    fn callee_arg_modes(&self, callee: &Term, n_args: usize) -> Vec<ParamMode> {
        let name = match callee {
            Term::Var { name } => name,
            _ => return vec![],
        };
        // First resort: same-module + builtin signatures.
        let ty = match self.globals.get(name) {
            Some(t) => t,
            // Fallback: a qualified `<mod>.<def>` callee whose
            // signature lives in a *different* module's defs (e.g. a
            // monomorphised intrinsic like `raw_buf.RawBuf_get__Int`
            // reached from a `main` in another module). Split on the
            // last `.` and resolve positionally against the
            // workspace-flat cross-module table. Without this the
            // qualified borrow callee would miss and every arg would
            // default to `Position::Consume`, mis-counting the borrow
            // arg and leaking the slab.
            None => match self.cross_module_callee_ty(name) {
                Some(t) => t,
                None => return vec![],
            },
        };
        let inner = strip_forall(ty);
        match inner {
            Type::Fn { param_modes, .. } => {
                if param_modes.is_empty() {
                    vec![ParamMode::Implicit; n_args]
                } else {
                    param_modes.clone()
                }
            }
            _ => vec![],
        }
    }

    /// Resolve a qualified `<mod>.<def>` callee name against the
    /// cross-module signature table. Splits on the last `.` so a def
    /// name carrying no dot (the monomorphic intrinsic symbols do not)
    /// resolves cleanly. Returns `None` for unqualified names or names
    /// whose module/def is absent from the table — the caller then
    /// falls back to the existing `vec![]` (all-consume) default.
    fn cross_module_callee_ty(&self, name: &str) -> Option<&'a Type> {
        let (module, def) = name.rsplit_once('.')?;
        self.cross_module_types.get(module)?.get(def)
    }

    fn record_binder(&mut self, name: &str, s: &BinderState) {
        let info = UniquenessInfo {
            uniqueness: classify(s.consume_count),
            consume_count: s.consume_count,
        };
        self.out
            .insert((self.def_name.clone(), name.to_string()), info);
    }
}

/// Conservative branch merge: for each binder name, the merged
/// `consume_count` is the **max** of the two branches. Max is the
/// right operator because the classification cares about the
/// worst-case path.
fn merge_states(into: &mut BTreeMap<String, BinderState>, other: &BTreeMap<String, BinderState>) {
    for (name, s) in other {
        let entry = into.entry(name.clone()).or_default();
        entry.consume_count = entry.consume_count.max(s.consume_count);
    }
}

fn collect_pattern_binders(p: &Pattern) -> Vec<String> {
    match p {
        Pattern::Wild | Pattern::Lit { .. } => vec![],
        Pattern::Var { name } => vec![name.clone()],
        Pattern::Ctor { fields, .. } => fields.iter().flat_map(collect_pattern_binders).collect(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ailang_core::ast::{FnDef, Literal};

    fn fn_simple(name: &str, params: Vec<&str>, body: Term) -> Def {
        Def::Fn(FnDef {
            name: name.into(),
            ty: Type::Fn {
                params: params.iter().map(|_| Type::int()).collect(),
                param_modes: vec![],
                ret: Box::new(Type::int()),
                ret_mode: ParamMode::Implicit,
                effects: vec![],
            },
            params: params.into_iter().map(|s| s.to_string()).collect(),
            body,
            suppress: vec![],
            doc: None,
        export: None,
        })
    }

    /// `(let x 0 0)` — `x` is unique, consume_count == 0.
    #[test]
    fn let_unused_binder_is_unique_zero_consume() {
        let body = Term::Let {
            name: "x".into(),
            value: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
            body: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
        };
        let m = Module {
            schema: ailang_core::SCHEMA.into(),
            name: "t".into(),
            kernel: false,
            imports: vec![],
            defs: vec![fn_simple("f", vec![], body)],
        };
        let table = infer_module(&m);
        let info = table.get(&("f".into(), "x".into())).copied().expect("x recorded");
        assert_eq!(info.uniqueness, Uniqueness::Unique);
        assert_eq!(info.consume_count, 0);
    }

    /// `(let x 0 x)` — `x` consumed once; classified Unique with
    /// consume_count == 1.
    #[test]
    fn let_singly_consumed_is_unique_one() {
        let body = Term::Let {
            name: "x".into(),
            value: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
            body: Box::new(Term::Var { name: "x".into() }),
        };
        let m = Module {
            schema: ailang_core::SCHEMA.into(),
            name: "t".into(),
            kernel: false,
            imports: vec![],
            defs: vec![fn_simple("f", vec![], body)],
        };
        let table = infer_module(&m);
        let info = table.get(&("f".into(), "x".into())).copied().expect("x recorded");
        assert_eq!(info.uniqueness, Uniqueness::Unique);
        assert_eq!(info.consume_count, 1);
    }

    /// `(let x 0 (let y 0 (app + x x)))` — `x` consumed twice in
    /// distinct positions of the same call → Shared, consume_count
    /// == 2.
    #[test]
    fn let_doubly_consumed_is_shared_two() {
        let body = Term::Let {
            name: "x".into(),
            value: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
            body: Box::new(Term::App {
                callee: Box::new(Term::Var { name: "+".into() }),
                args: vec![
                    Term::Var { name: "x".into() },
                    Term::Var { name: "x".into() },
                ],
                tail: false,
            }),
        };
        let m = Module {
            schema: ailang_core::SCHEMA.into(),
            name: "f".into(),
            kernel: false,
            imports: vec![],
            defs: vec![fn_simple("f", vec![], body)],
        };
        let table = infer_module(&m);
        let info = table.get(&("f".into(), "x".into())).copied().expect("x recorded");
        assert_eq!(info.uniqueness, Uniqueness::Shared);
        assert_eq!(info.consume_count, 2);
    }

    /// `(match scr ((pat-var x) x))` — `x` (a pattern binder) is
    /// consumed once → Unique with consume_count == 1.
    #[test]
    fn pattern_binder_is_recorded() {
        let body = Term::Match {
            scrutinee: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
            arms: vec![Arm {
                pat: Pattern::Var { name: "x".into() },
                body: Term::Var { name: "x".into() },
            }],
        };
        let m = Module {
            schema: ailang_core::SCHEMA.into(),
            name: "t".into(),
            kernel: false,
            imports: vec![],
            defs: vec![fn_simple("f", vec![], body)],
        };
        let table = infer_module(&m);
        let info = table.get(&("f".into(), "x".into())).copied().expect("x recorded");
        assert_eq!(info.uniqueness, Uniqueness::Unique);
        assert_eq!(info.consume_count, 1);
    }

    /// `(clone x)` is a Borrow — does NOT increment consume_count.
    /// Body: `(let x 0 (clone x))` → x consume_count == 0.
    #[test]
    fn clone_is_borrow_not_consume() {
        let body = Term::Let {
            name: "x".into(),
            value: Box::new(Term::Lit { lit: Literal::Int { value: 0 } }),
            body: Box::new(Term::Clone {
                value: Box::new(Term::Var { name: "x".into() }),
            }),
        };
        let m = Module {
            schema: ailang_core::SCHEMA.into(),
            name: "t".into(),
            kernel: false,
            imports: vec![],
            defs: vec![fn_simple("f", vec![], body)],
        };
        let table = infer_module(&m);
        let info = table.get(&("f".into(), "x".into())).copied().expect("x recorded");
        assert_eq!(info.uniqueness, Uniqueness::Unique);
        assert_eq!(info.consume_count, 0);
    }
}