//! Type system and inference engine for CJC.
//!
//! This crate provides the core type representation ([`Type`]), the type
//! environment ([`TypeEnv`]), Hindley-Milner style unification ([`unify`]),
//! constraint-based inference ([`inference::InferCtx`]), and a full type
//! checker ([`TypeChecker`]) that validates CJC programs.
//!
//! # Architecture
//!
//! The type system is organized around several key components:
//!
//! - **[`Type`]** -- The union of all CJC types, from primitives (`i32`, `f64`)
//!   through composites (`Tensor`, `Array`, `Tuple`) to user-defined aggregates
//!   (`Struct`, `Class`, `Record`, `Enum`) and quantum primitives.
//! - **[`TypeEnv`]** -- Scoped environment that tracks variable bindings,
//!   type definitions, trait implementations, and function signatures.
//! - **[`unify`] / [`unify_spanned`]** -- Robinson-style unification with
//!   occurs check, producing type variable bindings in a [`TypeSubst`] map.
//! - **[`inference::InferCtx`]** -- Constraint-based inference context that
//!   collects constraints during type checking and solves them via unification.
//! - **[`TypeChecker`]** -- Two-pass checker (register declarations, then
//!   type-check bodies) that drives the full pipeline.
//! - **[`EffectSet`]** -- Bitflag classification of function side effects
//!   (IO, allocation, GC, nondeterminism, mutation, capture).
//! - **[`effect_registry`]** -- Exhaustive registry mapping every builtin
//!   function name to its [`EffectSet`] flags.
//!
//! # Determinism Guarantees
//!
//! All internal collections use [`BTreeMap`] and [`BTreeSet`] rather than
//! hash-based containers, ensuring iteration order is deterministic across
//! runs and platforms.

pub mod inference;

use std::collections::BTreeMap;
use std::fmt;

// ── Core Type Representation ────────────────────────────────────

/// Unique type identifier used for internal type numbering.
pub type TypeId = usize;

/// Substitution map from type variable IDs to their resolved concrete types.
///
/// Built incrementally during [`unify`] calls. Uses [`BTreeMap`] for
/// deterministic iteration order.
pub type TypeSubst = BTreeMap<TypeVarId, Type>;

/// Substitution map from symbolic shape variable names to their resolved
/// concrete dimensions.
///
/// Built during [`unify_shape_dim`] and [`unify_shapes`] calls.
/// Uses [`BTreeMap`] for deterministic iteration order.
pub type ShapeSubst = BTreeMap<String, usize>;

/// The CJC type system representation.
///
/// Every CJC value has a corresponding `Type`. Types range from numeric
/// primitives through parameterized composites (tensors, arrays, tuples,
/// functions) to user-defined aggregates (structs, classes, records, enums)
/// and quantum simulation primitives.
///
/// Special variants:
/// - [`Type::Var`] -- Unresolved type variable used during inference.
/// - [`Type::Unresolved`] -- Named type not yet resolved (e.g., generic `T`).
/// - [`Type::Error`] -- Sentinel for error recovery; unifies with everything.
#[derive(Debug, Clone, PartialEq)]
pub enum Type {
    /// 32-bit signed integer.
    I32,
    /// 64-bit signed integer (default integer type in CJC).
    I64,
    /// Unsigned 8-bit byte.
    U8,
    /// 32-bit IEEE 754 floating-point number.
    F32,
    /// 64-bit IEEE 754 floating-point number (default float type in CJC).
    F64,
    /// Boolean (`true` / `false`).
    Bool,
    /// Heap-allocated UTF-8 string.
    Str,
    /// Unit type representing no value (e.g., from statements or `print`).
    Void,

    /// Owning byte buffer (heap-allocated `Vec<u8>` semantics).
    Bytes,

    /// Non-owning byte slice view (zero-copy, borrowed from a [`Bytes`](Type::Bytes) or literal).
    ByteSlice,

    /// Validated UTF-8 string view (zero-copy, borrowed from a [`Str`](Type::Str)).
    StrView,

    /// Dense tensor with a typed element and optional static shape.
    ///
    /// When `shape` is `Some`, dimension sizes may be [`ShapeDim::Known`]
    /// (compile-time constant) or [`ShapeDim::Symbolic`] (resolved at runtime).
    Tensor {
        /// Element type (e.g., [`Type::F64`]).
        elem: Box<Type>,
        /// Optional static shape for compile-time shape checking.
        shape: Option<Vec<ShapeDim>>,
    },

    /// Typed buffer (COW-backed contiguous memory region).
    Buffer {
        /// Element type stored in the buffer.
        elem: Box<Type>,
    },

    /// Fixed-length array with a typed element and compile-time length.
    ///
    /// A `len` of `0` indicates a dynamically-sized array (length unknown
    /// at type-check time).
    Array {
        /// Element type of the array.
        elem: Box<Type>,
        /// Static length (`0` = dynamic).
        len: usize,
    },

    /// Heterogeneous, fixed-length tuple.
    Tuple(Vec<Type>),

    /// User-defined struct -- mutable value type with named fields.
    Struct(StructType),

    /// User-defined class -- GC-managed reference type with named fields.
    Class(ClassType),

    /// User-defined record -- immutable value type with named fields and structural equality.
    Record(RecordType),

    /// User-defined enum / algebraic data type (value type, stack-allocatable).
    Enum(EnumType),

    /// Brain float 16-bit type.
    Bf16,

    /// IEEE 754 half-precision 16-bit float type.
    F16,

    /// Range type: `start..end` (half-open interval).
    ///
    /// **Status:** Type-system placeholder. No `Value::Range` exists yet;
    /// the runtime represents ranges via for-loop desugaring. This type is
    /// reserved for future range-literal support.
    Range {
        elem: Box<Type>,
    },

    /// Generic slice (non-owning view into an array).
    ///
    /// **Status:** Type-system placeholder. No `Value::Slice` exists yet;
    /// the runtime uses `ByteSlice` / `StrView` for concrete slice-like
    /// values. This type is reserved for future generic slice support.
    Slice {
        elem: Box<Type>,
    },

    /// Complex number type (f64 real + f64 imaginary).
    Complex,

    /// Compiled regex pattern type.
    Regex,

    /// Function type with typed parameter list and return type.
    Fn {
        /// Positional parameter types.
        params: Vec<Type>,
        /// Return type.
        ret: Box<Type>,
    },

    /// Deterministic ordered map from keys to values.
    ///
    /// Backed by [`BTreeMap`] at runtime for deterministic iteration.
    Map {
        /// Key type.
        key: Box<Type>,
        /// Value type.
        value: Box<Type>,
    },

    /// Sparse tensor in CSR/COO format.
    SparseTensor {
        /// Element type stored in the sparse tensor.
        elem: Box<Type>,
    },

    /// Tidy data view -- a lazy bitmask + column projection over a DataFrame.
    ///
    /// Operations on `TidyView` are deferred until materialization, allowing
    /// predicate pushdown and efficient chaining.
    TidyView,

    /// Grouped tidy data view -- a [`TidyView`](Type::TidyView) partitioned by key columns.
    ///
    /// Produced by `tidy_group_by`. Aggregation functions operate per-group.
    GroupedTidyView,

    // ── Quantum Primitives ──────────────────────────────────────
    //
    // These are first-class types for CJC's built-in quantum simulation
    // stack. Each maps to a concrete Rust type in `cjc-quantum`, wrapped
    // at runtime as `Value::QuantumState(Rc<RefCell<dyn Any>>)`.

    /// Quantum circuit (gate-level description, ≤ 26 qubits).
    QuantumCircuit,

    /// Statevector (full 2^n amplitude vector from circuit execution).
    QuantumStatevector,

    /// Matrix Product State (MPS) for 50+ qubit simulation.
    QuantumMps,

    /// Stabilizer / CHP state for Clifford circuits (1000+ qubits).
    QuantumStabilizer,

    /// Density matrix for mixed states and noise channels.
    QuantumDensity,

    /// Graph structure for QAOA optimization problems.
    QuantumGraph,

    /// Surface code / repetition code for quantum error correction.
    QuantumSurfaceCode,

    /// Unresolved type variable introduced during inference.
    ///
    /// Fresh variables are created by [`TypeEnv::fresh_var`] or
    /// [`inference::InferCtx::fresh_var`] and resolved to concrete types
    /// through [`unify`].
    Var(TypeVarId),

    /// A named type that has not yet been resolved against the type environment.
    ///
    /// Produced when the parser encounters a type parameter name like `T` in
    /// a generic definition. Resolved to a concrete type by
    /// [`substitute_type_params_concrete`] after inference.
    Unresolved(String),

    /// Error sentinel type used for error recovery during type checking.
    ///
    /// Unifies with every other type (see [`unify`]), preventing cascading
    /// errors after a single type mismatch.
    Error,
}

impl Type {
    /// Returns `true` if this type is a numeric type (integer, float, bf16, f16, or complex).
    pub fn is_numeric(&self) -> bool {
        matches!(self, Type::I32 | Type::I64 | Type::U8 | Type::F32 | Type::F64 | Type::Bf16 | Type::F16 | Type::Complex)
    }

    /// Returns `true` if this type is a floating-point type (`f32`, `f64`, `bf16`, or `f16`).
    pub fn is_float(&self) -> bool {
        matches!(self, Type::F32 | Type::F64 | Type::Bf16 | Type::F16)
    }

    /// Returns `true` if this type is an integer type (`i32`, `i64`, or `u8`).
    pub fn is_int(&self) -> bool {
        matches!(self, Type::I32 | Type::I64 | Type::U8)
    }

    /// Returns `true` if this type is a value type (stack-allocatable, no GC needed).
    ///
    /// Value types include all primitives, tensors, buffers, arrays, tuples,
    /// structs, records, enums, maps, ranges, slices, and sparse tensors.
    /// The only non-value type is [`Type::Class`], which requires GC.
    pub fn is_value_type(&self) -> bool {
        matches!(
            self,
            Type::I32
                | Type::I64
                | Type::U8
                | Type::F32
                | Type::F64
                | Type::Bf16
                | Type::F16
                | Type::Bool
                | Type::Str
                | Type::Void
                | Type::Bytes
                | Type::ByteSlice
                | Type::StrView
                | Type::Tensor { .. }
                | Type::Buffer { .. }
                | Type::Array { .. }
                | Type::Tuple(_)
                | Type::Struct(_)
                | Type::Record(_)
                | Type::Enum(_)
                | Type::Map { .. }
                | Type::SparseTensor { .. }
                | Type::Complex
                | Type::Regex
                | Type::Range { .. }
                | Type::Slice { .. }
        )
    }

    /// Returns true if this type is NoGC-safe (no hidden allocations).
    pub fn is_nogc_safe(&self) -> bool {
        matches!(
            self,
            Type::I32
                | Type::I64
                | Type::U8
                | Type::F32
                | Type::F64
                | Type::Bf16
                | Type::F16
                | Type::Bool
                | Type::Complex
                | Type::Void
                | Type::ByteSlice
                | Type::StrView
        )
    }

    /// Returns `true` if this type is a GC-managed reference type.
    ///
    /// Currently only [`Type::Class`] requires garbage collection.
    pub fn is_gc_type(&self) -> bool {
        matches!(self, Type::Class(_))
    }

    /// Returns `true` if this is the [`Type::Error`] sentinel.
    pub fn is_error(&self) -> bool {
        matches!(self, Type::Error)
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Type::I32 => write!(f, "i32"),
            Type::I64 => write!(f, "i64"),
            Type::U8 => write!(f, "u8"),
            Type::F32 => write!(f, "f32"),
            Type::F64 => write!(f, "f64"),
            Type::Bool => write!(f, "bool"),
            Type::Str => write!(f, "String"),
            Type::Void => write!(f, "void"),
            Type::Bytes => write!(f, "Bytes"),
            Type::ByteSlice => write!(f, "ByteSlice"),
            Type::StrView => write!(f, "StrView"),
            Type::Tensor { elem, shape } => {
                write!(f, "Tensor<{}", elem)?;
                if let Some(dims) = shape {
                    write!(f, ", [")?;
                    for (i, dim) in dims.iter().enumerate() {
                        if i > 0 {
                            write!(f, ", ")?;
                        }
                        write!(f, "{}", dim)?;
                    }
                    write!(f, "]")?;
                }
                write!(f, ">")
            }
            Type::Buffer { elem } => write!(f, "Buffer<{}>", elem),
            Type::Array { elem, len } => write!(f, "[{}; {}]", elem, len),
            Type::Tuple(elems) => {
                write!(f, "(")?;
                for (i, elem) in elems.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", elem)?;
                }
                write!(f, ")")
            }
            Type::Struct(s) => write!(f, "{}", s.name),
            Type::Class(c) => write!(f, "{}", c.name),
            Type::Record(r) => write!(f, "{}", r.name),
            Type::Enum(e) => write!(f, "{}", e.name),
            Type::Bf16 => write!(f, "bf16"),
            Type::F16 => write!(f, "f16"),
            Type::Range { elem } => write!(f, "Range<{}>", elem),
            Type::Slice { elem } => write!(f, "Slice<{}>", elem),
            Type::Complex => write!(f, "Complex"),
            Type::Regex => write!(f, "Regex"),
            Type::Fn { params, ret } => {
                write!(f, "fn(")?;
                for (i, p) in params.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", p)?;
                }
                write!(f, ") -> {}", ret)
            }
            Type::Map { key, value } => write!(f, "Map<{}, {}>", key, value),
            Type::SparseTensor { elem } => write!(f, "SparseTensor<{}>", elem),
            Type::TidyView => write!(f, "TidyView"),
            Type::GroupedTidyView => write!(f, "GroupedTidyView"),
            Type::QuantumCircuit => write!(f, "QuantumCircuit"),
            Type::QuantumStatevector => write!(f, "QuantumStatevector"),
            Type::QuantumMps => write!(f, "QuantumMps"),
            Type::QuantumStabilizer => write!(f, "QuantumStabilizer"),
            Type::QuantumDensity => write!(f, "QuantumDensity"),
            Type::QuantumGraph => write!(f, "QuantumGraph"),
            Type::QuantumSurfaceCode => write!(f, "QuantumSurfaceCode"),
            Type::Var(id) => write!(f, "T{}", id.0),
            Type::Unresolved(name) => write!(f, "?{}", name),
            Type::Error => write!(f, "<error>"),
        }
    }
}

// ── Shape Dimensions ────────────────────────────────────────────

/// A single dimension in a tensor shape, either statically known or symbolic.
///
/// Used within [`Type::Tensor`] shape vectors to enable compile-time shape
/// checking and symbolic shape unification.
#[derive(Debug, Clone, PartialEq)]
pub enum ShapeDim {
    /// A concrete, compile-time-known dimension size.
    Known(usize),
    /// A symbolic dimension variable (e.g., `N`, `batch`) resolved at runtime
    /// or unified with other symbolic/known dimensions via [`unify_shape_dim`].
    Symbolic(String),
}

impl fmt::Display for ShapeDim {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ShapeDim::Known(n) => write!(f, "{}", n),
            ShapeDim::Symbolic(name) => write!(f, "{}", name),
        }
    }
}

// ── Type Unification ────────────────────────────────────────────

/// Unify two types using Robinson-style unification, producing bindings for
/// type variables in `subst`.
///
/// Returns the unified (most specific) type on success, or an error message
/// describing the mismatch on failure.
///
/// # Behavior
///
/// - [`Type::Error`] unifies with everything (error recovery).
/// - [`Type::Var`] on either side is bound in `subst` (with occurs check).
/// - Composite types (Tensor, Array, Tuple, Fn, Enum, etc.) are unified
///   structurally by recursing into their components.
/// - Nominal types (Struct, Class, Record, Enum) require matching names.
///
/// # Arguments
///
/// * `a` -- Left-hand type.
/// * `b` -- Right-hand type.
/// * `subst` -- Mutable substitution map accumulating variable bindings.
///
/// # Errors
///
/// Returns `Err(String)` when types are incompatible (e.g., `i32` vs `f64`,
/// array length mismatch, function arity mismatch, or infinite type via
/// occurs check).
pub fn unify(a: &Type, b: &Type, subst: &mut TypeSubst) -> Result<Type, String> {
    // Error type unifies with everything (error recovery)
    if a.is_error() {
        return Ok(b.clone());
    }
    if b.is_error() {
        return Ok(a.clone());
    }

    match (a, b) {
        // Type variable on the left
        (Type::Var(id), _) => {
            if let Some(bound) = subst.get(id).cloned() {
                unify(&bound, b, subst)
            } else {
                // Occurs check: prevent infinite types
                if occurs_in(*id, b, subst) {
                    return Err(format!("infinite type: T{} occurs in {}", id.0, b));
                }
                subst.insert(*id, b.clone());
                Ok(b.clone())
            }
        }
        // Type variable on the right
        (_, Type::Var(id)) => {
            if let Some(bound) = subst.get(id).cloned() {
                unify(a, &bound, subst)
            } else {
                if occurs_in(*id, a, subst) {
                    return Err(format!("infinite type: T{} occurs in {}", id.0, a));
                }
                subst.insert(*id, a.clone());
                Ok(a.clone())
            }
        }
        // Identical primitives
        (Type::I32, Type::I32)
        | (Type::I64, Type::I64)
        | (Type::F32, Type::F32)
        | (Type::F64, Type::F64)
        | (Type::Bf16, Type::Bf16)
        | (Type::F16, Type::F16)
        | (Type::Complex, Type::Complex)
        | (Type::Bool, Type::Bool)
        | (Type::Str, Type::Str)
        | (Type::Void, Type::Void)
        | (Type::QuantumCircuit, Type::QuantumCircuit)
        | (Type::QuantumStatevector, Type::QuantumStatevector)
        | (Type::QuantumMps, Type::QuantumMps)
        | (Type::QuantumStabilizer, Type::QuantumStabilizer)
        | (Type::QuantumDensity, Type::QuantumDensity)
        | (Type::QuantumGraph, Type::QuantumGraph)
        | (Type::QuantumSurfaceCode, Type::QuantumSurfaceCode) => Ok(a.clone()),

        // Range: unify element types
        (Type::Range { elem: e1 }, Type::Range { elem: e2 }) => {
            let elem = unify(e1, e2, subst)?;
            Ok(Type::Range {
                elem: Box::new(elem),
            })
        }

        // Slice: unify element types
        (Type::Slice { elem: e1 }, Type::Slice { elem: e2 }) => {
            let elem = unify(e1, e2, subst)?;
            Ok(Type::Slice {
                elem: Box::new(elem),
            })
        }

        // Tensor: unify element types (shapes handled separately)
        (
            Type::Tensor { elem: e1, shape: s1 },
            Type::Tensor { elem: e2, shape: s2 },
        ) => {
            let elem = unify(e1, e2, subst)?;
            // For shape unification: if both have shapes, verify compatibility
            let shape = match (s1, s2) {
                (Some(sh1), Some(sh2)) => {
                    let mut shape_subst = ShapeSubst::new();
                    Some(unify_shapes(sh1, sh2, &mut shape_subst)?)
                }
                (Some(s), None) | (None, Some(s)) => Some(s.clone()),
                (None, None) => None,
            };
            Ok(Type::Tensor {
                elem: Box::new(elem),
                shape,
            })
        }

        // Buffer
        (Type::Buffer { elem: e1 }, Type::Buffer { elem: e2 }) => {
            let elem = unify(e1, e2, subst)?;
            Ok(Type::Buffer {
                elem: Box::new(elem),
            })
        }

        // Array
        (Type::Array { elem: e1, len: l1 }, Type::Array { elem: e2, len: l2 }) => {
            if l1 != l2 {
                return Err(format!("array length mismatch: {} vs {}", l1, l2));
            }
            let elem = unify(e1, e2, subst)?;
            Ok(Type::Array {
                elem: Box::new(elem),
                len: *l1,
            })
        }

        // Tuple
        (Type::Tuple(t1), Type::Tuple(t2)) => {
            if t1.len() != t2.len() {
                return Err(format!(
                    "tuple length mismatch: {} vs {}",
                    t1.len(),
                    t2.len()
                ));
            }
            let elems: Result<Vec<Type>, String> = t1
                .iter()
                .zip(t2.iter())
                .map(|(a, b)| unify(a, b, subst))
                .collect();
            Ok(Type::Tuple(elems?))
        }

        // Function
        (
            Type::Fn { params: p1, ret: r1 },
            Type::Fn { params: p2, ret: r2 },
        ) => {
            if p1.len() != p2.len() {
                return Err(format!(
                    "function arity mismatch: {} params vs {}",
                    p1.len(),
                    p2.len()
                ));
            }
            let params: Result<Vec<Type>, String> = p1
                .iter()
                .zip(p2.iter())
                .map(|(a, b)| unify(a, b, subst))
                .collect();
            let ret = unify(r1, r2, subst)?;
            Ok(Type::Fn {
                params: params?,
                ret: Box::new(ret),
            })
        }

        // Struct (nominal)
        (Type::Struct(s1), Type::Struct(s2)) => {
            if s1.name == s2.name {
                Ok(a.clone())
            } else {
                Err(format!(
                    "type mismatch: struct `{}` vs struct `{}`",
                    s1.name, s2.name
                ))
            }
        }

        // Class (nominal)
        (Type::Class(c1), Type::Class(c2)) => {
            if c1.name == c2.name {
                Ok(a.clone())
            } else {
                Err(format!(
                    "type mismatch: class `{}` vs class `{}`",
                    c1.name, c2.name
                ))
            }
        }

        // Record (nominal)
        (Type::Record(r1), Type::Record(r2)) => {
            if r1.name == r2.name {
                Ok(a.clone())
            } else {
                Err(format!(
                    "type mismatch: record `{}` vs record `{}`",
                    r1.name, r2.name
                ))
            }
        }

        // Enum (nominal — same name means same enum, unify variant fields pairwise)
        (Type::Enum(e1), Type::Enum(e2)) => {
            if e1.name == e2.name {
                // If both have concrete variant info, unify fields pairwise
                if e1.variants.len() == e2.variants.len() {
                    let mut unified_variants = Vec::new();
                    for (v1, v2) in e1.variants.iter().zip(e2.variants.iter()) {
                        if v1.name != v2.name || v1.fields.len() != v2.fields.len() {
                            return Err(format!(
                                "enum variant mismatch in `{}`: `{}` vs `{}`",
                                e1.name, v1.name, v2.name
                            ));
                        }
                        let mut unified_fields = Vec::new();
                        for (f1, f2) in v1.fields.iter().zip(v2.fields.iter()) {
                            unified_fields.push(unify(f1, f2, subst)?);
                        }
                        unified_variants.push(EnumVariant {
                            name: v1.name.clone(),
                            fields: unified_fields,
                        });
                    }
                    Ok(Type::Enum(EnumType {
                        name: e1.name.clone(),
                        type_params: e1.type_params.clone(),
                        variants: unified_variants,
                    }))
                } else {
                    Ok(a.clone())
                }
            } else {
                Err(format!(
                    "type mismatch: enum `{}` vs enum `{}`",
                    e1.name, e2.name
                ))
            }
        }

        _ => Err(format!("type mismatch: expected `{}`, found `{}`", a, b)),
    }
}

/// Span-aware unification: like `unify()` but emits a spanned `Diagnostic`
/// into `diag` on failure instead of returning `Err(String)`.
///
/// Returns `Ok(unified_type)` on success, or `Ok(Type::Error)` after emitting
/// a spanned error diagnostic on failure.  The caller never sees an `Err` —
/// error recovery proceeds with `Type::Error` which unifies with everything.
pub fn unify_spanned(
    a: &Type,
    b: &Type,
    subst: &mut TypeSubst,
    span: cjc_diag::Span,
    diag: &mut cjc_diag::DiagnosticBag,
) -> Type {
    match unify(a, b, subst) {
        Ok(ty) => ty,
        Err(msg) => {
            diag.emit(
                cjc_diag::Diagnostic::error(
                    "E0100",
                    msg,
                    span,
                )
                .with_hint(format!(
                    "expected `{}`, found `{}`",
                    a, b
                )),
            );
            Type::Error
        }
    }
}

/// Occurs check: does type variable `id` occur in type `ty`?
fn occurs_in(id: TypeVarId, ty: &Type, subst: &TypeSubst) -> bool {
    match ty {
        Type::Var(other) => {
            if *other == id {
                return true;
            }
            if let Some(bound) = subst.get(other) {
                occurs_in(id, bound, subst)
            } else {
                false
            }
        }
        Type::Tensor { elem, .. } => occurs_in(id, elem, subst),
        Type::Buffer { elem } => occurs_in(id, elem, subst),
        Type::Array { elem, .. } => occurs_in(id, elem, subst),
        Type::Tuple(elems) => elems.iter().any(|e| occurs_in(id, e, subst)),
        Type::Fn { params, ret } => {
            params.iter().any(|p| occurs_in(id, p, subst)) || occurs_in(id, ret, subst)
        }
        Type::Enum(e) => e
            .variants
            .iter()
            .any(|v| v.fields.iter().any(|f| occurs_in(id, f, subst))),
        _ => false,
    }
}

/// Apply a type substitution, replacing all bound [`Type::Var`] occurrences
/// with their concrete types from `subst`.
///
/// Recursively follows chains of substitutions (e.g., `T0 -> T1 -> i64`
/// resolves `T0` to `i64`). Non-variable types and unbound variables are
/// returned unchanged.
///
/// # Arguments
///
/// * `ty` -- The type to resolve.
/// * `subst` -- Substitution map from [`unify`] or manual construction.
///
/// # Returns
///
/// The fully resolved type with all known variables replaced.
pub fn apply_subst(ty: &Type, subst: &TypeSubst) -> Type {
    match ty {
        Type::Var(id) => {
            if let Some(bound) = subst.get(id) {
                // Recursively apply in case of chained substitutions
                apply_subst(bound, subst)
            } else {
                ty.clone()
            }
        }
        Type::Tensor { elem, shape } => Type::Tensor {
            elem: Box::new(apply_subst(elem, subst)),
            shape: shape.clone(),
        },
        Type::Buffer { elem } => Type::Buffer {
            elem: Box::new(apply_subst(elem, subst)),
        },
        Type::Array { elem, len } => Type::Array {
            elem: Box::new(apply_subst(elem, subst)),
            len: *len,
        },
        Type::Tuple(elems) => Type::Tuple(elems.iter().map(|e| apply_subst(e, subst)).collect()),
        Type::Fn { params, ret } => Type::Fn {
            params: params.iter().map(|p| apply_subst(p, subst)).collect(),
            ret: Box::new(apply_subst(ret, subst)),
        },
        Type::Enum(e) => Type::Enum(EnumType {
            name: e.name.clone(),
            type_params: e.type_params.clone(),
            variants: e
                .variants
                .iter()
                .map(|v| EnumVariant {
                    name: v.name.clone(),
                    fields: v.fields.iter().map(|f| apply_subst(f, subst)).collect(),
                })
                .collect(),
        }),
        _ => ty.clone(),
    }
}

/// Substitute [`Type::Unresolved`] names with concrete types based on a
/// name-to-type mapping.
///
/// After inferring concrete types for generic type parameters (e.g., `T=i64`
/// from `Some(42)`), this function rewrites the type so that `Option<T>`
/// becomes `Option<i64>`.
///
/// Recursively walks into composite types (Enum, Tuple, Array, Fn, Tensor,
/// Buffer, Map, Range, Slice, SparseTensor). Non-`Unresolved` leaf types
/// are returned unchanged.
///
/// # Arguments
///
/// * `ty` -- The type potentially containing [`Type::Unresolved`] leaves.
/// * `subst` -- Mapping from type parameter names to their inferred concrete types.
///
/// # Returns
///
/// A new type with all matching [`Type::Unresolved`] names replaced.
pub fn substitute_type_params_concrete(ty: &Type, subst: &BTreeMap<String, Type>) -> Type {
    match ty {
        Type::Unresolved(name) => {
            subst.get(name).cloned().unwrap_or_else(|| ty.clone())
        }
        Type::Enum(et) => Type::Enum(EnumType {
            name: et.name.clone(),
            type_params: et.type_params.clone(),
            variants: et
                .variants
                .iter()
                .map(|v| EnumVariant {
                    name: v.name.clone(),
                    fields: v
                        .fields
                        .iter()
                        .map(|f| substitute_type_params_concrete(f, subst))
                        .collect(),
                })
                .collect(),
        }),
        Type::Tuple(ts) => {
            Type::Tuple(ts.iter().map(|t| substitute_type_params_concrete(t, subst)).collect())
        }
        Type::Array { elem, len } => Type::Array {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
            len: *len,
        },
        Type::Fn { params, ret } => Type::Fn {
            params: params.iter().map(|p| substitute_type_params_concrete(p, subst)).collect(),
            ret: Box::new(substitute_type_params_concrete(ret, subst)),
        },
        Type::Tensor { elem, shape } => Type::Tensor {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
            shape: shape.clone(),
        },
        Type::Buffer { elem } => Type::Buffer {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
        },
        Type::Map { key, value } => Type::Map {
            key: Box::new(substitute_type_params_concrete(key, subst)),
            value: Box::new(substitute_type_params_concrete(value, subst)),
        },
        Type::Range { elem } => Type::Range {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
        },
        Type::Slice { elem } => Type::Slice {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
        },
        Type::SparseTensor { elem } => Type::SparseTensor {
            elem: Box::new(substitute_type_params_concrete(elem, subst)),
        },
        _ => ty.clone(),
    }
}

/// Walk a declared type (which may contain `Unresolved` type-param names) in
/// parallel with a concrete type inferred from an argument expression. When a
/// leaf `Unresolved("T")` is paired with a concrete type, record the binding
/// in `subst`.  This is a best-effort structural walk -- mismatches are silently
/// ignored (the type checker will catch them separately).
fn collect_unresolved_bindings(declared: &Type, concrete: &Type, subst: &mut BTreeMap<String, Type>) {
    match (declared, concrete) {
        (Type::Unresolved(name), _) if !concrete.is_error() => {
            subst.entry(name.clone()).or_insert_with(|| concrete.clone());
        }
        (Type::Tuple(ds), Type::Tuple(cs)) if ds.len() == cs.len() => {
            for (d, c) in ds.iter().zip(cs.iter()) {
                collect_unresolved_bindings(d, c, subst);
            }
        }
        (Type::Array { elem: d, .. }, Type::Array { elem: c, .. })
        | (Type::Buffer { elem: d }, Type::Buffer { elem: c })
        | (Type::Tensor { elem: d, .. }, Type::Tensor { elem: c, .. })
        | (Type::Range { elem: d }, Type::Range { elem: c })
        | (Type::Slice { elem: d }, Type::Slice { elem: c })
        | (Type::SparseTensor { elem: d }, Type::SparseTensor { elem: c }) => {
            collect_unresolved_bindings(d, c, subst);
        }
        (Type::Map { key: dk, value: dv }, Type::Map { key: ck, value: cv }) => {
            collect_unresolved_bindings(dk, ck, subst);
            collect_unresolved_bindings(dv, cv, subst);
        }
        (Type::Fn { params: dp, ret: dr }, Type::Fn { params: cp, ret: cr }) if dp.len() == cp.len() => {
            for (d, c) in dp.iter().zip(cp.iter()) {
                collect_unresolved_bindings(d, c, subst);
            }
            collect_unresolved_bindings(dr, cr, subst);
        }
        (Type::Enum(de), Type::Enum(ce)) if de.name == ce.name => {
            for (dv, cv) in de.variants.iter().zip(ce.variants.iter()) {
                for (df, cf) in dv.fields.iter().zip(cv.fields.iter()) {
                    collect_unresolved_bindings(df, cf, subst);
                }
            }
        }
        _ => {}
    }
}

// ── Shape Unification ──────────────────────────────────────────

/// Unify two shape dimensions, producing bindings for symbolic variables.
///
/// - Two [`ShapeDim::Known`] values must be equal.
/// - A [`ShapeDim::Symbolic`] paired with a [`ShapeDim::Known`] binds the
///   symbol to the concrete value in `subst` (or verifies consistency if
///   already bound).
/// - Two distinct [`ShapeDim::Symbolic`] variables are equated if both are
///   unbound.
///
/// # Errors
///
/// Returns `Err` on dimension mismatch or conflicting bindings.
pub fn unify_shape_dim(
    a: &ShapeDim,
    b: &ShapeDim,
    subst: &mut ShapeSubst,
) -> Result<ShapeDim, String> {
    match (a, b) {
        (ShapeDim::Known(n1), ShapeDim::Known(n2)) => {
            if n1 == n2 {
                Ok(ShapeDim::Known(*n1))
            } else {
                Err(format!(
                    "shape dimension mismatch: expected {}, found {}",
                    n1, n2
                ))
            }
        }
        (ShapeDim::Symbolic(s), ShapeDim::Known(n)) | (ShapeDim::Known(n), ShapeDim::Symbolic(s)) => {
            if let Some(&bound) = subst.get(s.as_str()) {
                if bound == *n {
                    Ok(ShapeDim::Known(*n))
                } else {
                    Err(format!(
                        "symbolic shape `{}` already bound to {}, cannot unify with {}",
                        s, bound, n
                    ))
                }
            } else {
                subst.insert(s.clone(), *n);
                Ok(ShapeDim::Known(*n))
            }
        }
        (ShapeDim::Symbolic(s1), ShapeDim::Symbolic(s2)) => {
            if s1 == s2 {
                Ok(ShapeDim::Symbolic(s1.clone()))
            } else {
                // If one is bound, use that
                match (subst.get(s1).copied(), subst.get(s2).copied()) {
                    (Some(n1), Some(n2)) => {
                        if n1 == n2 {
                            Ok(ShapeDim::Known(n1))
                        } else {
                            Err(format!(
                                "shape conflict: `{}` = {} but `{}` = {}",
                                s1, n1, s2, n2
                            ))
                        }
                    }
                    (Some(n), None) => {
                        subst.insert(s2.clone(), n);
                        Ok(ShapeDim::Known(n))
                    }
                    (None, Some(n)) => {
                        subst.insert(s1.clone(), n);
                        Ok(ShapeDim::Known(n))
                    }
                    (None, None) => {
                        // Both unbound: keep symbolic, they are equated
                        Ok(ShapeDim::Symbolic(s1.clone()))
                    }
                }
            }
        }
    }
}

/// Unify two shape vectors dimension by dimension.
///
/// Both shapes must have the same rank (number of dimensions). Each
/// dimension pair is unified via [`unify_shape_dim`].
///
/// # Errors
///
/// Returns `Err` on rank mismatch or any individual dimension mismatch.
pub fn unify_shapes(
    a: &[ShapeDim],
    b: &[ShapeDim],
    subst: &mut ShapeSubst,
) -> Result<Vec<ShapeDim>, String> {
    if a.len() != b.len() {
        return Err(format!(
            "rank mismatch: {}-D tensor vs {}-D tensor",
            a.len(),
            b.len()
        ));
    }
    a.iter()
        .zip(b.iter())
        .enumerate()
        .map(|(i, (da, db))| {
            unify_shape_dim(da, db, subst).map_err(|e| format!("dimension {}: {}", i, e))
        })
        .collect()
}

/// Compute the broadcast-compatible shape from two shapes following NumPy
/// broadcasting rules.
///
/// Shapes are right-aligned and padded with `Known(1)` on the left for the
/// shorter shape. Each dimension pair must either be equal or one of them
/// must be `1`. If either input is `None`, returns `None` (unknown shape).
///
/// # Errors
///
/// Returns `Err` if any dimension pair is incompatible for broadcasting.
pub fn broadcast_shapes(
    a: &Option<Vec<ShapeDim>>,
    b: &Option<Vec<ShapeDim>>,
) -> Result<Option<Vec<ShapeDim>>, String> {
    match (a, b) {
        (Some(sa), Some(sb)) => {
            let max_rank = sa.len().max(sb.len());
            let mut result = Vec::with_capacity(max_rank);

            // Right-align: pad shorter shape with Known(1) on the left
            for i in 0..max_rank {
                let da = if i < max_rank - sa.len() {
                    &ShapeDim::Known(1)
                } else {
                    &sa[i - (max_rank - sa.len())]
                };
                let db = if i < max_rank - sb.len() {
                    &ShapeDim::Known(1)
                } else {
                    &sb[i - (max_rank - sb.len())]
                };

                let dim = broadcast_dim(da, db, i)?;
                result.push(dim);
            }
            Ok(Some(result))
        }
        _ => Ok(None), // No shape info: pass through
    }
}

/// Broadcast a single dimension pair.
fn broadcast_dim(a: &ShapeDim, b: &ShapeDim, dim_idx: usize) -> Result<ShapeDim, String> {
    match (a, b) {
        (ShapeDim::Known(1), other) | (other, ShapeDim::Known(1)) => Ok(other.clone()),
        (ShapeDim::Known(n1), ShapeDim::Known(n2)) => {
            if n1 == n2 {
                Ok(ShapeDim::Known(*n1))
            } else {
                Err(format!(
                    "broadcast incompatible at dimension {}: {} vs {}",
                    dim_idx, n1, n2
                ))
            }
        }
        (ShapeDim::Symbolic(s), ShapeDim::Known(n)) | (ShapeDim::Known(n), ShapeDim::Symbolic(s)) => {
            if *n == 1 {
                Ok(ShapeDim::Symbolic(s.clone()))
            } else {
                // Symbolic could be 1 (broadcasts) or n (equal); conservatively allow
                Ok(ShapeDim::Known(*n))
            }
        }
        (ShapeDim::Symbolic(s1), ShapeDim::Symbolic(s2)) => {
            if s1 == s2 {
                Ok(ShapeDim::Symbolic(s1.clone()))
            } else {
                // Two different symbolic dims: conservatively keep first
                Ok(ShapeDim::Symbolic(s1.clone()))
            }
        }
    }
}

// ── Struct, Class, and Enum Type Info ────────────────────────────

/// User-defined struct type -- a mutable value type with named fields.
///
/// Structs are stack-allocatable and support field mutation. They use
/// nominal typing: two structs with the same fields but different names
/// are distinct types.
#[derive(Debug, Clone, PartialEq)]
pub struct StructType {
    /// The struct's declared name (e.g., `"Point"`).
    pub name: String,
    /// Generic type parameter names (e.g., `["T", "U"]`).
    pub type_params: Vec<String>,
    /// Ordered list of `(field_name, field_type)` pairs.
    pub fields: Vec<(String, Type)>,
}

/// User-defined class type -- a GC-managed reference type with named fields.
///
/// Classes are heap-allocated and garbage-collected. They are the only type
/// that returns `true` from [`Type::is_gc_type`].
#[derive(Debug, Clone, PartialEq)]
pub struct ClassType {
    /// The class's declared name.
    pub name: String,
    /// Generic type parameter names.
    pub type_params: Vec<String>,
    /// Ordered list of `(field_name, field_type)` pairs.
    pub fields: Vec<(String, Type)>,
}

/// User-defined record type -- an immutable value type with structural equality.
///
/// Records differ from structs in two ways: their fields cannot be mutated
/// after construction, and assigning to a record field produces a type error
/// (E0160).
#[derive(Debug, Clone, PartialEq)]
pub struct RecordType {
    /// The record's declared name.
    pub name: String,
    /// Generic type parameter names.
    pub type_params: Vec<String>,
    /// Ordered list of `(field_name, field_type)` pairs.
    pub fields: Vec<(String, Type)>,
}

/// User-defined enum type -- an algebraic data type (ADT) with named variants.
///
/// Enums are value types (stack-allocatable) and support pattern matching
/// via `match` expressions. Variant fields can contain generic type
/// parameters represented as [`Type::Unresolved`].
#[derive(Debug, Clone, PartialEq)]
pub struct EnumType {
    /// The enum's declared name (e.g., `"Option"`).
    pub name: String,
    /// Generic type parameter names (e.g., `["T"]` for `Option<T>`).
    pub type_params: Vec<String>,
    /// The enum's variants in declaration order.
    pub variants: Vec<EnumVariant>,
}

/// A single variant within an [`EnumType`].
///
/// Variants may carry zero or more positional fields (like tuple structs).
/// For example, `Some(T)` has one field, while `None` has zero.
#[derive(Debug, Clone, PartialEq)]
pub struct EnumVariant {
    /// The variant's name (e.g., `"Some"`, `"None"`).
    pub name: String,
    /// Positional field types (empty for unit variants like `None`).
    pub fields: Vec<Type>,
}

// ── Type Variables ──────────────────────────────────────────────

/// Unique identifier for a type variable, used during type inference.
///
/// Fresh IDs are allocated sequentially by [`TypeEnv::fresh_var`] and
/// [`inference::InferCtx::fresh_var`]. The inner `usize` is a monotonically
/// increasing counter.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct TypeVarId(pub usize);

// ── Traits / Typeclasses ────────────────────────────────────────

/// Definition of a trait (typeclass) with optional supertrait hierarchy and methods.
///
/// CJC has a built-in trait hierarchy (see [`builtin_trait_defs`]):
/// `Numeric` -> `Int` / `Float` -> `Differentiable`. User-defined traits
/// extend this system.
#[derive(Debug, Clone, PartialEq)]
pub struct TraitDef {
    /// The trait's name (e.g., `"Numeric"`, `"Float"`).
    pub name: String,
    /// Generic type parameter names for the trait itself.
    pub type_params: Vec<String>,
    /// Names of supertraits that implementors must also satisfy.
    pub super_traits: Vec<String>,
    /// Required method signatures.
    pub methods: Vec<MethodSig>,
}

/// Signature of a single method within a [`TraitDef`].
#[derive(Debug, Clone, PartialEq)]
pub struct MethodSig {
    /// Method name (e.g., `"zero"`, `"sqrt"`).
    pub name: String,
    /// Positional parameter types. `Type::Var(TypeVarId(0))` represents `Self`.
    pub params: Vec<Type>,
    /// Return type.
    pub ret: Type,
}

/// A trait implementation declaring that a concrete type satisfies a trait.
///
/// Registered in [`TypeEnv::trait_impls`] and queried by
/// [`TypeEnv::satisfies_trait`] during generic bound checking.
#[derive(Debug, Clone)]
pub struct TraitImpl {
    /// The trait being implemented (e.g., `"Numeric"`).
    pub trait_name: String,
    /// The concrete type implementing the trait (e.g., [`Type::F64`]).
    pub target_type: Type,
    /// Type arguments for parameterized traits (currently unused).
    pub type_args: Vec<Type>,
}

// ── Built-in Trait Hierarchy ────────────────────────────────────

/// Built-in trait name: all numeric types (`i32`, `i64`, `u8`, `f32`, `f64`, `bf16`, `f16`).
pub const TRAIT_NUMERIC: &str = "Numeric";
/// Built-in trait name: floating-point types (`f32`, `f64`, `bf16`, `f16`). Subtrait of `Numeric`.
pub const TRAIT_FLOAT: &str = "Float";
/// Built-in trait name: integer types (`i32`, `i64`, `u8`). Subtrait of `Numeric`.
pub const TRAIT_INT: &str = "Int";
/// Built-in trait name: types supporting automatic differentiation (`f32`, `f64`). Subtrait of `Float`.
pub const TRAIT_DIFFERENTIABLE: &str = "Differentiable";
/// Built-in trait name: types supporting ordering comparisons.
pub const TRAIT_ORD: &str = "Ord";
/// Built-in trait name: types that can be formatted for display.
pub const TRAIT_DISPLAY: &str = "Display";
/// Built-in trait name: types that support binary serialization via `cjc-snap`.
pub const TRAIT_SNAP: &str = "Snap";

/// Return the built-in trait definitions forming the CJC trait hierarchy.
///
/// ```text
///   Numeric
///   ├── Int
///   └── Float
///       └── Differentiable
///   Ord
///   Display
///   Snap
/// ```
///
/// Each [`TraitDef`] includes its supertrait relationships and required
/// method signatures. These are registered into [`TypeEnv::trait_defs`]
/// during [`TypeEnv::new`].
pub fn builtin_trait_defs() -> Vec<TraitDef> {
    vec![
        TraitDef {
            name: TRAIT_NUMERIC.into(),
            type_params: vec![],
            super_traits: vec![],
            methods: vec![
                MethodSig {
                    name: "zero".into(),
                    params: vec![],
                    ret: Type::Var(TypeVarId(0)), // Self
                },
                MethodSig {
                    name: "one".into(),
                    params: vec![],
                    ret: Type::Var(TypeVarId(0)),
                },
            ],
        },
        TraitDef {
            name: TRAIT_INT.into(),
            type_params: vec![],
            super_traits: vec![TRAIT_NUMERIC.into()],
            methods: vec![],
        },
        TraitDef {
            name: TRAIT_FLOAT.into(),
            type_params: vec![],
            super_traits: vec![TRAIT_NUMERIC.into()],
            methods: vec![
                MethodSig {
                    name: "sqrt".into(),
                    params: vec![Type::Var(TypeVarId(0))],
                    ret: Type::Var(TypeVarId(0)),
                },
                MethodSig {
                    name: "ln".into(),
                    params: vec![Type::Var(TypeVarId(0))],
                    ret: Type::Var(TypeVarId(0)),
                },
                MethodSig {
                    name: "exp".into(),
                    params: vec![Type::Var(TypeVarId(0))],
                    ret: Type::Var(TypeVarId(0)),
                },
            ],
        },
        TraitDef {
            name: TRAIT_DIFFERENTIABLE.into(),
            type_params: vec![],
            super_traits: vec![TRAIT_FLOAT.into()],
            methods: vec![],
        },
        TraitDef {
            name: TRAIT_ORD.into(),
            type_params: vec![],
            super_traits: vec![],
            methods: vec![],
        },
        TraitDef {
            name: TRAIT_DISPLAY.into(),
            type_params: vec![],
            super_traits: vec![],
            methods: vec![],
        },
        TraitDef {
            name: TRAIT_SNAP.into(),
            type_params: vec![],
            super_traits: vec![],
            methods: vec![],
        },
    ]
}

/// Return the built-in trait implementations mapping primitive types to
/// their satisfied traits.
///
/// Includes implementations of `Numeric`, `Int`, `Float`, `Differentiable`,
/// `Ord`, `Display`, and `Snap` for the appropriate primitive types.
pub fn builtin_trait_impls() -> Vec<TraitImpl> {
    vec![
        // i32: Numeric, Int
        TraitImpl { trait_name: TRAIT_NUMERIC.into(), target_type: Type::I32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_INT.into(), target_type: Type::I32, type_args: vec![] },
        // i64: Numeric, Int
        TraitImpl { trait_name: TRAIT_NUMERIC.into(), target_type: Type::I64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_INT.into(), target_type: Type::I64, type_args: vec![] },
        // f32: Numeric, Float, Differentiable
        TraitImpl { trait_name: TRAIT_NUMERIC.into(), target_type: Type::F32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_FLOAT.into(), target_type: Type::F32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DIFFERENTIABLE.into(), target_type: Type::F32, type_args: vec![] },
        // f64: Numeric, Float, Differentiable
        TraitImpl { trait_name: TRAIT_NUMERIC.into(), target_type: Type::F64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_FLOAT.into(), target_type: Type::F64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DIFFERENTIABLE.into(), target_type: Type::F64, type_args: vec![] },
        // Ord: i32, i64, f32, f64, u8, str, bool
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::I32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::I64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::U8, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::F32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::F64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::Str, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_ORD.into(), target_type: Type::Bool, type_args: vec![] },
        // Display: all primitives
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::I32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::I64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::U8, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::F32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::F64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::Bool, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_DISPLAY.into(), target_type: Type::Str, type_args: vec![] },
        // Snap: all value types (primitives + structs)
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::I32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::I64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::U8, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::F32, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::F64, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::Bool, type_args: vec![] },
        TraitImpl { trait_name: TRAIT_SNAP.into(), target_type: Type::Str, type_args: vec![] },
    ]
}

// ── Type Environment ────────────────────────────────────────────

/// Type checking environment -- the central registry for types, traits,
/// variables, and function signatures.
///
/// [`TypeEnv`] is a scoped environment that supports nested lexical scopes
/// via [`push_scope`](TypeEnv::push_scope) / [`pop_scope`](TypeEnv::pop_scope).
/// It is created with [`TypeEnv::new`], which pre-registers all built-in
/// types, traits, and quantum function signatures.
///
/// All internal maps use [`BTreeMap`] for deterministic iteration order.
pub struct TypeEnv {
    /// Named type definitions (e.g., `"i32" -> Type::I32`, `"Option" -> Type::Enum(...)`).
    pub type_defs: BTreeMap<String, Type>,
    /// Trait definitions keyed by trait name.
    pub trait_defs: BTreeMap<String, TraitDef>,
    /// All registered trait implementations.
    pub trait_impls: Vec<TraitImpl>,
    /// Variable scopes as a stack of scope frames.
    ///
    /// Each frame is a map from variable name to `(type, is_mutable)`.
    /// The innermost (last) frame is the current scope.
    scopes: Vec<BTreeMap<String, (Type, bool)>>,
    /// Compile-time constant definitions (name -> value type). Always immutable.
    pub const_defs: BTreeMap<String, Type>,
    /// Function signature table supporting multiple dispatch (overloading).
    ///
    /// Each function name maps to a list of [`FnSigEntry`] overloads.
    pub fn_sigs: BTreeMap<String, Vec<FnSigEntry>>,
    /// Monotonically increasing counter for generating fresh [`TypeVarId`]s.
    next_var: usize,
}

// ---------------------------------------------------------------------------
// Effect Classification
// ---------------------------------------------------------------------------

/// Bitflag effect classification for builtin functions.
///
/// Each flag represents a category of side effect. A function with `PURE` (no
/// flags set) is guaranteed to be deterministic and allocation-free.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct EffectSet {
    bits: u16,
}

impl Default for EffectSet {
    fn default() -> Self {
        Self::PURE
    }
}

impl EffectSet {
    /// No effects — pure, deterministic, no allocation.
    pub const PURE: Self = Self { bits: 0 };

    // Individual effect flags.
    pub const IO: u16 = 0b0000_0001; // print, file read/write
    pub const ALLOC: u16 = 0b0000_0010; // heap allocation (Rc/String/Vec/Tensor)
    pub const GC: u16 = 0b0000_0100; // triggers GC (gc_alloc, gc_collect)
    pub const NONDET: u16 = 0b0000_1000; // nondeterministic (randn, clock, hash-order)
    pub const MUTATES: u16 = 0b0001_0000; // mutates arguments (push, Map.insert)
    pub const ARENA_OK: u16 = 0b0010_0000; // result can safely live on arena
    pub const CAPTURES: u16 = 0b0100_0000; // may capture/store arguments beyond call

    /// Create an EffectSet from raw bits.
    pub const fn new(bits: u16) -> Self {
        Self { bits }
    }

    /// Check whether a specific flag is set.
    pub const fn has(&self, flag: u16) -> bool {
        self.bits & flag != 0
    }

    /// Return a new set with the given flag added.
    pub const fn with(self, flag: u16) -> Self {
        Self {
            bits: self.bits | flag,
        }
    }

    /// True if no flags are set (pure function).
    pub const fn is_pure(&self) -> bool {
        self.bits == 0
    }

    /// True if the function does NOT trigger GC.
    pub const fn is_nogc_safe(&self) -> bool {
        !self.has(Self::GC)
    }

    /// True if the function may allocate (heap or GC).
    pub const fn may_alloc(&self) -> bool {
        self.bits & (Self::ALLOC | Self::GC) != 0
    }

    /// Raw bits (for serialization / debug).
    pub const fn bits(&self) -> u16 {
        self.bits
    }
}

impl fmt::Display for EffectSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_pure() {
            return write!(f, "PURE");
        }
        let mut parts = Vec::new();
        if self.has(Self::IO) {
            parts.push("IO");
        }
        if self.has(Self::ALLOC) {
            parts.push("ALLOC");
        }
        if self.has(Self::GC) {
            parts.push("GC");
        }
        if self.has(Self::NONDET) {
            parts.push("NONDET");
        }
        if self.has(Self::MUTATES) {
            parts.push("MUTATES");
        }
        if self.has(Self::ARENA_OK) {
            parts.push("ARENA_OK");
        }
        if self.has(Self::CAPTURES) {
            parts.push("CAPTURES");
        }
        write!(f, "{}", parts.join("|"))
    }
}

// ---------------------------------------------------------------------------
// Effect Registry
// ---------------------------------------------------------------------------

pub mod effect_registry;

// ---------------------------------------------------------------------------
// Function Signature Entry
// ---------------------------------------------------------------------------

/// Entry in the function signature table, supporting generic type parameters
/// and multiple dispatch (overloading by arity/type).
///
/// Registered via [`TypeEnv::register_fn`] and looked up during
/// [`TypeChecker::check_fn_call`] for overload resolution.
#[derive(Debug, Clone)]
pub struct FnSigEntry {
    /// Function name (or qualified name like `"Tensor.zeros"`).
    pub name: String,
    /// Generic type parameters as `(name, trait_bounds)` pairs.
    ///
    /// For example, `[("T", ["Numeric", "Float"])]` means type parameter `T`
    /// must implement both `Numeric` and `Float`.
    pub type_params: Vec<(String, Vec<String>)>,
    /// Positional parameters as `(param_name, param_type)` pairs.
    pub params: Vec<(String, Type)>,
    /// Return type.
    pub ret: Type,
    /// Whether this function is declared as `@nogc` (no GC allocations allowed).
    pub is_nogc: bool,
    /// Effect classification for this function.
    ///
    /// Used by the NoGC verifier (`cjc-mir/nogc_verify`), escape analysis
    /// (`cjc-mir/escape`), and the optimizer to reason about side effects.
    pub effects: EffectSet,
}

impl TypeEnv {
    pub fn new() -> Self {
        let mut env = Self {
            type_defs: BTreeMap::new(),
            trait_defs: BTreeMap::new(),
            trait_impls: Vec::new(),
            scopes: vec![BTreeMap::new()],
            fn_sigs: BTreeMap::new(),
            next_var: 0,
            const_defs: BTreeMap::new(),
        };

        // Register built-in types
        env.type_defs.insert("i32".into(), Type::I32);
        env.type_defs.insert("i64".into(), Type::I64);
        env.type_defs.insert("f32".into(), Type::F32);
        env.type_defs.insert("f64".into(), Type::F64);
        env.type_defs.insert("bf16".into(), Type::Bf16);
        env.type_defs.insert("f16".into(), Type::F16);
        env.type_defs.insert("Complex".into(), Type::Complex);
        env.type_defs.insert("bool".into(), Type::Bool);
        env.type_defs.insert("String".into(), Type::Str);

        // Register prelude enum types: Option<T> and Result<T, E>
        env.type_defs.insert(
            "Option".into(),
            Type::Enum(EnumType {
                name: "Option".into(),
                type_params: vec!["T".into()],
                variants: vec![
                    EnumVariant {
                        name: "Some".into(),
                        fields: vec![Type::Unresolved("T".into())],
                    },
                    EnumVariant {
                        name: "None".into(),
                        fields: vec![],
                    },
                ],
            }),
        );
        env.type_defs.insert(
            "Result".into(),
            Type::Enum(EnumType {
                name: "Result".into(),
                type_params: vec!["T".into(), "E".into()],
                variants: vec![
                    EnumVariant {
                        name: "Ok".into(),
                        fields: vec![Type::Unresolved("T".into())],
                    },
                    EnumVariant {
                        name: "Err".into(),
                        fields: vec![Type::Unresolved("E".into())],
                    },
                ],
            }),
        );

        // Register Error struct (for error metadata)
        env.type_defs.insert(
            "Error".into(),
            Type::Struct(StructType {
                name: "Error".into(),
                type_params: vec![],
                fields: vec![
                    ("span".into(), Type::I64),
                    ("message".into(), Type::Str),
                ],
            }),
        );

        // Register built-in traits
        for def in builtin_trait_defs() {
            env.trait_defs.insert(def.name.clone(), def);
        }

        // Register built-in impls
        env.trait_impls = builtin_trait_impls();

        // Register bf16 trait impls
        env.trait_impls.push(TraitImpl {
            trait_name: TRAIT_NUMERIC.into(),
            target_type: Type::Bf16,
            type_args: vec![],
        });
        env.trait_impls.push(TraitImpl {
            trait_name: TRAIT_FLOAT.into(),
            target_type: Type::Bf16,
            type_args: vec![],
        });

        // Register f16 trait impls
        env.trait_impls.push(TraitImpl {
            trait_name: TRAIT_NUMERIC.into(),
            target_type: Type::F16,
            type_args: vec![],
        });
        env.trait_impls.push(TraitImpl {
            trait_name: TRAIT_FLOAT.into(),
            target_type: Type::F16,
            type_args: vec![],
        });

        // Register prelude variant constructors as functions
        // Some(v) -> Option<T>
        env.fn_sigs.entry("Some".into()).or_default().push(FnSigEntry {
            name: "Some".into(),
            type_params: vec![("T".into(), vec![])],
            params: vec![("_0".into(), Type::Unresolved("T".into()))],
            ret: Type::Enum(EnumType {
                name: "Option".into(),
                type_params: vec!["T".into()],
                variants: vec![
                    EnumVariant { name: "Some".into(), fields: vec![Type::Unresolved("T".into())] },
                    EnumVariant { name: "None".into(), fields: vec![] },
                ],
            }),
            is_nogc: false,
            effects: EffectSet::default(),
        });
        // None -> Option<T>
        env.fn_sigs.entry("None".into()).or_default().push(FnSigEntry {
            name: "None".into(),
            type_params: vec![("T".into(), vec![])],
            params: vec![],
            ret: Type::Enum(EnumType {
                name: "Option".into(),
                type_params: vec!["T".into()],
                variants: vec![
                    EnumVariant { name: "Some".into(), fields: vec![Type::Unresolved("T".into())] },
                    EnumVariant { name: "None".into(), fields: vec![] },
                ],
            }),
            is_nogc: false,
            effects: EffectSet::default(),
        });
        // Ok(v) -> Result<T, E>
        env.fn_sigs.entry("Ok".into()).or_default().push(FnSigEntry {
            name: "Ok".into(),
            type_params: vec![("T".into(), vec![]), ("E".into(), vec![])],
            params: vec![("_0".into(), Type::Unresolved("T".into()))],
            ret: Type::Enum(EnumType {
                name: "Result".into(),
                type_params: vec!["T".into(), "E".into()],
                variants: vec![
                    EnumVariant { name: "Ok".into(), fields: vec![Type::Unresolved("T".into())] },
                    EnumVariant { name: "Err".into(), fields: vec![Type::Unresolved("E".into())] },
                ],
            }),
            is_nogc: false,
            effects: EffectSet::default(),
        });
        // Err(e) -> Result<T, E>
        env.fn_sigs.entry("Err".into()).or_default().push(FnSigEntry {
            name: "Err".into(),
            type_params: vec![("T".into(), vec![]), ("E".into(), vec![])],
            params: vec![("_0".into(), Type::Unresolved("E".into()))],
            ret: Type::Enum(EnumType {
                name: "Result".into(),
                type_params: vec!["T".into(), "E".into()],
                variants: vec![
                    EnumVariant { name: "Ok".into(), fields: vec![Type::Unresolved("T".into())] },
                    EnumVariant { name: "Err".into(), fields: vec![Type::Unresolved("E".into())] },
                ],
            }),
            is_nogc: false,
            effects: EffectSet::default(),
        });

        // ── Quantum Primitive Types ─────────────────────────────────
        env.type_defs.insert("QuantumCircuit".into(), Type::QuantumCircuit);
        env.type_defs.insert("QuantumStatevector".into(), Type::QuantumStatevector);
        env.type_defs.insert("QuantumMps".into(), Type::QuantumMps);
        env.type_defs.insert("QuantumStabilizer".into(), Type::QuantumStabilizer);
        env.type_defs.insert("QuantumDensity".into(), Type::QuantumDensity);
        env.type_defs.insert("QuantumGraph".into(), Type::QuantumGraph);
        env.type_defs.insert("QuantumSurfaceCode".into(), Type::QuantumSurfaceCode);

        // ── Quantum Builtin Function Signatures ─────────────────────
        register_quantum_builtins(&mut env);

        env
    }

    pub fn fresh_var(&mut self) -> TypeVarId {
        let id = TypeVarId(self.next_var);
        self.next_var += 1;
        id
    }

    pub fn push_scope(&mut self) {
        self.scopes.push(BTreeMap::new());
    }

    pub fn pop_scope(&mut self) {
        self.scopes.pop();
    }

    /// Define an immutable variable in the current scope.
    pub fn define_var(&mut self, name: &str, ty: Type) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.insert(name.to_string(), (ty, false));
        }
    }

    /// Define a mutable variable in the current scope.
    pub fn define_var_mut(&mut self, name: &str, ty: Type) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.insert(name.to_string(), (ty, true));
        }
    }

    pub fn lookup_var(&self, name: &str) -> Option<&Type> {
        for scope in self.scopes.iter().rev() {
            if let Some((ty, _)) = scope.get(name) {
                return Some(ty);
            }
        }
        None
    }

    /// Look up a variable and return its (type, is_mutable) pair.
    pub fn lookup_var_entry(&self, name: &str) -> Option<(&Type, bool)> {
        for scope in self.scopes.iter().rev() {
            if let Some((ty, mutable)) = scope.get(name) {
                return Some((ty, *mutable));
            }
        }
        None
    }

    /// Returns true if the named variable exists AND is mutable.
    pub fn is_var_mutable(&self, name: &str) -> bool {
        self.lookup_var_entry(name).map(|(_, m)| m).unwrap_or(false)
    }

    pub fn resolve_type_name(&self, name: &str) -> Option<Type> {
        self.type_defs.get(name).cloned()
    }

    /// Check if a type satisfies a trait bound.
    pub fn satisfies_trait(&self, ty: &Type, trait_name: &str) -> bool {
        // Error type satisfies all traits (for error recovery)
        if ty.is_error() {
            return true;
        }

        // Check direct implementations
        for imp in &self.trait_impls {
            if imp.trait_name == trait_name && self.types_match(&imp.target_type, ty) {
                return true;
            }
        }

        // Check sub-trait relationships
        if let Some(trait_def) = self.trait_defs.get(trait_name) {
            // If the type satisfies a sub-trait, it also satisfies super-traits
            for sub_trait_name in self.trait_defs.keys().cloned().collect::<Vec<_>>() {
                if sub_trait_name == trait_name {
                    continue;
                }
                if let Some(sub_def) = self.trait_defs.get(&sub_trait_name) {
                    if sub_def.super_traits.contains(&trait_name.to_string()) {
                        // If type implements the sub-trait, it implements this trait
                        for imp in &self.trait_impls {
                            if imp.trait_name == sub_trait_name
                                && self.types_match(&imp.target_type, ty)
                            {
                                return true;
                            }
                        }
                    }
                }
            }
            // Also check if the trait has super_traits that the type satisfies
            let _ = trait_def;
        }

        false
    }

    /// Check if two types match (simple structural equality for v1).
    pub fn types_match(&self, a: &Type, b: &Type) -> bool {
        if a.is_error() || b.is_error() {
            return true;
        }
        match (a, b) {
            (Type::I32, Type::I32)
            | (Type::I64, Type::I64)
            | (Type::F32, Type::F32)
            | (Type::F64, Type::F64)
            | (Type::Bool, Type::Bool)
            | (Type::Str, Type::Str)
            | (Type::Void, Type::Void) => true,
            (Type::Tensor { elem: e1, .. }, Type::Tensor { elem: e2, .. }) => {
                self.types_match(e1, e2)
            }
            (Type::Buffer { elem: e1 }, Type::Buffer { elem: e2 }) => self.types_match(e1, e2),
            (
                Type::Array {
                    elem: e1, len: l1, ..
                },
                Type::Array {
                    elem: e2, len: l2, ..
                },
            ) => l1 == l2 && self.types_match(e1, e2),
            (Type::Struct(s1), Type::Struct(s2)) => s1.name == s2.name,
            (Type::Class(c1), Type::Class(c2)) => c1.name == c2.name,
            (Type::Enum(e1), Type::Enum(e2)) => e1.name == e2.name,
            (Type::Bf16, Type::Bf16) => true,
            (Type::F16, Type::F16) => true,
            (Type::Range { elem: e1 }, Type::Range { elem: e2 }) => self.types_match(e1, e2),
            (Type::Slice { elem: e1 }, Type::Slice { elem: e2 }) => self.types_match(e1, e2),
            (
                Type::Fn {
                    params: p1,
                    ret: r1,
                },
                Type::Fn {
                    params: p2,
                    ret: r2,
                },
            ) => {
                p1.len() == p2.len()
                    && p1.iter().zip(p2).all(|(a, b)| self.types_match(a, b))
                    && self.types_match(r1, r2)
            }
            (Type::Map { key: k1, value: v1 }, Type::Map { key: k2, value: v2 }) => {
                self.types_match(k1, k2) && self.types_match(v1, v2)
            }
            (Type::SparseTensor { elem: e1 }, Type::SparseTensor { elem: e2 }) => {
                self.types_match(e1, e2)
            }
            (Type::QuantumCircuit, Type::QuantumCircuit)
            | (Type::QuantumStatevector, Type::QuantumStatevector)
            | (Type::QuantumMps, Type::QuantumMps)
            | (Type::QuantumStabilizer, Type::QuantumStabilizer)
            | (Type::QuantumDensity, Type::QuantumDensity)
            | (Type::QuantumGraph, Type::QuantumGraph)
            | (Type::QuantumSurfaceCode, Type::QuantumSurfaceCode) => true,
            (Type::Var(a), Type::Var(b)) => a == b,
            _ => false,
        }
    }

    /// Check shapes are compatible for matmul: [M, K] x [K, N] -> [M, N].
    pub fn check_matmul_shapes(
        &self,
        a_shape: &Option<Vec<ShapeDim>>,
        b_shape: &Option<Vec<ShapeDim>>,
    ) -> Result<Option<Vec<ShapeDim>>, String> {
        match (a_shape, b_shape) {
            (Some(a), Some(b)) => {
                if a.len() != 2 || b.len() != 2 {
                    return Err(format!(
                        "matmul requires 2D tensors, got {}D and {}D",
                        a.len(),
                        b.len()
                    ));
                }
                // Check inner dimensions match
                match (&a[1], &b[0]) {
                    (ShapeDim::Known(k1), ShapeDim::Known(k2)) => {
                        if k1 != k2 {
                            return Err(format!(
                                "matmul inner dimension mismatch: {} vs {}",
                                k1, k2
                            ));
                        }
                    }
                    (ShapeDim::Symbolic(s1), ShapeDim::Symbolic(s2)) => {
                        if s1 != s2 {
                            return Err(format!(
                                "matmul inner dimension mismatch: {} vs {}",
                                s1, s2
                            ));
                        }
                    }
                    _ => {} // Mixed known/symbolic: allow (runtime check)
                }
                Ok(Some(vec![a[0].clone(), b[1].clone()]))
            }
            _ => Ok(None), // No shape info: pass through
        }
    }

    /// Check if a concrete type satisfies all listed trait bounds.
    pub fn check_bounds(&self, ty: &Type, bounds: &[String]) -> bool {
        bounds.iter().all(|b| self.satisfies_trait(ty, b))
    }

    /// Check trait bounds for all type parameters at a generic function call site.
    ///
    /// After unification resolves type variables to concrete types, this method
    /// verifies that each inferred type satisfies the trait bounds declared on
    /// the corresponding type parameter.
    ///
    /// Returns a list of violations as `(type_param_name, inferred_type, bounds)`.
    /// An empty list means all bounds are satisfied. Unresolved type variables
    /// (`Type::Var`) are skipped (non-breaking: the check is conservative).
    pub fn check_generic_call_bounds(
        &self,
        type_params: &[(String, Vec<String>)],
        param_var_map: &BTreeMap<String, TypeVarId>,
        subst: &TypeSubst,
    ) -> Vec<(String, Type, Vec<String>)> {
        let mut violations = Vec::new();
        for (tp_name, bounds) in type_params {
            if bounds.is_empty() {
                continue;
            }
            if let Some(&var_id) = param_var_map.get(tp_name) {
                let inferred = apply_subst(&Type::Var(var_id), subst);
                // Skip unresolved type variables — can't check bounds yet
                if matches!(inferred, Type::Var(_)) {
                    continue;
                }
                if !self.check_bounds(&inferred, bounds) {
                    violations.push((tp_name.clone(), inferred, bounds.clone()));
                }
            }
        }
        violations
    }

    /// Register a function signature.
    pub fn register_fn(&mut self, entry: FnSigEntry) {
        self.fn_sigs
            .entry(entry.name.clone())
            .or_default()
            .push(entry);
    }

    /// Register a struct type.
    pub fn register_struct(&mut self, name: &str, st: StructType) {
        self.type_defs.insert(name.to_string(), Type::Struct(st));
    }

    /// Register a class type.
    pub fn register_class(&mut self, name: &str, ct: ClassType) {
        self.type_defs.insert(name.to_string(), Type::Class(ct));
    }

    /// Register a record type (immutable value type).
    pub fn register_record(&mut self, name: &str, rt: RecordType) {
        self.type_defs.insert(name.to_string(), Type::Record(rt));
    }

    /// Register an enum type.
    pub fn register_enum_type(&mut self, name: &str, et: EnumType) {
        self.type_defs.insert(name.to_string(), Type::Enum(et));
    }

    /// Look up which enum a variant belongs to (by variant name).
    /// Returns the enum type if found.
    pub fn lookup_variant_enum(&self, variant_name: &str) -> Option<&EnumType> {
        for ty in self.type_defs.values() {
            if let Type::Enum(et) = ty {
                if et.variants.iter().any(|v| v.name == variant_name) {
                    return Some(et);
                }
            }
        }
        None
    }
}

impl Default for TypeEnv {
    fn default() -> Self {
        Self::new()
    }
}

// ── Quantum Builtin Registration ────────────────────────────────

/// Helper: register a single quantum builtin signature.
fn q_sig(env: &mut TypeEnv, name: &str, params: Vec<(&str, Type)>, ret: Type) {
    let alloc_effects = EffectSet::new(EffectSet::ALLOC);
    env.fn_sigs.entry(name.into()).or_default().push(FnSigEntry {
        name: name.into(),
        type_params: vec![],
        params: params.into_iter().map(|(n, t)| (n.to_string(), t)).collect(),
        ret,
        is_nogc: true,
        effects: alloc_effects,
    });
}

/// Register all quantum builtin function signatures so the type checker
/// can validate quantum programs without requiring `Any` annotations.
fn register_quantum_builtins(env: &mut TypeEnv) {
    use Type::*;

    // === Circuit (≤ 26 qubits) ===
    q_sig(env, "qubits", vec![("n", I64)], QuantumCircuit);
    // Single-qubit gates: (circuit, qubit) -> circuit
    for name in &["q_h", "q_x", "q_y", "q_z", "q_s", "q_t"] {
        q_sig(env, name, vec![("circuit", QuantumCircuit), ("qubit", I64)], QuantumCircuit);
    }
    // Parameterized single-qubit gates: (circuit, qubit, angle) -> circuit
    for name in &["q_rx", "q_ry", "q_rz"] {
        q_sig(env, name, vec![("circuit", QuantumCircuit), ("qubit", I64), ("theta", F64)], QuantumCircuit);
    }
    // Two-qubit gates: (circuit, q1, q2) -> circuit
    for name in &["q_cx", "q_cnot", "q_cz", "q_swap"] {
        q_sig(env, name, vec![("circuit", QuantumCircuit), ("a", I64), ("b", I64)], QuantumCircuit);
    }
    // Three-qubit: (circuit, c1, c2, target) -> circuit
    for name in &["q_toffoli", "q_ccx"] {
        q_sig(env, name, vec![
            ("circuit", QuantumCircuit), ("c1", I64), ("c2", I64), ("target", I64),
        ], QuantumCircuit);
    }
    // Execute / measure
    q_sig(env, "q_run", vec![("circuit", QuantumCircuit)], QuantumStatevector);
    q_sig(env, "q_measure", vec![("circuit", QuantumCircuit), ("seed", I64)],
        Array { elem: Box::new(I64), len: 0 });
    q_sig(env, "q_probs", vec![("circuit", QuantumCircuit)],
        Array { elem: Box::new(F64), len: 0 });
    q_sig(env, "q_sample", vec![("circuit", QuantumCircuit), ("n_shots", I64), ("seed", I64)],
        Array { elem: Box::new(I64), len: 0 });
    q_sig(env, "q_amplitudes", vec![("circuit", QuantumCircuit)],
        Array { elem: Box::new(Complex), len: 0 });
    q_sig(env, "q_n_qubits", vec![("circuit", QuantumCircuit)], I64);
    q_sig(env, "q_n_gates", vec![("circuit", QuantumCircuit)], I64);

    // === MPS (Matrix Product States, 50+ qubits) ===
    q_sig(env, "mps_new", vec![("n_qubits", I64), ("max_bond", I64)], QuantumMps);
    // Single-qubit MPS gates: (mps, qubit) -> mps
    for name in &["mps_h", "mps_x"] {
        q_sig(env, name, vec![("mps", QuantumMps), ("qubit", I64)], QuantumMps);
    }
    q_sig(env, "mps_ry", vec![("mps", QuantumMps), ("qubit", I64), ("theta", F64)], QuantumMps);
    q_sig(env, "mps_cnot", vec![("mps", QuantumMps), ("control", I64), ("target", I64)], QuantumMps);
    q_sig(env, "mps_z_expectation", vec![("mps", QuantumMps), ("qubit", I64)], F64);
    q_sig(env, "mps_energy", vec![("mps", QuantumMps), ("hamiltonian", Str)], F64);
    q_sig(env, "mps_memory", vec![("mps", QuantumMps)], I64);

    // === VQE ===
    q_sig(env, "vqe_heisenberg", vec![
        ("n_qubits", I64), ("max_bond", I64), ("lr", F64),
        ("iters", I64), ("seed", I64),
    ], F64);
    q_sig(env, "vqe_full_heisenberg", vec![
        ("n_qubits", I64), ("max_bond", I64), ("lr", F64),
        ("iters", I64), ("seed", I64),
    ], F64);

    // === QAOA ===
    q_sig(env, "qaoa_graph_cycle", vec![("n_vertices", I64)], QuantumGraph);
    q_sig(env, "qaoa_maxcut", vec![
        ("graph", QuantumGraph), ("max_bond", I64), ("p_layers", I64),
        ("lr", F64), ("iters", I64), ("seed", I64),
    ], Array { elem: Box::new(F64), len: 0 }); // returns [energy, cut_value]

    // === Stabilizer (1000+ qubits) ===
    q_sig(env, "stabilizer_new", vec![("n_qubits", I64)], QuantumStabilizer);
    for name in &["stabilizer_h", "stabilizer_s", "stabilizer_x", "stabilizer_y", "stabilizer_z"] {
        q_sig(env, name, vec![("state", QuantumStabilizer), ("qubit", I64)], QuantumStabilizer);
    }
    q_sig(env, "stabilizer_cnot", vec![
        ("state", QuantumStabilizer), ("control", I64), ("target", I64),
    ], QuantumStabilizer);
    q_sig(env, "stabilizer_measure", vec![
        ("state", QuantumStabilizer), ("qubit", I64), ("seed", I64),
    ], I64);
    q_sig(env, "stabilizer_n_qubits", vec![("state", QuantumStabilizer)], I64);

    // === Density Matrix ===
    q_sig(env, "density_new", vec![("n_qubits", I64)], QuantumDensity);
    q_sig(env, "density_gate", vec![
        ("dm", QuantumDensity), ("gate_name", Str), ("qubit", I64),
    ], QuantumDensity);
    q_sig(env, "density_cnot", vec![
        ("dm", QuantumDensity), ("control", I64), ("target", I64),
    ], QuantumDensity);
    q_sig(env, "density_depolarize", vec![
        ("dm", QuantumDensity), ("qubit", I64), ("p", F64),
    ], QuantumDensity);
    q_sig(env, "density_dephase", vec![
        ("dm", QuantumDensity), ("qubit", I64), ("p", F64),
    ], QuantumDensity);
    q_sig(env, "density_amplitude_damp", vec![
        ("dm", QuantumDensity), ("qubit", I64), ("gamma", F64),
    ], QuantumDensity);
    q_sig(env, "density_trace", vec![("dm", QuantumDensity)], F64);
    q_sig(env, "density_purity", vec![("dm", QuantumDensity)], F64);
    q_sig(env, "density_entropy", vec![("dm", QuantumDensity)], F64);
    q_sig(env, "density_probs", vec![("dm", QuantumDensity)],
        Array { elem: Box::new(F64), len: 0 });

    // === DMRG ===
    q_sig(env, "dmrg_ising", vec![
        ("n_qubits", I64), ("max_bond", I64), ("sweeps", I64), ("tol", F64),
    ], F64);
    q_sig(env, "dmrg_heisenberg", vec![
        ("n_qubits", I64), ("max_bond", I64), ("sweeps", I64), ("tol", F64),
    ], F64);

    // === QEC (Quantum Error Correction) ===
    q_sig(env, "qec_repetition_code", vec![("distance", I64)], QuantumSurfaceCode);
    q_sig(env, "qec_surface_code", vec![("distance", I64)], QuantumSurfaceCode);
    q_sig(env, "qec_syndrome", vec![
        ("state", QuantumStabilizer), ("code", QuantumSurfaceCode), ("seed", I64),
    ], Array { elem: Box::new(I64), len: 0 });
    q_sig(env, "qec_decode", vec![
        ("syndrome", Array { elem: Box::new(I64), len: 0 }),
        ("code", QuantumSurfaceCode),
    ], Array { elem: Box::new(I64), len: 0 });
    q_sig(env, "qec_logical_error_rate", vec![
        ("distance", I64), ("error_rate", F64), ("rounds", I64), ("seed", I64),
    ], F64);

    // === QML (Quantum Machine Learning) ===
    // qml_train has complex args — registered with Any-like flexible types
    // qml_predict similarly — both use array args
    q_sig(env, "qml_predict", vec![
        ("n_qubits", I64), ("layers", I64), ("n_classes", I64), ("max_bond", I64),
        ("params", Array { elem: Box::new(F64), len: 0 }),
        ("input", Array { elem: Box::new(F64), len: 0 }),
    ], I64);
}

// ── Type Checker ────────────────────────────────────────────────

use cjc_ast::*;
use cjc_diag::{Diagnostic, DiagnosticBag};

fn to_diag_span(span: cjc_ast::Span) -> cjc_diag::Span {
    cjc_diag::Span::new(span.start, span.end)
}

/// P2-3: Returns true if an expression is a compile-time constant (pure literal).
/// Only literal values are accepted as const initializers.
// ── Effect Typing Helpers ─────────────────────────────────────────

/// Valid effect annotation names.
const VALID_EFFECTS: &[&str] = &["pure", "io", "alloc", "gc", "nondet", "mutates", "arena_ok", "captures"];

/// Simple Levenshtein distance for suggesting closest valid effect name.
fn levenshtein(a: &str, b: &str) -> usize {
    let m = a.len();
    let n = b.len();
    let mut prev: Vec<usize> = (0..=n).collect();
    let mut curr = vec![0; n + 1];
    for (i, ca) in a.chars().enumerate() {
        curr[0] = i + 1;
        for (j, cb) in b.chars().enumerate() {
            let cost = if ca == cb { 0 } else { 1 };
            curr[j + 1] = (prev[j] + cost)
                .min(prev[j + 1] + 1)
                .min(curr[j] + 1);
        }
        std::mem::swap(&mut prev, &mut curr);
    }
    prev[n]
}

/// Convert effect annotation names (e.g. "pure", "io", "alloc") to an EffectSet.
/// Returns the EffectSet and a list of error messages for unknown effect names.
fn effects_from_names(names: &[String]) -> (EffectSet, Vec<String>) {
    let mut bits: u16 = 0;
    let mut unknown = Vec::new();
    for name in names {
        match name.as_str() {
            "pure" => {} // pure = no flags
            "io" => bits |= EffectSet::IO,
            "alloc" => bits |= EffectSet::ALLOC,
            "gc" => bits |= EffectSet::GC,
            "nondet" => bits |= EffectSet::NONDET,
            "mutates" => bits |= EffectSet::MUTATES,
            "arena_ok" => bits |= EffectSet::ARENA_OK,
            "captures" => bits |= EffectSet::CAPTURES,
            _ => {
                // Find closest valid effect for suggestion
                let mut best = ("", usize::MAX);
                for &valid in VALID_EFFECTS {
                    let d = levenshtein(name, valid);
                    if d < best.1 {
                        best = (valid, d);
                    }
                }
                let msg = if best.1 <= 2 {
                    format!("unknown effect `{name}`, did you mean `{}`?", best.0)
                } else {
                    format!("unknown effect `{name}`; valid effects: {}", VALID_EFFECTS.join(", "))
                };
                unknown.push(msg);
            }
        }
    }
    (EffectSet::new(bits), unknown)
}

/// Return the name of an effect flag constant.
fn effect_flag_name(flag: u16) -> &'static str {
    match flag {
        EffectSet::IO => "io",
        EffectSet::ALLOC => "alloc",
        EffectSet::GC => "gc",
        EffectSet::NONDET => "nondet",
        EffectSet::MUTATES => "mutates",
        EffectSet::ARENA_OK => "arena_ok",
        EffectSet::CAPTURES => "captures",
        _ => "unknown",
    }
}

/// Return the set of effect flags that `actual` has but `allowed` does not.
fn effect_violations(actual: &EffectSet, allowed: &EffectSet) -> Vec<u16> {
    let mut violations = Vec::new();
    let excess = actual.bits() & !allowed.bits();
    for &flag in &[
        EffectSet::IO,
        EffectSet::ALLOC,
        EffectSet::GC,
        EffectSet::NONDET,
        EffectSet::MUTATES,
        EffectSet::ARENA_OK,
        EffectSet::CAPTURES,
    ] {
        if excess & flag != 0 {
            violations.push(flag);
        }
    }
    violations
}

/// Recursively collect all function call names from a block.
fn collect_calls_in_block(block: &cjc_ast::Block, calls: &mut Vec<(String, cjc_ast::Span)>) {
    for stmt in &block.stmts {
        collect_calls_in_stmt(stmt, calls);
    }
    if let Some(ref expr) = block.expr {
        collect_calls_in_expr(expr, calls);
    }
}

fn collect_calls_in_stmt(stmt: &cjc_ast::Stmt, calls: &mut Vec<(String, cjc_ast::Span)>) {
    use cjc_ast::StmtKind;
    match &stmt.kind {
        StmtKind::Expr(e) => collect_calls_in_expr(e, calls),
        StmtKind::Let(l) => {
            collect_calls_in_expr(&l.init, calls);
        }
        StmtKind::Return(Some(e)) => collect_calls_in_expr(e, calls),
        StmtKind::Return(None) => {}
        StmtKind::If(if_stmt) => {
            collect_calls_in_if_stmt(if_stmt, calls);
        }
        StmtKind::While(w) => {
            collect_calls_in_expr(&w.condition, calls);
            collect_calls_in_block(&w.body, calls);
        }
        StmtKind::For(f) => {
            match &f.iter {
                cjc_ast::ForIter::Range { start, end, .. } => {
                    collect_calls_in_expr(start, calls);
                    collect_calls_in_expr(end, calls);
                }
                cjc_ast::ForIter::Expr(e) => {
                    collect_calls_in_expr(e, calls);
                }
            }
            collect_calls_in_block(&f.body, calls);
        }
        StmtKind::NoGcBlock(b) => {
            collect_calls_in_block(b, calls);
        }
        StmtKind::Break | StmtKind::Continue => {}
    }
}

fn collect_calls_in_if_stmt(if_stmt: &cjc_ast::IfStmt, calls: &mut Vec<(String, cjc_ast::Span)>) {
    collect_calls_in_expr(&if_stmt.condition, calls);
    collect_calls_in_block(&if_stmt.then_block, calls);
    if let Some(ref else_branch) = if_stmt.else_branch {
        match else_branch {
            cjc_ast::ElseBranch::ElseIf(nested) => {
                collect_calls_in_if_stmt(nested, calls);
            }
            cjc_ast::ElseBranch::Else(block) => {
                collect_calls_in_block(block, calls);
            }
        }
    }
}

fn collect_calls_in_expr(expr: &cjc_ast::Expr, calls: &mut Vec<(String, cjc_ast::Span)>) {
    use cjc_ast::ExprKind;
    match &expr.kind {
        ExprKind::Call { callee, args } => {
            // Extract function name
            match &callee.kind {
                ExprKind::Ident(id) => {
                    calls.push((id.name.clone(), expr.span));
                }
                ExprKind::Field { object: _, name } => {
                    // Method call — use just the method name for now
                    calls.push((name.name.clone(), expr.span));
                }
                _ => {}
            }
            // Recurse into callee and args
            collect_calls_in_expr(callee, calls);
            for arg in args {
                collect_calls_in_expr(&arg.value, calls);
            }
        }
        ExprKind::Binary { left, right, .. } => {
            collect_calls_in_expr(left, calls);
            collect_calls_in_expr(right, calls);
        }
        ExprKind::Unary { operand, .. } => {
            collect_calls_in_expr(operand, calls);
        }
        ExprKind::ArrayLit(elems) => {
            for e in elems {
                collect_calls_in_expr(e, calls);
            }
        }
        ExprKind::TupleLit(elems) => {
            for e in elems {
                collect_calls_in_expr(e, calls);
            }
        }
        ExprKind::Index { object, index } => {
            collect_calls_in_expr(object, calls);
            collect_calls_in_expr(index, calls);
        }
        ExprKind::MultiIndex { object, indices } => {
            collect_calls_in_expr(object, calls);
            for idx in indices {
                collect_calls_in_expr(idx, calls);
            }
        }
        ExprKind::Field { object, .. } => {
            collect_calls_in_expr(object, calls);
        }
        ExprKind::Block(b) => {
            collect_calls_in_block(b, calls);
        }
        ExprKind::Lambda { body, .. } => {
            collect_calls_in_expr(body, calls);
        }
        ExprKind::IfExpr { condition, then_block, else_branch } => {
            collect_calls_in_expr(condition, calls);
            collect_calls_in_block(then_block, calls);
            if let Some(ref eb) = else_branch {
                match eb {
                    cjc_ast::ElseBranch::ElseIf(nested) => {
                        // Treat nested if as having its own condition + blocks
                        collect_calls_in_if_stmt(nested, calls);
                    }
                    cjc_ast::ElseBranch::Else(block) => {
                        collect_calls_in_block(block, calls);
                    }
                }
            }
        }
        ExprKind::Assign { target, value } | ExprKind::CompoundAssign { target, value, .. } => {
            collect_calls_in_expr(target, calls);
            collect_calls_in_expr(value, calls);
        }
        ExprKind::Pipe { left, right } => {
            collect_calls_in_expr(left, calls);
            collect_calls_in_expr(right, calls);
        }
        ExprKind::FStringLit(segments) => {
            for (_lit, interp) in segments {
                if let Some(e) = interp {
                    collect_calls_in_expr(e, calls);
                }
            }
        }
        ExprKind::StructLit { fields, .. } => {
            for fi in fields {
                collect_calls_in_expr(&fi.value, calls);
            }
        }
        ExprKind::TensorLit { rows } => {
            for row in rows {
                for e in row {
                    collect_calls_in_expr(e, calls);
                }
            }
        }
        ExprKind::Match { scrutinee, arms } => {
            collect_calls_in_expr(scrutinee, calls);
            for arm in arms {
                collect_calls_in_expr(&arm.body, calls);
            }
        }
        ExprKind::Try(inner) => {
            collect_calls_in_expr(inner, calls);
        }
        ExprKind::VariantLit { fields, .. } => {
            for e in fields {
                collect_calls_in_expr(e, calls);
            }
        }
        ExprKind::Cast { expr, .. } => {
            collect_calls_in_expr(expr, calls);
        }
        // Literals and identifiers — no calls to collect
        ExprKind::IntLit(_) | ExprKind::FloatLit(_) | ExprKind::StringLit(_)
        | ExprKind::BoolLit(_) | ExprKind::NaLit | ExprKind::Ident(_) | ExprKind::ByteStringLit(_)
        | ExprKind::ByteCharLit(_) | ExprKind::RawStringLit(_) | ExprKind::RawByteStringLit(_)
        | ExprKind::RegexLit { .. } | ExprKind::Col(_) => {}
    }
}

fn is_const_expr(expr: &cjc_ast::Expr) -> bool {
    use cjc_ast::ExprKind;
    match &expr.kind {
        ExprKind::IntLit(_)
        | ExprKind::FloatLit(_)
        | ExprKind::BoolLit(_)
        | ExprKind::NaLit
        | ExprKind::StringLit(_)
        | ExprKind::RawStringLit(_)
        | ExprKind::ByteCharLit(_) => true,
        ExprKind::Unary { op: cjc_ast::UnaryOp::Neg, operand } => is_const_expr(operand),
        _ => false,
    }
}

pub struct TypeChecker {
    pub env: TypeEnv,
    pub diagnostics: DiagnosticBag,
    /// Current source filename for multi-file diagnostics.
    /// When set, all emitted diagnostics will include this filename.
    pub current_filename: Option<String>,
}

impl TypeChecker {
    pub fn new() -> Self {
        Self {
            env: TypeEnv::new(),
            diagnostics: DiagnosticBag::new(),
            current_filename: None,
        }
    }

    /// Create a type checker for a specific source file (multi-file mode).
    pub fn new_with_filename(filename: impl Into<String>) -> Self {
        Self {
            env: TypeEnv::new(),
            diagnostics: DiagnosticBag::new(),
            current_filename: Some(filename.into()),
        }
    }

    pub fn check_program(&mut self, program: &Program) {
        // First pass: register all type and function declarations
        for decl in &program.declarations {
            self.register_decl(decl);
        }

        // Second pass: type-check function bodies
        for decl in &program.declarations {
            self.check_decl(decl);
        }
    }

    fn register_decl(&mut self, decl: &Decl) {
        match &decl.kind {
            DeclKind::Struct(s) => {
                let fields: Vec<(String, Type)> = s
                    .fields
                    .iter()
                    .map(|f| (f.name.name.clone(), self.resolve_type_expr(&f.ty)))
                    .collect();
                let st = StructType {
                    name: s.name.name.clone(),
                    type_params: s.type_params.iter().map(|p| p.name.name.clone()).collect(),
                    fields,
                };
                self.env.register_struct(&s.name.name, st);
            }
            DeclKind::Class(c) => {
                let fields: Vec<(String, Type)> = c
                    .fields
                    .iter()
                    .map(|f| (f.name.name.clone(), self.resolve_type_expr(&f.ty)))
                    .collect();
                let ct = ClassType {
                    name: c.name.name.clone(),
                    type_params: c.type_params.iter().map(|p| p.name.name.clone()).collect(),
                    fields,
                };
                self.env.register_class(&c.name.name, ct);
            }
            DeclKind::Record(r) => {
                let fields: Vec<(String, Type)> = r
                    .fields
                    .iter()
                    .map(|f| (f.name.name.clone(), self.resolve_type_expr(&f.ty)))
                    .collect();
                let rt = RecordType {
                    name: r.name.name.clone(),
                    type_params: r.type_params.iter().map(|p| p.name.name.clone()).collect(),
                    fields,
                };
                self.env.register_record(&r.name.name, rt);
            }
            DeclKind::Fn(f) => {
                let params: Vec<(String, Type)> = f
                    .params
                    .iter()
                    .map(|p| (p.name.name.clone(), self.resolve_type_expr(&p.ty)))
                    .collect();
                let ret = f
                    .return_type
                    .as_ref()
                    .map(|t| self.resolve_type_expr(t))
                    .unwrap_or(Type::Void);
                let type_params: Vec<(String, Vec<String>)> = f
                    .type_params
                    .iter()
                    .map(|tp| {
                        (
                            tp.name.name.clone(),
                            tp.bounds.iter().map(|b| self.type_expr_name(b)).collect(),
                        )
                    })
                    .collect();
                let effects = if let Some(ref ann) = f.effect_annotation {
                    let (eset, unknowns) = effects_from_names(ann);
                    for msg in unknowns {
                        self.diagnostics.emit(
                            Diagnostic::error("E0410", msg, to_diag_span(decl.span))
                        );
                    }
                    eset
                } else {
                    EffectSet::default()
                };
                self.env.register_fn(FnSigEntry {
                    name: f.name.name.clone(),
                    type_params,
                    params,
                    ret,
                    is_nogc: f.is_nogc,
                    effects,
                });
            }
            DeclKind::Trait(t) => {
                let methods: Vec<MethodSig> = t
                    .methods
                    .iter()
                    .map(|m| MethodSig {
                        name: m.name.name.clone(),
                        params: m.params.iter().map(|p| self.resolve_type_expr(&p.ty)).collect(),
                        ret: m
                            .return_type
                            .as_ref()
                            .map(|t| self.resolve_type_expr(t))
                            .unwrap_or(Type::Void),
                    })
                    .collect();
                let def = TraitDef {
                    name: t.name.name.clone(),
                    type_params: t.type_params.iter().map(|p| p.name.name.clone()).collect(),
                    super_traits: t.super_traits.iter().map(|s| self.type_expr_name(s)).collect(),
                    methods,
                };
                self.env.trait_defs.insert(t.name.name.clone(), def);
            }
            DeclKind::Impl(i) => {
                if let Some(ref trait_ref) = i.trait_ref {
                    let trait_name = self.type_expr_name(trait_ref);
                    let target = self.resolve_type_expr(&i.target);
                    self.env.trait_impls.push(TraitImpl {
                        trait_name,
                        target_type: target,
                        type_args: vec![],
                    });
                }
                // Register methods
                for method in &i.methods {
                    let target_name = self.type_expr_name(&i.target);
                    let params: Vec<(String, Type)> = method
                        .params
                        .iter()
                        .map(|p| (p.name.name.clone(), self.resolve_type_expr(&p.ty)))
                        .collect();
                    let ret = method
                        .return_type
                        .as_ref()
                        .map(|t| self.resolve_type_expr(t))
                        .unwrap_or(Type::Void);
                    let qualified_name = format!("{}.{}", target_name, method.name.name);
                    let effects = if let Some(ref ann) = method.effect_annotation {
                        let (eset, unknowns) = effects_from_names(ann);
                        for msg in unknowns {
                            self.diagnostics.emit(
                                Diagnostic::error("E0410", msg, to_diag_span(decl.span))
                            );
                        }
                        eset
                    } else {
                        EffectSet::default()
                    };
                    self.env.register_fn(FnSigEntry {
                        name: qualified_name,
                        type_params: vec![],
                        params,
                        ret,
                        is_nogc: method.is_nogc,
                        effects,
                    });
                }
            }
            DeclKind::Enum(e) => {
                self.register_enum(e);
            }
            _ => {}
        }
    }

    /// Register an enum type and its variant constructors in the type environment.
    fn register_enum(&mut self, e: &cjc_ast::EnumDecl) {
        let type_params: Vec<String> =
            e.type_params.iter().map(|p| p.name.name.clone()).collect();
        let variants: Vec<EnumVariant> = e
            .variants
            .iter()
            .map(|v| EnumVariant {
                name: v.name.name.clone(),
                fields: v.fields.iter().map(|f| self.resolve_type_expr(f)).collect(),
            })
            .collect();
        let enum_type = EnumType {
            name: e.name.name.clone(),
            type_params: type_params.clone(),
            variants: variants.clone(),
        };
        self.env
            .type_defs
            .insert(e.name.name.clone(), Type::Enum(enum_type.clone()));

        // Register each variant as a constructor function
        let ret_type = Type::Enum(enum_type);
        for variant in &variants {
            let params: Vec<(String, Type)> = variant
                .fields
                .iter()
                .enumerate()
                .map(|(i, f)| (format!("_{}", i), f.clone()))
                .collect();
            let tp: Vec<(String, Vec<String>)> =
                type_params.iter().map(|t| (t.clone(), vec![])).collect();
            self.env.register_fn(FnSigEntry {
                name: variant.name.clone(),
                type_params: tp,
                params,
                ret: ret_type.clone(),
                is_nogc: false,
                effects: EffectSet::default(),
            });
        }
    }

    fn check_decl(&mut self, decl: &Decl) {
        match &decl.kind {
            DeclKind::Fn(f) => self.check_fn(f),
            DeclKind::Let(l) => {
                self.check_let(l, decl.span);
            }
            DeclKind::Impl(i) => {
                self.check_impl(i, decl.span);
            }
            // P2-3: Const expressions — type-check and register in env.
            DeclKind::Const(c) => {
                self.check_const_decl(c, decl.span);
            }
            _ => {}
        }
    }

    /// P2-3: Type-check a compile-time constant declaration.
    ///
    /// Evaluates only pure literal expressions; emits E0400 if the
    /// initializer is not a compile-time constant.
    fn check_const_decl(&mut self, c: &cjc_ast::ConstDecl, span: cjc_ast::Span) {
        let declared_ty = self.resolve_type_expr(&c.ty);
        let init_ty = self.check_expr(&c.value);

        // Verify the initializer is a compile-time constant (pure literal).
        if !is_const_expr(&c.value) {
            self.diagnostics.emit(
                Diagnostic::error(
                    "E0400",
                    format!(
                        "const `{}` initializer is not a compile-time constant expression",
                        c.name.name
                    ),
                    to_diag_span(span),
                )
                .with_hint("const initializers must be literal values (integer, float, bool, or string)"),
            );
        }

        // Type-check the declared type vs initializer type.
        if !declared_ty.is_error() && !init_ty.is_error() && !self.env.types_match(&declared_ty, &init_ty) {
            self.diagnostics.emit(
                Diagnostic::error(
                    "E0401",
                    format!(
                        "const `{}` type mismatch: declared `{}`, initialized with `{}`",
                        c.name.name, declared_ty, init_ty
                    ),
                    to_diag_span(span),
                ),
            );
        }

        // Register in const_defs AND as an immutable var in the top-level scope.
        self.env.const_defs.insert(c.name.name.clone(), declared_ty.clone());
        self.env.define_var(&c.name.name, declared_ty);
    }

    /// Type-check an `impl` block.
    ///
    /// When a `trait_ref` is present (`impl Trait for Type`), enforces:
    /// 1. The trait must be defined (known to the type environment).
    /// 2. All required trait methods must be present in the impl.
    /// 3. Duplicate impls for the same (trait, type) pair are rejected.
    ///
    /// Also type-checks all method bodies regardless of whether a trait is named.
    fn check_impl(&mut self, i: &cjc_ast::ImplDecl, span: cjc_ast::Span) {
        let target_name = self.type_expr_name(&i.target);

        if let Some(ref trait_ref) = i.trait_ref {
            let trait_name = self.type_expr_name(trait_ref);

            // (1) Verify the trait is defined.
            if !self.env.trait_defs.contains_key(&trait_name) {
                self.diagnostics.emit(
                    Diagnostic::error(
                        "E0200",
                        format!("trait `{}` is not defined", trait_name),
                        to_diag_span(span),
                    )
                    .with_hint(format!(
                        "declare the trait with `trait {} {{ ... }}`",
                        trait_name
                    )),
                );
                // Still type-check method bodies for good diagnostics
                for method in &i.methods {
                    self.check_fn(method);
                }
                return;
            }

            // (2) Check for duplicate impls: same trait for same type.
            let impl_count = self.env.trait_impls.iter().filter(|imp| {
                imp.trait_name == trait_name
                    && self.env.types_match(&imp.target_type, &self.env.resolve_type_name(&target_name).unwrap_or(Type::Unresolved(target_name.clone())))
            }).count();
            if impl_count > 1 {
                self.diagnostics.emit(
                    Diagnostic::error(
                        "E0201",
                        format!(
                            "duplicate implementation of trait `{}` for type `{}`",
                            trait_name, target_name
                        ),
                        to_diag_span(span),
                    )
                    .with_hint("only one impl of a trait per type is allowed"),
                );
            }

            // (3) Verify all required trait methods are present.
            let required_methods: Vec<String> = self.env
                .trait_defs
                .get(&trait_name)
                .map(|def| def.methods.iter().map(|m| m.name.clone()).collect())
                .unwrap_or_default();

            let provided_methods: Vec<String> =
                i.methods.iter().map(|m| m.name.name.clone()).collect();

            for required in &required_methods {
                if !provided_methods.contains(required) {
                    self.diagnostics.emit(
                        Diagnostic::error(
                            "E0202",
                            format!(
                                "impl of `{}` for `{}` is missing method `{}`",
                                trait_name, target_name, required
                            ),
                            to_diag_span(span),
                        )
                        .with_hint(format!(
                            "add `fn {}(...)` to satisfy the trait requirement",
                            required
                        )),
                    );
                }
            }
        }

        // Type-check all method bodies.
        for method in &i.methods {
            self.check_fn(method);
        }
    }

    fn check_fn(&mut self, f: &FnDecl) {
        self.env.push_scope();

        // Bind parameters
        for param in &f.params {
            let ty = self.resolve_type_expr(&param.ty);
            self.env.define_var(&param.name.name, ty);
        }

        // Bind type parameters
        for tp in &f.type_params {
            let var = self.env.fresh_var();
            self.env.define_var(&tp.name.name, Type::Var(var));
        }

        // Check nogc constraints
        if f.is_nogc {
            self.check_nogc_block(&f.body);
        }

        // Check body
        let body_type = self.check_block(&f.body);

        // Return type: infer from body if not annotated
        if let Some(ref ret_ty_expr) = f.return_type {
            let expected_ret = self.resolve_type_expr(ret_ty_expr);
            if !expected_ret.is_error()
                && !body_type.is_error()
                && expected_ret != Type::Void
                && !self.env.types_match(&body_type, &expected_ret)
            {
                self.diagnostics.emit(
                    Diagnostic::error(
                        "E0103",
                        format!(
                            "mismatched return type: expected `{}`, found `{}`",
                            expected_ret, body_type
                        ),
                        to_diag_span(f.body.span),
                    )
                    .with_hint(format!("function `{}` should return `{}`", f.name.name, expected_ret)),
                );
            }
        }
        // When return_type is None, the function's return type is inferred
        // from the body's tail expression (body_type). This is used by the
        // register_decl phase which defaults to Void, but at checking time
        // we just validate the body is well-typed.

        // Effect checking: verify body effects are subset of declared effects
        if let Some(ref declared_effects) = f.effect_annotation {
            self.check_effect_annotation(f, declared_effects);
        }

        self.env.pop_scope();
    }

    /// Verify that a function body's computed effects are a subset of the declared effects.
    fn check_effect_annotation(&mut self, f: &FnDecl, declared_effects: &[String]) {
        use crate::effect_registry::builtin_effects;

        let (declared_set, unknowns) = effects_from_names(declared_effects);
        for msg in unknowns {
            self.diagnostics.emit(
                Diagnostic::error("E0410", msg, to_diag_span(f.name.span))
            );
        }

        // Collect all function calls in the body
        let mut called_fns = Vec::new();
        collect_calls_in_block(&f.body, &mut called_fns);

        let registry = builtin_effects();

        // Check each call against declared effects
        for (call_name, call_span) in &called_fns {
            let call_effects = if let Some(effects) = registry.get(call_name.as_str()) {
                *effects
            } else {
                // User-defined function: look up from registered functions
                // For now, assume user-defined functions called from an effect-annotated
                // function must themselves be effect-compatible. We check if they have
                // effect annotations too (conservative: if unknown, assume any-effect).
                if let Some(sigs) = self.env.fn_sigs.get(call_name) {
                    sigs.first().map(|s| s.effects).unwrap_or(EffectSet::PURE)
                } else {
                    // Unknown function — conservatively skip (will be caught at call-site type checking)
                    continue;
                }
            };

            // Check: computed effects must be subset of declared effects
            let violations = effect_violations(&call_effects, &declared_set);
            if !violations.is_empty() {
                let violation_names: Vec<&str> = violations.iter().map(|v| effect_flag_name(*v)).collect();
                self.diagnostics.emit(
                    Diagnostic::error(
                        "E4002",
                        format!(
                            "function `{}` declared as `/ {}` but calls `{}` which has {} effects",
                            f.name.name,
                            declared_effects.join(" + "),
                            call_name,
                            violation_names.join(", "),
                        ),
                        to_diag_span(*call_span),
                    )
                    .with_hint(format!(
                        "remove `{}` from the body or widen the effect annotation",
                        call_name,
                    )),
                );
            }
        }
    }

    fn check_let(&mut self, l: &LetStmt, span: cjc_ast::Span) {
        let init_type = self.check_expr(&l.init);

        if let Some(ref ty_expr) = l.ty {
            let declared = self.resolve_type_expr(ty_expr);
            if !declared.is_error()
                && !init_type.is_error()
                && !self.env.types_match(&declared, &init_type)
            {
                self.diagnostics.emit(
                    Diagnostic::error(
                        "E0104",
                        format!(
                            "type mismatch: declared `{}`, initialized with `{}`",
                            declared, init_type
                        ),
                        to_diag_span(span),
                    )
                    .with_hint(format!(
                        "change the type annotation to `{}` or fix the initializer",
                        init_type
                    )),
                );
            }
            if l.mutable {
                self.env.define_var_mut(&l.name.name, declared);
            } else {
                self.env.define_var(&l.name.name, declared);
            }
        } else if l.mutable {
            self.env.define_var_mut(&l.name.name, init_type);
        } else {
            self.env.define_var(&l.name.name, init_type);
        }
    }

    fn check_block(&mut self, block: &Block) -> Type {
        self.env.push_scope();

        for stmt in &block.stmts {
            self.check_stmt(stmt);
        }

        let result = if let Some(ref expr) = block.expr {
            self.check_expr(expr)
        } else {
            Type::Void
        };

        self.env.pop_scope();
        result
    }

    fn check_stmt(&mut self, stmt: &Stmt) {
        match &stmt.kind {
            StmtKind::Let(l) => self.check_let(l, stmt.span),
            StmtKind::Expr(e) => {
                self.check_expr(e);
            }
            StmtKind::Return(e) => {
                if let Some(expr) = e {
                    self.check_expr(expr);
                }
            }
            StmtKind::If(if_stmt) => self.check_if(if_stmt),
            StmtKind::While(w) => {
                let cond_type = self.check_expr(&w.condition);
                if !cond_type.is_error() && cond_type != Type::Bool {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0105",
                        format!("while condition must be `bool`, found `{}`", cond_type),
                        to_diag_span(stmt.span),
                    ));
                }
                self.check_block(&w.body);
            }
            StmtKind::For(f) => {
                // Type-check the iterator expressions
                match &f.iter {
                    cjc_ast::ForIter::Range { start, end } => {
                        self.check_expr(start);
                        self.check_expr(end);
                    }
                    cjc_ast::ForIter::Expr(expr) => {
                        self.check_expr(expr);
                    }
                }
                self.check_block(&f.body);
            }
            StmtKind::NoGcBlock(block) => {
                self.check_nogc_block(block);
                self.check_block(block);
            }
            // Break/Continue are validated at parse time (loop depth check).
            // No type checking needed — they produce no values.
            StmtKind::Break | StmtKind::Continue => {}
        }
    }

    fn check_if(&mut self, if_stmt: &IfStmt) {
        let cond_type = self.check_expr(&if_stmt.condition);
        if !cond_type.is_error() && cond_type != Type::Bool {
            self.diagnostics.emit(Diagnostic::error(
                "E0105",
                format!("if condition must be `bool`, found `{}`", cond_type),
                to_diag_span(if_stmt.condition.span),
            ));
        }
        self.check_block(&if_stmt.then_block);
        if let Some(ref else_branch) = if_stmt.else_branch {
            match else_branch {
                ElseBranch::ElseIf(elif) => self.check_if(elif),
                ElseBranch::Else(block) => {
                    self.check_block(block);
                }
            }
        }
    }

    fn check_nogc_block(&mut self, block: &Block) {
        // Check that no GC types are allocated in the block
        for stmt in &block.stmts {
            if let StmtKind::Let(ref l) = stmt.kind {
                let ty = self.check_expr(&l.init);
                if ty.is_gc_type() {
                    self.diagnostics.emit(
                        Diagnostic::error(
                            "E0201",
                            format!(
                                "cannot allocate GC-managed type `{}` in `nogc` context",
                                ty
                            ),
                            to_diag_span(stmt.span),
                        )
                        .with_label(
                            to_diag_span(stmt.span),
                            "GC allocation not allowed here",
                        )
                        .with_hint(
                            "use value types (struct, Tensor, Buffer) in nogc zones; \
                             move GC-managed objects outside the nogc block",
                        ),
                    );
                }
            }
        }
    }

    /// Static exhaustiveness checking for match expressions.
    /// Supports: enum types, bool types.
    /// Emits a diagnostic if not all cases are covered and no wildcard/binding exists.
    fn check_match_exhaustiveness(
        &mut self,
        scrutinee_type: &Type,
        arms: &[MatchArm],
        span: cjc_ast::Span,
    ) {
        // Check for bool exhaustiveness
        if matches!(scrutinee_type, Type::Bool) {
            self.check_bool_exhaustiveness(arms, span);
            return;
        }

        let enum_type = match scrutinee_type {
            Type::Enum(e) => e,
            _ => return, // Only check exhaustiveness for enum and bool types
        };

        // Collect all variant names for the enum under scrutiny
        let all_variant_names: Vec<String> = enum_type
            .variants
            .iter()
            .map(|v| v.name.clone())
            .collect();

        // Collect covered variant names and check for wildcard/binding patterns.
        // A `Binding(name)` is treated as a variant pattern (not a wildcard) if
        // `name` matches one of the enum's declared variant names.  This handles
        // the common case where bare enum variant names (e.g., `Red`, `Green`) are
        // parsed as `Binding` rather than `Variant` by the parser when written
        // without parentheses.
        let mut covered_variants: Vec<String> = Vec::new();
        let mut has_wildcard = false;

        for arm in arms {
            match &arm.pattern.kind {
                PatternKind::Wildcard => {
                    has_wildcard = true;
                }
                PatternKind::Binding(name) => {
                    if all_variant_names.contains(&name.name) {
                        // Name matches an enum variant — treat as variant pattern.
                        covered_variants.push(name.name.clone());
                    } else {
                        // Pure binding variable — covers all remaining arms.
                        has_wildcard = true;
                    }
                }
                PatternKind::Variant {
                    variant, ..
                } => {
                    covered_variants.push(variant.name.clone());
                }
                _ => {}
            }
        }

        if has_wildcard {
            return; // Wildcard covers everything
        }

        // Check for missing variants
        let missing: Vec<&str> = enum_type
            .variants
            .iter()
            .filter(|v| !covered_variants.contains(&v.name))
            .map(|v| v.name.as_str())
            .collect();

        if !missing.is_empty() {
            self.diagnostics.emit(
                Diagnostic::error(
                    "E0130",
                    format!(
                        "non-exhaustive match on enum `{}`: missing variant(s) {}",
                        enum_type.name,
                        missing
                            .iter()
                            .map(|n| format!("`{}`", n))
                            .collect::<Vec<_>>()
                            .join(", ")
                    ),
                    to_diag_span(span),
                )
                .with_hint("add the missing variant(s) or a wildcard `_` arm"),
            );
        }
    }

    /// Exhaustiveness checking for match on bool.
    fn check_bool_exhaustiveness(
        &mut self,
        arms: &[MatchArm],
        span: cjc_ast::Span,
    ) {
        let mut has_true = false;
        let mut has_false = false;
        let mut has_wildcard = false;

        for arm in arms {
            match &arm.pattern.kind {
                PatternKind::Wildcard => { has_wildcard = true; }
                PatternKind::Binding(_) => { has_wildcard = true; }
                PatternKind::LitBool(true) => {
                    has_true = true;
                }
                PatternKind::LitBool(false) => {
                    has_false = true;
                }
                _ => {}
            }
        }

        if has_wildcard || (has_true && has_false) {
            return;
        }

        let mut missing = Vec::new();
        if !has_true { missing.push("`true`"); }
        if !has_false { missing.push("`false`"); }

        if !missing.is_empty() {
            self.diagnostics.emit(
                Diagnostic::error(
                    "E0131",
                    format!(
                        "non-exhaustive match on `bool`: missing {}",
                        missing.join(" and ")
                    ),
                    to_diag_span(span),
                )
                .with_hint("add the missing arm(s) or a wildcard `_` arm"),
            );
        }
    }

    pub fn check_expr(&mut self, expr: &Expr) -> Type {
        match &expr.kind {
            ExprKind::IntLit(_) => Type::I64,
            ExprKind::FloatLit(_) => Type::F64,
            ExprKind::StringLit(_) => Type::Str,
            ExprKind::ByteStringLit(_) => Type::ByteSlice,
            ExprKind::ByteCharLit(_) => Type::U8,
            ExprKind::RawStringLit(_) => Type::Str,
            ExprKind::RawByteStringLit(_) => Type::ByteSlice,
            ExprKind::FStringLit(segments) => {
                // P2-5: Type-check each interpolated expression; result type is always Str.
                for (_lit, interp) in segments {
                    if let Some(e) = interp {
                        self.check_expr(e);
                    }
                }
                Type::Str
            }
            ExprKind::RegexLit { .. } => Type::Regex,
            ExprKind::TensorLit { rows } => {
                // Type-check all elements, infer element type as f64
                for row in rows {
                    for expr in row {
                        self.check_expr(expr);
                    }
                }
                Type::Tensor {
                    elem: Box::new(Type::F64),
                    shape: None,
                }
            }
            ExprKind::BoolLit(_) => Type::Bool,
            // NA has no concrete type — treat as Void (compatible with any nullable context)
            ExprKind::NaLit => Type::Void,
            ExprKind::Ident(id) => {
                if let Some(ty) = self.env.lookup_var(&id.name) {
                    ty.clone()
                } else {
                    self.diagnostics.emit(
                        Diagnostic::error(
                            "E0106",
                            format!("undefined variable `{}`", id.name),
                            to_diag_span(id.span),
                        )
                        .with_hint("check for typos or declare the variable with `let`"),
                    );
                    Type::Error
                }
            }
            ExprKind::Binary { op, left, right } => {
                let lt = self.check_expr(left);
                let rt = self.check_expr(right);
                self.check_binary_op(*op, &lt, &rt, expr.span)
            }
            ExprKind::Unary { op, operand } => {
                let t = self.check_expr(operand);
                match op {
                    UnaryOp::Neg => {
                        if t.is_numeric() || t.is_error() {
                            t
                        } else {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0107",
                                format!("cannot negate type `{}`", t),
                                to_diag_span(expr.span),
                            ));
                            Type::Error
                        }
                    }
                    UnaryOp::Not => {
                        if t == Type::Bool || t.is_error() {
                            Type::Bool
                        } else {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0107",
                                format!("cannot apply `!` to type `{}`", t),
                                to_diag_span(expr.span),
                            ));
                            Type::Error
                        }
                    }
                    UnaryOp::BitNot => {
                        if t.is_int() || t.is_error() {
                            t
                        } else {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0107",
                                format!("cannot apply `~` to type `{}`; bitwise NOT requires an integer", t),
                                to_diag_span(expr.span),
                            ));
                            Type::Error
                        }
                    }
                }
            }
            ExprKind::Call { callee, args } => {
                // Simple function call resolution
                if let ExprKind::Ident(id) = &callee.kind {
                    self.check_fn_call(&id.name, args, expr.span)
                } else if let ExprKind::Field { object, name } = &callee.kind {
                    let obj_ty = self.check_expr(object);
                    let qualified = format!("{}.{}", obj_ty, name.name);
                    self.check_fn_call(&qualified, args, expr.span)
                } else {
                    let _callee_ty = self.check_expr(callee);
                    for arg in args {
                        self.check_expr(&arg.value);
                    }
                    Type::Error
                }
            }
            ExprKind::Field { object, name } => {
                let obj_ty = self.check_expr(object);
                match &obj_ty {
                    Type::Struct(st) => {
                        if let Some((_, ft)) = st.fields.iter().find(|(n, _)| n == &name.name) {
                            ft.clone()
                        } else {
                            self.diagnostics.emit(
                                Diagnostic::error(
                                    "E0108",
                                    format!(
                                        "no field `{}` on struct `{}`",
                                        name.name, st.name
                                    ),
                                    to_diag_span(name.span),
                                )
                                .with_hint(format!(
                                    "available fields: {}",
                                    st.fields
                                        .iter()
                                        .map(|(n, _)| n.as_str())
                                        .collect::<Vec<_>>()
                                        .join(", ")
                                )),
                            );
                            Type::Error
                        }
                    }
                    Type::Class(ct) => {
                        if let Some((_, ft)) = ct.fields.iter().find(|(n, _)| n == &name.name) {
                            ft.clone()
                        } else {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0108",
                                format!("no field `{}` on class `{}`", name.name, ct.name),
                                to_diag_span(name.span),
                            ));
                            Type::Error
                        }
                    }
                    Type::Record(rt) => {
                        if let Some((_, ft)) = rt.fields.iter().find(|(n, _)| n == &name.name) {
                            ft.clone()
                        } else {
                            self.diagnostics.emit(
                                Diagnostic::error(
                                    "E0108",
                                    format!(
                                        "no field `{}` on record `{}`",
                                        name.name, rt.name
                                    ),
                                    to_diag_span(name.span),
                                )
                                .with_hint(format!(
                                    "available fields: {}",
                                    rt.fields
                                        .iter()
                                        .map(|(n, _)| n.as_str())
                                        .collect::<Vec<_>>()
                                        .join(", ")
                                )),
                            );
                            Type::Error
                        }
                    }
                    Type::Tensor { .. } => {
                        match name.name.as_str() {
                            "shape" => Type::Array {
                                elem: Box::new(Type::I64),
                                len: 0,
                            },
                            "buffer" => Type::Buffer {
                                elem: Box::new(Type::F64),
                            },
                            _ => {
                                // Could be a method — return Error for now
                                Type::Error
                            }
                        }
                    }
                    Type::Error => Type::Error,
                    _ => {
                        self.diagnostics.emit(Diagnostic::error(
                            "E0108",
                            format!("cannot access field `{}` on type `{}`", name.name, obj_ty),
                            to_diag_span(name.span),
                        ));
                        Type::Error
                    }
                }
            }
            ExprKind::Index { object, index } => {
                let obj_ty = self.check_expr(object);
                let _idx_ty = self.check_expr(index);
                match &obj_ty {
                    Type::Array { elem, .. } => *elem.clone(),
                    Type::Tensor { elem, .. } => *elem.clone(),
                    Type::Buffer { elem } => *elem.clone(),
                    Type::Error => Type::Error,
                    _ => {
                        self.diagnostics.emit(Diagnostic::error(
                            "E0109",
                            format!("cannot index into type `{}`", obj_ty),
                            to_diag_span(expr.span),
                        ));
                        Type::Error
                    }
                }
            }
            ExprKind::MultiIndex { object, indices } => {
                let obj_ty = self.check_expr(object);
                for idx in indices {
                    self.check_expr(idx);
                }
                match &obj_ty {
                    Type::Tensor { elem, .. } => *elem.clone(),
                    Type::Error => Type::Error,
                    _ => {
                        self.diagnostics.emit(Diagnostic::error(
                            "E0109",
                            format!("multi-index not supported on type `{}`", obj_ty),
                            to_diag_span(expr.span),
                        ));
                        Type::Error
                    }
                }
            }
            ExprKind::Assign { target, value } => {
                let _tt = self.check_expr(target);
                let _vt = self.check_expr(value);
                // P0-3: Enforce immutable binding — emit E0150 if target is a
                // bare identifier that was declared without `mut`.
                if let ExprKind::Ident(id) = &target.kind {
                    if let Some((_ty, is_mut)) = self.env.lookup_var_entry(&id.name) {
                        if !is_mut {
                            self.diagnostics.emit(
                                Diagnostic::error(
                                    "E0150",
                                    format!(
                                        "cannot assign to immutable variable `{}`",
                                        id.name
                                    ),
                                    to_diag_span(expr.span),
                                )
                                .with_hint(format!(
                                    "consider making `{}` mutable: `let mut {} = ...`",
                                    id.name, id.name
                                )),
                            );
                        }
                    }
                }
                // LH10: Enforce record immutability — field assignment on records is a type error.
                if let ExprKind::Field { object, name: field_name } = &target.kind {
                    let obj_ty = self.check_expr(object);
                    if let Type::Record(ref rt) = obj_ty {
                        self.diagnostics.emit(
                            Diagnostic::error(
                                "E0160",
                                format!(
                                    "cannot assign to field `{}` of record `{}`: records are immutable",
                                    field_name.name, rt.name
                                ),
                                to_diag_span(expr.span),
                            )
                            .with_hint("consider using a struct instead of a record if you need mutability".to_string()),
                        );
                    }
                }
                Type::Void
            }
            ExprKind::Pipe { left, right } => {
                let _lt = self.check_expr(left);
                // Pipe semantics: result type is the result of the right expression
                self.check_expr(right)
            }
            ExprKind::Block(block) => self.check_block(block),
            ExprKind::StructLit { name, fields } => {
                if let Some(ty) = self.env.resolve_type_name(&name.name) {
                    // Check field types for both struct and record types
                    let type_fields = match &ty {
                        Type::Struct(ref st) => Some(&st.fields),
                        Type::Record(ref rt) => Some(&rt.fields),
                        _ => None,
                    };
                    if let Some(defined_fields) = type_fields {
                        for field in fields {
                            let ft = self.check_expr(&field.value);
                            if let Some((_, expected)) =
                                defined_fields.iter().find(|(n, _)| n == &field.name.name)
                            {
                                if !ft.is_error()
                                    && !expected.is_error()
                                    && !self.env.types_match(&ft, expected)
                                {
                                    self.diagnostics.emit(Diagnostic::error(
                                        "E0110",
                                        format!(
                                            "field `{}` type mismatch: expected `{}`, found `{}`",
                                            field.name.name, expected, ft
                                        ),
                                        to_diag_span(field.span),
                                    ));
                                }
                            }
                        }
                    }
                    ty
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0111",
                        format!("undefined type `{}`", name.name),
                        to_diag_span(name.span),
                    ));
                    Type::Error
                }
            }
            ExprKind::ArrayLit(elems) => {
                if elems.is_empty() {
                    Type::Array {
                        elem: Box::new(Type::Error),
                        len: 0,
                    }
                } else {
                    let first_type = self.check_expr(&elems[0]);
                    for elem in elems.iter().skip(1) {
                        let et = self.check_expr(elem);
                        if !et.is_error()
                            && !first_type.is_error()
                            && !self.env.types_match(&et, &first_type)
                        {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0112",
                                format!(
                                    "array element type mismatch: expected `{}`, found `{}`",
                                    first_type, et
                                ),
                                to_diag_span(elem.span),
                            ));
                        }
                    }
                    Type::Array {
                        elem: Box::new(first_type),
                        len: elems.len(),
                    }
                }
            }
            ExprKind::Col(_) => {
                // Data DSL column reference — returns an opaque Expr type
                Type::Unresolved("ColumnExpr".into())
            }
            ExprKind::Lambda { params, body } => {
                self.env.push_scope();
                let param_types: Vec<Type> = params
                    .iter()
                    .map(|p| {
                        let ty = self.resolve_type_expr(&p.ty);
                        self.env.define_var(&p.name.name, ty.clone());
                        ty
                    })
                    .collect();
                let ret_type = self.check_expr(body);
                self.env.pop_scope();
                Type::Fn {
                    params: param_types,
                    ret: Box::new(ret_type),
                }
            }
            ExprKind::Match { scrutinee, arms } => {
                let scrut_ty = self.check_expr(scrutinee);
                // Check exhaustiveness for enum types
                self.check_match_exhaustiveness(&scrut_ty, arms, expr.span);
                // Type of a match = type of the first arm body (simplified)
                if arms.is_empty() {
                    Type::Void
                } else {
                    let first = self.check_expr(&arms[0].body);
                    for arm in arms.iter().skip(1) {
                        self.check_expr(&arm.body);
                    }
                    first
                }
            }
            ExprKind::TupleLit(elems) => {
                let types: Vec<Type> = elems.iter().map(|e| self.check_expr(e)).collect();
                Type::Tuple(types)
            }
            ExprKind::Try(inner) => {
                // `expr?` — inner must be Result<T, E>, expression type is T
                let inner_ty = self.check_expr(inner);
                match &inner_ty {
                    Type::Enum(e) if e.name == "Result" => {
                        // Extract T from Result<T, E>
                        if let Some(ok_variant) = e.variants.iter().find(|v| v.name == "Ok") {
                            if let Some(t) = ok_variant.fields.first() {
                                t.clone()
                            } else {
                                Type::Void
                            }
                        } else {
                            Type::Error
                        }
                    }
                    Type::Error => Type::Error,
                    _ => {
                        self.diagnostics.emit(
                            Diagnostic::error(
                                "E0120",
                                format!(
                                    "`?` operator requires `Result<T, E>`, found `{}`",
                                    inner_ty
                                ),
                                to_diag_span(expr.span),
                            )
                            .with_hint("the `?` operator can only be applied to Result types"),
                        );
                        Type::Error
                    }
                }
            }
            ExprKind::VariantLit {
                enum_name,
                variant,
                fields,
            } => {
                // Type-check variant literal by looking up enum and variant
                let field_types: Vec<Type> = fields.iter().map(|f| self.check_expr(f)).collect();

                // Try to find the enum via explicit enum_name or variant lookup
                let enum_name_str = enum_name.as_ref().map(|id| id.name.as_str());
                let found_enum = if let Some(name) = enum_name_str {
                    self.env.resolve_type_name(name)
                } else {
                    // Look up by variant name
                    self.env
                        .lookup_variant_enum(&variant.name)
                        .map(|e| Type::Enum(e.clone()))
                };

                if let Some(Type::Enum(et)) = found_enum {
                    if let Some(v) = et.variants.iter().find(|v| v.name == variant.name) {
                        if v.fields.len() != field_types.len() {
                            self.diagnostics.emit(Diagnostic::error(
                                "E0121",
                                format!(
                                    "variant `{}` expects {} field(s), found {}",
                                    variant.name,
                                    v.fields.len(),
                                    field_types.len()
                                ),
                                to_diag_span(expr.span),
                            ));
                            Type::Error
                        } else {
                            // Build a substitution map from type params to concrete
                            // types by matching variant field types against argument
                            // types.  For example, `Some(42)` matches field
                            // `Unresolved("T")` against `i64`, producing {T -> i64}.
                            if et.type_params.is_empty() {
                                Type::Enum(et)
                            } else {
                                let mut subst: BTreeMap<String, Type> = BTreeMap::new();
                                for (field_ty, arg_ty) in v.fields.iter().zip(field_types.iter()) {
                                    collect_unresolved_bindings(field_ty, arg_ty, &mut subst);
                                }
                                substitute_type_params_concrete(&Type::Enum(et), &subst)
                            }
                        }
                    } else {
                        self.diagnostics.emit(Diagnostic::error(
                            "E0122",
                            format!("no variant `{}` in enum `{}`", variant.name, et.name),
                            to_diag_span(variant.span),
                        ));
                        Type::Error
                    }
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0122",
                        format!("unknown enum variant `{}`", variant.name),
                        to_diag_span(variant.span),
                    ));
                    Type::Error
                }
            }
            ExprKind::CompoundAssign { op, target, value } => {
                let tt = self.check_expr(target);
                let vt = self.check_expr(value);
                // Type-check the binary operation (target op value)
                let _result_ty = self.check_binary_op(*op, &tt, &vt, expr.span);
                // Check mutability like a regular assignment
                if let ExprKind::Ident(id) = &target.kind {
                    if let Some((_ty, is_mut)) = self.env.lookup_var_entry(&id.name) {
                        if !is_mut {
                            self.diagnostics.emit(
                                Diagnostic::error(
                                    "E0150",
                                    format!(
                                        "cannot assign to immutable variable `{}`",
                                        id.name
                                    ),
                                    to_diag_span(expr.span),
                                )
                                .with_hint(format!(
                                    "consider making `{}` mutable: `let mut {} = ...`",
                                    id.name, id.name
                                )),
                            );
                        }
                    }
                }
                Type::Void
            }
            ExprKind::IfExpr { condition, then_block, else_branch } => {
                let cond_type = self.check_expr(condition);
                if !cond_type.is_error() && cond_type != Type::Bool {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0105",
                        format!("if condition must be `bool`, found `{}`", cond_type),
                        to_diag_span(condition.span),
                    ));
                }
                let then_ty = self.check_block(then_block);
                if let Some(ref eb) = else_branch {
                    match eb {
                        ElseBranch::ElseIf(elif) => {
                            self.check_if(elif);
                        }
                        ElseBranch::Else(block) => {
                            self.check_block(block);
                        }
                    }
                }
                then_ty
            }
            ExprKind::Cast { expr, target_type } => {
                self.check_expr(expr);
                // Return the target type
                match target_type.name.as_str() {
                    "f64" | "float" | "Float" => Type::F64,
                    "i64" | "int" | "Int" => Type::I64,
                    "bool" | "Bool" => Type::Bool,
                    "String" | "string" => Type::Str,
                    _ => Type::Void,
                }
            }
        }
    }

    fn check_binary_op(
        &mut self,
        op: BinOp,
        lt: &Type,
        rt: &Type,
        span: cjc_ast::Span,
    ) -> Type {
        if lt.is_error() || rt.is_error() {
            return Type::Error;
        }

        match op {
            BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Mod | BinOp::Pow => {
                if lt.is_numeric() && self.env.types_match(lt, rt) {
                    lt.clone()
                } else if lt.is_numeric() && rt.is_numeric() {
                    self.diagnostics.emit(
                        Diagnostic::error(
                            "E0101",
                            format!(
                                "cannot apply `{}` to `{}` and `{}`",
                                op, lt, rt
                            ),
                            to_diag_span(span),
                        )
                        .with_hint(format!(
                            "both operands must have the same numeric type; \
                             consider casting one to `{}`",
                            lt
                        )),
                    );
                    Type::Error
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0101",
                        format!("cannot apply `{}` to `{}` and `{}`", op, lt, rt),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
            BinOp::Eq | BinOp::Ne => {
                if self.env.types_match(lt, rt) {
                    Type::Bool
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0102",
                        format!(
                            "cannot compare `{}` and `{}` for equality",
                            lt, rt
                        ),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
            BinOp::Lt | BinOp::Gt | BinOp::Le | BinOp::Ge => {
                if lt.is_numeric() && self.env.types_match(lt, rt) {
                    Type::Bool
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0102",
                        format!("cannot compare `{}` and `{}`", lt, rt),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
            BinOp::And | BinOp::Or => {
                if *lt == Type::Bool && *rt == Type::Bool {
                    Type::Bool
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0101",
                        format!(
                            "`{}` requires `bool` operands, found `{}` and `{}`",
                            op, lt, rt
                        ),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
            BinOp::BitAnd | BinOp::BitOr | BinOp::BitXor | BinOp::Shl | BinOp::Shr => {
                if lt.is_int() && self.env.types_match(lt, rt) {
                    lt.clone()
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0101",
                        format!(
                            "cannot apply `{}` to `{}` and `{}`; bitwise ops require integer operands",
                            op, lt, rt
                        ),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
            BinOp::Match | BinOp::NotMatch => {
                // LHS should be ByteSlice/String/StrView, RHS should be Regex
                if *rt == Type::Regex {
                    Type::Bool
                } else {
                    self.diagnostics.emit(Diagnostic::error(
                        "E0101",
                        format!(
                            "`{}` requires a Regex on the right side, found `{}`",
                            op, rt
                        ),
                        to_diag_span(span),
                    ));
                    Type::Error
                }
            }
        }
    }

    fn check_fn_call(&mut self, name: &str, args: &[CallArg], span: cjc_ast::Span) -> Type {
        // Check argument types
        let arg_types: Vec<Type> = args.iter().map(|a| self.check_expr(&a.value)).collect();

        // Look up function
        if let Some(sigs) = self.env.fn_sigs.get(name).cloned() {
            // Find matching signature with type unification
            for sig in &sigs {
                if sig.params.len() != arg_types.len() {
                    continue;
                }

                // If function has type parameters, attempt unification
                if !sig.type_params.is_empty() {
                    let mut subst = TypeSubst::new();

                    // Create fresh type variables for each type parameter
                    let mut param_var_map: BTreeMap<String, TypeVarId> = BTreeMap::new();
                    for (tp_name, _bounds) in &sig.type_params {
                        let var_id = self.env.fresh_var();
                        param_var_map.insert(tp_name.clone(), var_id);
                    }

                    // Replace type parameter names in param types with fresh vars
                    let sig_param_types: Vec<Type> = sig
                        .params
                        .iter()
                        .map(|(_, ty)| self.substitute_type_params(ty, &param_var_map))
                        .collect();
                    let sig_ret = self.substitute_type_params(&sig.ret, &param_var_map);

                    // Try to unify each argument type with the parameter type
                    let mut ok = true;
                    for (arg_ty, param_ty) in arg_types.iter().zip(sig_param_types.iter()) {
                        if unify(arg_ty, param_ty, &mut subst).is_err() {
                            ok = false;
                            break;
                        }
                    }

                    if ok {
                        // Check that all inferred types satisfy their trait bounds.
                        let violations = self.env.check_generic_call_bounds(
                            &sig.type_params,
                            &param_var_map,
                            &subst,
                        );
                        if violations.is_empty() {
                            return apply_subst(&sig_ret, &subst);
                        }
                        for (tp_name, inferred, bounds) in &violations {
                            self.diagnostics.emit(
                                Diagnostic::error(
                                    "E6001",
                                    format!(
                                        "trait bound not satisfied: type `{}` does not implement `{}` (required by type parameter `{}`)",
                                        inferred,
                                        bounds.join(" + "),
                                        tp_name
                                    ),
                                    to_diag_span(span),
                                )
                                .with_hint(format!(
                                    "type parameter `{}` requires: `{}`",
                                    tp_name,
                                    bounds.join(" + ")
                                )),
                            );
                        }
                    }
                } else {
                    // Non-generic function: verify argument types with span-aware unification.
                    let mut all_ok = true;
                    for ((_, param_ty), arg_ty) in sig.params.iter().zip(arg_types.iter()) {
                        let mut subst_local = TypeSubst::new();
                        let result = unify_spanned(
                            arg_ty,
                            param_ty,
                            &mut subst_local,
                            to_diag_span(span),
                            &mut self.diagnostics,
                        );
                        if result.is_error() && !arg_ty.is_error() {
                            all_ok = false;
                        }
                    }
                    if all_ok || arg_types.iter().any(|t| t.is_error()) {
                        return sig.ret.clone();
                    }
                    return Type::Error;
                }
            }
            self.diagnostics.emit(
                Diagnostic::error(
                    "E0113",
                    format!(
                        "no matching overload for `{}` with {} argument(s)",
                        name,
                        arg_types.len()
                    ),
                    to_diag_span(span),
                )
                .with_hint(format!(
                    "available overloads: {}",
                    sigs.iter()
                        .map(|s| format!("{}({})", s.name, s.params.len()))
                        .collect::<Vec<_>>()
                        .join(", ")
                )),
            );
            Type::Error
        } else {
            // Built-in function check
            match name {
                "print" => Type::Void,
                "mean" | "sum" | "max" | "min" | "count" => {
                    // Data DSL aggregation functions
                    Type::F64
                }
                "bf16_to_f32" => Type::F32,
                "f32_to_bf16" => Type::Bf16,
                "f16_to_f64" => Type::F64,
                "f64_to_f16" => Type::F16,
                "f16_to_f32" => Type::F32,
                "f32_to_f16" => Type::F16,
                _ => {
                    self.diagnostics.emit(
                        Diagnostic::error(
                            "E0114",
                            format!("undefined function `{}`", name),
                            to_diag_span(span),
                        )
                        .with_hint("check for typos or define the function"),
                    );
                    Type::Error
                }
            }
        }
    }

    /// Replace type parameter names (like "T") with fresh type variables in a type.
    fn substitute_type_params(&self, ty: &Type, map: &BTreeMap<String, TypeVarId>) -> Type {
        match ty {
            Type::Unresolved(name) => {
                if let Some(&var_id) = map.get(name) {
                    Type::Var(var_id)
                } else {
                    ty.clone()
                }
            }
            Type::Tensor { elem, shape } => Type::Tensor {
                elem: Box::new(self.substitute_type_params(elem, map)),
                shape: shape.clone(),
            },
            Type::Buffer { elem } => Type::Buffer {
                elem: Box::new(self.substitute_type_params(elem, map)),
            },
            Type::Array { elem, len } => Type::Array {
                elem: Box::new(self.substitute_type_params(elem, map)),
                len: *len,
            },
            Type::Tuple(elems) => {
                Type::Tuple(elems.iter().map(|e| self.substitute_type_params(e, map)).collect())
            }
            Type::Fn { params, ret } => Type::Fn {
                params: params.iter().map(|p| self.substitute_type_params(p, map)).collect(),
                ret: Box::new(self.substitute_type_params(ret, map)),
            },
            Type::Enum(e) => Type::Enum(EnumType {
                name: e.name.clone(),
                type_params: e.type_params.clone(),
                variants: e
                    .variants
                    .iter()
                    .map(|v| EnumVariant {
                        name: v.name.clone(),
                        fields: v
                            .fields
                            .iter()
                            .map(|f| self.substitute_type_params(f, map))
                            .collect(),
                    })
                    .collect(),
            }),
            _ => ty.clone(),
        }
    }

    fn resolve_type_expr(&self, ty: &TypeExpr) -> Type {
        match &ty.kind {
            TypeExprKind::Named { name, args } => {
                let base_name = &name.name;
                match base_name.as_str() {
                    "i32" => Type::I32,
                    "i64" => Type::I64,
                    "u8" => Type::U8,
                    "f32" => Type::F32,
                    "f64" => Type::F64,
                    "bool" => Type::Bool,
                    "String" => Type::Str,
                    "Bytes" => Type::Bytes,
                    "ByteSlice" => Type::ByteSlice,
                    "StrView" => Type::StrView,
                    "Regex" => Type::Regex,
                    "Tensor" => {
                        let elem = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::F64
                            }
                        } else {
                            Type::F64
                        };
                        let shape = if args.len() > 1 {
                            if let TypeArg::Shape(dims) = &args[1] {
                                Some(
                                    dims.iter()
                                        .map(|d| match d {
                                            cjc_ast::ShapeDim::Name(n) => {
                                                ShapeDim::Symbolic(n.name.clone())
                                            }
                                            cjc_ast::ShapeDim::Lit(v) => {
                                                ShapeDim::Known(*v as usize)
                                            }
                                        })
                                        .collect(),
                                )
                            } else {
                                None
                            }
                        } else {
                            None
                        };
                        Type::Tensor {
                            elem: Box::new(elem),
                            shape,
                        }
                    }
                    "Buffer" => {
                        let elem = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::F64
                            }
                        } else {
                            Type::F64
                        };
                        Type::Buffer {
                            elem: Box::new(elem),
                        }
                    }
                    "bf16" => Type::Bf16,
                    "f16" => Type::F16,
                    "Complex" => Type::Complex,
                    "Range" => {
                        let elem = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::I64
                            }
                        } else {
                            Type::I64
                        };
                        Type::Range {
                            elem: Box::new(elem),
                        }
                    }
                    "Slice" => {
                        let elem = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::F64
                            }
                        } else {
                            Type::F64
                        };
                        Type::Slice {
                            elem: Box::new(elem),
                        }
                    }
                    "Option" => {
                        let inner = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::Unresolved("T".into())
                            }
                        } else {
                            Type::Unresolved("T".into())
                        };
                        Type::Enum(EnumType {
                            name: "Option".into(),
                            type_params: vec!["T".into()],
                            variants: vec![
                                EnumVariant { name: "Some".into(), fields: vec![inner] },
                                EnumVariant { name: "None".into(), fields: vec![] },
                            ],
                        })
                    }
                    "Result" => {
                        let ok_ty = if !args.is_empty() {
                            if let TypeArg::Type(t) = &args[0] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::Unresolved("T".into())
                            }
                        } else {
                            Type::Unresolved("T".into())
                        };
                        let err_ty = if args.len() > 1 {
                            if let TypeArg::Type(t) = &args[1] {
                                self.resolve_type_expr(t)
                            } else {
                                Type::Unresolved("E".into())
                            }
                        } else {
                            Type::Unresolved("E".into())
                        };
                        Type::Enum(EnumType {
                            name: "Result".into(),
                            type_params: vec!["T".into(), "E".into()],
                            variants: vec![
                                EnumVariant { name: "Ok".into(), fields: vec![ok_ty] },
                                EnumVariant { name: "Err".into(), fields: vec![err_ty] },
                            ],
                        })
                    }
                    _ => {
                        if let Some(ty) = self.env.type_defs.get(base_name) {
                            ty.clone()
                        } else {
                            Type::Unresolved(base_name.clone())
                        }
                    }
                }
            }
            TypeExprKind::Array { elem, size } => {
                let elem_ty = self.resolve_type_expr(elem);
                let len = if let ExprKind::IntLit(v) = size.kind {
                    v as usize
                } else {
                    0
                };
                Type::Array {
                    elem: Box::new(elem_ty),
                    len,
                }
            }
            TypeExprKind::Tuple(elems) => {
                Type::Tuple(elems.iter().map(|e| self.resolve_type_expr(e)).collect())
            }
            TypeExprKind::Fn { params, ret } => Type::Fn {
                params: params.iter().map(|p| self.resolve_type_expr(p)).collect(),
                ret: Box::new(self.resolve_type_expr(ret)),
            },
            TypeExprKind::ShapeLit(_) => Type::Error, // Shape literals only valid in type args
        }
    }

    fn type_expr_name(&self, ty: &TypeExpr) -> String {
        match &ty.kind {
            TypeExprKind::Named { name, .. } => name.name.clone(),
            _ => "<unknown>".into(),
        }
    }
}

impl Default for TypeChecker {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builtin_types() {
        let env = TypeEnv::new();
        assert_eq!(env.resolve_type_name("i32"), Some(Type::I32));
        assert_eq!(env.resolve_type_name("f64"), Some(Type::F64));
        assert_eq!(env.resolve_type_name("bool"), Some(Type::Bool));
    }

    #[test]
    fn test_trait_satisfaction() {
        let env = TypeEnv::new();
        assert!(env.satisfies_trait(&Type::I32, "Numeric"));
        assert!(env.satisfies_trait(&Type::I32, "Int"));
        assert!(!env.satisfies_trait(&Type::I32, "Float"));

        assert!(env.satisfies_trait(&Type::F64, "Numeric"));
        assert!(env.satisfies_trait(&Type::F64, "Float"));
        assert!(env.satisfies_trait(&Type::F64, "Differentiable"));
        assert!(!env.satisfies_trait(&Type::F64, "Int"));

        assert!(!env.satisfies_trait(&Type::Bool, "Numeric"));
    }

    #[test]
    fn test_type_display() {
        assert_eq!(format!("{}", Type::I32), "i32");
        assert_eq!(format!("{}", Type::F64), "f64");
        assert_eq!(
            format!(
                "{}",
                Type::Tensor {
                    elem: Box::new(Type::F32),
                    shape: Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)])
                }
            ),
            "Tensor<f32, [3, 4]>"
        );
    }

    #[test]
    fn test_matmul_shape_check() {
        let env = TypeEnv::new();

        // Valid: [3, 4] x [4, 5] -> [3, 5]
        let result = env.check_matmul_shapes(
            &Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]),
            &Some(vec![ShapeDim::Known(4), ShapeDim::Known(5)]),
        );
        assert!(result.is_ok());
        let shape = result.unwrap().unwrap();
        assert_eq!(shape, vec![ShapeDim::Known(3), ShapeDim::Known(5)]);

        // Invalid: [3, 4] x [5, 6] -> error
        let result = env.check_matmul_shapes(
            &Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]),
            &Some(vec![ShapeDim::Known(5), ShapeDim::Known(6)]),
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_scope() {
        let mut env = TypeEnv::new();
        env.define_var("x", Type::I32);
        assert_eq!(env.lookup_var("x"), Some(&Type::I32));

        env.push_scope();
        env.define_var("y", Type::F64);
        assert_eq!(env.lookup_var("x"), Some(&Type::I32));
        assert_eq!(env.lookup_var("y"), Some(&Type::F64));

        env.pop_scope();
        assert_eq!(env.lookup_var("x"), Some(&Type::I32));
        assert_eq!(env.lookup_var("y"), None);
    }

    #[test]
    fn test_types_match() {
        let env = TypeEnv::new();
        assert!(env.types_match(&Type::I32, &Type::I32));
        assert!(!env.types_match(&Type::I32, &Type::I64));
        assert!(env.types_match(&Type::Error, &Type::I32)); // Error matches everything
    }

    // ── Unification Tests ────────────────────────────────────────

    #[test]
    fn test_unify_concrete_types() {
        let mut subst = TypeSubst::new();
        assert_eq!(unify(&Type::I32, &Type::I32, &mut subst).unwrap(), Type::I32);
        assert_eq!(unify(&Type::F64, &Type::F64, &mut subst).unwrap(), Type::F64);
        assert_eq!(unify(&Type::Bool, &Type::Bool, &mut subst).unwrap(), Type::Bool);
        assert_eq!(unify(&Type::Str, &Type::Str, &mut subst).unwrap(), Type::Str);
    }

    #[test]
    fn test_unify_var_binds() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        let result = unify(&Type::Var(var0), &Type::F64, &mut subst).unwrap();
        assert_eq!(result, Type::F64);
        assert_eq!(subst.get(&var0), Some(&Type::F64));
    }

    #[test]
    fn test_unify_var_right_side() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        let result = unify(&Type::I32, &Type::Var(var0), &mut subst).unwrap();
        assert_eq!(result, Type::I32);
        assert_eq!(subst.get(&var0), Some(&Type::I32));
    }

    #[test]
    fn test_unify_mismatch_fails() {
        let mut subst = TypeSubst::new();
        assert!(unify(&Type::I32, &Type::F64, &mut subst).is_err());
        assert!(unify(&Type::Bool, &Type::Str, &mut subst).is_err());
    }

    #[test]
    fn test_unify_tensor_elem() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        let a = Type::Tensor {
            elem: Box::new(Type::Var(var0)),
            shape: None,
        };
        let b = Type::Tensor {
            elem: Box::new(Type::F32),
            shape: None,
        };
        let result = unify(&a, &b, &mut subst).unwrap();
        assert_eq!(
            result,
            Type::Tensor {
                elem: Box::new(Type::F32),
                shape: None,
            }
        );
        assert_eq!(subst.get(&var0), Some(&Type::F32));
    }

    #[test]
    fn test_unify_fn_types() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        let a = Type::Fn {
            params: vec![Type::Var(var0)],
            ret: Box::new(Type::Var(var0)),
        };
        let b = Type::Fn {
            params: vec![Type::I64],
            ret: Box::new(Type::I64),
        };
        let result = unify(&a, &b, &mut subst).unwrap();
        assert_eq!(
            result,
            Type::Fn {
                params: vec![Type::I64],
                ret: Box::new(Type::I64),
            }
        );
    }

    #[test]
    fn test_unify_error_recovery() {
        let mut subst = TypeSubst::new();
        assert_eq!(unify(&Type::Error, &Type::I32, &mut subst).unwrap(), Type::I32);
        assert_eq!(unify(&Type::F64, &Type::Error, &mut subst).unwrap(), Type::F64);
    }

    #[test]
    fn test_apply_subst_basic() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        subst.insert(var0, Type::F32);
        assert_eq!(apply_subst(&Type::Var(var0), &subst), Type::F32);
        assert_eq!(apply_subst(&Type::I32, &subst), Type::I32);
    }

    #[test]
    fn test_apply_subst_nested() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        subst.insert(var0, Type::F32);
        let tensor = Type::Tensor {
            elem: Box::new(Type::Var(var0)),
            shape: None,
        };
        assert_eq!(
            apply_subst(&tensor, &subst),
            Type::Tensor {
                elem: Box::new(Type::F32),
                shape: None,
            }
        );
    }

    #[test]
    fn test_apply_subst_chained() {
        let mut subst = TypeSubst::new();
        let var0 = TypeVarId(0);
        let var1 = TypeVarId(1);
        subst.insert(var0, Type::Var(var1));
        subst.insert(var1, Type::I64);
        assert_eq!(apply_subst(&Type::Var(var0), &subst), Type::I64);
    }

    #[test]
    fn test_check_bounds_numeric() {
        let env = TypeEnv::new();
        assert!(env.check_bounds(&Type::I32, &["Numeric".into()]));
        assert!(env.check_bounds(&Type::F64, &["Numeric".into()]));
        assert!(!env.check_bounds(&Type::Bool, &["Numeric".into()]));
    }

    #[test]
    fn test_check_bounds_differentiable() {
        let env = TypeEnv::new();
        assert!(env.check_bounds(&Type::F64, &["Differentiable".into()]));
        assert!(env.check_bounds(&Type::F32, &["Differentiable".into()]));
        assert!(!env.check_bounds(&Type::I32, &["Differentiable".into()]));
    }

    #[test]
    fn test_check_bounds_multiple() {
        let env = TypeEnv::new();
        assert!(env.check_bounds(&Type::F64, &["Numeric".into(), "Float".into()]));
        assert!(!env.check_bounds(&Type::I32, &["Numeric".into(), "Float".into()]));
    }

    // ── Shape Unification Tests ────────────────────────────────

    #[test]
    fn test_unify_shape_known_match() {
        let mut subst = ShapeSubst::new();
        let result = unify_shape_dim(&ShapeDim::Known(3), &ShapeDim::Known(3), &mut subst);
        assert_eq!(result.unwrap(), ShapeDim::Known(3));
    }

    #[test]
    fn test_unify_shape_known_mismatch() {
        let mut subst = ShapeSubst::new();
        let result = unify_shape_dim(&ShapeDim::Known(3), &ShapeDim::Known(4), &mut subst);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("mismatch"));
    }

    #[test]
    fn test_unify_shape_symbolic_binds() {
        let mut subst = ShapeSubst::new();
        let result = unify_shape_dim(
            &ShapeDim::Symbolic("N".into()),
            &ShapeDim::Known(5),
            &mut subst,
        );
        assert_eq!(result.unwrap(), ShapeDim::Known(5));
        assert_eq!(subst.get("N"), Some(&5));
    }

    #[test]
    fn test_unify_shape_symbolic_conflict() {
        let mut subst = ShapeSubst::new();
        subst.insert("N".into(), 3);
        let result = unify_shape_dim(
            &ShapeDim::Symbolic("N".into()),
            &ShapeDim::Known(5),
            &mut subst,
        );
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("already bound"));
    }

    #[test]
    fn test_unify_shapes_rank_mismatch() {
        let mut subst = ShapeSubst::new();
        let a = vec![ShapeDim::Known(3), ShapeDim::Known(4)];
        let b = vec![ShapeDim::Known(3)];
        let result = unify_shapes(&a, &b, &mut subst);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("rank mismatch"));
    }

    // ── Broadcasting Tests ─────────────────────────────────────

    #[test]
    fn test_broadcast_same_shape() {
        let a = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let b = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let result = broadcast_shapes(&a, &b).unwrap().unwrap();
        assert_eq!(result, vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
    }

    #[test]
    fn test_broadcast_scalar_to_matrix() {
        let a = Some(vec![ShapeDim::Known(1)]);
        let b = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let result = broadcast_shapes(&a, &b).unwrap().unwrap();
        assert_eq!(result, vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
    }

    #[test]
    fn test_broadcast_row_to_matrix() {
        let a = Some(vec![ShapeDim::Known(1), ShapeDim::Known(4)]);
        let b = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let result = broadcast_shapes(&a, &b).unwrap().unwrap();
        assert_eq!(result, vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
    }

    #[test]
    fn test_broadcast_col_to_matrix() {
        let a = Some(vec![ShapeDim::Known(3), ShapeDim::Known(1)]);
        let b = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let result = broadcast_shapes(&a, &b).unwrap().unwrap();
        assert_eq!(result, vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
    }

    #[test]
    fn test_broadcast_incompatible() {
        let a = Some(vec![ShapeDim::Known(3), ShapeDim::Known(5)]);
        let b = Some(vec![ShapeDim::Known(3), ShapeDim::Known(4)]);
        let result = broadcast_shapes(&a, &b);
        assert!(result.is_err());
    }

    #[test]
    fn test_broadcast_symbolic_with_one() {
        let a = Some(vec![ShapeDim::Symbolic("N".into())]);
        let b = Some(vec![ShapeDim::Known(1)]);
        let result = broadcast_shapes(&a, &b).unwrap().unwrap();
        assert_eq!(result, vec![ShapeDim::Symbolic("N".into())]);
    }

    #[test]
    fn test_broadcast_none_passthrough() {
        assert_eq!(broadcast_shapes(&None, &None).unwrap(), None);
        let a = Some(vec![ShapeDim::Known(3)]);
        assert_eq!(broadcast_shapes(&a, &None).unwrap(), None);
    }

    #[test]
    fn test_value_vs_gc_type() {
        assert!(Type::I32.is_value_type());
        assert!(Type::F64.is_value_type());
        assert!(Type::Struct(StructType {
            name: "Foo".into(),
            type_params: vec![],
            fields: vec![]
        })
        .is_value_type());

        assert!(Type::Class(ClassType {
            name: "Bar".into(),
            type_params: vec![],
            fields: vec![]
        })
        .is_gc_type());

        assert!(!Type::I32.is_gc_type());
        assert!(!Type::Class(ClassType {
            name: "Bar".into(),
            type_params: vec![],
            fields: vec![]
        })
        .is_value_type());
    }
}
