AILang/crates/ailang-core/src/ast.rs

//! AST nodes for the AILang language.
//!
//! **The canonical schema lives in `design/contracts/data-model.md`**;
//! this module is the Rust-side projection of it. When the two drift,
//! `crates/ailang-core/tests/design_schema_drift.rs` fires.
//!
//! The serde attributes carry the schema: field renames (`as`, `type`,
//! `fn`, `paramTypes`, `retType`), enum tags (`kind`, `t`, `k`, `p`),
//! and `skip_serializing_if` predicates that keep the canonical-JSON
//! representation backwards compatible across schema extensions.
//!
//! The entry type is [`Module`]. The two helpers [`def_name`] and
//! [`def_kind`] are intended for tools (`ail diff`, `ail manifest`) that
//! consume a [`Def`] without going through method calls.
//!
//! This module does **not** typecheck, evaluate, or hash anything — see
//! [`crate::canonical`] for canonical bytes, [`crate::hash`] for content
//! hashes, and the `ailang-check` crate for typechecking.

use serde::{Deserialize, Serialize};

/// A complete AILang translation unit.
///
/// Carries the schema tag (always [`crate::SCHEMA`] for this version),
/// a module name, the list of imports, and the list of top-level
/// definitions. A module is the smallest unit that can be loaded,
/// hashed, and emitted as LLVM IR.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Module {
    /// Schema identifier; must equal [`crate::SCHEMA`] at load time.
    pub schema: String,
    /// Module name. By convention matches the file stem
    /// (`<name>.ail.json`); the [`crate::workspace`] loader enforces
    /// this on import.
    pub name: String,
    /// Imports of other modules. Resolved by the workspace loader as
    /// `<root_dir>/<module>.ail.json`.
    #[serde(default)]
    pub imports: Vec<Import>,
    /// Top-level definitions in declaration order.
    pub defs: Vec<Def>,
}

/// An import statement.
///
/// `module` is the bare module name (no path, no extension). `alias`
/// renames it for use in the body; absent means the module is referred
/// to under its own name.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Import {
    /// Imported module name, resolved relative to the entry file's
    /// directory.
    pub module: String,
    /// Optional rename (`import foo as bar`). Serialized as `as` in
    /// JSON; omitted when absent.
    #[serde(rename = "as", default, skip_serializing_if = "Option::is_none")]
    pub alias: Option<String>,
}

/// A top-level definition — the unit that gets a content hash.
///
/// Discriminated by the JSON `kind` tag: `"fn"`, `"const"`, `"type"`,
/// `"class"`, or `"instance"`. Use [`def_name`] / [`def_kind`] (or the
/// [`Def::name`] method) to inspect generically without matching every
/// variant.
///
/// `Class` and `Instance` are additive — fixtures that predate
/// the typeclass layer never produce these tags, so their
/// canonical-JSON bytes (and therefore their content hashes) are
/// unchanged. The `Def`-level match exhaustiveness in downstream
/// crates is the only caller that has to acknowledge the variants.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub enum Def {
    /// Function definition; see [`FnDef`].
    Fn(FnDef),
    /// Constant definition; see [`ConstDef`].
    Const(ConstDef),
    /// Type (ADT) definition; see [`TypeDef`].
    Type(TypeDef),
    /// Typeclass declaration; see [`ClassDef`].
    Class(ClassDef),
    /// Instance declaration; see [`InstanceDef`].
    Instance(InstanceDef),
}

impl Def {
    /// Source-level name of the definition, regardless of kind.
    ///
    /// For `Def::Instance` the "name" is the class being instantiated;
    /// instances are not separately named at the source level. The
    /// workspace registry keys instances by `(class, type-hash)`, so
    /// the class name alone is the closest analogue of a "name".
    pub fn name(&self) -> &str {
        match self {
            Def::Fn(f) => &f.name,
            Def::Const(c) => &c.name,
            Def::Type(t) => &t.name,
            Def::Class(c) => &c.name,
            Def::Instance(i) => &i.class,
        }
    }
}

/// External helper: name of a definition.
///
/// Equivalent to [`Def::name`], exposed as a free function so tools
/// like `ail diff` that hold a [`Def`] node can read its name without
/// importing the inherent-method namespace.
pub fn def_name(def: &Def) -> &str {
    def.name()
}

/// External helper: discriminator tag of a definition (`"fn"`,
/// `"const"`, `"type"`, `"class"`, `"instance"`).
///
/// Identical to the `kind` field in the JSON representation. Useful
/// for tools that group or filter definitions by kind.
pub fn def_kind(def: &Def) -> &'static str {
    match def {
        Def::Fn(_) => "fn",
        Def::Const(_) => "const",
        Def::Type(_) => "type",
        Def::Class(_) => "class",
        Def::Instance(_) => "instance",
    }
}

/// An algebraic data type definition.
///
/// A `TypeDef` introduces a named type with one or more constructors.
/// Monomorphic types (`vars` empty) and parameterised types (`vars`
/// non-empty) share the same node — see the field docs for the
/// schema-compatibility note.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TypeDef {
    /// Type name (capitalised by convention).
    pub name: String,
    /// Type parameters for parameterised ADTs. A monomorphic ADT has
    /// `vars` empty and the field is **omitted** when empty,
    /// preserving the canonical-JSON hash of every existing module on
    /// disk. See the regression test
    /// `iter13a_schema_extension_preserves_pre_13a_hashes` in
    /// [`crate::hash`].
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub vars: Vec<String>,
    /// Constructors in declaration order. Pattern-match arms in
    /// `ailang-check` are validated against this list.
    pub ctors: Vec<Ctor>,
    /// Optional source-level documentation string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub doc: Option<String>,
    /// Opt-in `(drop-iterative)` annotation. When `true`,
    /// codegen emits `drop_<m>_<T>` with an iterative worklist body
    /// instead of the recursive cascade — chosen by the LLM-author when
    /// the type is expected to form long chains (millions of cells)
    /// that would overflow the C stack on free.
    ///
    /// Serialised as `"drop-iterative": true` (kebab-case) when set;
    /// the field is omitted when `false` so canonical-JSON hashes of
    /// every fixture that does not opt in remain bit-stable. See the
    /// regression test `iter18e_drop_iterative_default_preserves_hashes`
    /// in [`crate::hash`].
    #[serde(
        default,
        rename = "drop-iterative",
        skip_serializing_if = "is_false"
    )]
    pub drop_iterative: bool,
}


/// A single constructor of a [`TypeDef`].
///
/// `fields` is the constructor's positional argument list as types;
/// nullary constructors have an empty list.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Ctor {
    /// Constructor name (capitalised by convention; unique within its
    /// `TypeDef`).
    pub name: String,
    /// Positional field types. Empty for nullary constructors.
    #[serde(default)]
    pub fields: Vec<Type>,
}

/// A function definition.
///
/// `ty` is the full function type (including effect set and any
/// `Forall` quantifier for top-level polymorphism). `params` are the
/// names bound in `body`, in order.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FnDef {
    /// Function name.
    pub name: String,
    /// Declared function type. Top-level polymorphism is opt-in via
    /// [`Type::Forall`].
    #[serde(rename = "type")]
    pub ty: Type,
    /// Parameter names, in the order they appear in `ty.params`.
    pub params: Vec<String>,
    /// Function body.
    pub body: Term,
    /// Optional source-level documentation string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub doc: Option<String>,
    /// Iter embedding-abi-m1: when `Some(sym)`, this fn is the
    /// embedding boundary. Codegen's `Target::StaticLib` mode
    /// additionally emits an externally-visible C entrypoint
    /// `@<sym>` (signature frozen as of M3 —
    /// design/contracts/embedding-abi.md) forwarding
    /// to `@ail_<module>_<fn>`. The symbol is author-chosen and
    /// decoupled from the `ail_<module>_<def>` mangling so the
    /// M3-frozen ABI survives module/fn refactors.
    ///
    /// Serialised `skip_serializing_if = "Option::is_none"` so every
    /// pre-M1 fixture's canonical-JSON hash stays bit-identical —
    /// the same additive-schema pattern as [`FnDef::doc`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub export: Option<String>,
    /// Structured-diagnostic suppressions opted into for this
    /// fn. Each entry names a diagnostic code and an author-asserted
    /// reason. The only current consumer is `over-strict-mode`, but
    /// the mechanism is generic across codes. `because` must be
    /// non-empty — the typechecker emits `empty-suppress-reason`
    /// (Error) otherwise.
    ///
    /// Serialised with `skip_serializing_if = "Vec::is_empty"` so
    /// every pre-suppress fixture's canonical-JSON hash stays
    /// bit-identical. The same additive-schema pattern is used by
    /// [`TypeDef::vars`] and [`Type::Con::args`].
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub suppress: Vec<Suppress>,
}

/// One entry in [`FnDef::suppress`]. Marks a structured diagnostic
/// the author has consciously decided to allow on this def, with a
/// mandatory reason.
///
/// `because` is non-empty by schema rule — the typechecker emits
/// `empty-suppress-reason` (Error severity) when it is empty or
/// whitespace-only, and a wrong/unknown `code` simply matches no
/// diagnostic and therefore suppresses nothing (the original
/// diagnostic still fires unmasked).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Suppress {
    /// The diagnostic code being suppressed (e.g.
    /// `"over-strict-mode"`). Matched against
    /// [`crate::SCHEMA`]-side codes; an unknown code suppresses
    /// nothing but is not itself an error (the diagnostic registry
    /// is open-set).
    pub code: String,
    /// The author's stated reason. Must be non-empty — the
    /// typechecker emits `empty-suppress-reason` (Error) otherwise.
    pub because: String,
}

/// A typeclass declaration (narrative in
/// `design/contracts/typeclasses.md`).
///
/// Single-parameter, multi-method, optional-default, optional-superclass
/// typeclass. The `param` is a single string — multi-param classes are
/// rejected by the typeclass design (concrete-types-only, kind `*`),
/// and the schema enforces it by shape (`param: String`, not
/// `Vec<String>`).
///
/// `superclass`, when present, MUST have its `type` field equal to
/// `param`. The check is enforced via the `InvalidSuperclassParam`
/// diagnostic; the schema does not encode it.
///
/// All optional fields are omitted from canonical JSON when absent /
/// empty, so future schema evolution can land here without disturbing
/// existing hashes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassDef {
    /// Class name (capitalised by convention). Bare — symmetric to
    /// `TypeDef.name`, the field is the defining-site context, not
    /// a reference. Cross-module class references live in
    /// `InstanceDef.class`, `Constraint.class`, and
    /// `SuperclassRef.class`; those fields carry the canonical form
    /// (bare for same-module, `<module>.<Class>` for cross-module).
    pub name: String,
    /// Single type-parameter name.
    pub param: String,
    /// Optional single-superclass relation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub superclass: Option<SuperclassRef>,
    /// Methods of the class, in declaration order.
    pub methods: Vec<ClassMethod>,
    /// Optional source-level documentation string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub doc: Option<String>,
}

/// Reference to a superclass relation in [`ClassDef`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuperclassRef {
    /// Superclass name in canonical form: bare for a same-module
    /// class, `<module>.<Class>` for a cross-module class.
    /// Symmetric to `Type::Con.name`'s canonical-form rule.
    pub class: String,
    /// Type the superclass is applied to. MUST equal the parent
    /// `ClassDef.param` (validated by typecheck — schema does not
    /// enforce).
    #[serde(rename = "type")]
    pub type_: String,
}

/// One method declared in a [`ClassDef`].
///
/// `default` is `None` when the method is abstract-required (every
/// instance must specify it); `Some(body)` when the method has a
/// default body (instances may inherit or override).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassMethod {
    /// Method name.
    pub name: String,
    /// Full method signature including any mode annotations. The
    /// class parameter appears as a [`Type::Var`] inside this signature.
    #[serde(rename = "type")]
    pub ty: Type,
    /// Default body. `None` ⇒ abstract-required; `Some(body)` ⇒
    /// default-with-fallback.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default: Option<Term>,
}

/// An instance declaration (narrative in
/// `design/contracts/typeclasses.md`).
///
/// `class` is the name of the class being instantiated. `type_` is the
/// concrete type expression the class is applied to — never the class
/// param. `methods` contains bodies for the required methods plus any
/// overrides of default-bearing methods.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InstanceDef {
    /// Class being instantiated, in canonical form: bare for a
    /// same-module class, `<module>.<Class>` for a cross-module
    /// class. Symmetric to `Type::Con.name`'s canonical-form rule.
    pub class: String,
    /// Concrete type the class is applied to.
    #[serde(rename = "type")]
    pub type_: Type,
    /// Method bodies in declaration order.
    pub methods: Vec<InstanceMethod>,
    /// Optional source-level documentation string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub doc: Option<String>,
}

/// One method body in an [`InstanceDef`].
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InstanceMethod {
    /// Method name (must match a method in the corresponding class).
    pub name: String,
    /// Method body. The class's declared method type with the class
    /// param substituted to the instance type is the body's expected
    /// type.
    pub body: Term,
}

/// A class constraint on a polymorphic function (narrative in
/// `design/contracts/typeclasses.md`). `(class, type)` pair where
/// `class` is a class name and `type` is a `Type` expression —
/// typically a single `Type::Var` (e.g. `(Eq, a)` for
/// `Eq a => ...`). Concrete-type constraints are legal schema-wise
/// but fired as `no-instance` at typecheck time if no matching
/// registry entry exists.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Constraint {
    /// Class name in canonical form: bare for a same-module class,
    /// `<module>.<Class>` for a cross-module class. Symmetric to
    /// `Type::Con.name`'s canonical-form rule.
    pub class: String,
    /// Type the class is applied to.
    #[serde(rename = "type")]
    pub type_: Type,
}

/// A constant (top-level binding to a value).
///
/// Differs from a zero-arg [`FnDef`] in that it is evaluated once and
/// has no parameter list; the codegen emits it as a global.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConstDef {
    /// Constant name.
    pub name: String,
    /// Declared type of the value.
    #[serde(rename = "type")]
    pub ty: Type,
    /// The value expression. Must be pure (no effects).
    pub value: Term,
    /// Optional source-level documentation string.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub doc: Option<String>,
}

/// An expression node ("Term" = computation that produces a value).
///
/// The JSON discriminator is the `t` field. Each variant's doc
/// describes its concrete-syntax intent; the typing rules live in
/// `ailang-check`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "t", rename_all = "lowercase")]
pub enum Term {
    /// Literal value (see [`Literal`]).
    Lit { lit: Literal },
    /// Variable reference (parameter, local, top-level def, or
    /// imported alias).
    Var { name: String },
    /// Function application. `callee` is evaluated to a function value;
    /// `args` are evaluated left-to-right.
    ///
    /// `tail` marks this call as occurring in tail position (see
    /// `design/contracts/tail-calls.md`). When set, codegen lowers
    /// the call as `musttail call`. The flag defaults to `false` and
    /// is omitted during canonical-JSON serialisation when unset, so
    /// pre-tail-flag fixtures keep bit-identical hashes.
    App {
        #[serde(rename = "fn")]
        callee: Box<Term>,
        args: Vec<Term>,
        #[serde(default, skip_serializing_if = "is_false")]
        tail: bool,
    },
    /// Let-binding: `value` is evaluated and bound to `name` in `body`.
    Let {
        name: String,
        value: Box<Term>,
        body: Box<Term>,
    },
    /// Local recursive let-binding. Always fn-shaped:
    /// the bound name `name` is recursively visible inside `body`.
    /// Eliminated by `crate::desugar` before typecheck — lifted to a
    /// synthetic top-level fn when `body` does not capture any name
    /// from the enclosing lexical scope. The on-disk schema gains the
    /// `"t": "letrec"` tag; pre-existing fixtures hash bit-identically
    /// because the variant is additive.
    LetRec {
        name: String,
        #[serde(rename = "type")]
        ty: Type,
        params: Vec<String>,
        body: Box<Term>,
        #[serde(rename = "in")]
        in_term: Box<Term>,
    },
    /// If-expression. Both branches must have the same type.
    If {
        cond: Box<Term>,
        then: Box<Term>,
        #[serde(rename = "else")]
        else_: Box<Term>,
    },
    /// Effect operation invocation (e.g. `do io/print_str "hi"`). The
    /// `op` is resolved at typecheck against `Env::effect_ops` (an
    /// `IndexMap<String, EffectOpSig>`) and lowered by a literal
    /// `match` in codegen (`lower_effect_op`); there is no
    /// effect-handler table and no link-time resolution.
    ///
    /// See [`Term::App`] for the `tail` field semantics.
    Do {
        op: String,
        args: Vec<Term>,
        #[serde(default, skip_serializing_if = "is_false")]
        tail: bool,
    },
    /// Constructor application. `type_name` binds the ADT, `ctor` the
    /// variant. Example: `Some(42)` ->
    /// `{ "t": "ctor", "type": "Option", "ctor": "Some",
    /// "args": [{"t":"lit","lit":{"kind":"int","value":42}}] }`.
    Ctor {
        #[serde(rename = "type")]
        type_name: String,
        ctor: String,
        #[serde(default)]
        args: Vec<Term>,
    },
    /// Pattern matching over a value. Arms are tried top-to-bottom;
    /// exhaustiveness is checked by `ailang-check`.
    Match {
        scrutinee: Box<Term>,
        arms: Vec<Arm>,
    },
    /// Anonymous function. Captures any free variables of
    /// `body` from the enclosing scope. Param/return types are
    /// declared inline so the typechecker stays HM-monomorphic on
    /// the inferred shape.
    Lam {
        params: Vec<String>,
        #[serde(rename = "paramTypes")]
        param_tys: Vec<Type>,
        #[serde(rename = "retType")]
        ret_ty: Box<Type>,
        #[serde(default)]
        effects: Vec<String>,
        body: Box<Term>,
    },
    /// Sequencing. `lhs` is evaluated for its effects and its
    /// result discarded; `rhs` is the value of the whole expression.
    /// Equivalent to `let _ = lhs in rhs`, but with a dedicated node so
    /// pretty-print and diagnostics read cleanly.
    Seq {
        lhs: Box<Term>,
        rhs: Box<Term>,
    },
    /// Explicit RC clone. Codegen emits
    /// `call void @ailang_rc_inc(ptr %v)` before returning `%v` under
    /// `--alloc=rc`. The variant is additive: `(clone X)` round-trips
    /// through every fixture that does not use the tag without its
    /// hash changing.
    Clone {
        value: Box<Term>,
    },
    /// Explicit reuse-as hint. Wraps an allocating `body` (typically
    /// `Term::Ctor`, also `Term::Lam`) and names a `source` term
    /// whose memory slot the body's allocation should reuse.
    /// Conventionally `source` is a `Term::Var { name }` — the
    /// linearity check rejects anything else with
    /// `reuse-as-source-not-bare-var`. The body must be allocating;
    /// non-allocating bodies are rejected at typecheck with
    /// `reuse-as-non-allocating-body`. Codegen lowers as in-place
    /// rewrite under `--alloc=rc`. The variant is additive —
    /// fixtures that do not use the tag keep their canonical-JSON
    /// hash.
    #[serde(rename = "reuse-as")]
    ReuseAs {
        source: Box<Term>,
        body: Box<Term>,
    },
    /// loop-recur iter 1: a strict iteration block. `binders`
    /// declares one or more loop parameters (name, type, init),
    /// evaluated in order on loop entry; `body` is evaluated with
    /// all binders in scope. The loop's value is `body`'s value on
    /// the iteration that exits via a non-`recur` branch. Strictly
    /// additive: pre-existing fixtures hash bit-identically because
    /// none carry the `"t":"loop"` tag. `binders` has no
    /// `skip_serializing_if` — the field is part of the shape.
    /// Typecheck (loop-recur.2): `synth`
    /// types each binder init, the loop's type is the body's type;
    /// `recur` arity/type is checked positionally via `loop_stack`
    /// and `verify_loop_body` enforces `recur`-in-tail-position; a
    /// lambda capturing a binder is rejected
    /// (`loop-binder-captured-by-lambda`, loop-recur.tidy). Codegen
    /// (loop-recur.3): binders lower to entry-block allocas reached
    /// from a `loop.header` block; `recur` stores + back-edges; the
    /// loop-carried SSA / phi form is produced by `clang -O2`
    /// mem2reg (NOT hand-emitted phi). No totality claim — an
    /// infinite loop is legal. See
    /// `docs/specs/2026-05-17-loop-recur.md`.
    Loop {
        binders: Vec<LoopBinder>,
        body: Box<Term>,
    },
    /// loop-recur iter 1: re-enter the lexically innermost enclosing
    /// `Term::Loop`, rebinding its binders positionally to `args`.
    /// Transfers control (no fall-through); valid only in tail
    /// position of its enclosing loop — enforced at typecheck in
    /// iter 2 (`RecurNotInTailPosition`). Additive `"t":"recur"`
    /// tag; pre-existing fixtures hash bit-identically.
    Recur {
        args: Vec<Term>,
    },
}

/// One arm of a [`Term::Match`].
///
/// `pat` is the pattern; `body` is evaluated when the pattern matches,
/// with any pattern-bound variables in scope.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Arm {
    /// Pattern to match the scrutinee against.
    pub pat: Pattern,
    /// Body to evaluate on a successful match.
    pub body: Term,
}

/// loop-recur iter 1: one binder of a [`Term::Loop`]. The
/// `(name, type, init)` triple is the Form-A surface vocabulary
/// `(NAME TYPE INIT)`; it is a nested field of `Term::Loop`, not a
/// first-class `Term` (loop binders cannot escape the enclosing
/// loop). `recur` rebinds these positionally per iteration. The
/// `ty` JSON field is `"type"`, matching the spec schema.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoopBinder {
    /// The binder's lexical name. Within the enclosing `Term::Loop`,
    /// a `Term::Var { name }` resolves to this binding.
    pub name: String,
    /// The binder's declared type.
    #[serde(rename = "type")]
    pub ty: Type,
    /// Initial value, evaluated once on loop entry in scope of the
    /// outer environment plus already-declared binders of the same
    /// `Term::Loop` (declaration order).
    pub init: Term,
}

/// A match pattern.
///
/// The JSON discriminator is the `p` field. Patterns are linear: each
/// pattern variable may appear at most once.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "p", rename_all = "lowercase")]
pub enum Pattern {
    /// `_` — binds nothing, matches everything.
    Wild,
    /// `x` — binds the value to a name.
    Var { name: String },
    /// Match on a literal.
    Lit { lit: Literal },
    /// Match on a constructor with sub-patterns for its fields.
    Ctor {
        ctor: String,
        #[serde(default)]
        fields: Vec<Pattern>,
    },
}

/// Private serde helper: emits a `u64` as a 16-character lowercase
/// hex JSON string and parses the same shape back. The
/// 16-character invariant covers the full `u64` range zero-padded
/// (`format!("{:016x}", 0u64) == "0000000000000000"`,
/// `format!("{:016x}", u64::MAX) == "ffffffffffffffff"`). Used by
/// [`Literal::Float`] so float bit patterns flow through the
/// canonical-JSON *string* path — guaranteeing bit stability across
/// `serde_json` versions and surfacing NaN / ±Inf, which JSON
/// numbers cannot represent.
mod hex_u64 {
    use serde::{de::Error as DeError, Deserialize, Deserializer, Serializer};

    pub fn serialize<S: Serializer>(value: &u64, s: S) -> Result<S::Ok, S::Error> {
        s.serialize_str(&format!("{:016x}", value))
    }

    pub fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
        let s = <&str>::deserialize(d)?;
        if s.len() != 16 {
            return Err(D::Error::custom(format!(
                "Float bits: expected 16 hex chars, got {}",
                s.len()
            )));
        }
        u64::from_str_radix(s, 16).map_err(D::Error::custom)
    }
}

/// A literal value.
///
/// The JSON discriminator is the `kind` field. `Unit` carries no
/// payload and serializes as `{"kind":"unit"}`. `Float` carries a
/// `u64` IEEE-754 binary64 bit pattern, serialized via a private
/// `hex_u64` helper as a 16-character lowercase hex string —
/// canonical-JSON bytes therefore stay bit-stable for hashing and
/// can represent NaN / ±Inf (which JSON numbers cannot).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub enum Literal {
    /// Signed 64-bit integer.
    Int { value: i64 },
    /// Boolean.
    Bool { value: bool },
    /// UTF-8 string.
    Str { value: String },
    /// The unit value `()`.
    Unit,
    /// IEEE-754 binary64 (LLVM `double`) carried as its bit pattern.
    /// The private `hex_u64` helper module above documents why the
    /// field is routed through the JSON string path.
    Float {
        #[serde(with = "hex_u64")]
        bits: u64,
    },
}

/// A type expression.
///
/// The JSON discriminator is the `k` field. The four variants are
/// type constructor application ([`Type::Con`]), function type
/// ([`Type::Fn`]), type variable ([`Type::Var`]), and universal
/// quantifier ([`Type::Forall`] — only valid at the top level of an
/// [`FnDef`] / [`ConstDef`] type).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "k", rename_all = "lowercase")]
pub enum Type {
    /// Type-constructor application: `Int`, `Bool`, user ADT names,
    /// and parameterised forms like `List<Int>`.
    Con {
        name: String,
        /// Type arguments. For non-parameterised uses (`Int`, `Bool`,
        /// non-parameterised user ADTs) this stays empty and is
        /// **omitted** during serialization, so the canonical-JSON
        /// hash of every pre-existing module remains bit-identical.
        /// See the regression test in [`crate::hash`].
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        args: Vec<Type>,
    },
    /// Function type with optional effect annotation. `effects` is a
    /// set; equality compares it modulo order (see the [`PartialEq`]
    /// impl below).
    ///
    /// `param_modes` and `ret_mode` carry the `(borrow T)` /
    /// `(own T)` wrappers from the surface form. They are metadata
    /// on `Type::Fn`, not new `Type` variants — so unification,
    /// occurs, apply, and every other `Type` match-arm keeps working
    /// unchanged. `param_modes` is omitted from canonical JSON when
    /// every entry is `Implicit`; `ret_mode` is omitted when it is
    /// `Implicit`, so pre-mode-annotation fixtures hash
    /// bit-identically. Full contract in
    /// `design/contracts/memory-model.md`.
    Fn {
        params: Vec<Type>,
        #[serde(default, skip_serializing_if = "all_implicit")]
        param_modes: Vec<ParamMode>,
        ret: Box<Type>,
        #[serde(default, skip_serializing_if = "ParamMode::is_implicit")]
        ret_mode: ParamMode,
        #[serde(default)]
        effects: Vec<String>,
    },
    /// Type variable. During checking, names with the prefix `$m`
    /// denote checker-internal metavariables; source-level names
    /// cannot collide because identifiers may not start with `$`.
    Var {
        name: String,
    },
    /// Universal quantifier (top-level polymorphism only). `body`
    /// is the quantified type; `vars` are the bound type-variable
    /// names, instantiated fresh at each use site.
    Forall {
        vars: Vec<String>,
        /// Class constraints quantified together with `vars`. Empty
        /// for unconstrained polymorphic types; serialised with
        /// `skip_serializing_if = "Vec::is_empty"` so existing
        /// canonical-JSON bytes stay bit-identical.
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        constraints: Vec<Constraint>,
        body: Box<Type>,
    },
}

impl Type {
    /// Convenience constructor for `Int` (no args).
    pub fn int() -> Type {
        Type::Con { name: "Int".into(), args: vec![] }
    }
    /// Convenience constructor for `Bool` (no args).
    pub fn bool_() -> Type {
        Type::Con { name: "Bool".into(), args: vec![] }
    }
    /// Convenience constructor for `Unit` (no args).
    pub fn unit() -> Type {
        Type::Con { name: "Unit".into(), args: vec![] }
    }
    /// Convenience constructor for `Str` (no args).
    pub fn str_() -> Type {
        Type::Con { name: "Str".into(), args: vec![] }
    }
    /// Convenience constructor for `Float` (no args). IEEE-754
    /// binary64. Parallel to [`Type::int`] / [`Type::bool_`].
    pub fn float() -> Type {
        Type::Con { name: "Float".into(), args: vec![] }
    }

    /// Build a `Type::Fn` with all parameter modes set to
    /// `ParamMode::Implicit` and `ret_mode` set to `Implicit`. This
    /// is the form every typechecker / desugar / codegen site that
    /// synthesises a fn-type should use, so that newly inferred
    /// fn-types retain pre-mode-annotation canonical-JSON bytes.
    pub fn fn_implicit(params: Vec<Type>, ret: Type, effects: Vec<String>) -> Type {
        let n = params.len();
        Type::Fn {
            params,
            param_modes: vec![ParamMode::Implicit; n],
            ret: Box::new(ret),
            ret_mode: ParamMode::Implicit,
            effects,
        }
    }
}

/// Per-parameter / return mode marker on a [`Type::Fn`]. Full
/// contract lives in `design/contracts/memory-model.md`.
///
/// `Implicit` is the legacy state for fn-types that were constructed
/// before the borrow/own surface annotations existed. Semantically,
/// `Implicit ≡ Own`; the distinction exists only so pre-annotation
/// JSON fixtures continue to serialize without a `"mode"` wrapper
/// and therefore keep their canonical-JSON hash.
///
/// `Own` and `Borrow` are author-asserted: the surface form
/// `(own T)` / `(borrow T)` round-trips through this enum.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ParamMode {
    /// Unannotated / back-compat. Treated as `Own` by the typechecker.
    #[default]
    Implicit,
    /// `(own T)` — caller transfers ownership; callee consumes.
    Own,
    /// `(borrow T)` — caller retains ownership; callee may not consume.
    Borrow,
}

impl ParamMode {
    /// Used by the `skip_serializing_if` predicate on
    /// [`Type::Fn::ret_mode`].
    pub fn is_implicit(&self) -> bool {
        matches!(self, ParamMode::Implicit)
    }
}

/// Serde helper for [`Type::Fn::param_modes`]. Returns `true` when
/// every entry is [`ParamMode::Implicit`] (or when the list is
/// empty), so canonical JSON omits the field for any fn-type without
/// explicit `(borrow)` / `(own)` annotations and pre-annotation
/// fixtures keep bit-identical hashes.
fn all_implicit(modes: &[ParamMode]) -> bool {
    modes.iter().all(|m| m.is_implicit())
}

/// Equality of [`ParamMode`] for the purposes of `Type` equality.
/// `Implicit` and `Own` are treated as the same mode; `Borrow` is
/// distinct. This keeps pre-annotation fixtures (whose fn-types
/// serialize `Implicit`) compatible with newly-written fixtures
/// that mark the same fn-type explicitly with `(own T)`.
fn mode_eq(a: &ParamMode, b: &ParamMode) -> bool {
    match (a, b) {
        (ParamMode::Borrow, ParamMode::Borrow) => true,
        (ParamMode::Borrow, _) | (_, ParamMode::Borrow) => false,
        // Implicit and Own are interchangeable.
        _ => true,
    }
}

/// Equality of two `param_modes` slices, robust to the
/// "elided when all-implicit" representation used by typechecker /
/// desugar / codegen sites that construct fn-types with
/// `param_modes: vec![]`. Both slices are normalised to "implicit
/// padding to match the longer one"; equality then proceeds
/// element-wise via [`mode_eq`].
fn mode_slices_eq(a: &[ParamMode], b: &[ParamMode]) -> bool {
    let n = a.len().max(b.len());
    for i in 0..n {
        let x = a.get(i).copied().unwrap_or(ParamMode::Implicit);
        let y = b.get(i).copied().unwrap_or(ParamMode::Implicit);
        if !mode_eq(&x, &y) {
            return false;
        }
    }
    true
}

impl PartialEq for Type {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (
                Type::Con { name: a, args: aa },
                Type::Con { name: b, args: ba },
            ) => a == b && aa == ba,
            (
                Type::Fn {
                    params: ap,
                    param_modes: apm,
                    ret: ar,
                    ret_mode: arm,
                    effects: ae,
                },
                Type::Fn {
                    params: bp,
                    param_modes: bpm,
                    ret: br,
                    ret_mode: brm,
                    effects: be,
                },
            ) => {
                ap == bp
                    && ar == br
                    && mode_slices_eq(apm, bpm)
                    && mode_eq(arm, brm)
                    && {
                        let mut a = ae.clone();
                        let mut b = be.clone();
                        a.sort();
                        b.sort();
                        a == b
                    }
            }
            (Type::Var { name: a }, Type::Var { name: b }) => a == b,
            (
                Type::Forall { vars: a, constraints: _, body: ab },
                Type::Forall { vars: b, constraints: _, body: bb },
            ) => a == b && ab == bb,
            _ => false,
        }
    }
}
impl Eq for Type {}

/// Serde helper for `#[serde(skip_serializing_if = "is_false")]`.
///
/// Used by [`Term::App::tail`] and [`Term::Do::tail`] so the `tail`
/// flag is omitted from the canonical JSON whenever it is false,
/// preserving bit-identical hashes for every fixture that does not
/// carry the flag.
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_false(b: &bool) -> bool {
    !*b
}

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

    /// loop-recur iter 1: pin the canonical-bytes shape of a
    /// `Term::Loop` with one binder: `binders` stays present in
    /// canonical JSON (no `skip_serializing_if`).
    #[test]
    fn term_loop_one_binder_serialises_with_explicit_binders_field() {
        let t = Term::Loop {
            binders: vec![LoopBinder {
                name: "i".into(),
                ty: Type::int(),
                init: Term::Lit {
                    lit: Literal::Int { value: 0 },
                },
            }],
            body: Box::new(Term::Var { name: "i".into() }),
        };
        let bytes = serde_json::to_string(&t).expect("serialise");
        assert_eq!(
            bytes,
            r#"{"t":"loop","binders":[{"name":"i","type":{"k":"con","name":"Int"},"init":{"t":"lit","lit":{"kind":"int","value":0}}}],"body":{"t":"var","name":"i"}}"#,
        );
        let back: Term = serde_json::from_str(&bytes).expect("deserialise");
        match back {
            Term::Loop { binders, body } => {
                assert_eq!(binders.len(), 1);
                assert_eq!(binders[0].name, "i");
                match *body {
                    Term::Var { name } => assert_eq!(name, "i"),
                    other => panic!("body mismatch: {other:?}"),
                }
            }
            other => panic!("variant mismatch: {other:?}"),
        }
    }

    /// loop-recur iter 1: round-trip a `Term::Recur` through JSON.
    /// Pins `{ "t": "recur", "args": [...] }`.
    #[test]
    fn term_recur_round_trips_through_json() {
        let t = Term::Recur {
            args: vec![Term::Lit {
                lit: Literal::Int { value: 1 },
            }],
        };
        let bytes = serde_json::to_string(&t).expect("serialise");
        assert_eq!(
            bytes,
            r#"{"t":"recur","args":[{"t":"lit","lit":{"kind":"int","value":1}}]}"#,
        );
        let back: Term = serde_json::from_str(&bytes).expect("deserialise");
        match back {
            Term::Recur { args } => assert_eq!(args.len(), 1),
            other => panic!("variant mismatch: {other:?}"),
        }
    }
}