earthsci-ast 0.6.0

//! The operator registry: the single source of truth for **which `op` strings
//! exist** and **how many `args` each one takes** (esm-spec §4.2 / §4.3).
//!
//! Before this module the crate had no arity check anywhere — neither
//! [`crate::validate`] nor [`crate::structural`] — and the schema deliberately
//! leaves `args` at `minItems: 0`, because a *schema* cannot express "`atan2`
//! takes exactly two arguments but `+` takes any number". The result was that a
//! malformed-but-schema-valid node reached the evaluators, where the two of them
//! disagreed about what it meant: the per-cell oracle
//! ([`crate::simulate_array::eval`]) folded `-(3,1,1)` to `NaN` while the
//! vectorized overlay ([`crate::simulate_array::vectorized`]) left-folded it to
//! `1.0`, and `atan2` with one argument indexed `args[1]` and **panicked**.
//!
//! The registry closes that class of bug at the source: an op/arity pair that
//! this table rejects never reaches an evaluator, so the evaluators are only
//! ever asked to agree on nodes that are *legal*, and for those they agree by
//! construction.
//!
//! # The two tiers (esm-spec §4.2)
//!
//! - **Evaluable core (closed).** Everything in [`arity_of`]. Adding one is a
//!   spec change. A core op with the wrong arity is a hard error
//!   ([`OpError::Arity`]).
//! - **Rewrite-target (open).** Everything else — the sugar ops
//!   `grad`/`div`/`laplacian`/`curl`/`∇`/`integral`, a *spatial* `D`, and any
//!   user-defined op. These have **no evaluator**; they MUST be eliminated by a
//!   rewrite rule (§9.6) before evaluation. Reaching an evaluator with one still
//!   present is [`OpError::Unlowered`] → diagnostic `unlowered_operator`.
//!
//! A **misspelled** op (`"expp"`) is, formally, an open-tier op for which no
//! rewrite rule exists — so it lands in the same `unlowered_operator` bucket
//! rather than being silently evaluated to `NaN`. That is the spec's own answer
//! (§9.6.3 constraint 6: "any node whose `op` is not in the evaluable-core set
//! … is rejected with diagnostic `unlowered_operator`"), and it is why a typo is
//! now loud instead of quietly poisoning a trajectory with `NaN`.

use crate::types::{Expr, ExpressionNode};

/// How many `args` an evaluable-core operator accepts.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Arity {
    /// Exactly `n` arguments (`atan2` → 2, `ifelse` → 3, `const` → 0).
    Exact(usize),
    /// Between `lo` and `hi` inclusive (`-` → 1 or 2: negation or subtraction).
    Between(usize, usize),
    /// At least `n` (`+`/`*` are n-ary; `min`/`max` are n-ary with `n >= 2`).
    AtLeast(usize),
    /// Any arity, including zero. Used by the ops whose operands do not live in
    /// `args` at all (`faq`/`makearray` carry them in `expr`/`values`;
    /// `fn` defers to the closed-function registry's own signature check).
    Any,
}

impl Arity {
    /// Does `n` satisfy this arity?
    #[must_use]
    pub fn admits(self, n: usize) -> bool {
        match self {
            Self::Exact(k) => n == k,
            Self::Between(lo, hi) => n >= lo && n <= hi,
            Self::AtLeast(k) => n >= k,
            Self::Any => true,
        }
    }

    /// Human-readable expectation, for diagnostics.
    #[must_use]
    pub fn describe(self) -> String {
        match self {
            Self::Exact(k) => format!("exactly {k}"),
            Self::Between(lo, hi) => format!("{lo} or {hi}"),
            Self::AtLeast(k) => format!("at least {k}"),
            Self::Any => "any number of".to_string(),
        }
    }
}

/// What is wrong with an operator node.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OpError {
    /// A known evaluable-core op applied to the wrong number of arguments.
    Arity {
        /// The operator name.
        op: String,
        /// How many arguments it was given.
        got: usize,
        /// What the spec says it takes.
        expected: String,
    },
    /// A rewrite-target op (sugar, spatial `D`, user op, or a typo) reached
    /// evaluation without being lowered. Diagnostic code `unlowered_operator`.
    Unlowered {
        /// The operator name.
        op: String,
    },
    /// A `makearray` whose `regions` are ragged (regions disagreeing on rank)
    /// or carry an inverted bound pair. Diagnostic code
    /// `makearray_region_inverted` (esm-spec §4.3.2).
    MakearrayRegion {
        /// What is wrong with the regions.
        reason: String,
    },
    /// A `broadcast` node whose `fn` is absent, or names something that is not
    /// a SCALAR operator (esm-spec §4.3.4: "the `fn` value must name a scalar
    /// operator"). Diagnostic code `invalid_broadcast_fn`.
    ///
    /// An `fn`-arity mismatch is reported as [`OpError::Arity`] against the
    /// `fn` NAME instead, because that is literally the same defect the bare
    /// node would have: `broadcast(fn = "min", [x])` is wrong for exactly the
    /// reason `{"op": "min", "args": [x]}` is wrong.
    BroadcastFn {
        /// The offending `fn` value, or `None` when the field is absent.
        fn_name: Option<String>,
        /// Why it is rejected.
        reason: String,
    },
}

impl std::fmt::Display for OpError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Arity { op, got, expected } => {
                write!(f, "operator '{op}' takes {expected} argument(s), got {got}")
            }
            Self::Unlowered { op } => write!(
                f,
                "operator '{op}' has no evaluator and was not lowered by any rewrite rule"
            ),
            Self::MakearrayRegion { reason } | Self::BroadcastFn { reason, .. } => {
                write!(f, "{reason}")
            }
        }
    }
}

/// Validate a `makearray` node's `regions` block (esm-spec §4.3.2).
///
/// Two things can be wrong with it, and both used to **panic** the per-cell
/// evaluator rather than being reported:
///
/// 1. **Ragged regions.** `eval_makearray` took its rank from `regions[0]` and
///    then indexed `lo[d]`/`hi[d]` for every `d` of *every* region, so a second
///    region of higher rank ran off the end of the bounding box.
/// 2. **An inverted bound pair.** The bounding-box extent `hi - lo + 1` was cast
///    `as usize`, so a reversed region such as `[5, 2]` produced an extent of
///    `-2`, wrapped to `usize::MAX - 1`, and blew up in `ArrayD::zeros` with a
///    capacity overflow.
///
/// The spec draws the line precisely: `stop == start - 1` is the canonical
/// **empty** region — it is what a metaparameter-folded interior region
/// (`[2, N-1]` at `N = 2`) legitimately produces, contributes no cells, and MUST
/// load cleanly. Anything *further* inverted (`stop < start - 1`) is an
/// authoring error and is rejected with `makearray_region_inverted`.
///
/// # Errors
///
/// [`OpError::MakearrayRegion`] naming the offending region.
pub fn check_makearray_regions(node: &ExpressionNode) -> Result<(), OpError> {
    let Some(regions) = node.regions.as_deref() else {
        return Ok(());
    };
    let bad = |reason: String| Err(OpError::MakearrayRegion { reason });

    if let Some(values) = node.values.as_deref()
        && values.len() != regions.len()
    {
        return bad(format!(
            "makearray has {} region(s) but {} value(s) — they must correspond one-to-one",
            regions.len(),
            values.len()
        ));
    }

    let Some(first) = regions.first() else {
        return Ok(());
    };
    let ndim = first.len();
    for (i, region) in regions.iter().enumerate() {
        if region.len() != ndim {
            return bad(format!(
                "makearray region {i} has rank {} but region 0 has rank {ndim} — every region \
                 must address the same number of dimensions",
                region.len()
            ));
        }
        for (d, pair) in region.iter().enumerate() {
            // A still-symbolic (unfolded) bound has no orderable value yet, so
            // there is nothing to check: §9.7.6 folding runs before evaluation
            // and this same check re-runs on the folded node.
            let Some([start, stop]) = crate::types::region_bounds(pair) else {
                continue;
            };
            // `stop == start - 1` is the legal empty spelling; only a further
            // inversion is an error.
            if stop < start - 1 {
                return bad(format!(
                    "makearray region {i} dimension {d} has inverted bounds [{start}, {stop}]: \
                     stop < start - 1 (the empty spelling [{start}, {}] is legal, anything \
                     further inverted is not)",
                    start - 1
                ));
            }
        }
    }
    Ok(())
}

/// The arity of an evaluable-core operator, or `None` if `op` is not in the
/// closed core set (esm-spec §4.2) — i.e. it is an open-tier rewrite target.
///
/// The unary elementary functions are grouped; every arity here is the one the
/// spec's §4.2 tables state.
#[must_use]
pub fn arity_of(op: &str) -> Option<Arity> {
    let a = match op {
        // --- Arithmetic (§4.2). `+`/`*` are n-ary; `-` is unary OR binary;
        // `/` and `^` are strictly binary.
        "+" | "*" => Arity::AtLeast(1),
        "-" => Arity::Between(1, 2),
        // `pow` is the word spelling of `^` — the same binary power, evaluated
        // by the same kernel — as the Julia and Python registries list it.
        "/" | "^" | "pow" => Arity::Exact(2),
        // `neg` is the canonical unary negation `canonicalize.rs` emits.
        "neg" => Arity::Exact(1),
        // The engine-internal precision-boundary marker
        // (`crate::precision_infer::MARKER_OP`), inserted after the typed
        // parse and never authored: a transparent unary wrapper naming the
        // element type its operand is evaluated at.
        crate::precision_infer::MARKER_OP => Arity::Exact(1),

        // --- Elementary functions (§4.2). All unary except `atan2` and
        // `min`/`max`.
        "exp" | "log" | "log10" | "sqrt" | "abs" | "sign" | "floor" | "ceil" | "sin" | "cos"
        | "tan" | "asin" | "acos" | "atan" | "sinh" | "cosh" | "tanh" | "asinh" | "acosh"
        | "atanh" => Arity::Exact(1),
        "atan2" => Arity::Exact(2),
        // "n-ary with arity >= 2 … Conforming bindings MUST reject `min`/`max`
        // nodes with fewer than two arguments." (§4.2)
        "min" | "max" => Arity::AtLeast(2),

        // --- Conditionals / logic (§4.2).
        "ifelse" => Arity::Exact(3),
        "==" | "!=" | "<" | "<=" | ">" | ">=" => Arity::Exact(2),
        "and" | "or" => Arity::AtLeast(2),
        "not" => Arity::Exact(1),

        // --- Calculus (§4.2 "Arity of `D`"). This table is keyed by op STRING,
        // so it states the STRUCTURAL time derivative's contract: `wrt: "t"` (or
        // an absent `wrt`) is evaluable-core, consumed by system assembly, and
        // STRICTLY UNARY. A SPATIAL `D` is a rewrite target — it never reaches
        // an arity check here (`check_node` short-circuits it to `Unlowered`
        // first via `is_rewrite_target_derivative`), so it may carry trailing
        // auxiliary boundary operands (esm-spec §4.2 "Arity of `D`").
        "D" | "ic" => Arity::Exact(1),

        // --- Events (§4.2).
        "Pre" => Arity::Exact(1),

        // --- Inline constants (§4.2): "`args` MUST be empty `[]`".
        "const" => Arity::Exact(0),
        // Nullary boolean literals — `true` is an always-true join/`filter`
        // predicate (§4.3); `false` is its counterpart, 0 in the evaluators'
        // 1.0 / 0.0 boolean encoding.
        "true" | "false" => Arity::Exact(0),

        // --- Closed-registry invocation (§4.2 / §9.2). The registry checks the
        // per-function signature itself (`unknown_closed_function`), so arity is
        // open here.
        "fn" => Arity::Any,
        // Lowered at load time (§4.5 / §9.5). `enum` is `[enum_name, symbol]`;
        // `table_lookup` carries its inputs in `axes`, so "`args` MUST be empty".
        "enum" => Arity::Exact(2),
        "table_lookup" => Arity::Exact(0),
        // Expanded at load time (§9.6); its operands live in `bindings`.
        "apply_expression_template" => Arity::Any,

        // --- Array / tensor (§4.3). `faq` and `makearray` carry their
        // real operands in `expr` / `values`, and `args` is conventionally the
        // (possibly empty) operand list, so their arity is open.
        "faq" => Arity::Any,
        "makearray" => Arity::Any,
        // "`args[0]` is the array; `args[1..]` are the index expressions."
        "index" => Arity::AtLeast(1),
        "broadcast" => Arity::AtLeast(1),
        "reshape" | "transpose" => Arity::Exact(1),
        "concat" => Arity::AtLeast(1),

        // --- Relational / value-invention & geometry (§4.3 companions). These
        // are BUILD-TIME ops, evaluated by `value_invention.rs` (not by the
        // per-cell oracle), so they are known-and-legal here even though
        // `eval_op` has no arm for them. Their operand shapes are checked by
        // that evaluator.
        "skolem" | "rank" | "distinct" => Arity::Any,
        "argmin" | "argmax" => Arity::Any,
        "intersect_polygon" | "polygon_intersection_area" => Arity::Exact(2),

        _ => return None,
    };
    Some(a)
}

/// Is `op` an evaluable-core operator (esm-spec §4.2)?
#[must_use]
pub fn is_core_op(op: &str) -> bool {
    arity_of(op).is_some()
}

/// Is `name` a **built-in function head** — always legal as a bare leaf in
/// VARIABLE position even though no variable declares it (esm-spec §4.2's
/// named elementary functions and reductions, `ifelse`, and the event op
/// `Pre`)?
///
/// This is the excuse list the crate's two undefined-variable gates share:
/// [`crate::structural`]'s reference checks and the array compiler's
/// free-variable check ([`crate::simulate_array`]) skip these names instead of
/// reporting an undefined variable. Each gate used to carry its own copy of
/// the list, with a doc comment asking the other to stay in sync; the registry
/// holds the one copy now.
///
/// Every name here is an evaluable-core op, but deliberately NOT every
/// identifier-shaped core op is here: the operator spellings `neg`, `ln`,
/// `and`, `or` and `not` (and the form ops `D`, `ic`, `const`, `true`, …) have
/// never been excused as variable names, and widening the excuse set would
/// silently pass a genuinely undeclared variable that happens to carry one of
/// those names. `builtin_function_names_are_core_and_nothing_widened` pins
/// both directions.
#[must_use]
pub fn is_builtin_function_name(name: &str) -> bool {
    matches!(
        name,
        "exp"
            | "log"
            | "log10"
            | "sqrt"
            | "abs"
            | "sign"
            | "sin"
            | "cos"
            | "tan"
            | "asin"
            | "acos"
            | "atan"
            | "atan2"
            | "sinh"
            | "cosh"
            | "tanh"
            | "asinh"
            | "acosh"
            | "atanh"
            | "min"
            | "max"
            | "floor"
            | "ceil"
            | "ifelse"
            | "Pre"
    )
}

/// Is `op` a **scalar operator** — one whose meaning is a POINTWISE map from
/// scalar operands to a scalar result (esm-spec §4.2 "arithmetic, elementary
/// functions, comparisons")?
///
/// One set answers two questions the format asks in two different places,
/// because they are the same question:
///
/// * **What may a `broadcast` node's `fn` name?** §4.3.4: "the `fn` value must
///   name a scalar operator". That is what makes the `broadcast` contract
///   stateable in one line —
///
///   > `{"op": "broadcast", "fn": F, "args": A}` means what `{"op": F, "args": A}`
///   > means, applied element-wise
///
///   — and [`check_broadcast_fn`] enforces exactly it.
/// * **Does element alignment descend into a node's operands?** An ARRAY-LEVEL
///   `a * b` means "multiply corresponding elements", so the operands must first
///   be brought onto a common index space (§4.3.4). Every other op consumes its
///   operands WHOLE and defines its own operand contract — a `faq` and a
///   `makearray` name their axes, an `index` gathers, the shape ops restructure,
///   the relational and geometry kernels return arrays of an unrelated shape —
///   so alignment must not reach inside them. [`is_elementwise_node`] lifts this
///   predicate to nodes.
///
/// Keeping ONE list rather than two is deliberate: issues #100 and #101 each
/// arrived carrying its own copy, the two copies agreed on all forty names, and
/// two hand-maintained spellings of one question is precisely how they would
/// eventually stop agreeing.
///
/// Everything whose meaning is NOT pointwise is excluded, and each exclusion is
/// load-bearing:
///
/// * the array/tensor ops (`faq`, `makearray`, `index`, `reshape`,
///   `transpose`, `concat`, and `broadcast` itself) — they RESHAPE, so
///   "apply element-wise" is not defined for them (and a self-referential
///   `fn: "broadcast"` would recurse forever in the evaluators);
/// * the closed-registry invocation `fn` — its callee lives in `name`, not in
///   the `broadcast` node's `fn` field, so it can never be spelled here;
/// * the form/lowering ops (`D`, `ic`, `Pre`, `const`, `true`, `enum`,
///   `table_lookup`, `apply_expression_template`) — consumed by system
///   assembly or a lowering pass, never evaluated element-wise;
/// * the build-time relational and geometry ops (`skolem`, `rank`, `distinct`,
///   `argmin`, `argmax`, `intersect_polygon`, `polygon_intersection_area`).
///
/// Every name here is also in [`arity_of`]: this is a SUBSET of the evaluable
/// core, never an extension of it (pinned by `every_scalar_operator_has_an_arity`
/// and `the_scalar_operator_set_partitions_the_core`), which is the invariant
/// [`check_broadcast_fn`] relies on when it derives the `fn`'s legal arity. In
/// particular the registry has no alias mechanism and no alias entries — `pow`,
/// `**`, `power` are NOT spellings of `^`, `=` is not a spelling of `==`, and
/// none of them is an operator in this format at all (they are open-tier names
/// that [`crate::simulate_array`] rejects as `unlowered_operator`). Adding one
/// here would invent an operator rather than classify one;
/// `the_registry_has_no_operator_aliases` pins that.
#[must_use]
pub fn is_scalar_operator(op: &str) -> bool {
    matches!(
        op,
        // Arithmetic (§4.2), including the canonical unary negation.
        "+" | "-" | "*" | "/" | "^" | "neg"
        // Elementary functions (§4.2).
        | "exp" | "log" | "log10" | "sqrt" | "abs" | "sign" | "floor" | "ceil"
        | "sin" | "cos" | "tan" | "asin" | "acos" | "atan"
        | "sinh" | "cosh" | "tanh" | "asinh" | "acosh" | "atanh"
        | "atan2" | "min" | "max"
        // Comparisons, logic and the pointwise conditional (§4.2).
        | "==" | "!=" | "<" | "<=" | ">" | ">=" | "and" | "or" | "not" | "ifelse"
    )
}

/// Does element alignment see THROUGH this node to its `args` — i.e. are the
/// node's operands in elementwise position (esm-spec §4.3.4)?
///
/// This is [`is_scalar_operator`] lifted from op NAMES to NODES, plus the one
/// node whose elementwise-ness lives in a FIELD rather than in its name:
/// `broadcast` applies the scalar operator named in `fn` element by element, so
/// `{"op":"broadcast","fn":"*","args":[w2,z1]}` MEANS `{"op":"*","args":[w2,z1]}`
/// and its operands must align exactly as the bare spelling's do. Treating
/// `broadcast` as opaque made the two spellings disagree — the bare form aligned
/// by index-set name while the `broadcast` form silently kept the positional
/// flatten, which is the very divergence issue #100 exists to close.
///
/// The node-level / name-level split is exactly why [`is_scalar_operator`] must
/// keep EXCLUDING `broadcast`: a `broadcast` NODE is elementwise in its `args`,
/// but the STRING `"broadcast"` is not a kernel name, and any caller handed it
/// as one — as an `fn`, or as a re-dispatch target in the evaluators — recurses
/// forever. Ask the node question of a node and the name question of a name.
///
/// The `fn` guard is deliberately conservative: alignment descends only when
/// `fn` names a scalar operator the evaluators actually have a kernel for. A
/// `broadcast` carrying some other `fn` has no elementwise meaning to preserve —
/// it is now rejected outright (issue #101) — so it keeps the legacy positional
/// lowering rather than being reinterpreted. An ABSENT `fn` likewise keeps the
/// historical `+` classification: since #101 such a node is a hard
/// `invalid_broadcast_fn` error ([`check_broadcast_fn`]) and never reaches an
/// evaluator, so this arm decides only how an UNVALIDATED document is aligned,
/// and it keeps aligning it the way it always did.
#[must_use]
pub fn is_elementwise_node(node: &ExpressionNode) -> bool {
    if node.op == "broadcast" {
        return is_scalar_operator(node.broadcast_fn.as_deref().unwrap_or("+"));
    }
    is_scalar_operator(&node.op)
}

/// Validate a `broadcast` node's `fn` field (esm-spec §4.3.4).
///
/// `broadcast` is the one op whose OPERATOR is data: the node says `broadcast`
/// but the arithmetic it performs is named by a sibling string field. Nothing
/// checked that string, so three distinct defects were accepted silently and
/// all three produced a plausible wrong number rather than a diagnostic
/// (issue #101):
///
/// 1. **A missing `fn`.** All three evaluators defaulted it to `"+"`. There is
///    no default in the spec, and inventing one turns a truncated node into a
///    silent sum.
/// 2. **An `fn` that names no scalar operator** — a typo (`"not_a_real_op"`) or
///    a non-pointwise op (`"faq"`). §4.3.4 requires a scalar operator, so
///    this is the exact analogue of the §9.1 rule for `fn`-NODE names.
/// 3. **An `fn`/`args` arity mismatch.** `broadcast(fn = "sin", [a, b])` and
///    `broadcast(fn = "min", [x])` are wrong for precisely the reason the bare
///    `sin(a, b)` / `min(x)` nodes are wrong, so they get the SAME
///    [`OpError::Arity`] diagnostic, reported against the `fn` name.
///
/// Note what rule (3) does NOT reject: a one-operand `broadcast` whose `fn` is
/// genuinely n-ary or unary. `broadcast(fn = "+", [x])` and
/// `broadcast(fn = "-", [x])` stay legal, because `+(x)` and `-(x)` are legal
/// bare nodes (identity and negation respectively). That is the whole point of
/// deriving the rule from [`arity_of`] rather than hard-coding "≥ 2".
///
/// # Errors
///
/// [`OpError::BroadcastFn`] for (1) and (2); [`OpError::Arity`] for (3).
pub fn check_broadcast_fn(node: &ExpressionNode) -> Result<(), OpError> {
    let Some(fn_name) = node.broadcast_fn.as_deref() else {
        return Err(OpError::BroadcastFn {
            fn_name: None,
            reason: "'broadcast' is missing its 'fn': the field names the scalar operator to \
                     apply element-wise and has no default (esm-spec §4.3.4)"
                .to_string(),
        });
    };
    // One `else` covers both "not in the registry at all" and "in the registry
    // but not pointwise" — from the file's point of view they are the same
    // defect: `fn` does not name a scalar operator.
    let Some(arity) = arity_of(fn_name).filter(|_| is_scalar_operator(fn_name)) else {
        return Err(OpError::BroadcastFn {
            fn_name: Some(fn_name.to_string()),
            reason: format!(
                "'broadcast' fn '{fn_name}' does not name a scalar operator (esm-spec §4.3.4)"
            ),
        });
    };
    if !arity.admits(node.args.len()) {
        return Err(OpError::Arity {
            op: fn_name.to_string(),
            got: node.args.len(),
            expected: arity.describe(),
        });
    }
    Ok(())
}

/// The `wrt` value that marks the STRUCTURAL time derivative. A `D` node
/// carrying this `wrt` — or carrying no `wrt` at all, which esm-spec §4.2
/// defines to MEAN this — is the structural tier, not a rewrite target.
pub const STRUCTURAL_DERIVATIVE_WRT: &str = "t";

/// The axis a `D` node differentiates along, with esm-spec §4.2's default
/// applied: **an absent `wrt` MEANS `t`**.
///
/// Every consumer of a `D` node reads its axis through this one function. The
/// default is a property of the FORMAT, not of any one pipeline, and spelling
/// it out per call site is what let it be applied on some paths and dropped on
/// others (EarthSciAST#407).
#[must_use]
pub fn derivative_wrt(node: &ExpressionNode) -> &str {
    node.wrt.as_deref().unwrap_or(STRUCTURAL_DERIVATIVE_WRT)
}

/// Is this node a **rewrite-target** `D` — a derivative whose `wrt` names a
/// SPATIAL axis rather than the time variable (esm-spec §4.2 / §9.6.8)?
///
/// `D` is the one op whose tier — and therefore whose operand contract — depends
/// on a FIELD of the node rather than on its name. A spatial `D` has no
/// evaluator (it MUST be lowered to a stencil by a discretization rule), which
/// is exactly why it MAY carry TRAILING AUXILIARY OPERANDS after `args[0]`: the
/// per-face boundary/halo values the rule binds as ordinary §9.6.1 wildcards and
/// consumes. The STRUCTURAL time derivative (`wrt: "t"`, or absent — an absent
/// `wrt` means `t`) is not a rewrite target and stays strictly unary.
///
/// This is the single home of that predicate; both `check_node` and the
/// dimensional checker (`crate::units`) read it rather than re-deriving it.
#[must_use]
pub fn is_rewrite_target_derivative(node: &ExpressionNode) -> bool {
    node.op == "D" && derivative_wrt(node) != STRUCTURAL_DERIVATIVE_WRT
}

/// Classify one operator node: `Ok(())` if it is an evaluable-core op with a
/// legal arity, otherwise the reason it may not reach an evaluator.
///
/// A `D` node is the one op whose tier depends on a *field* rather than its
/// name: `wrt: "t"` is the structural time derivative (core), while a spatial
/// `wrt` is a rewrite target (§4.2). `check_node` is where that distinction
/// lives, so every caller inherits it.
///
/// # Errors
///
/// [`OpError::Unlowered`] for an open-tier op (sugar, spatial `D`, user op, or a
/// misspelling); [`OpError::Arity`] for a core op with the wrong argument count.
pub fn check_node(node: &ExpressionNode) -> Result<(), OpError> {
    let op = node.op.as_str();

    // A spatial `D` — or any `D` carrying a non-time `wrt` — is a rewrite
    // target, not the structural time derivative.
    if is_rewrite_target_derivative(node) {
        return Err(OpError::Unlowered { op: op.to_string() });
    }

    let Some(arity) = arity_of(op) else {
        // Open tier: the named sugar ops and anything unrecognised (including a
        // typo) share one diagnostic — `unlowered_operator` (§9.6.3 c.6).
        return Err(OpError::Unlowered { op: op.to_string() });
    };

    if !arity.admits(node.args.len()) {
        return Err(OpError::Arity {
            op: op.to_string(),
            got: node.args.len(),
            expected: arity.describe(),
        });
    }

    // `makearray` carries its real structure in `regions`/`values`, not `args`,
    // so arity alone says nothing about whether it is well-formed.
    if op == "makearray" {
        check_makearray_regions(node)?;
    }
    // `broadcast` likewise: `Arity::AtLeast(1)` says nothing about the OPERATOR
    // it applies, which lives in the `fn` field (esm-spec §4.3.4).
    if op == "broadcast" {
        check_broadcast_fn(node)?;
    }
    Ok(())
}

/// Walk an expression tree and return the first operator node that may not
/// reach an evaluator.
///
/// The walk uses [`ExpressionNode::for_each_child`], the crate's one-walker
/// keystone, so it descends into every expression-bearing sidecar field
/// (`expr`, `values`, `filter`, `key`, `axes`, `bindings`, …) and not merely
/// `args` — a hand-rolled `args`-only walk would miss exactly the
/// `aggregate.expr` bodies where these malformed nodes hide.
///
/// # Errors
///
/// The first [`OpError`] encountered, in pre-order.
pub fn check_expr(expr: &Expr) -> Result<(), OpError> {
    match expr {
        Expr::Number(_) | Expr::Integer(_) | Expr::Variable(_) => Ok(()),
        Expr::Operator(node) => {
            check_node(node)?;
            let mut first_err: Option<OpError> = None;
            node.for_each_child(&mut |child| {
                if first_err.is_none()
                    && let Err(e) = check_expr(child)
                {
                    first_err = Some(e);
                }
            });
            match first_err {
                Some(e) => Err(e),
                None => Ok(()),
            }
        }
    }
}

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

    fn node(json: serde_json::Value) -> Expr {
        serde_json::from_value(json).expect("fixture parses")
    }

    #[test]
    fn spec_arities_are_pinned() {
        assert_eq!(arity_of("atan2"), Some(Arity::Exact(2)));
        assert_eq!(arity_of("min"), Some(Arity::AtLeast(2)));
        assert_eq!(arity_of("max"), Some(Arity::AtLeast(2)));
        assert_eq!(arity_of("/"), Some(Arity::Exact(2)));
        assert_eq!(arity_of("^"), Some(Arity::Exact(2)));
        assert_eq!(arity_of("-"), Some(Arity::Between(1, 2)));
        assert_eq!(arity_of("+"), Some(Arity::AtLeast(1)));
        assert_eq!(arity_of("ifelse"), Some(Arity::Exact(3)));
        assert_eq!(arity_of("const"), Some(Arity::Exact(0)));
        assert_eq!(arity_of("expp"), None, "a typo is not a core op");
    }

    /// R1: the node that used to panic with an index-out-of-bounds.
    #[test]
    fn underapplied_atan2_is_an_arity_error() {
        let e = check_expr(&node(serde_json::json!({"op": "atan2", "args": [1.0]})))
            .expect_err("atan2/1 must be rejected");
        assert!(matches!(e, OpError::Arity { ref op, got: 1, .. } if op == "atan2"));
        let e = check_expr(&node(serde_json::json!({"op": "atan2", "args": []})))
            .expect_err("atan2/0 must be rejected");
        assert!(matches!(e, OpError::Arity { got: 0, .. }));
    }

    /// R2: the empty `neg` that panicked on the vectorized path.
    #[test]
    fn empty_neg_is_an_arity_error() {
        assert!(matches!(
            check_expr(&node(serde_json::json!({"op": "neg", "args": []}))),
            Err(OpError::Arity { got: 0, .. })
        ));
    }

    /// R3: the exact nodes on which the two evaluators disagreed. None of them
    /// is legal, so neither evaluator ever sees one again.
    #[test]
    fn the_divergent_arities_are_all_rejected() {
        for j in [
            serde_json::json!({"op": "-", "args": [3.0, 1.0, 1.0]}),
            serde_json::json!({"op": "/", "args": [8.0, 2.0, 2.0]}),
            serde_json::json!({"op": "^", "args": [2.0, 3.0, 4.0]}),
            serde_json::json!({"op": "min", "args": [5.0]}),
            serde_json::json!({"op": "max", "args": [5.0]}),
            serde_json::json!({"op": "and", "args": [5.0]}),
        ] {
            let e = check_expr(&node(j.clone()));
            assert!(
                matches!(e, Err(OpError::Arity { .. })),
                "{j} must be an arity error, got {e:?}"
            );
        }
    }

    #[test]
    fn legal_arities_pass() {
        for j in [
            serde_json::json!({"op": "+", "args": ["a", "b", "c"]}),
            serde_json::json!({"op": "-", "args": ["a"]}),
            serde_json::json!({"op": "-", "args": ["a", "b"]}),
            serde_json::json!({"op": "min", "args": ["a", "b", "c"]}),
            serde_json::json!({"op": "atan2", "args": ["a", "b"]}),
            serde_json::json!({"op": "ifelse", "args": ["c", "a", "b"]}),
            serde_json::json!({"op": "D", "args": ["u"], "wrt": "t"}),
        ] {
            check_expr(&node(j.clone())).unwrap_or_else(|e| panic!("{j} should be legal: {e}"));
        }
    }

    /// R7: a misspelled op is loud, not a silent `NaN`.
    #[test]
    fn a_typo_is_an_unlowered_operator() {
        let e = check_expr(&node(serde_json::json!({"op": "expp", "args": ["x"]})))
            .expect_err("a typo must be rejected");
        assert!(matches!(e, OpError::Unlowered { ref op } if op == "expp"));
    }

    #[test]
    fn sugar_and_spatial_d_are_unlowered() {
        for j in [
            serde_json::json!({"op": "grad", "args": ["u"], "dim": "x"}),
            serde_json::json!({"op": "laplacian", "args": ["u"]}),
            serde_json::json!({"op": "D", "args": ["u"], "wrt": "x"}),
        ] {
            assert!(
                matches!(check_expr(&node(j.clone())), Err(OpError::Unlowered { .. })),
                "{j} must be unlowered"
            );
        }
    }

    /// R5: the two `makearray` region shapes that panicked the evaluator.
    #[test]
    fn ragged_and_inverted_makearray_regions_are_rejected() {
        // Ragged: region 1 has rank 2, region 0 has rank 1.
        let ragged = node(serde_json::json!({
            "op": "makearray",
            "regions": [[[1, 3]], [[1, 3], [1, 2]]],
            "values": [0.0, 1.0],
            "args": []
        }));
        assert!(
            matches!(check_expr(&ragged), Err(OpError::MakearrayRegion { .. })),
            "a ragged regions list must be rejected, not panic"
        );

        // Inverted: [5, 2] has stop < start - 1 (extent would be -2).
        let inverted = node(serde_json::json!({
            "op": "makearray",
            "regions": [[[5, 2]]],
            "values": [0.0],
            "args": []
        }));
        assert!(
            matches!(check_expr(&inverted), Err(OpError::MakearrayRegion { .. })),
            "an inverted region must be rejected, not overflow ArrayD::zeros"
        );
    }

    /// The legal EMPTY spelling `stop == start - 1` (§4.3.2) — what a
    /// metaparameter-folded interior region produces at its minimum extent —
    /// must still load cleanly. This is the boundary the inverted check must
    /// not overshoot.
    #[test]
    fn the_empty_region_spelling_is_legal() {
        let empty = node(serde_json::json!({
            "op": "makearray",
            "regions": [[[2, 1]], [[1, 1]]],
            "values": [0.0, 1.0],
            "args": []
        }));
        check_expr(&empty).expect("[2, 1] is the canonical empty region and is legal");
    }

    /// The walk must descend into sidecar fields, not just `args` — a malformed
    /// node hiding in an `aggregate.expr` body is the whole point.
    #[test]
    fn the_walk_descends_into_sidecar_fields() {
        let agg = node(serde_json::json!({
            "op": "faq",
            "output_idx": ["i"],
            "expr": {"op": "atan2", "args": [1.0]},
            "args": ["u"]
        }));
        assert!(
            matches!(check_expr(&agg), Err(OpError::Arity { .. })),
            "an arity error inside `aggregate.expr` must be found"
        );
    }

    // -----------------------------------------------------------------------
    // Issues #100 / #101 — the scalar-operator set. ONE list answers both
    // "what may a `broadcast.fn` name?" and "does element alignment descend
    // through this node?", so both branches' pins run against it.
    // -----------------------------------------------------------------------

    /// Every op this registry knows, partitioned by [`is_scalar_operator`].
    ///
    /// The list is exhaustive over [`arity_of`] as of this commit. Adding an op
    /// to `arity_of` without deciding which side it falls on will fail
    /// `the_scalar_operator_set_partitions_the_core`, which is the point: both
    /// `broadcast.fn` legality and array-level element alignment (esm-spec
    /// §4.3.4) are defined by this partition, and a new op silently defaulting
    /// to "not scalar" is a decision, not an accident.
    const SCALAR: &[&str] = &[
        "+", "-", "*", "/", "^", "neg", "exp", "log", "log10", "sqrt", "abs", "sign", "floor",
        "ceil", "sin", "cos", "tan", "asin", "acos", "atan", "sinh", "cosh", "tanh", "asinh",
        "acosh", "atanh", "atan2", "min", "max", "ifelse", "==", "!=", "<", "<=", ">", ">=", "and",
        "or", "not",
    ];
    const NOT_SCALAR: &[&str] = &[
        "D",
        "ic",
        "Pre",
        "const",
        "true",
        "false",
        "fn",
        "enum",
        "table_lookup",
        "apply_expression_template",
        "faq",
        "makearray",
        "index",
        "broadcast",
        "reshape",
        "transpose",
        "concat",
        "skolem",
        "rank",
        "distinct",
        "argmin",
        "argmax",
        "intersect_polygon",
        "polygon_intersection_area",
    ];

    /// [`is_scalar_operator`] may only CLASSIFY registry ops, never invent them.
    #[test]
    fn the_scalar_operator_set_partitions_the_core() {
        for op in SCALAR {
            assert!(is_core_op(op), "`{op}` is scalar but not a core op");
            assert!(is_scalar_operator(op), "`{op}` should be a scalar operator");
        }
        for op in NOT_SCALAR {
            assert!(is_core_op(op), "`{op}` is listed but not a core op");
            assert!(
                !is_scalar_operator(op),
                "`{op}` should not be a scalar operator"
            );
        }
    }

    /// Every scalar operator must have an entry in `arity_of` — the invariant
    /// `check_broadcast_fn` relies on when it derives the `fn`'s legal arity.
    #[test]
    fn every_scalar_operator_has_an_arity() {
        for op in SCALAR {
            assert!(is_scalar_operator(op), "`{op}` must be a scalar operator");
            assert!(
                arity_of(op).is_some(),
                "scalar operator `{op}` has no arity entry"
            );
        }
    }

    /// The built-in-function-head excuse set, pinned exactly: the named
    /// elementary functions and reductions, `ifelse`, and `Pre`. Every name is
    /// a core op (the set only CLASSIFIES registry ops), and the
    /// identifier-shaped core spellings the undefined-variable gates have
    /// never excused stay excluded — widening the set would silently pass a
    /// genuinely undeclared variable carrying one of those names.
    #[test]
    fn builtin_function_names_are_core_and_nothing_widened() {
        const BUILTIN: &[&str] = &[
            "exp", "log", "log10", "sqrt", "abs", "sign", "sin", "cos", "tan", "asin", "acos",
            "atan", "atan2", "sinh", "cosh", "tanh", "asinh", "acosh", "atanh", "min", "max",
            "floor", "ceil", "ifelse", "Pre",
        ];
        for name in BUILTIN {
            assert!(
                is_builtin_function_name(name),
                "`{name}` must be an excused built-in function head"
            );
            assert!(is_core_op(name), "`{name}` is excused but not a core op");
        }
        for name in [
            "neg", "ln", "and", "or", "not", "D", "ic", "const", "true", "fn", "faq", "index",
            "expp", "t",
        ] {
            assert!(
                !is_builtin_function_name(name),
                "`{name}` must not be excused as a built-in function head"
            );
        }
    }

    /// `pow` and `false` are evaluable-core ops: `pow` is the word spelling of
    /// `^` and `false` the nullary counterpart of `true`, as the Julia and
    /// Python registries list them. Neither is a `broadcast.fn` — the §4.3.4
    /// scalar-operator list names `^`, not `pow`, and a literal has no operands.
    ///
    /// `**`, `power` and `=` are still not operators in this format: they are
    /// open-tier names, rejected as `unlowered_operator` when a model carrying
    /// one is built. Adding them to [`is_scalar_operator`] so array-level
    /// alignment covers them would invent operators the spec's §4.2 table does
    /// not have, and would not make them runnable either.
    #[test]
    fn pow_and_false_are_core_and_the_other_spellings_are_not() {
        assert_eq!(arity_of("pow"), Some(Arity::Exact(2)));
        assert_eq!(arity_of("false"), Some(Arity::Exact(0)));
        for op in ["pow", "false"] {
            assert!(
                !is_scalar_operator(op),
                "`{op}` must not be a scalar operator"
            );
        }
        for op in ["**", "power", "="] {
            assert!(!is_core_op(op), "`{op}` must not be a registry op");
            assert!(
                !is_scalar_operator(op),
                "`{op}` must not be a scalar operator"
            );
            assert!(
                matches!(
                    check_node(&ExpressionNode {
                        op: op.to_string(),
                        args: vec![Expr::Variable("z".into()), Expr::Number(2.0)],
                        ..Default::default()
                    }),
                    Err(OpError::Unlowered { .. })
                ),
                "`{op}` must be reported as an open-tier / unlowered operator"
            );
        }
    }

    /// The exclusions are the load-bearing half of `is_scalar_operator`: none of
    /// these has a pointwise meaning, so none may be spelled as a `broadcast.fn`.
    /// `broadcast` itself is in [`NOT_SCALAR`] — allowing it would make the
    /// evaluators' `fn` re-dispatch recurse forever. The two trailing names are
    /// not registry ops at all: an open-tier rewrite target and a typo.
    #[test]
    fn non_pointwise_ops_are_not_scalar_operators() {
        for op in NOT_SCALAR.iter().copied().chain(["grad", "not_a_real_op"]) {
            assert!(
                !is_scalar_operator(op),
                "`{op}` must NOT be usable as a `broadcast` fn"
            );
        }
    }

    /// The headline bug: `fn: "not_a_real_op"` was accepted and then silently
    /// discarded, so the node evaluated to its operand.
    #[test]
    fn unregistered_broadcast_fn_is_rejected() {
        let e = check_expr(&node(serde_json::json!({
            "op": "broadcast", "fn": "not_a_real_op", "args": ["x"]
        })))
        .expect_err("an unregistered `fn` must be rejected");
        assert!(
            matches!(e, OpError::BroadcastFn { fn_name: Some(ref f), .. } if f == "not_a_real_op"),
            "got {e:?}"
        );
    }

    /// A MISSING `fn` used to default to `+` in all three evaluators. There is
    /// no default in the spec.
    #[test]
    fn missing_broadcast_fn_is_rejected() {
        let e = check_expr(&node(
            serde_json::json!({"op": "broadcast", "args": ["x", "y"]}),
        ))
        .expect_err("a `broadcast` with no `fn` must be rejected");
        assert!(
            matches!(e, OpError::BroadcastFn { fn_name: None, .. }),
            "got {e:?}"
        );
    }

    /// A non-scalar op named as `fn` is rejected for the same reason a typo is:
    /// §4.3.4 requires a SCALAR operator.
    #[test]
    fn non_scalar_broadcast_fn_is_rejected() {
        for f in ["faq", "index", "makearray", "broadcast", "fn", "D"] {
            let e = check_expr(&node(serde_json::json!({
                "op": "broadcast", "fn": f, "args": ["x"]
            })))
            .expect_err("a non-scalar `fn` must be rejected");
            assert!(
                matches!(e, OpError::BroadcastFn { fn_name: Some(ref g), .. } if g == f),
                "fn `{f}`: got {e:?}"
            );
        }
    }

    /// `broadcast(fn = F, args)` is legal exactly when `{op: F, args}` is —
    /// which is also what it now MEANS. The arity rule is therefore derived
    /// from `arity_of(F)`, not hard-coded.
    #[test]
    fn broadcast_fn_arity_mirrors_the_bare_node() {
        let legal = |f: &str, n: usize| {
            let args: Vec<serde_json::Value> =
                (0..n).map(|i| serde_json::json!(format!("x{i}"))).collect();
            check_expr(&node(
                serde_json::json!({"op": "broadcast", "fn": f, "args": args}),
            ))
        };
        // n-ary ops fold to the identity on one operand — `+(x)` is `x` — so a
        // one-operand broadcast of them stays legal.
        legal("+", 1).expect("`+` is n-ary from 1");
        legal("*", 3).expect("`*` is n-ary");
        // `-` is unary OR binary: negation or subtraction, exactly as bare.
        legal("-", 1).expect("unary `-` is negation");
        legal("-", 2).expect("binary `-` is subtraction");
        assert!(matches!(legal("-", 3), Err(OpError::Arity { got: 3, .. })));
        // Strictly binary.
        assert!(matches!(legal("/", 1), Err(OpError::Arity { ref op, got: 1, .. }) if op == "/"));
        assert!(matches!(
            legal("atan2", 3),
            Err(OpError::Arity { got: 3, .. })
        ));
        // `min`/`max` are n-ary with n >= 2 (§4.2 "MUST reject … fewer than
        // two"), so a ONE-operand `broadcast(fn="min")` is an arity error —
        // the same verdict the bare `{"op":"min","args":[x]}` node gets.
        assert!(
            matches!(legal("min", 1), Err(OpError::Arity { ref op, got: 1, .. }) if op == "min")
        );
        legal("min", 2).expect("`min` of two is fine");
        // Unary elementary functions: exactly one. A 2-arg `broadcast(fn="sin")`
        // used to resolve to `BinCode::Unknown` and evaluate to NaN in silence.
        legal("sin", 1).expect("`sin` is unary");
        assert!(
            matches!(legal("sin", 2), Err(OpError::Arity { ref op, got: 2, .. }) if op == "sin")
        );
        legal("neg", 1).expect("`neg` is unary");
        assert!(matches!(
            legal("neg", 2),
            Err(OpError::Arity { got: 2, .. })
        ));
        legal("ifelse", 3).expect("`ifelse` is ternary");
        assert!(matches!(
            legal("ifelse", 2),
            Err(OpError::Arity { got: 2, .. })
        ));
    }

    /// The in-tree authored idiom (`canonicalize.rs`, `tests/display/
    /// structural_ops.json`): a ONE-operand `broadcast` of a unary function.
    #[test]
    fn one_operand_unary_broadcast_stays_legal() {
        check_expr(&node(serde_json::json!({
            "op": "broadcast", "fn": "sin", "args": ["u"]
        })))
        .expect("`broadcast(fn=\"sin\", [u])` is the authored unary spelling");
        check_expr(&node(serde_json::json!({
            "op": "broadcast", "fn": "exp", "args": ["A"]
        })))
        .expect("`broadcast(fn=\"exp\", [A])` is pinned by tests/display/structural_ops.json");
    }

    /// [`is_elementwise_node`] sees through a `broadcast` to its `args`, because
    /// `broadcast(fn: F, args)` MEANS `{op: F, args}` — so the two spellings of
    /// one array-level expression align their operands identically (issue #100).
    /// The `fn` guard keeps that conservative: a `broadcast` whose `fn` is not a
    /// scalar operator has no elementwise meaning to preserve.
    #[test]
    fn is_elementwise_node_sees_through_broadcast() {
        let bcast = |f: Option<&str>| ExpressionNode {
            op: "broadcast".to_string(),
            broadcast_fn: f.map(str::to_string),
            args: vec![Expr::Variable("a".into()), Expr::Variable("b".into())],
            ..Default::default()
        };
        assert!(is_elementwise_node(&bcast(Some("*"))));
        assert!(is_elementwise_node(&bcast(Some("atan2"))));
        // An absent `fn` keeps the historical `+` classification; the node
        // itself is now rejected by `check_broadcast_fn` (issue #101).
        assert!(is_elementwise_node(&bcast(None)));
        // Not a scalar operator: stay out.
        assert!(!is_elementwise_node(&bcast(Some("pow"))));
        assert!(!is_elementwise_node(&bcast(Some("faq"))));

        // The NODE question and the NAME question differ on exactly one input:
        // a `broadcast` node is elementwise in its args, but the string
        // "broadcast" is not a legal kernel name.
        assert!(!is_scalar_operator("broadcast"));

        // Non-`broadcast` nodes are classified by name alone, as before.
        let named = |op: &str| ExpressionNode {
            op: op.to_string(),
            args: vec![Expr::Variable("a".into())],
            ..Default::default()
        };
        assert!(is_elementwise_node(&named("*")));
        assert!(!is_elementwise_node(&named("index")));
        assert!(!is_elementwise_node(&named("intersect_polygon")));
    }
}