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

//! Type substitution + unification helpers for monomorphisation.
//!
//! Free functions extracted from `lib.rs` during the 18g tidy split.
//! `unify_for_subst` walks declared params against actual arg types and
//! binds `Type::Var`s; `apply_subst_to_type` specialises a polymorphic
//! type under that binding; `qualify_local_types_codegen` rewrites bare
//! ADT names into `module.Type` form when a sig crosses an import
//! boundary. (The mir.1b switch deleted the codegen-side type-mirror,
//! the only caller of the former `derive_substitution` entry; its
//! callers in `match_lower` build the substitution directly via
//! `unify_for_subst`.)

use ailang_core::ast::*;
use std::collections::{BTreeMap, BTreeSet};

use super::{CodegenError, Result};

/// Walks `param` and `arg` in parallel, treating any `Type::Var { name }`
/// on the param side whose name is in `vars` as an unknown to be bound
/// in `subst`. Identical concrete shapes pass through; structural
/// mismatches yield an internal error.
pub(crate) fn unify_for_subst(
    param: &Type,
    arg: &Type,
    vars: &BTreeSet<&str>,
    subst: &mut BTreeMap<String, Type>,
) -> Result<()> {
    // A `$u`-prefixed var is a
    // synth-only wildcard the checker emits for nullary
    // ctors of a parameterised ADT (e.g. `Nil : List<$u>`). It
    // carries no real constraint — accept without binding so a
    // sibling arg can pin the type var instead. Without this,
    // `Cons(Int, Nil)` synth would unify `a = Int` (from head) and
    // then `a = $u` (from tail's recursive `List<a>` slot) and
    // falsely error.
    //
    // 15g-aux: the early-return must accept `$u` on **either** side.
    // `$u` enters in arg position from synth, but the prev-binding
    // recursion below (`unify_for_subst(&prev, arg, ...)`) can swap
    // a `$u` onto the param side when a previously-bound type is
    // unified against a fresher arg whose roles differ. Reduced
    // repro: `length [Left 1, Right 10]` — `a` first binds to
    // `Either<Int, $u>` from `Left 1`, then a recursive unification
    // against `Either<$u, Int>` from `Right 10` lands `$u` in the
    // param-pos[1] slot. Symmetric early-return is correct because
    // `$u` is a synth-only wildcard regardless of which side carries
    // it after the prev-binding swap.
    if let Type::Var { name } = arg {
        if name.starts_with("$u") {
            return Ok(());
        }
    }
    if let Type::Var { name } = param {
        if name.starts_with("$u") {
            return Ok(());
        }
    }
    match (param, arg) {
        (Type::Var { name }, _) if vars.contains(name.as_str()) => {
            if let Some(prev) = subst.get(name).cloned() {
                // the previously-bound type may be more
                // concrete than `arg` (e.g. `prev = List<Int>` from a
                // sibling binding, `arg = List<$u>` from a synth-
                // wildcard nullary ctor). Use recursive unification
                // instead of strict equality so the inner `$u`
                // wildcard matches `Int`. The previous strict-
                // equality check rejected such overlaps as bogus
                // duplicate bindings.
                return unify_for_subst(&prev, arg, vars, subst);
            }
            subst.insert(name.clone(), arg.clone());
            Ok(())
        }
        (
            Type::Con { name: pn, args: pa },
            Type::Con { name: an, args: aa },
        ) if pn == an && pa.len() == aa.len() => {
            for (p, a) in pa.iter().zip(aa.iter()) {
                unify_for_subst(p, a, vars, subst)?;
            }
            Ok(())
        }
        (
            Type::Fn { params: pp, ret: pr, .. },
            Type::Fn { params: ap, ret: ar, .. },
        ) => {
            if pp.len() != ap.len() {
                return Err(CodegenError::Internal(
                    "monomorphisation: fn arity mismatch in arg".into(),
                ));
            }
            for (p, a) in pp.iter().zip(ap.iter()) {
                unify_for_subst(p, a, vars, subst)?;
            }
            unify_for_subst(pr, ar, vars, subst)
        }
        (Type::Var { name: pn }, Type::Var { name: an }) if pn == an => Ok(()),
        _ => Err(CodegenError::Internal(format!(
            "monomorphisation: cannot match param `{}` to arg `{}`",
            ailang_core::pretty::type_to_string(param),
            ailang_core::pretty::type_to_string(arg),
        ))),
    }
}

/// rewrites bare `Type::Con` references that resolve against
/// `owner_local_types` into qualified `module.Type` form. Mirrors
/// `ailang_check::qualify_local_types` and complements prep.1's
/// `qualify_workspace_types` (which qualifies consumer-side bare
/// cross-module refs). Used when the codegen pulls a polymorphic
/// fn signature across the import boundary; without this the
/// substitution derived from the call site's qualified args
/// (e.g. `std_maybe.Maybe<Int>`) would fail to unify against the
/// owner-local signature (`Maybe<a>`).
pub(crate) fn qualify_local_types_codegen(
    t: &Type,
    owner_module: &str,
    owner_local_types: &BTreeSet<String>,
) -> Type {
    match t {
        Type::Con { name, args } => {
            // The first two branches share the body `name.clone()` but
            // express semantically distinct reasons (already qualified;
            // primitive needs no qualification). Combining them with
            // `||` would obscure why each disqualifies the name.
            #[allow(clippy::if_same_then_else)]
            let qualified = if name.contains('.') {
                name.clone()
            } else if ailang_core::primitives::is_primitive_name(name) {
                name.clone()
            } else if owner_local_types.contains(name) {
                format!("{owner_module}.{name}")
            } else {
                name.clone()
            };
            Type::Con {
                name: qualified,
                args: args
                    .iter()
                    .map(|a| qualify_local_types_codegen(a, owner_module, owner_local_types))
                    .collect(),
            }
        }
        Type::Fn { params, ret, effects, param_modes, ret_mode } => Type::Fn {
            params: params
                .iter()
                .map(|p| qualify_local_types_codegen(p, owner_module, owner_local_types))
                .collect(),
            ret: Box::new(qualify_local_types_codegen(ret, owner_module, owner_local_types)),
            effects: effects.clone(),
            param_modes: param_modes.clone(),
            ret_mode: *ret_mode,
        },
        Type::Forall { vars, constraints, body } => Type::Forall {
            vars: vars.clone(),
            constraints: constraints.clone(),
            body: Box::new(qualify_local_types_codegen(body, owner_module, owner_local_types)),
        },
        Type::Var { .. } => t.clone(),
    }
}

/// substitute rigid type vars in `t` according to `subst`.
/// Used to specialise the type of a polymorphic def for a given
/// instantiation.
pub(crate) fn apply_subst_to_type(t: &Type, subst: &BTreeMap<String, Type>) -> Type {
    match t {
        Type::Var { name } => subst.get(name).cloned().unwrap_or_else(|| t.clone()),
        Type::Con { name, args } => Type::Con {
            name: name.clone(),
            args: args.iter().map(|a| apply_subst_to_type(a, subst)).collect(),
        },
        Type::Fn { params, ret, effects, param_modes, ret_mode } => {
            let new_params: Vec<Type> =
                params.iter().map(|p| apply_subst_to_type(p, subst)).collect();
            let new_ret = apply_subst_to_type(ret, subst);
            // spec 0062: a polymorphic (borrow a) specialised onto a
            // value type becomes (own value-type) — borrow-over-value
            // is forbidden and is a no-op for unboxed types (no RC).
            let coerce = |ty: &Type, m: &ParamMode| -> ParamMode {
                if matches!(m, ParamMode::Borrow) {
                    if let Type::Con { name, args } = ty {
                        if args.is_empty() && ailang_core::primitives::is_value_type(name) {
                            return ParamMode::Own;
                        }
                    }
                }
                *m
            };
            let new_param_modes: Vec<ParamMode> = new_params
                .iter()
                .zip(param_modes.iter())
                .map(|(t, m)| coerce(t, m))
                .collect();
            let new_ret_mode = coerce(&new_ret, ret_mode);
            Type::Fn {
                params: new_params,
                ret: Box::new(new_ret),
                effects: effects.clone(),
                param_modes: new_param_modes,
                ret_mode: new_ret_mode,
            }
        }
        Type::Forall { vars, constraints, body } => {
            // Inner forall shadows: don't substitute re-bound names.
            let inner: BTreeMap<String, Type> = subst
                .iter()
                .filter(|(k, _)| !vars.contains(k))
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect();
            Type::Forall {
                vars: vars.clone(),
                constraints: constraints.clone(),
                body: Box::new(apply_subst_to_type(body, &inner)),
            }
        }
    }
}

// iter 23.4: `apply_subst_to_term` and `descriptor_for_subst` were
// the codegen-side body-substitution and symbol-mangling helpers used
// by `lower_polymorphic_call` / `emit_specialised_fn`. Both are gone:
// the typecheck-time mono pass synthesises every monomorphic body
// (via `ailang_check::substitute_rigids_in_term`) and produces
// surface-named mono symbols (via `ailang_check::mono::mono_symbol_n`).