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

//! AST nodes for the AILang language.
//!
//! Every type in this module is the in-memory mirror of a node in the
//! AILang JSON schema documented in `docs/DESIGN.md`. 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"`, or
/// `"type"`. Use [`def_name`] / [`def_kind`] (or the [`Def::name`]
/// method) to inspect generically without matching every variant.
#[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),
}

impl Def {
    /// Source-level name of the definition, regardless of kind.
    pub fn name(&self) -> &str {
        match self {
            Def::Fn(f) => &f.name,
            Def::Const(c) => &c.name,
            Def::Type(t) => &t.name,
        }
    }
}

/// 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"`).
///
/// 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",
    }
}

/// 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 (Iter 13a parameterised-ADT support). A
    /// monomorphic ADT has `vars` empty and is serialized identically
    /// to the pre-13a schema (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>,
    /// Iter 18e: 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 pre-18e fixture 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>,
}

/// 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.
    ///
    /// Iter 14e: `tail` marks this call as occurring in tail position
    /// (per Decision 8). 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-14e
    /// 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 (Iter 16b.1). 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 print "hi"`). The `op` is
    /// resolved against the effect-handler table at link time.
    ///
    /// Iter 14e: 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 (Iter 8b). 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 (Iter 10). `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>,
    },
    /// Iter 18c.1: explicit RC clone. Lowers identically to its inner
    /// term in 18c.1; in 18c.3 the codegen will emit
    /// `call void @ailang_rc_inc(ptr %v)` before returning `%v` under
    /// `--alloc=rc`. The variant is additive: `(clone X)` round-trips
    /// through every pre-18c.1 fixture without their hashes changing
    /// because none of them use the new tag.
    Clone {
        value: Box<Term>,
    },
    /// Iter 18d.1: 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`. Lowers as identity in 18d.1
    /// (returns `body`'s `(ssa, ty)`, ignores `source`); 18d.2 will
    /// lower this as in-place rewrite under `--alloc=rc`. The variant
    /// is additive — pre-18d fixtures keep their canonical-JSON hash
    /// because none of them use the new tag.
    #[serde(rename = "reuse-as")]
    ReuseAs {
        source: Box<Term>,
        body: Box<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,
}

/// 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>,
    },
}

/// A literal value.
///
/// The JSON discriminator is the `kind` field. `Unit` carries no
/// payload and serializes as `{"kind":"unit"}`.
#[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,
}

/// 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 (Iter 13a). For pre-13a 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).
    ///
    /// Iter 18a (Decision 10): `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`. Pre-18a fixtures therefore hash
    /// bit-identically.
    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>,
        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![] }
    }

    /// Iter 18a: 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-18a 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,
        }
    }
}

/// Iter 18a (Decision 10): per-parameter / return mode marker on a
/// [`Type::Fn`].
///
/// `Implicit` is the legacy state for fn-types that were constructed
/// before the borrow/own surface annotations existed. Semantically,
/// `Implicit ≡ Own` throughout the 18-series; the distinction exists
/// only so pre-18a 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, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ParamMode {
    /// Pre-18a / unannotated. Treated as `Own` by the typechecker.
    Implicit,
    /// `(own T)` — caller transfers ownership; callee consumes.
    Own,
    /// `(borrow T)` — caller retains ownership; callee may not consume.
    Borrow,
}

impl Default for ParamMode {
    fn default() -> Self {
        ParamMode::Implicit
    }
}

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

/// Iter 18a: 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-18a
/// fixtures keep bit-identical hashes.
fn all_implicit(modes: &[ParamMode]) -> bool {
    modes.iter().all(|m| m.is_implicit())
}

/// Iter 18a: equality of [`ParamMode`] for the purposes of `Type`
/// equality. `Implicit` and `Own` are treated as the same mode
/// throughout the 18-series; `Borrow` is distinct. This keeps
/// pre-18a fixtures (whose fn-types serialize `Implicit`) compatible
/// with newly-written 18a 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,
    }
}

/// Iter 18a: 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, body: ab },
                Type::Forall { vars: b, 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`] (Iter 14e) so the
/// `tail` flag is omitted from the canonical JSON whenever it is false,
/// preserving bit-identical hashes for every pre-14e definition.
#[allow(clippy::trivially_copy_pass_by_ref)]
fn is_false(b: &bool) -> bool {
    !*b
}