AILang/crates/ailang-codegen/src/dropmono.rs

//! Per-monomorph drop-fn collection for polymorphic ADTs (leg C of
//! the #55-cutover drop-soundness family).
//!
//! Background. Pre-cutover, the per-ADT drop fn was emitted exactly
//! once per `Def::Type`, keyed by the *declared* field types. For a
//! polymorphic ADT (`Box a`, `Pair a b`) a ctor field typed at a
//! type-var `a` failed `synth::llvm_type` and fell back to `ptr`,
//! emitting an `ailang_rc_dec` on it. At a value-type monomorph
//! (`Box Int`) that field is an inline `i64`, so `rc_dec(7)`
//! dereferences the scalar and SIGSEGVs.
//!
//! Fix. Emit one drop fn per (polymorphic-ADT, concrete instantiation)
//! and decide dec-vs-skip on the *substituted* field type. A
//! value-type field (`Int`/`Bool`/`Float`/`Unit`) is skipped (inline
//! scalar, no RC); a heap field is dec'd through its own per-monomorph
//! drop symbol so the cascade composes.
//!
//! Symbol naming. A concrete-monomorphic ADT (`vars` empty: IntList,
//! Ordering) keeps its un-suffixed `drop_<m>_<T>` symbol byte-for-byte
//! (the `ir_snapshot` goldens pin these). A polymorphic ADT
//! instantiation gets a mono suffix mangled the same way
//! `ailang_check::mono::mono_symbol_n` mangles fn symbols
//! (`drop_<m>_Pair__Int_Int`; compound args hash-route to stay
//! bounded). Intrinsic-storage types (RawBuf) are polymorphic but use
//! the flat intrinsic drop path — they are NOT suffixed (their drop is
//! element-type independent and the golden pins `drop_<m>_RawBuf`).
//!
//! This module owns the workspace-global collection
//! ([`collect_drop_monos`]) and the shared suffix decision
//! ([`DropAdtMeta`]); the emission and call-site manglers in `drop.rs`
//! / `match_lower.rs` consult the same [`DropAdtMeta`] so a definition
//! and its callers always agree on the symbol (a mismatch is a link
//! error).

use ailang_core::ast::{Def, Module, Term, Type, TypeDef};
use ailang_mir::{Callee, MArg, MNewArg, MTerm, MirWorkspace};
use std::collections::{BTreeMap, BTreeSet};

use super::subst::{apply_subst_to_type, qualify_local_types_codegen};

/// A canonical key for an ADT: `(owner_module, bare_type_name)`.
pub(crate) type AdtKey = (String, String);

/// Workspace-global drop-monomorphisation metadata, built once before
/// the per-module codegen loop and shared (by reference) with every
/// `Emitter`.
#[derive(Debug, Default)]
pub(crate) struct DropAdtMeta {
    /// ADTs that take a per-monomorph drop-symbol suffix: declared
    /// `vars` non-empty AND not intrinsic-storage. A `Type::Con` whose
    /// resolved key is in this set, and whose `args` are non-empty,
    /// gets a `__<suffix>` appended to its `drop_`/`partial_drop_`
    /// symbol.
    suffixed: BTreeSet<AdtKey>,
    /// For each suffixed ADT, the concrete arg-tuples it is
    /// instantiated at across the whole workspace, keyed by a canonical
    /// hash of the arg-tuple (`Type` is not `Ord`, so the tuple itself
    /// cannot key a `BTreeSet`; the hash gives a deterministic,
    /// dedup-stable key). The emission loop reads the values to emit one
    /// drop fn per instantiation in the owning module.
    monos: BTreeMap<AdtKey, BTreeMap<String, Vec<Type>>>,
}

/// Canonical, deterministic key for an arg-tuple (used to dedup
/// instantiations that `Type`'s non-`Ord`-ness blocks from a set key).
fn args_key(args: &[Type]) -> String {
    args.iter()
        .map(ailang_core::canonical::type_hash)
        .collect::<Vec<_>>()
        .join("|")
}

impl DropAdtMeta {
    /// Is this ADT (resolved to `key`) suffixed (poly + non-intrinsic)?
    pub(crate) fn is_suffixed(&self, key: &AdtKey) -> bool {
        self.suffixed.contains(key)
    }

    /// The concrete instantiations of `key` to emit in its owner
    /// module. Empty for monomorphic / intrinsic / never-instantiated
    /// ADTs.
    pub(crate) fn instantiations(&self, key: &AdtKey) -> Vec<Vec<Type>> {
        self.monos
            .get(key)
            .map(|s| s.values().cloned().collect())
            .unwrap_or_default()
    }

    /// The mono symbol suffix for an instantiation `args`, mangled the
    /// same way `mono_symbol_n` mangles fn type-args. `None` when the
    /// type is not suffixed (monomorphic / intrinsic) or has no args —
    /// caller keeps the un-suffixed base symbol. The leading `__`
    /// joiner is included so callers concatenate directly.
    pub(crate) fn suffix_for(&self, key: &AdtKey, args: &[Type]) -> Option<String> {
        if args.is_empty() || !self.is_suffixed(key) {
            return None;
        }
        // Reuse the fn-symbol mangler: `mono_symbol_n("", &args)` would
        // prepend an empty base + joiner. We want only the per-arg
        // suffix joined by `__`, prefixed with the `__` that separates
        // it from the drop base. Mirror `mono_symbol_n`'s join exactly.
        let parts: Vec<String> = args
            .iter()
            .map(ailang_check::mono::type_mono_suffix)
            .collect();
        Some(format!("__{}", parts.join("__")))
    }
}

/// Resolve a (possibly qualified) `Type::Con` name to its canonical
/// `(owner_module, bare_name)` key, using `import_map` for the
/// qualified case and `current_module` for the bare case.
pub(crate) fn resolve_adt_key(
    name: &str,
    current_module: &str,
    import_map: &BTreeMap<String, String>,
) -> AdtKey {
    if name.matches('.').count() == 1 {
        let (prefix, suffix) = name.split_once('.').expect("checked");
        let owner = import_map
            .get(prefix)
            .map(|s| s.as_str())
            .unwrap_or(prefix);
        (owner.to_string(), suffix.to_string())
    } else {
        (current_module.to_string(), name.to_string())
    }
}

/// Does this `Def::Type` use the flat intrinsic-storage drop path
/// (its `new` op is an `(intrinsic)`, not a real term-ctor body)?
/// Mirror of the gate in `lib.rs`'s drop-fn emission loop.
fn is_intrinsic_storage(module: &Module, td: &TypeDef) -> bool {
    module.defs.iter().any(|d| {
        matches!(d, Def::Fn(f)
            if f.name == "new"
                && matches!(f.body, Term::Intrinsic)
                && fn_returns_type_name(f, &td.name))
    })
}

/// Does fn `f` return a `Type::Con` named `tname` (outermost)? Mirror
/// of `lib.rs::fn_returns_type` (kept local to avoid widening that
/// fn's visibility).
fn fn_returns_type_name(f: &ailang_core::ast::FnDef, tname: &str) -> bool {
    fn outer(t: &Type, tname: &str) -> bool {
        match t {
            Type::Con { name, .. } => name == tname,
            Type::Forall { body, .. } => outer(body, tname),
            _ => false,
        }
    }
    if let Type::Fn { ret, .. } = &f.ty {
        outer(ret, tname)
    } else if let Type::Forall { body, .. } = &f.ty {
        if let Type::Fn { ret, .. } = body.as_ref() {
            outer(ret, tname)
        } else {
            false
        }
    } else {
        false
    }
}

/// Build the workspace-global drop-monomorphisation table: which ADTs
/// are suffixed, and which concrete instantiations each is used at.
pub(crate) fn collect_drop_monos(mir: &MirWorkspace) -> DropAdtMeta {
    let mut meta = DropAdtMeta::default();

    // Per-module set of locally-declared ADT names. Used to qualify a
    // bare type-con reference (`List` in `std_list`'s body) to its
    // canonical `module.Type` form *at the module where it appears*, so
    // a stored arg / field resolves to the same owner regardless of
    // which module later emits a drop fn referencing it. Without this,
    // a `Maybe (List Int)` collected in `std_list` would store a bare
    // `List`, and the drop fn emitted in `std_maybe` (Maybe's owner)
    // would resolve that bare `List` to `drop_std_maybe_List` (wrong
    // owner) — an undefined-symbol link error.
    let mut local_types: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();

    // 1. Determine the suffixed set + record each ADT's declared
    //    type-vars and ctor field types (under their owner module),
    //    so the fixpoint below can expand nested field instantiations.
    //    `adt_decl[(m, T)] = (vars, ctor_field_types_flattened)`.
    let mut adt_decl: BTreeMap<AdtKey, (Vec<String>, Vec<Type>)> = BTreeMap::new();
    for (mname, mir_module) in &mir.modules {
        let m = &mir_module.ast;
        for def in &m.defs {
            if let Def::Type(td) = def {
                let key = (mname.clone(), td.name.clone());
                let fields: Vec<Type> =
                    td.ctors.iter().flat_map(|c| c.fields.clone()).collect();
                adt_decl.insert(key.clone(), (td.vars.clone(), fields));
                local_types
                    .entry(mname.clone())
                    .or_default()
                    .insert(td.name.clone());
                if !td.vars.is_empty() && !is_intrinsic_storage(m, td) {
                    meta.suffixed.insert(key);
                }
            }
        }
    }

    // 2. Seed the instantiation set by walking every MIR body / const
    //    and collecting every `Type` that appears. Each walked type is
    //    qualified against its module's local types before collection,
    //    so stored args carry canonical `module.Type` names.
    let mut seeds: Vec<(AdtKey, Vec<Type>)> = Vec::new();
    for (mname, mir_module) in &mir.modules {
        let m = &mir_module.ast;
        let import_map = build_import_map(m);
        let owner_local = local_types.get(mname).cloned().unwrap_or_default();
        let mut sink = |t: &Type| {
            let q = qualify_local_types_codegen(t, mname, &owner_local);
            collect_instantiations(&q, mname, &import_map, &meta.suffixed, &mut seeds);
        };
        for d in &mir_module.defs {
            walk_term_types(&d.body, &mut sink);
        }
        for c in &mir_module.consts {
            walk_term_types(&c.body, &mut sink);
        }
        // Also walk every fn signature type appearing in the AST defs —
        // an Own-returned polymorphic ADT whose ctor is built in a
        // callee still needs its drop emitted in the caller's module
        // when the let-close decs it. The body walk catches the ctor
        // site; the signature walk is belt-and-braces for ret types
        // that surface only through a call.
        for def in &m.defs {
            if let Def::Fn(f) = def {
                walk_type(&f.ty, &mut sink);
            }
        }
    }

    // 3. Fixpoint: an instantiation `T<args>` needs each of its ctor
    //    fields, substituted by {var -> arg}, to also be emitted (a
    //    `Pair (Box Int) Int` field-drops `Box Int`). Expand until no
    //    new (key, args) pairs appear. Dedup on `(AdtKey, args-hash)`
    //    since `Type` is not `Ord`.
    let mut worklist: Vec<(AdtKey, Vec<Type>)> = seeds;
    let mut seen: BTreeSet<(AdtKey, String)> = BTreeSet::new();
    while let Some((key, args)) = worklist.pop() {
        let ak = args_key(&args);
        if !seen.insert((key.clone(), ak.clone())) {
            continue;
        }
        meta.monos
            .entry(key.clone())
            .or_default()
            .insert(ak, args.clone());
        let Some((vars, fields)) = adt_decl.get(&key) else {
            continue;
        };
        if vars.len() != args.len() {
            // Defensive: arity mismatch means our resolution missed;
            // skip rather than mis-substitute.
            continue;
        }
        let subst: BTreeMap<String, Type> =
            vars.iter().cloned().zip(args.iter().cloned()).collect();
        // Field types are written in the owner module's local
        // namespace; resolve against the owner's import map and qualify
        // the owner's own bare type-cons (e.g. a `List a` self-field, or
        // a sibling ADT) so the nested instantiation keys on the
        // canonical owner. The substituted-in args are already
        // qualified from the seed pass.
        let owner = key.0.clone();
        let import_map = mir
            .modules
            .get(&owner)
            .map(|mm| build_import_map(&mm.ast))
            .unwrap_or_default();
        let owner_local = local_types.get(&owner).cloned().unwrap_or_default();
        for f in fields {
            let concrete = apply_subst_to_type(f, &subst);
            let concrete = qualify_local_types_codegen(&concrete, &owner, &owner_local);
            collect_instantiations(
                &concrete,
                &owner,
                &import_map,
                &meta.suffixed,
                &mut worklist,
            );
        }
    }

    meta
}

/// Build a module's import map (alias|name -> actual module), mirroring
/// the per-module loop in `lower_workspace_inner`.
fn build_import_map(m: &Module) -> BTreeMap<String, String> {
    let mut import_map: BTreeMap<String, String> = BTreeMap::new();
    for imp in &m.imports {
        let key = imp.alias.clone().unwrap_or_else(|| imp.module.clone());
        import_map.insert(key, imp.module.clone());
    }
    if m.name != "prelude" {
        import_map
            .entry("prelude".to_string())
            .or_insert_with(|| "prelude".to_string());
    }
    import_map
}

/// Collect, from a single `Type`, every suffixed-ADT instantiation
/// (recursing into args).
fn collect_instantiations(
    t: &Type,
    current_module: &str,
    import_map: &BTreeMap<String, String>,
    suffixed: &BTreeSet<AdtKey>,
    out: &mut Vec<(AdtKey, Vec<Type>)>,
) {
    match t {
        Type::Con { name, args } => {
            if !args.is_empty() {
                let key = resolve_adt_key(name, current_module, import_map);
                if suffixed.contains(&key) {
                    out.push((key, args.clone()));
                }
            }
            for a in args {
                collect_instantiations(a, current_module, import_map, suffixed, out);
            }
        }
        Type::Fn { params, ret, .. } => {
            for p in params {
                collect_instantiations(p, current_module, import_map, suffixed, out);
            }
            collect_instantiations(ret, current_module, import_map, suffixed, out);
        }
        Type::Forall { body, .. } => {
            collect_instantiations(body, current_module, import_map, suffixed, out)
        }
        Type::Var { .. } => {}
    }
}

/// Apply `sink` to every `Type` reachable from a `Type`.
fn walk_type<F: FnMut(&Type)>(t: &Type, sink: &mut F) {
    sink(t);
}

/// Apply `sink` to every `Type` carried by `term` and its sub-terms.
/// `sink` itself recurses into the type structure
/// ([`collect_instantiations`]), so this walker only has to surface
/// every type-bearing node once.
fn walk_term_types<F: FnMut(&Type)>(term: &MTerm, sink: &mut F) {
    // The node's own static type.
    let ty = term.ty();
    sink(&ty);
    match term {
        MTerm::Lit { .. }
        | MTerm::Var { .. }
        | MTerm::Str { .. }
        | MTerm::Intrinsic { .. }
        | MTerm::Recur { .. } => {}
        MTerm::App { callee, args, .. } => {
            walk_callee_types(callee, sink);
            walk_args(args, sink);
        }
        MTerm::Do { args, .. } => walk_args(args, sink),
        MTerm::Ctor { args, .. } => walk_args(args, sink),
        MTerm::New { elem, args, .. } => {
            if let Some(e) = elem {
                sink(e);
            }
            for a in args {
                if let MNewArg::Value(v) = a {
                    walk_term_types(v, sink);
                } else if let MNewArg::Type(t) = a {
                    sink(t);
                }
            }
        }
        MTerm::Let { init, body, .. } => {
            walk_term_types(init, sink);
            walk_term_types(body, sink);
        }
        MTerm::LetRec { sig, body, in_term, .. } => {
            sink(sig);
            walk_term_types(body, sink);
            walk_term_types(in_term, sink);
        }
        MTerm::If { cond, then, else_, .. } => {
            walk_term_types(cond, sink);
            walk_term_types(then, sink);
            walk_term_types(else_, sink);
        }
        MTerm::Match { scrutinee, arms, .. } => {
            walk_term_types(scrutinee, sink);
            for arm in arms {
                walk_term_types(&arm.body, sink);
            }
        }
        MTerm::Lam { param_tys, ret_ty, body, .. } => {
            for p in param_tys {
                sink(p);
            }
            sink(ret_ty);
            walk_term_types(body, sink);
        }
        MTerm::Seq { lhs, rhs, .. } => {
            walk_term_types(lhs, sink);
            walk_term_types(rhs, sink);
        }
        MTerm::Clone { value, .. } => walk_term_types(value, sink),
        MTerm::ReuseAs { source, body, .. } => {
            walk_term_types(source, sink);
            walk_term_types(body, sink);
        }
        MTerm::Loop { binders, body, .. } => {
            for b in binders {
                sink(&b.ty);
                walk_term_types(&b.init, sink);
            }
            walk_term_types(body, sink);
        }
    }
}

fn walk_args<F: FnMut(&Type)>(args: &[MArg], sink: &mut F) {
    for a in args {
        walk_term_types(&a.term, sink);
    }
}

fn walk_callee_types<F: FnMut(&Type)>(callee: &Callee, sink: &mut F) {
    match callee {
        Callee::Static { sig, .. } | Callee::Builtin { sig, .. } => sink(sig),
        Callee::Indirect(inner) => walk_term_types(inner, sink),
    }
}