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

//! Iter 22b.3: workspace monomorphisation pass.
//!
//! Slots into the build pipeline after [`crate::lift_letrecs`] and
//! before `ailang_codegen::lower_workspace_with_alloc`.
//!
//! ## What Task 1 delivers
//!
//! Skeleton only. [`monomorphise_workspace`] is a no-op pass with
//! two branches:
//!
//! - Class-free workspaces (no [`Def::Class`] / [`Def::Instance`]
//!   anywhere) take the early-out and the input workspace is
//!   cloned unchanged.
//! - Class-present workspaces fall through to a second
//!   `Ok(ws.clone())`, dead-equivalent today but the placeholder
//!   for the real fixpoint body that later tasks will fill in.
//!
//! Both branches return a byte-identical workspace clone, so
//! `module_hash` is preserved end-to-end through the pass at this
//! stage of the iteration.
//!
//! ## Planned (later tasks in iter 22b.3)
//!
//! See `docs/plans/2026-05-09-22b3-monomorphisation.md` for the
//! full task list. The eventual contract is for the pass to
//! consume the workspace's typeclass [`Registry`] and the per-
//! module class-method tables (already populated by
//! [`crate::check_in_workspace`]) and produce a workspace with:
//!
//! 1. Synthesised [`Def::Fn`] entries — one per unique
//!    `(class, method, type-hash)` triple observed at any
//!    class-method call site — appended to the [`Registry`]'s
//!    `defining_module` for that instance.
//! 2. Rewritten call sites — every `Term::Var { name }` whose
//!    `name` resolves through [`Env::class_methods`] is replaced
//!    by the corresponding mono-symbol name (qualified with the
//!    instance's `defining_module` iff that module differs from
//!    the calling module).
//!
//! Pre-existing `Def::Class` and `Def::Instance` entries are
//! preserved verbatim — codegen ignores them, but downstream
//! tooling (e.g. `ail describe`) may still consult them.
//!
//! ## Symbol-hashing invariant (eventual-state)
//!
//! Once the synthesis body lands, synthesised FnDefs will be
//! post-typecheck artefacts. They will NOT enter
//! `CheckedModule.symbols` (built from the original module at
//! typecheck time and used by `ail diff` / `ail manifest`). Same
//! convention as [`crate::lift_letrecs`]. Documented here ahead of
//! the implementation so the constraint is visible to anyone
//! extending the skeleton.

use ailang_core::ast::{Arm, ClassDef, ConstDef, Def, FnDef as AstFnDef, InstanceDef, Pattern, Term, Type};
use ailang_core::workspace::Workspace;
use crate::{CtorRef, Result};
use indexmap::IndexMap;
use std::collections::{BTreeMap, BTreeSet};

/// Iter 22b.3: workspace-wide monomorphisation pass entry. See the
/// module-level doc for the architecture and contract.
///
/// Pre-condition: `ws` has been typechecked (`check_workspace(ws)`
/// returned no errors) and lifted (`lift_letrecs` per module). The
/// pass does not perform new type checking — it queries types via
/// `synth` on already-typechecked bodies.
pub fn monomorphise_workspace(ws: &Workspace) -> Result<Workspace> {
    // Fast path — no class / instance defs anywhere → nothing to do.
    // The pass also has no targets when there are class defs but no
    // instance defs (no callable methods at concrete types), but
    // the body walks would still be a wasted traversal; the
    // class-free check is the cheap way to opt out.
    if !workspace_has_typeclasses(ws) {
        return Ok(ws.clone());
    }

    let mut ws_owned: Workspace = ws.clone();
    let env = build_workspace_env(&ws_owned);
    let class_index = build_class_index(&ws_owned);

    let mut synthesised: BTreeSet<(String, String, String)> = BTreeSet::new();

    // Fixpoint: keep collecting until a round adds nothing new.
    // Each round walks every fn body in every module — including
    // bodies appended by the previous round, which is what makes
    // the loop close on chained class-method calls (e.g.
    // `Eq.ne x y = not (eq x y)`).
    loop {
        let new_targets = collect_targets_workspace_wide(&ws_owned, &env)?;
        // Dedup within a round (multiple call sites of `show` at
        // the same type produce N copies of the same MonoTarget)
        // AND across rounds (already-synthesised keys filtered out).
        // Stable iteration order — preserve first-seen target so
        // any later debugging round-trips to a deterministic output.
        let mut seen_this_round: BTreeSet<(String, String, String)> = BTreeSet::new();
        let mut new: Vec<MonoTarget> = Vec::new();
        for t in new_targets {
            let k = mono_target_key(&t);
            if synthesised.contains(&k) || !seen_this_round.insert(k) {
                continue;
            }
            new.push(t);
        }
        if new.is_empty() {
            break;
        }

        // Synthesise each new target; append its FnDef to the
        // target's `defining_module`. Mark the key as synthesised
        // so the next round won't re-collect.
        //
        // Iter 22b.3 uses ws_owned.registry as the source of
        // truth for ClassDef + InstanceDef lookups. Both come
        // from the un-modified workspace registry — synthesis
        // never mutates the registry.
        for t in &new {
            let key = mono_target_key(t);
            let registry_key = (t.class.clone(), ailang_core::canonical::type_hash(&t.type_));
            let entry = ws_owned
                .registry
                .entries
                .get(&registry_key)
                .ok_or_else(|| {
                    crate::CheckError::Internal(format!(
                        "monomorphise_workspace: target `{} {}` has no registry entry",
                        t.class,
                        ailang_core::pretty::type_to_string(&t.type_),
                    ))
                })?;
            let class_def = class_index.get(&t.class).ok_or_else(|| {
                crate::CheckError::Internal(format!(
                    "monomorphise_workspace: class `{}` not found",
                    t.class
                ))
            })?;
            let f = synthesise_mono_fn(t, class_def, &entry.instance)?;
            let target_module = ws_owned
                .modules
                .get_mut(&t.defining_module)
                .ok_or_else(|| {
                    crate::CheckError::Internal(format!(
                        "monomorphise_workspace: defining module `{}` missing",
                        t.defining_module
                    ))
                })?;
            target_module.defs.push(Def::Fn(f));
            synthesised.insert(key);
        }
    }

    // Phase 3: rewrite call sites in every fn / const body. Walk in
    // the same pre-order as collect_residuals_ordered so the cursor
    // and the per-callsite target list align position-by-position.
    let env = build_workspace_env(&ws_owned); // re-build: ws_owned has new defs.
    let module_names: Vec<String> = ws_owned.modules.keys().cloned().collect();
    for mname in &module_names {
        // Per-module ctor_index overlay — see
        // [`apply_per_module_ctor_index_overlay`] for rationale.
        // Mirrors the overlay applied in
        // `collect_targets_workspace_wide` so the rewrite-phase
        // residual replay produces the same residual list as the
        // collection-phase walk (cursor alignment depends on it).
        let mut env_mod = env.clone();
        apply_per_module_ctor_index_overlay(&mut env_mod, &ws_owned, mname);
        let n_defs = ws_owned.modules[mname].defs.len();
        for i in 0..n_defs {
            let ordered: Vec<Option<MonoTarget>> = {
                let m = &ws_owned.modules[mname];
                let d = &m.defs[i];
                match d {
                    Def::Fn(f) => collect_residuals_ordered(f, mname, &env_mod)?,
                    Def::Const(c) => {
                        let pseudo = const_as_pseudo_fn(c);
                        collect_residuals_ordered(&pseudo, mname, &env_mod)?
                    }
                    _ => continue,
                }
            };
            let m = ws_owned.modules.get_mut(mname).unwrap();
            let d = &mut m.defs[i];
            let mut cursor = 0usize;
            let mut locals: BTreeSet<String> = BTreeSet::new();
            match d {
                Def::Fn(f) => {
                    // Top-level fn params shadow class-method names too.
                    for p in &f.params {
                        locals.insert(p.clone());
                    }
                    rewrite_class_method_calls(
                        &mut f.body,
                        &env.class_methods,
                        mname,
                        &ordered,
                        &mut cursor,
                        &mut locals,
                    );
                }
                Def::Const(c) => {
                    rewrite_class_method_calls(
                        &mut c.value,
                        &env.class_methods,
                        mname,
                        &ordered,
                        &mut cursor,
                        &mut locals,
                    );
                }
                _ => {}
            }
        }
    }

    Ok(ws_owned)
}

/// Iter 22b.3: walk every fn / const body in the workspace,
/// returning the union of [`collect_mono_targets`] outputs. Const
/// bodies are wrapped in a synthetic zero-arg `FnDef` for the
/// collection call — the residual gathering only depends on body
/// shape, so the wrapping is harmless.
fn collect_targets_workspace_wide(
    ws: &Workspace,
    env: &crate::Env,
) -> Result<Vec<MonoTarget>> {
    let mut out: Vec<MonoTarget> = Vec::new();
    for (mname, m) in &ws.modules {
        // Per-module ctor_index overlay — see
        // [`apply_per_module_ctor_index_overlay`] for rationale. The
        // env arriving here is workspace-flat (via `build_check_env`);
        // synth's `Pattern::Ctor` resolution requires per-module
        // shape to keep the qualified-type-name comparison intact
        // for cross-module ctor patterns.
        let mut env_mod = env.clone();
        apply_per_module_ctor_index_overlay(&mut env_mod, ws, mname);
        for d in &m.defs {
            match d {
                Def::Fn(f) => {
                    out.extend(collect_mono_targets(f, mname, &env_mod)?);
                }
                Def::Const(c) => {
                    let pseudo = const_as_pseudo_fn(c);
                    out.extend(collect_mono_targets(&pseudo, mname, &env_mod)?);
                }
                _ => {}
            }
        }
    }
    Ok(out)
}

/// Iter 22b.3.6: wrap a [`ConstDef`] in a synthetic zero-arg
/// [`AstFnDef`] for residual-collection / rewrite purposes. Const
/// bodies have no parameter list, so the residual gathering only
/// depends on body shape — the wrapping is a structural adapter,
/// not a semantic conversion. Two callers (Phase 1 collection and
/// Phase 3 rewrite) need the same shape; sharing the constructor
/// keeps them in lockstep.
fn const_as_pseudo_fn(c: &ConstDef) -> AstFnDef {
    AstFnDef {
        name: c.name.clone(),
        ty: c.ty.clone(),
        params: Vec::new(),
        body: c.value.clone(),
        doc: None,
        suppress: Vec::new(),
    }
}

/// Iter 22b.3: workspace-wide `class-name -> ClassDef` index.
/// Used by the fixpoint to look up the matching class definition
/// when synthesising a fn.
fn build_class_index(ws: &Workspace) -> BTreeMap<String, ClassDef> {
    let mut idx = BTreeMap::new();
    for m in ws.modules.values() {
        for d in &m.defs {
            if let Def::Class(c) = d {
                idx.insert(c.name.clone(), c.clone());
            }
        }
    }
    idx
}

/// Iter 22b.3: returns `true` iff any module in `ws` declares at
/// least one [`Def::Class`] or [`Def::Instance`]. Cheap workspace
/// scan; the early-out keeps class-free workspaces byte-identical
/// through the pass.
fn workspace_has_typeclasses(ws: &Workspace) -> bool {
    ws.modules.values().any(|m| {
        m.defs.iter().any(|d| matches!(d, Def::Class(_) | Def::Instance(_)))
    })
}

/// Iter 22b.3: deterministic mono-symbol name for a `(method,
/// type)` pair. Primitive types (`Int`, `Bool`, `Str`, `Unit`)
/// produce `<method>__<surface-name>` for diagnostic and
/// ABI legibility. All other types — parameterised cons,
/// user-defined ADTs, function types — fall to
/// `<method>__<8-hex-prefix-of-canonical-type-hash>`. The hash
/// route ensures uniqueness without requiring a flattened
/// surface form for arbitrarily nested types.
///
/// Separator choice: `__` (double underscore) rather than the spec's
/// recommended `#` — `#` terminates LLVM IR global identifiers (and
/// breaks C-ABI symbol names on most targets), so the produced name
/// has to survive codegen unaltered. The double underscore is
/// already in use elsewhere in the codegen pipeline as a descriptor
/// separator (see `emit_specialised_fn`), keeping the convention
/// uniform. Iter 22b.3.7 documents the substitution at the e2e
/// gate; Tasks 2/4 unit-tests are updated in lockstep.
///
/// Determinism: `ailang_core::canonical::type_hash` is the same
/// function `workspace::build_registry` uses to key
/// [`Registry::entries`], so a registry-key match implies a
/// `mono_symbol` match.
pub fn mono_symbol(method: &str, ty: &Type) -> String {
    if let Some(prim) = primitive_surface_name(ty) {
        return format!("{method}__{prim}");
    }
    let full_hash = ailang_core::canonical::type_hash(ty);
    // 8-hex prefix is enough for low-collision keying across the
    // workspace; the full hash remains in the registry for
    // disambiguation should one ever be needed.
    format!("{method}__{}", &full_hash[..8])
}

/// Returns the surface name iff `ty` is a zero-arity primitive
/// `Type::Con`. Used by [`mono_symbol`] to gate the human-readable
/// form. The match is intentionally narrow: `Int<args>` (malformed
/// but parser-accepting) is treated as compound, so it falls to
/// the hash form. The primitive-set itself lives in
/// [`ailang_core::primitives::primitive_surface_name`].
fn primitive_surface_name(ty: &Type) -> Option<&'static str> {
    match ty {
        Type::Con { name, args } if args.is_empty() => {
            ailang_core::primitives::primitive_surface_name(name)
        }
        _ => None,
    }
}

/// Iter 22b.3: one synthesisation request — a fully-concrete
/// `(class, method, type)` triple observed at a class-method call
/// site, plus the registry's `defining_module` for the matching
/// `InstanceDef` (the synthesised fn lives there). Equality of
/// targets is compared via [`mono_target_key`] — `(class,
/// method, type-hash)`.
#[derive(Debug, Clone)]
pub struct MonoTarget {
    pub class: String,
    pub method: String,
    pub type_: Type,
    pub defining_module: String,
}

/// Iter 22b.3: dedup key for a [`MonoTarget`] — the same triple
/// `Registry::entries` uses for instance lookup, plus the method
/// name. Two targets with the same key produce the same
/// synthesised symbol and the same body. Consumed by the workspace
/// fixpoint (Task 5) and the rewrite walker (Task 6); declared
/// here in Task 3 so the dedup contract lives next to `MonoTarget`.
#[allow(dead_code)]
pub(crate) fn mono_target_key(t: &MonoTarget) -> (String, String, String) {
    (
        t.class.clone(),
        t.method.clone(),
        ailang_core::canonical::type_hash(&t.type_),
    )
}

/// Thin wrapper over [`crate::build_check_env`]. The mono pass needs
/// the same workspace-flat `Env` shape as `check_in_workspace`, so both
/// share one source of truth. Per-fn entry points
/// (`collect_mono_targets`, `collect_residuals_ordered`) clone the env
/// and apply per-fn overlay (current_module, globals from
/// module_globals, imports from module_imports, rigid_vars) plus the
/// per-module ctor/type overlay
/// (see [`apply_per_module_ctor_index_overlay`]).
pub fn build_workspace_env(ws: &Workspace) -> crate::Env {
    crate::build_check_env(ws)
}

/// Iter 22c (sibling of commit 5c5180f for `env.ctor_index`):
/// rebuild `env.types` AND `env.ctor_index` to contain ONLY the
/// types and ctors declared in `module_name`'s `Def::Type`s,
/// mirroring `check_in_workspace`'s per-module overlay at
/// `crates/ailang-check/src/lib.rs:1257-1287`.
///
/// `build_workspace_env` (via `build_check_env`) leaves both fields
/// workspace-flat. Two synth paths need the per-module shape:
///
///   * `Pattern::Ctor` resolution (lib.rs:2486-2526) does a
///     local-first `ctor_index` lookup followed by an imports-fallback
///     that produces a *qualified* `resolved_type_name` (`Mod.Type`).
///     With a flat `ctor_index`, a cross-module ctor pattern resolves
///     locally and yields a bare type name, mismatching the qualified
///     scrutinee type and firing `CheckError::PatternTypeMismatch`.
///   * `Term::Ctor` synth (lib.rs:1962-2026) does a local-first
///     `env.types` lookup followed by an imports-fallback that
///     produces a qualified `result_type_name`. With a flat `env.types`,
///     a `Term::Ctor` against a prelude (or imported) type resolves
///     locally and yields a bare scrutinee type — which then mismatches
///     the qualified `Pattern::Ctor` `resolved_type_name` produced by
///     the per-module-cleared `ctor_index`.
///
/// Both fields must be cleared together so the two synth paths agree
/// on bare-vs-qualified naming. Mirrors the pattern in
/// `check_in_workspace`, which clears and rebuilds both for the
/// same reason.
///
/// Caller-contract assumption: the workspace has already typechecked,
/// so duplicate type or ctor names within a single module cannot
/// occur (the fail-fast `DuplicateType` / `DuplicateCtor` diagnostics
/// in `check_in_workspace` would have surfaced).
fn apply_per_module_ctor_index_overlay(env: &mut crate::Env, ws: &Workspace, module_name: &str) {
    env.types.clear();
    env.ctor_index.clear();
    if let Some(m) = ws.modules.get(module_name) {
        for d in &m.defs {
            if let Def::Type(td) = d {
                for c in &td.ctors {
                    env.ctor_index.insert(
                        c.name.clone(),
                        CtorRef { type_name: td.name.clone() },
                    );
                }
                env.types.insert(td.name.clone(), td.clone());
            }
        }
    }
}

/// Iter 22b.3: re-run [`crate::synth`] on `f`'s body to recover
/// the per-fn residual class constraints. Filter to fully-
/// concrete residuals (the only ones eligible for monomorphisation),
/// look up each `(class, type-hash)` in the registry to recover
/// the `defining_module`, and return the list. Var-shaped or
/// metavar-shaped residuals are silently skipped — the
/// 22b.2 typecheck pass has already fired
/// `MissingConstraint`/`NoInstance` for any that should not exist
/// at this point.
pub fn collect_mono_targets(
    f: &AstFnDef,
    module_name: &str,
    env: &crate::Env,
) -> Result<Vec<MonoTarget>> {
    // Build the per-def env exactly as `check_fn` does — install
    // rigid vars, set the current module, etc. This mirrors
    // `crate::check_fn` minus the diagnostic emission.
    let (rigids, inner_ty): (Vec<String>, Type) = match &f.ty {
        Type::Forall { vars, constraints: _, body } => (vars.clone(), (**body).clone()),
        other => (vec![], other.clone()),
    };
    let (param_tys, _ret_ty, _eff): (Vec<Type>, Type, Vec<String>) = match &inner_ty {
        Type::Fn { params, ret, effects, .. } => {
            (params.clone(), (**ret).clone(), effects.clone())
        }
        _ => return Ok(Vec::new()),
    };

    let mut env = env.clone();
    for v in &rigids {
        env.rigid_vars.insert(v.clone());
    }
    env.current_module = module_name.to_string();
    // Iter 22c (sibling of commit 5c5180f): seed `env.globals` from
    // the current module's fns so `synth`'s `Term::Var` lookup
    // (lib.rs:1678) resolves bare same-module references — most
    // importantly self-recursive top-level fns. Mirrors
    // `check_in_workspace` (lib.rs:1135-1139). Top-level fn names
    // are only per-module-unique (workspace-wide collisions are
    // legal), so seeding must be scoped to the current module —
    // unlike the workspace-wide flat `env.types` / `env.ctor_index`
    // tables that 5c5180f populated in `build_workspace_env`.
    if let Some(g) = env.module_globals.get(module_name).cloned() {
        for (n, t) in g {
            env.globals.insert(n, t);
        }
    }
    // Iter 22c (sibling of commit 13b36cc): seed `env.imports` from
    // the current module's import list so `synth`'s qualified-var
    // path (lib.rs:1697) resolves `Mod.fn` references. Mirrors
    // `check_in_workspace` (lib.rs:1147-1152). Per-module because
    // aliases collide across modules.
    if let Some(im) = env.module_imports.get(module_name).cloned() {
        env.imports = im;
    }

    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 effects: BTreeSet<String> = BTreeSet::new();
    let mut subst = crate::Subst::default();
    let mut counter: u32 = 0;
    let mut residuals: Vec<crate::ResidualConstraint> = Vec::new();
    crate::synth(
        &f.body,
        &env,
        &mut locals,
        &mut effects,
        &f.name,
        &mut subst,
        &mut counter,
        &mut residuals,
    )?;

    // Filter residuals to fully-concrete ones; look up
    // defining_module via the registry.
    let mut out: Vec<MonoTarget> = Vec::new();
    for r in residuals {
        let r_ty = subst.apply(&r.type_);
        if !crate::is_fully_concrete(&r_ty) {
            continue;
        }
        let key = (r.class.clone(), ailang_core::canonical::type_hash(&r_ty));
        let entry = match env.workspace_registry.entries.get(&key) {
            Some(e) => e,
            None => continue, // no-instance — Task 10 of 22b.2 fires; skip silently here.
        };
        out.push(MonoTarget {
            class: r.class.clone(),
            method: r.method.clone(),
            type_: r_ty,
            defining_module: entry.defining_module.clone(),
        });
    }
    Ok(out)
}

/// Iter 22b.3: produce a `FnDef` for a single (target, class,
/// instance) triple. The synthesised fn:
///
/// 1. has name [`mono_symbol`]`(target.method, target.type_)`,
/// 2. has type = the class's method type with the class param
///    substituted to `target.type_` (rigid-var substitution via
///    [`crate::substitute_rigids`]); the result is a plain
///    `Type::Fn` with no `Forall`,
/// 3. has params + body taken from the instance's matching
///    `InstanceMethod.body` if present and Lam-shaped; from the
///    class's `default` body if the instance omits the method;
///    or directly from the body if the method has no params
///    (zero-arg method types).
///
/// Errors (all [`crate::CheckError::Internal`] — caller-contract
/// violations that cannot occur after typecheck has succeeded):
///
/// - The instance omits the method AND the class has no `default`
///   — unreachable in practice because
///   `workspace::build_registry`'s `MissingMethod` check fires at
///   load.
/// - The class itself has no method by that name — also caught at
///   load.
/// - The method type is `(args) -> ret` with `args` non-empty but
///   the resolved body is not [`Term::Lam`] — schema-shape
///   mismatch enforced by 22b.2's instance-method check.
pub fn synthesise_mono_fn(
    target: &MonoTarget,
    class_def: &ClassDef,
    instance: &InstanceDef,
) -> Result<AstFnDef> {
    // Locate the class method declaration.
    let class_method = class_def
        .methods
        .iter()
        .find(|m| m.name == target.method)
        .ok_or_else(|| {
            crate::CheckError::Internal(format!(
                "synthesise_mono_fn: class `{}` has no method `{}`",
                class_def.name, target.method
            ))
        })?;

    // Build the substitution `param := target.type_` and apply it
    // to the method type. The result is a concrete `Type::Fn` (no
    // `Forall`).
    let mut mapping: BTreeMap<String, Type> = BTreeMap::new();
    mapping.insert(class_def.param.clone(), target.type_.clone());
    let concrete_method_ty = crate::substitute_rigids(&class_method.ty, &mapping);

    // Resolve the body — instance override first, then class default.
    let body_term: Term = match instance.methods.iter().find(|im| im.name == target.method) {
        Some(im) => im.body.clone(),
        None => class_method.default.clone().ok_or_else(|| {
            crate::CheckError::Internal(format!(
                "synthesise_mono_fn: instance `{} {}` omits method `{}` and class has no \
                 default (registry-build should have rejected this)",
                target.class,
                ailang_core::pretty::type_to_string(&target.type_),
                target.method,
            ))
        })?,
    };

    // Decide the (params, body) shape based on whether the method
    // type takes positional args.
    let method_has_params = matches!(
        &class_method.ty,
        Type::Fn { params, .. } if !params.is_empty()
    );
    let (params, body): (Vec<String>, Term) = if method_has_params {
        match body_term {
            Term::Lam { params, body, .. } => (params, *body),
            other => {
                return Err(crate::CheckError::Internal(format!(
                    "synthesise_mono_fn: method `{}` has positional params but body is not \
                     a Lam: {:?}",
                    target.method, other
                )));
            }
        }
    } else {
        (Vec::new(), body_term)
    };

    Ok(AstFnDef {
        name: mono_symbol(&target.method, &target.type_),
        ty: concrete_method_ty,
        params,
        body,
        doc: None,
        suppress: Vec::new(),
    })
}

/// Iter 22b.3: rewrite every class-method call site in `body` to
/// the corresponding mono symbol, using `ordered_targets` —
/// the per-call-site resolved targets collected by a parallel
/// synth-replay run on the same body.
///
/// The walker increments a positional counter at each
/// class-method-named `Term::Var` it encounters that is NOT
/// shadowed by a local binding. The counter must match the order
/// in which `synth` pushes residuals (pre-order AST walk; child
/// evaluation order matches the `match` arms in [`crate::synth`]).
/// When this invariant holds, `ordered_targets[idx]` is the target
/// resolved at the i-th class-method call site.
///
/// `caller_module`: name of the module enclosing this body. If
/// the target's `defining_module` differs, the rewritten name is
/// qualified `<defining_module>.<mono_symbol>` (cross-module
/// resolution path); otherwise unqualified.
///
/// `locals` mirrors `synth`'s lookup-precedence rule: a `Term::Var`
/// whose `name` is in `locals` resolves to the local binding and
/// `synth` pushes NO residual for it, so the walker must NOT
/// advance the cursor either. Each binding-form pushes its binders
/// before recursing into the relevant scope and pops them after.
///
/// `ordered_targets[idx]`:
/// - `Some(t)` — concrete instance found; rewrite the name.
/// - `None` — residual was non-concrete or had no instance; cursor
///   still advances (to keep alignment) but the name is left
///   unchanged.
fn rewrite_class_method_calls(
    body: &mut Term,
    class_methods: &BTreeMap<String, crate::ClassMethodEntry>,
    caller_module: &str,
    ordered_targets: &[Option<MonoTarget>],
    cursor: &mut usize,
    locals: &mut BTreeSet<String>,
) {
    match body {
        Term::Var { name } => {
            if class_methods.contains_key(name) {
                // Mirror synth's lookup-precedence rule: if `name` is
                // shadowed by a local, synth resolves to the local
                // and pushes no residual — we must not advance.
                if locals.contains(name) {
                    return;
                }
                if let Some(slot) = ordered_targets.get(*cursor) {
                    if let Some(t) = slot {
                        let sym = mono_symbol(&t.method, &t.type_);
                        let new_name = if t.defining_module == caller_module {
                            sym
                        } else {
                            format!("{}.{}", t.defining_module, sym)
                        };
                        *name = new_name;
                    }
                    // else: residual was non-concrete → leave name unchanged.
                }
                *cursor += 1;
            }
        }
        Term::App { callee, args, .. } => {
            rewrite_class_method_calls(callee, class_methods, caller_module, ordered_targets, cursor, locals);
            for a in args {
                rewrite_class_method_calls(a, class_methods, caller_module, ordered_targets, cursor, locals);
            }
        }
        Term::Let { name, value, body } => {
            // synth: `value` is in the outer scope; `body` sees `name`.
            rewrite_class_method_calls(value, class_methods, caller_module, ordered_targets, cursor, locals);
            let inserted = locals.insert(name.clone());
            rewrite_class_method_calls(body, class_methods, caller_module, ordered_targets, cursor, locals);
            if inserted {
                locals.remove(name);
            }
        }
        Term::LetRec { name, params, body, in_term, .. } => {
            // synth: `name` is in scope in BOTH `body` (with `params`
            // also in scope) AND `in_term` (without `params`).
            let name_inserted = locals.insert(name.clone());
            // Body scope: name + params.
            let mut params_inserted: Vec<String> = Vec::new();
            for p in params {
                if locals.insert(p.clone()) {
                    params_inserted.push(p.clone());
                }
            }
            rewrite_class_method_calls(body, class_methods, caller_module, ordered_targets, cursor, locals);
            for p in &params_inserted {
                locals.remove(p);
            }
            // In-clause scope: name only.
            rewrite_class_method_calls(in_term, class_methods, caller_module, ordered_targets, cursor, locals);
            if name_inserted {
                locals.remove(name);
            }
        }
        Term::If { cond, then, else_ } => {
            rewrite_class_method_calls(cond, class_methods, caller_module, ordered_targets, cursor, locals);
            rewrite_class_method_calls(then, class_methods, caller_module, ordered_targets, cursor, locals);
            rewrite_class_method_calls(else_, class_methods, caller_module, ordered_targets, cursor, locals);
        }
        Term::Do { args, .. } => {
            for a in args {
                rewrite_class_method_calls(a, class_methods, caller_module, ordered_targets, cursor, locals);
            }
        }
        Term::Ctor { args, .. } => {
            for a in args {
                rewrite_class_method_calls(a, class_methods, caller_module, ordered_targets, cursor, locals);
            }
        }
        Term::Match { scrutinee, arms } => {
            // synth: scrutinee is in outer scope; each arm's body sees
            // its pattern's binders.
            rewrite_class_method_calls(scrutinee, class_methods, caller_module, ordered_targets, cursor, locals);
            for Arm { pat, body } in arms {
                let binders = pattern_binders(pat);
                let mut inserted: Vec<String> = Vec::new();
                for b in &binders {
                    if locals.insert(b.clone()) {
                        inserted.push(b.clone());
                    }
                }
                rewrite_class_method_calls(body, class_methods, caller_module, ordered_targets, cursor, locals);
                for b in &inserted {
                    locals.remove(b);
                }
            }
        }
        Term::Lam { params, body, .. } => {
            // synth: params are in scope in `body`.
            let mut inserted: Vec<String> = Vec::new();
            for p in params {
                if locals.insert(p.clone()) {
                    inserted.push(p.clone());
                }
            }
            rewrite_class_method_calls(body, class_methods, caller_module, ordered_targets, cursor, locals);
            for p in &inserted {
                locals.remove(p);
            }
        }
        Term::Seq { lhs, rhs } => {
            rewrite_class_method_calls(lhs, class_methods, caller_module, ordered_targets, cursor, locals);
            rewrite_class_method_calls(rhs, class_methods, caller_module, ordered_targets, cursor, locals);
        }
        Term::Clone { value } => {
            rewrite_class_method_calls(value, class_methods, caller_module, ordered_targets, cursor, locals);
        }
        Term::ReuseAs { source, body } => {
            rewrite_class_method_calls(source, class_methods, caller_module, ordered_targets, cursor, locals);
            rewrite_class_method_calls(body, class_methods, caller_module, ordered_targets, cursor, locals);
        }
        Term::Lit { .. } => {}
    }
}

/// Iter 22b.3.6: collect the variable binders introduced by a
/// match pattern. Used by the walker's `Term::Match` arm to extend
/// `locals` for each arm body. Mirrors `Pattern`'s shape:
/// - [`Pattern::Wild`] / [`Pattern::Lit`] bind nothing,
/// - [`Pattern::Var`] binds its `name`,
/// - [`Pattern::Ctor`] recurses through its `fields`.
fn pattern_binders(pat: &Pattern) -> Vec<String> {
    let mut out: Vec<String> = Vec::new();
    fn rec(pat: &Pattern, out: &mut Vec<String>) {
        match pat {
            Pattern::Wild | Pattern::Lit { .. } => {}
            Pattern::Var { name } => out.push(name.clone()),
            Pattern::Ctor { fields, .. } => {
                for f in fields {
                    rec(f, out);
                }
            }
        }
    }
    rec(pat, &mut out);
    out
}

/// Iter 22b.3: traversal-ordered residual collection — used by
/// the rewrite walker to align cursor positions. Unlike
/// [`collect_mono_targets`], this includes non-concrete residuals
/// as `None`, preserving one-entry-per-callsite alignment.
pub(crate) fn collect_residuals_ordered(
    f: &AstFnDef,
    module_name: &str,
    env: &crate::Env,
) -> Result<Vec<Option<MonoTarget>>> {
    let (rigids, inner_ty): (Vec<String>, Type) = match &f.ty {
        Type::Forall { vars, constraints: _, body } => (vars.clone(), (**body).clone()),
        other => (vec![], other.clone()),
    };
    let param_tys: Vec<Type> = match &inner_ty {
        Type::Fn { params, .. } => params.clone(),
        _ => return Ok(Vec::new()),
    };

    let mut env = env.clone();
    for v in &rigids {
        env.rigid_vars.insert(v.clone());
    }
    env.current_module = module_name.to_string();
    // Iter 22c: same `env.globals` seeding as `collect_mono_targets`
    // — see the comment there for rationale. Both fns re-run `synth`
    // and must agree with the main check path's per-module env shape.
    if let Some(g) = env.module_globals.get(module_name).cloned() {
        for (n, t) in g {
            env.globals.insert(n, t);
        }
    }
    // Iter 22c: same `env.imports` seeding as `collect_mono_targets`
    // — see the comment there for rationale.
    if let Some(im) = env.module_imports.get(module_name).cloned() {
        env.imports = im;
    }

    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 effects: BTreeSet<String> = BTreeSet::new();
    let mut subst = crate::Subst::default();
    let mut counter: u32 = 0;
    let mut residuals: Vec<crate::ResidualConstraint> = Vec::new();
    crate::synth(
        &f.body,
        &env,
        &mut locals,
        &mut effects,
        &f.name,
        &mut subst,
        &mut counter,
        &mut residuals,
    )?;

    let mut out: Vec<Option<MonoTarget>> = Vec::new();
    for r in residuals {
        let r_ty = subst.apply(&r.type_);
        if !crate::is_fully_concrete(&r_ty) {
            out.push(None);
            continue;
        }
        let key = (r.class.clone(), ailang_core::canonical::type_hash(&r_ty));
        let entry = match env.workspace_registry.entries.get(&key) {
            Some(e) => e,
            None => {
                out.push(None);
                continue;
            }
        };
        out.push(Some(MonoTarget {
            class: r.class.clone(),
            method: r.method.clone(),
            type_: r_ty,
            defining_module: entry.defining_module.clone(),
        }));
    }
    Ok(out)
}