Skip to main content

earthsci_ast/simulate_array/
mod.rs

1//! Native array runtime for `faq`, `makearray`, `index`, `reshape`,
2//! `transpose`, `concat`, and `broadcast` expression nodes (gt-oxr).
3//!
4//! This is the crate's one evaluator: [`crate::problem::esm_problem`] builds it
5//! for every document, scalar or gridded — its tape under `native` / `xla` and
6//! its per-cell oracle under `interpreter` (API_SPEC §5.8). The solver
7//! plumbing and run vocabulary around it live in [`crate::simulate`].
8//!
9//! ## Approach
10//!
11//! The flat state vector that diffsol consumes is a contiguous
12//! concatenation of per-variable blocks. Each array variable occupies a
13//! column-major-ordered block sized by its inferred shape; scalar
14//! variables occupy a single slot. Shape inference walks every `index`
15//! call and every `faq` `ranges` dict to compute per-variable, per-
16//! dimension bounds.
17//!
18//! At RHS evaluation time the interpreter wraps the flat state slice into
19//! [`ndarray::ArrayD`] views (one per variable), binds `faq` loop
20//! indices into a context, and evaluates each equation's body expression
21//! into a [`Value`] — either `Scalar(f64)` or `Array(ArrayD<f64>)`. For
22//! array-producing operators (`reshape`, `transpose`, `concat`,
23//! `broadcast`, `makearray`) the whole array is materialised as an
24//! intermediate so downstream `index` extractions can select any element.
25//!
26//! Column-major ordering is the convention used by the Julia sibling and
27//! reflected in the cross-language conformance fixtures (e.g.
28//! `faq_11_reshape_roundtrip.esm`).
29//!
30//! ## Subsystems
31//!
32//! The runtime is split along its subsystem seams:
33//!
34//! * [`compile`] — the build path: array-op/spatial file detection, subsystem
35//!   mounting (esm-spec §4.6), ragged keyed-factor scoping, the staged
36//!   `from_model` lowering, and shape-inference / LHS-parsing helpers.
37//! * `eval` — the per-cell oracle interpreter (the correctness reference),
38//!   including the geometry leaf ops and standalone [`eval_expression`].
39//! * [`vectorized`] — the vectorized whole-array stencil overlay (ess-bdm):
40//!   shifted-slice `index` gathers, broadcast kernels, einsum folds.
41//! * [`rhs`] — per-call RHS evaluation: the zero-allocation scratch/buffer
42//!   pool (ess-mro), observed-rule materialization, and the rule driver.
43//! * [`driver`] — diffsol solver plumbing (`solve`/`solve_inspect`,
44//!   `debug_*` entry points) and the external refreshable forcing channel
45//!   (PR-1, ess-14f.7) exposed via `forcing_handle`.
46//! * [`layout`] — column-major flat↔multi index/array conversion helpers.
47//!
48//! This file keeps the shared data model: [`Value`], the per-variable
49//! [`VarShape`] table, the compiled rule types ([`RhsRule`],
50//! [`AlgebraicRule`], [`ContractDim`]), the compiled model [`ArrayCompiled`],
51//! the kernel buffer [`Pool`], and the evaluation context [`EvalCtx`].
52
53// This runtime is compiled for wasm too (EarthSciAST-akz): it reaches
54// s2geometry only through the already-wasm-safe `crate::geometry` API (planar
55// clips work; spherical/geodesic returns a runtime `GeometryError` stub on
56// wasm), and its solver is the diffsol/Faer path, which is pure Rust — so no
57// native-only dependency remains, and planar / geometry-free PDEs run in the
58// browser.
59#![allow(
60    clippy::type_complexity,
61    clippy::collapsible_if,
62    clippy::needless_range_loop
63)]
64
65mod compile;
66mod cse;
67mod driver;
68mod eval;
69mod layout;
70pub(crate) mod range;
71mod rhs;
72pub mod tape;
73mod vectorized;
74
75// `inline_tests` consumes `eval_buildtime_field` on every target; only the
76// native-only `area_faq` needs the scoped form, so it is gated to avoid an
77// unused-import warning on wasm.
78pub(crate) use compile::eval_buildtime_field;
79#[cfg(not(target_arch = "wasm32"))]
80pub(crate) use compile::eval_buildtime_field_in_scope;
81pub use compile::{
82    file_has_array_ops, file_has_spatial_model, run_value_invention,
83    run_value_invention_with_params,
84};
85// The ONE free-variable gate (CONFORMANCE_SPEC §5.23), shared with the build
86// pipeline: `crate::prepare` runs the same check the compile path runs, so the
87// two routes cannot disagree about which names a document declares.
88#[cfg(not(target_arch = "wasm32"))]
89pub(crate) use compile::check_free_variables;
90pub(crate) use compile::data_feed_source;
91#[cfg(not(target_arch = "wasm32"))]
92pub(crate) use compile::inline_parameter_arrays;
93pub(crate) use compile::{model_tree_any, parse_subsystem_model};
94// The build pipeline reads the authored model, as the single-model route does,
95// so it applies the same stand-ins for `flatten`'s rewrites.
96#[cfg(not(target_arch = "wasm32"))]
97pub(crate) use compile::{apply_flatten_rewrites, mount_subsystems, resolve_model_self_references};
98#[cfg(not(target_arch = "wasm32"))]
99pub(crate) use eval::eval_observed_taped;
100#[cfg(test)]
101pub(crate) use eval::per_cell_cells;
102pub(crate) use eval::{
103    ONE_CELL_NOTE, StopAtFirstCell, eval_observed_recurrence, eval_scalar_expression,
104    per_cell_walk_refused, per_cell_walks,
105};
106#[cfg(not(target_arch = "wasm32"))]
107pub(crate) use tape::StandaloneError;
108#[cfg(not(target_arch = "wasm32"))]
109pub(crate) use vectorized::take_overlay_decline;
110// Read only by `crate::expression`'s tests.
111pub(crate) use eval::GEOMETRY_CLIP_CODE;
112#[cfg(test)]
113pub(crate) use eval::check_scalar_evaluable;
114pub use eval::{
115    eval_expression, eval_expression_with_extents, eval_expression_with_extents_and_consts,
116    take_const_array_oob,
117};
118// The scalar-op leaf kernel is defined once here (backs the per-cell oracle and
119// the vectorized overlay); re-exported crate-wide so a caller outside the
120// runtime that needs one operator's numeric meaning (the expression
121// simplifier) routes through the SAME definition.
122pub(crate) use eval::{apply_binary, eval_expression_with_extents_and_consts_shared};
123pub use rhs::RhsScratch;
124
125use compile::*;
126use cse::*;
127use eval::*;
128use layout::*;
129use rhs::*;
130use vectorized::*;
131
132use crate::faq::{ReduceKind, empty_derived_extents};
133use crate::types::{Expr, ExpressionNode, IndexSet, RangeSpec};
134use crate::value_invention::BoundaryKind;
135use indexmap::IndexMap;
136use ndarray::{ArrayD, IxDyn};
137use rustc_hash::FxBuildHasher;
138use smallvec::SmallVec;
139use std::cell::RefCell;
140use std::collections::{HashMap, HashSet};
141use std::rc::Rc;
142
143/// Stack-inlined index vectors for per-axis kernel bookkeeping. Grid rank stays
144/// ≤ 4 in practice, so these never touch the heap — a precondition for the
145/// zero-allocation steady-state RHS (ess-mro).
146type DimI = SmallVec<[i64; 4]>;
147type DimU = SmallVec<[usize; 4]>;
148
149/// Fast, fixed-seed loop-index binding map for the hot per-cell interpreter.
150/// `EvalCtx::loop_binds` holds the current output/contraction index values and
151/// is probed FIRST on every variable reference during the per-cell tree walk
152/// ([`lookup_variable`]), so it is the highest-frequency hash in the evaluator —
153/// a coupled level-set solve spent ~19% of its samples in std's SipHash. The
154/// keys are trusted internal index names (`i`, `j`, `k`, …), so SipHash's DoS
155/// resistance is unused cost; `FxBuildHasher` is a fixed seed (so *more*
156/// deterministic than the randomized `RandomState` it replaces). `loop_binds` is
157/// only ever built inline at each `EvalCtx` construction and probed/inserted/
158/// removed by key — never iterated in an order-affecting way — so switching the
159/// hasher is byte-identical. (The borrowed `state_arrays`/`observed_arrays` name
160/// maps still use std `HashMap`; swapping their hasher too is a natural
161/// follow-up but crosses the public `eval_expression`/`forcing_handle` boundary.)
162type IdxMap = HashMap<String, i64, FxBuildHasher>;
163
164/// Fast, fixed-seed name→array map for the hot per-cell interpreter's variable
165/// resolution ([`lookup_variable`] probes `state_arrays` then `observed_arrays`
166/// on every non-index reference). Same rationale as [`IdxMap`]: trusted internal
167/// variable names, no order-affecting iteration, byte-identical results. The
168/// public boundary ([`eval_expression`]'s `inputs`, `forcing_handle`,
169/// `BuildInspection::setup_arrays`) stays on std `HashMap`; conversion happens
170/// there (a shallow clone of the small FAQ/coordinate input maps).
171pub(crate) type ArrMap = HashMap<String, ArrayD<f64>, FxBuildHasher>;
172
173/// Stack-inlined operand buffer for the per-node scalar evaluator. An operator's
174/// arity is ≤ 4 in practice (most are binary), so evaluating `node.args` into
175/// this never touches the heap — removing the per-node `Vec<Value>` temporary
176/// that dominated allocation in the per-cell interpreter profile.
177type ValVec = SmallVec<[Value; 4]>;
178
179// ============================================================================
180// Value type: scalar or dynamic-rank ndarray.
181// ============================================================================
182
183/// A runtime value carried through the array-aware interpreter.
184///
185/// Scalars and whole arrays are first-class so operators like `reshape`,
186/// `transpose`, `concat`, and `broadcast` can produce array-typed
187/// intermediates that later `index` calls sample from.
188#[derive(Debug, Clone)]
189pub enum Value {
190    Scalar(f64),
191    /// Boxed so the whole-array payload lives on the heap and `Value` stays
192    /// ~16 bytes: the per-cell oracle moves a `Value` at every AST node of
193    /// every grid cell, and the vast majority are `Scalar`s — an unboxed
194    /// `ArrayD` (~64 B) would make every one of those scalar moves a large
195    /// memcpy (knot #6). Constructing an `Array` now heap-allocates, but the
196    /// hot zero-allocation RHS paths (scalar ODE, vectorized whole-array) do
197    /// not construct per-cell `Array` values, so they are unaffected.
198    Array(Box<ArrayD<f64>>),
199}
200
201impl Value {
202    fn as_scalar(&self) -> Option<f64> {
203        match self {
204            Value::Scalar(v) => Some(*v),
205            Value::Array(a) if a.ndim() == 0 => Some(a[IxDyn(&[])]),
206            _ => None,
207        }
208    }
209}
210
211// ============================================================================
212// Array model: shape information per variable + compiled RHS rules.
213// ============================================================================
214
215/// Per-variable shape/origin description.
216#[derive(Debug, Clone)]
217pub struct VarShape {
218    /// Dimension extents. Empty vec means scalar.
219    pub shape: Vec<usize>,
220    /// Per-dimension origin (1-based indices per schema convention).
221    pub origin: Vec<i64>,
222    /// Flat offset in the state vector.
223    pub flat_offset: usize,
224}
225
226/// One state variable's declared default (esm-spec §6.3): one value
227/// broadcast over every cell (or none declared), or inline array data.
228#[derive(Debug, Clone)]
229pub(crate) enum StateDefault {
230    /// The same value in every cell, or `None` when the variable declares no
231    /// default.
232    Scalar(Option<f64>),
233    /// One value per cell, ROW-major over the variable's shape (the authored
234    /// nesting's order), which the slots' column-major order gathers from.
235    Field(Vec<f64>),
236}
237
238/// One contracted (reduction) index's loop bound in a `faq`
239/// einsum. Either a static inclusive interval, or a **ragged** bound whose
240/// upper limit `offsets[of…]` is gathered per output tuple at eval time
241/// (RFC `semiring-faq-unified-ir` §5.2 — variable-valence / unstructured-mesh
242/// reductions). The lower bound of a ragged dim is implicitly `1`.
243#[derive(Debug, Clone)]
244enum ContractDim {
245    /// Static inclusive `[lo, hi]` (interval / categorical index sets).
246    Static(i64, i64),
247    /// Ragged `[1, offsets[of…]]` — `offsets` names the per-parent length
248    /// factor; `of` names the parent index variables that address it.
249    Ragged { offsets: String, of: Vec<String> },
250    /// Derived `[1, n]` — `from_faq` names the FAQ producer node whose
251    /// materialized output sizes this contraction, resolved at eval time by
252    /// [`derived_extent`]: the value-invention engine's distinct-set cardinality
253    /// (RFC §6.1 / §5.5) if it materialized this producer, else the
254    /// `intersect_polygon` clip ring's distinct-vertex count (RFC §8.1).
255    Derived { from_faq: String },
256}
257
258impl ContractDim {
259    /// Build a contracted dim from a resolved range spec: a [`RangeSpec::RaggedDyn`]
260    /// becomes [`ContractDim::Ragged`], a [`RangeSpec::DerivedDyn`] becomes
261    /// [`ContractDim::Derived`]; anything else falls back to its static
262    /// `[lo, hi]` bounds (`[0, 0]` — an empty reduction — if unresolved).
263    fn from_range(spec: &RangeSpec) -> Self {
264        if let Some((offsets, of)) = spec.ragged() {
265            return ContractDim::Ragged {
266                offsets: offsets.to_string(),
267                of: of.to_vec(),
268            };
269        }
270        if let Some(from_faq) = spec.derived() {
271            return ContractDim::Derived {
272                from_faq: from_faq.to_string(),
273            };
274        }
275        let r = spec.bounds().unwrap_or([0, 0]);
276        ContractDim::Static(r[0], r[1])
277    }
278
279    /// Resolve to a concrete inclusive `(lo, hi)` range under the current loop
280    /// binds. A ragged dim gathers its parent index value(s) from `ctx` and
281    /// reads `offsets[parent…]`; a derived dim reads the materialized ring's
282    /// vertex count from the runtime registry; an empty bound (`lo > hi`, e.g.
283    /// an isolated cell with zero neighbours, or a disjoint clip) yields no
284    /// contraction tuples, so the reduction returns the additive identity 0̄.
285    fn concrete(&self, ctx: &EvalCtx) -> (i64, i64) {
286        match self {
287            ContractDim::Static(lo, hi) => (*lo, *hi),
288            ContractDim::Ragged { offsets, of } => (1, ragged_upper_bound(offsets, of, ctx)),
289            ContractDim::Derived { from_faq } => (1, derived_extent(from_faq, ctx)),
290        }
291    }
292
293    /// The cell-independent `(lo, hi)` of a static dim, or `None` for a
294    /// ragged/derived dim (whose bound varies per output tuple and so must be
295    /// resolved via [`Self::concrete`] each cell). Lets the caller hoist the
296    /// range derivation out of the per-cell reduction loop when all dims are
297    /// static.
298    fn static_bound(&self) -> Option<(i64, i64)> {
299        match self {
300            ContractDim::Static(lo, hi) => Some((*lo, *hi)),
301            ContractDim::Ragged { .. } | ContractDim::Derived { .. } => None,
302        }
303    }
304}
305
306/// An equation rule compiled for runtime RHS evaluation.
307#[derive(Debug, Clone)]
308enum RhsRule {
309    /// Scalar derivative `D(var) = body` — `var` is a 0-D state variable.
310    Scalar { slot: usize, body: Box<Expr> },
311    /// Indexed scalar derivative `D(var[i1, i2, ...]) = body` with all
312    /// indices concrete. Writes to a single flat slot.
313    IndexedScalar { slot: usize, body: Box<Expr> },
314    /// Array-op derivative. The body expression is evaluated once per tuple
315    /// of `output_idx` values (the tuple drawn from `output_ranges`) and the
316    /// resulting scalar is written into `var_name[idx...]`.
317    /// If `contract_names` is non-empty the body also contains contracted
318    /// (reduction) indices that are unrolled at eval time and combined via the
319    /// semiring's ⊕ (`reduce`), resolved once at build time.
320    ArrayLoop {
321        var_name: String,
322        output_idx_names: Vec<String>,
323        output_ranges: Vec<(i64, i64)>,
324        lhs_idx_exprs: Vec<Expr>,
325        body: Box<Expr>,
326        contract_names: Vec<String>,
327        /// Per-contracted-index loop bounds. A [`ContractDim::Ragged`] dim is
328        /// expanded to its dynamic `[1, offsets[of…]]` extent per output tuple.
329        contract_dims: Vec<ContractDim>,
330        reduce: ReduceKind,
331        /// Optional `filter` predicate (§5.3): combinations for which it is
332        /// false contribute the additive identity 0̄. `None` ⇒ no gating.
333        filter: Option<Box<Expr>>,
334    },
335}
336
337/// Instrumentation for one `evaluate_rhs` call: how the spatial array-op
338/// derivatives were evaluated. The load-bearing field for the
339/// no-scalarization contract (ess-bdm) is [`RhsStats::kernel_ops`] — the
340/// number of AST-node evaluations performed by the **vectorized** whole-array
341/// path. It is a function of the discretized RHS *expression* only, so it is
342/// **independent of the grid size N**: the same stencil evaluated on a 4-cell
343/// and an 8-cell grid visits the same number of array kernels (one shifted
344/// slice per neighbour, one broadcast per arithmetic node — not `O(N)` scalar
345/// sub-expressions). The per-cell oracle, by contrast, re-walks the body once
346/// per grid cell, so its `kernel_ops` scales with N. A test asserts the
347/// vectorized count is N-independent (criterion 3 of ess-bdm).
348#[derive(Debug, Clone, Default)]
349pub struct RhsStats {
350    /// AST-node evaluations performed by the vectorized (whole-array) path.
351    /// N-independent for a fixed discretized RHS.
352    pub kernel_ops: usize,
353    /// Number of array-op derivative rules evaluated vectorized (shifted-slice
354    /// stencils + region-materialized boundary makearrays, no per-cell loop).
355    pub vectorized_rules: usize,
356    /// Number of array-op derivative rules that fell back to the per-cell
357    /// oracle (general semiring contraction, non-affine indexing, periodic
358    /// wrap, …). The vectorized path is a verified-equivalent overlay; the
359    /// per-cell path remains the correctness reference.
360    pub scalar_rules: usize,
361    /// Number of array-shaped *observed* (algebraic) rules materialized via the
362    /// vectorized whole-array overlay rather than the per-cell oracle. For a
363    /// coupled model with a time/space-varying behaviour stack this is the
364    /// dominant per-step cost, so keeping it vectorized (not `obs_scalar_rules`)
365    /// is what makes the observed materialization N-independent.
366    pub obs_vectorized_rules: usize,
367    /// Number of array-shaped observed rules that fell back to the per-cell
368    /// oracle (a body op the vectorizer does not yet cover, a non-unit origin,
369    /// or a forced-scalar reference run).
370    pub obs_scalar_rules: usize,
371    /// Step 3b: number of rules (observed + RHS) executed through the compiled
372    /// tape's fast executor on this call: every rule of the model, since a
373    /// tape exists only when every rule lowered. Zero on the legacy
374    /// interpreter path.
375    pub taped_rules: usize,
376}
377
378/// Eliminated algebraic-variable definition. Evaluated once per RHS call
379/// into a transient ndarray (or scalar) that the `observed_values` map
380/// exposes to downstream expressions.
381///
382/// The body is an `Arc<Expr>`, not a `Box<Expr>`: the driver derives several
383/// cadence-filtered rule lists from `observed_rules` (CONST / DISCRETE /
384/// CONTINUOUS subsets, the per-closure RHS/Jacobian captures, the
385/// output-pass dependency cones), and the expanded discretization bodies are
386/// by far the model's largest allocation (~1 GiB for `simpleclimate.esm` at
387/// the production grid). Cloning a rule now shares the immutable body
388/// instead of deep-copying it, collapsing what used to be five simultaneous
389/// full copies at solve time into one. Sharing also keeps every `Expr` node
390/// address stable for the whole run, which the address-keyed CSE overlay
391/// ([`CseRt`] / `AddrClasses`) relies on; bodies are never mutated after
392/// compile, so no copy-on-write can move them mid-solve.
393#[derive(Debug, Clone)]
394enum AlgebraicRule {
395    /// `var := body` — pure scalar algebraic.
396    ///
397    /// `declared_shape` is the defined variable's declared extents, when it
398    /// declares a shape that resolves to static sizes. A body that evaluates
399    /// to a scalar is broadcast over it (esm-spec §4.3.4), so the variable
400    /// keeps its declared rank however the value came out.
401    Scalar {
402        var: String,
403        body: std::sync::Arc<Expr>,
404        declared_shape: Option<Vec<usize>>,
405    },
406    /// `var[i...] := body` — array algebraic defined via a faq over
407    /// the full shape of `var`.
408    ArrayLoop {
409        var: String,
410        output_idx_names: Vec<String>,
411        output_ranges: Vec<(i64, i64)>,
412        body: std::sync::Arc<Expr>,
413    },
414    /// `var[i…] := body`, where `body` reads `var` itself at a strictly earlier
415    /// position along axis `axis` — a CAUSAL SELF-REFERENCE (esm-spec §4.3.1.1,
416    /// `docs/content/rfcs/causal-self-reference-recurrence.md`).
417    ///
418    /// Structurally an [`AlgebraicRule::ArrayLoop`] plus the recurrence axis,
419    /// but a SEPARATE variant on purpose: its cells are not independent, so it
420    /// must never reach the whole-array overlay or any other path that
421    /// evaluates cells out of order or in a batch (CONFORMANCE_SPEC §5.19.2).
422    /// The interpreter evaluates it with `sweep_recurrence`; the tape with an
423    /// ordered `Sweep` that runs its cell body one cell at a time in the same
424    /// order. Being a distinct variant means every `match` over the rule
425    /// kinds has to say what it does with a recurrence instead of silently
426    /// inheriting a reordering path.
427    Recurrence {
428        var: String,
429        output_idx_names: Vec<String>,
430        output_ranges: Vec<(i64, i64)>,
431        body: std::sync::Arc<Expr>,
432        /// Position within `output_idx_names` of the axis the recurrence folds
433        /// along. Swept ascending as the OUTERMOST loop.
434        axis: usize,
435        /// Largest lag any self-read takes along `axis`, derived from the reads
436        /// themselves and never declared (RFC §2.1). Reported for
437        /// observability; no evaluation rule depends on it.
438        max_lag: i64,
439        /// Whether every self-read was proved strictly earlier statically.
440        /// See [`RecurrenceInfo::lag_proven`].
441        lag_proven: bool,
442    },
443}
444
445/// The partially materialized array of an [`AlgebraicRule::Recurrence`], in
446/// scope for the body's causal self-reads while the sweep is in progress.
447///
448/// Interior-mutable for the same reason [`EvalCtx::derived_rings`] is: the
449/// driver publishes cell `i` while the evaluation context that reads it at cell
450/// `i + c` holds a shared borrow of the scope.
451///
452/// `written` is not redundant with the axis bounds. A read is admitted only when
453/// the cell it names has ALREADY been published, which the bounds alone cannot
454/// decide (a body may compute an index from a contracted symbol). An unwritten
455/// or out-of-range read is fail-closed (CONFORMANCE_SPEC §5.19.4) — never the
456/// §5.5.5 zero ghost, and never a bare NaN that a `max(x, 0)` in the body would
457/// launder back into a plausible number.
458pub(super) struct RecurScope<'a> {
459    /// The variable being defined; an `index` gather on this name resolves here.
460    pub(super) name: &'a str,
461    /// Extents of the cell frame.
462    pub(super) shape: DimU,
463    /// Lower bound of each axis of the cell frame.
464    pub(super) origin: DimI,
465    /// Cell values, column-major over `shape`.
466    cells: RefCell<Vec<f64>>,
467    /// Which cells have been published.
468    written: RefCell<Vec<bool>>,
469}
470
471impl<'a> RecurScope<'a> {
472    pub(super) fn new(name: &'a str, shape: DimU, origin: DimI) -> Self {
473        let total = shape.iter().copied().product::<usize>().max(1);
474        Self {
475            name,
476            shape,
477            origin,
478            cells: RefCell::new(vec![0.0; total]),
479            written: RefCell::new(vec![false; total]),
480        }
481    }
482
483    /// Publish one cell (its 1-based multi-index) so later cells may read it.
484    ///
485    /// Rounded to the ACTIVE working precision (esm-spec §11.3.1) before it is
486    /// stored, not only when it is read back. The recurrence's carried state is
487    /// a cell of the variable being defined, so it carries that variable's
488    /// `element_type` — and for a `Float32` document that has to hold at every
489    /// step of the fold, exactly as it does in a `real*4` reference
490    /// implementation. Storing a binary64 partial and rounding only on the way
491    /// out would run the fold at a precision the document did not declare and
492    /// silently beat the reference it is being checked against.
493    pub(super) fn publish(&self, tuple: &[i64], v: f64) {
494        let flat = layout::multi_to_flat_col_major(tuple, &self.shape, &self.origin);
495        let mut cells = self.cells.borrow_mut();
496        if flat < cells.len() {
497            cells[flat] = crate::precision::active().round(v);
498            self.written.borrow_mut()[flat] = true;
499        }
500    }
501
502    /// Read one published cell, or `None` when the position is outside the cell
503    /// frame or has not been published yet.
504    pub(super) fn read(&self, raw: &[i64]) -> Option<f64> {
505        if raw.len() != self.shape.len() {
506            return None;
507        }
508        for (d, &one_based) in raw.iter().enumerate() {
509            let lo = self.origin[d];
510            let hi = lo + self.shape[d] as i64 - 1;
511            if one_based < lo || one_based > hi {
512                return None;
513            }
514        }
515        let flat = layout::multi_to_flat_col_major(raw, &self.shape, &self.origin);
516        if !*self.written.borrow().get(flat)? {
517            return None;
518        }
519        self.cells.borrow().get(flat).copied()
520    }
521
522    /// The finished array, as an owned ndarray in the evaluator's layout.
523    pub(super) fn finish(&self) -> ArrayD<f64> {
524        layout::col_major_to_arrayd(&self.cells.borrow(), &self.shape)
525    }
526}
527
528/// Build/run observability record — the Rust mirror of the Julia binding's
529/// `BuildInspection` (`build_evaluator(…; inspect=BuildInspection())`). Pass
530/// one to a EsmProblem built with [`crate::problem::ProblemOptions::inspect`] (or
531/// [`ArrayCompiled::solve_inspect`]) and the run fills it with named
532/// build-time products that are otherwise internal to the runtime:
533///
534/// * `setup_arrays` — the STATE-FREE observed arrays, materialized once at the
535///   initial state through the official observed machinery
536///   (the observed-materialization pass): the per-pair regrid geometry (`A_ij`, its
537///   row-sums `A_j`, the normalized weights `W_ij` — RFC §8.1 / esm-spec
538///   §8.6.1), const-op mesh factors and their bare-name aliases, and every
539///   other build-once array observed whose transitive references reach no
540///   state, no `t`, and no external forcing channel. This is the official
541///   inspection surface for conformance runners that gate per-pair regridding
542///   values (CONFORMANCE_SPEC §5.8) and for §6.6.5 assertions that target a
543///   state-free array observed directly (the MPAS `div_flux` max/min). A
544///   state- or time-dependent observed is deliberately ABSENT (its build-time
545///   snapshot would be wrong at any later time), so a reader errors rather
546///   than consuming a stale field.
547/// * `observed_exprs` — every observed rule's resolved body expression (post
548///   subsystem mounting, range resolution, and keyed-factor scoping), exactly
549///   as the runtime evaluates it.
550///
551/// Filling the record never changes the run: the returned
552/// [`crate::simulate::Solution`] is identical with or without a sink (the
553/// capture is one extra read-only observed-materialization pass at `t0`).
554#[derive(Debug, Clone, Default)]
555pub struct BuildInspection {
556    /// State-free observed arrays materialized at the initial state, keyed by
557    /// (possibly flattening-namespaced) observed name.
558    pub setup_arrays: HashMap<String, ArrayD<f64>>,
559    /// Resolved observed body expressions, keyed like `setup_arrays`.
560    pub observed_exprs: HashMap<String, Expr>,
561    /// Resolved SCALAR parameter values (model defaults with any overrides
562    /// applied), keyed by (flattened) parameter name. Load-time constants, so a
563    /// build-time cellwise evaluation (a §6.6.5 analytic `reference`,
564    /// coordinate-expression `ic` seeding) may bind them into scope — STATE
565    /// stays out. The observed-assertion form already binds them (a state-free
566    /// observed is materialized into `setup_arrays` with these values); `params`
567    /// exposes the same map for the reference / `ic` positions.
568    pub params: HashMap<String, f64>,
569    /// One entry per recognized CAUSAL SELF-REFERENCE (esm-spec §4.3.1.1), so
570    /// the interpretation a document's self-read was given is auditable without
571    /// being authored twice. This is what stands in for a declared `recur`
572    /// annotation (RFC §2.1): the recurrence, its axis and its derived maximum
573    /// lag are reported, not required.
574    pub recurrences: Vec<RecurrenceInfo>,
575}
576
577/// A recognized causal self-reference, as reported to a [`BuildInspection`].
578#[derive(Debug, Clone, PartialEq, Eq)]
579pub struct RecurrenceInfo {
580    /// The variable the recurrence defines.
581    pub var: String,
582    /// The output index symbol the sweep folds along, ascending.
583    pub axis: String,
584    /// Largest lag any self-read takes along `axis`, DERIVED from the reads.
585    pub max_lag: i64,
586    /// Whether every self-read was proved strictly earlier statically. When
587    /// `false` at least one read's lag depends on a symbol whose range admits
588    /// `0` or less, so the runtime's fail-closed read (never a value for an
589    /// unpublished cell) is what rules those cells out — typically because a
590    /// guard in the body means they are never evaluated.
591    pub lag_proven: bool,
592}
593
594/// Which compiler this compiled model is serving (API_SPEC §5.8).
595///
596/// The array runtime has three evaluation tiers under it — the tape, the
597/// whole-array overlay and the per-cell oracle — and the historical routing
598/// walks down them silently. A named compiler fixes which of them may run, so
599/// that what ran is a property of the NAME and not of the input.
600#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
601pub(crate) enum RuntimeMode {
602    /// The historical routing: the tape where it lowers, the overlay and the
603    /// per-cell oracle beneath it. Every entry point that is not
604    /// [`crate::problem::esm_problem`] — `compile_array`, the `debug_*` seams,
605    /// the adapters — builds this.
606    #[default]
607    Legacy,
608    /// [`crate::Compiler::Native`]: the tape and nothing under it. Every rule
609    /// is taped (construction refused otherwise), so the build-time and
610    /// output-time observed passes are served from the tape rather than from
611    /// the overlay.
612    Native,
613    /// [`crate::Compiler::Interpreter`]: no tape and no overlay — the per-cell
614    /// oracle everywhere, as the reference the other compilers are checked
615    /// against.
616    Interpreter,
617    /// [`crate::Compiler::Xla`]: the specialty compiler that needs a heavy
618    /// external dependency. The tape is lowered to a StableHLO computation by
619    /// `tape::xla_emit` and run through PJRT (`crate::xla_runtime::CompiledRhs`),
620    /// so the right-hand side — and the
621    /// finite-difference Jacobian built out of it — are the compiled
622    /// executable, not the tape's own interpreter.
623    ///
624    /// It is as STRICT as [`Self::Native`] and strict twice over: every rule
625    /// must lower to the tape (or construction refuses it), and the whole tape
626    /// must lower to XLA (or construction refuses that, naming the rule). The
627    /// passes the tape serves under `native` — the build-time materialization,
628    /// the per-segment seed, the inspection snapshot and the observeds
629    /// reported at output times — are served from the tape here too, because
630    /// the emitted program's only output is `du`. That is one evaluator, not
631    /// two: the same tape both feeds the emitter and answers those passes.
632    Xla,
633}
634
635/// Compiled, parameter-sweep-ready ODE model for array-op models.
636pub struct ArrayCompiled {
637    /// Which compiler this model serves. See [`RuntimeMode`]; set by
638    /// [`crate::problem::esm_problem`] right after the build, and
639    /// [`RuntimeMode::Legacy`] for every other entry point.
640    pub(crate) runtime_mode: RuntimeMode,
641    /// The XLA executable this model's right-hand side runs on under
642    /// [`RuntimeMode::Xla`], installed by
643    /// [`crate::problem::esm_problem`] at CONSTRUCTION — so a model the
644    /// emitter cannot lower is a build refusal rather than a surprise on the
645    /// first step — and empty under every other mode.
646    ///
647    /// A cell rather than a plain field because the program is emitted from a
648    /// finished [`ArrayCompiled`]: the emitter reads the model it is going to
649    /// be installed on. `Rc` so the per-segment RHS and Jacobian closures can
650    /// each hold the ONE executable; compiling it is the expensive step and
651    /// there must never be a second.
652    #[cfg(feature = "xla")]
653    pub(crate) xla_rhs: std::cell::OnceCell<Rc<crate::xla_runtime::CompiledRhs>>,
654    /// Every state spelling an `operator_compose` renaming match DELETED,
655    /// mapped onto the survivor (issue #230). Carried from
656    /// `FlattenMetadata::merged_variable_renames` by
657    /// [`ArrayCompiled::from_flattened`] so an override key naming a state the
658    /// merge moved resolves instead of silently designating nothing. Empty on
659    /// the `from_model` path, which has no coupling to have renamed anything.
660    merged_renames: HashMap<String, String>,
661    /// Every state variable's slot range — its base slot in the flat state
662    /// vector and its extents — in slot order. This is the whole state
663    /// layout: a slot's name (`"u[2,3]"`, `"s"`) is computed from it
664    /// (`layout::slot_name`, `layout::lookup_slot`) rather than stored.
665    var_shapes: IndexMap<String, VarShape>,
666    /// Every slot's name, parallel to the flat state vector, built the first
667    /// time [`Self::state_variable_names`] is asked. Nothing inside the build
668    /// or a solve's setup reads it; it is the public surface's spelling.
669    state_names: std::cell::OnceCell<Vec<String>>,
670    /// The same names qualified by `namespace`, built on first ask
671    /// ([`Self::qualified_state_names`]).
672    qualified_state_names: std::cell::OnceCell<Vec<String>>,
673    /// Each state variable's default, parallel to `var_shapes`.
674    state_defaults: Vec<StateDefault>,
675    param_names: Vec<String>,
676    param_index: HashMap<String, usize>,
677    param_defaults: Vec<Option<f64>>,
678    /// Algebraic variables eliminated from the state vector. Stored as
679    /// observed definitions evaluated at each RHS call in order (no cross-
680    /// dependency support for v1 — fixtures don't need it).
681    observed_rules: Vec<AlgebraicRule>,
682    /// Per-state RHS rules.
683    rhs_rules: Vec<RhsRule>,
684    /// Identity of `rhs_rules` + `observed_rules` for the scratch tables keyed
685    /// by their AST node addresses ([`RhsCall::rule_set`]).
686    rule_set: u64,
687    /// Number of flat state slots.
688    n_states: usize,
689    /// External refreshable forcing-array channel (PR-1, ess-14f.7): the live
690    /// runtime input a discrete-cadence loader's regridded field lands in, read
691    /// by the RHS each step. Keyed by variable name; a forcing-fed variable
692    /// resolves here (see [`lookup_variable`]) when it is bound in no other
693    /// channel. `Rc<RefCell<…>>` for the same reason [`RhsScratch`] is a
694    /// `RefCell` — diffsol's RHS is `Fn`, so the buffer needs interior
695    /// mutability — *plus* an `Rc` so a segmented driver (the future R-3
696    /// example-harness) can hold a clone and refresh entries *between* segments
697    /// while the captured closure reads the same buffer. Empty for every model
698    /// with no loader forcing, so the scalar-`p` path is byte-identical.
699    ///
700    /// Feasibility-gate verdict (the bead's declarative-or-fail opener): no
701    /// existing runtime channel suffices for a *refreshable external
702    /// forcing-array*. diffsol's `p` slice (`p.as_slice()`) is scalar-typed and
703    /// shape-less — fine for scalar forcings (which keep going through
704    /// `p`/`set_params`), awkward for fields; `state_arrays` is the integrator
705    /// state `y` (refilled from the solver, not a free input); `observed_arrays`
706    /// is a pure function of state, cleared and recomputed each call;
707    /// `derived_rings` is interior-mutable but built *fresh per RHS call*
708    /// (intra-evaluation FAQ-geometry scratch, wrong lifetime and overwritten by
709    /// `intersect_polygon` producers). The gap is real, so the channel is added
710    /// — as a runtime *binding*, not a new engine primitive (no faq, no
711    /// scalarizer arm, no `Discrete` `VariableType`). The optional typed
712    /// `ModelVariable.refresh` field (plan PR-2) is deferred: forcing resolves
713    /// by name at runtime and does not need it.
714    forcing: Rc<RefCell<HashMap<String, ArrayD<f64>>>>,
715    /// Every DATA-FED parameter routed to [`Self::forcing`] at classification:
716    /// the name the runtime looks up, and the `data_sources` key its `update`
717    /// names (esm-spec §5.4/§8.5). Read by
718    /// [`crate::data_fed::refuse_unbound`] at construction to answer whether
719    /// anything actually bound each one — a forcing nothing bound must be a
720    /// refusal, never a run at the parameter's `default` (esm-spec §9.6.6
721    /// `data_source_unbound`, CONFORMANCE_SPEC §5.46). Empty for every model
722    /// with no data feed.
723    data_fed: Vec<(String, String)>,
724    /// How many times [`Self::forcing_handle`] has handed out the buffer: a
725    /// host holding a handle can have written it since, which is what
726    /// [`driver::FieldIcMemo`] checks.
727    forcing_generation: std::cell::Cell<u64>,
728    /// Deferred scoped-reference / array `ic` equations (esm-spec §11.4.1),
729    /// classified out of the equation list by [`Self::from_model`] (single-model
730    /// path) or carried from [`crate::flatten::FlattenedSystem::field_ics`]
731    /// (coupled path). Each entry is `(target_state, rhs)`: at `u0` build time
732    /// [`Self::simulate`] resolves the target's grid cells and folds the initial
733    /// field — a provider-served loaded field, a broadcast constant, or a
734    /// coordinate expression over grid-geometry aggregates — into the flat state
735    /// vector cell-by-cell (DESIGN pde_simulation_pipeline §2 R2).
736    field_ics: Vec<(String, Expr)>,
737    /// The `faq`-valued initialization equations (esm-spec §6.2), each
738    /// `(target_state, faq)`, in document order. At `u0` build time each
739    /// assigns the cells of its ranges, read against the initial state seeded
740    /// so far, except a cell the caller's `u0` names
741    /// ([`Self::seed_initialization_faqs`]).
742    init_faqs: Vec<(String, Expr)>,
743    /// The STATE-FREE observed definitions a field `ic` RHS reads, in
744    /// dependency order (esm-spec §6.6.5 "Build-time evaluation scope").
745    ///
746    /// A state-free observed — one whose defining expression closes over
747    /// parameters, inline `const` data, provider-served fields and other
748    /// state-free observeds — is resolvable before the simulation runs, so
749    /// §11.4.1 admits it as an `ic` right-hand side alongside a loaded field,
750    /// a constant and a coordinate expression. Captured at compile (before the
751    /// observed bodies are moved into the rules) and materialized lazily at
752    /// `u0` build time, in the resolved parameter scope, by
753    /// [`Self::resolve_field_ics`]. Empty for a document with no field `ic`.
754    ic_scope: compile::IcScope,
755    /// Document-scoped index-set registry, kept so `ic` RHS coordinate
756    /// expressions (whose `faq` ranges may still carry `{ "from": <set> }`
757    /// references on the flattened path) resolve at `u0` build time exactly as
758    /// equation expressions do at compile time.
759    index_sets: HashMap<String, IndexSet>,
760    /// The model's CONST-ARRAY registry (CONFORMANCE_SPEC §5.5.5): the
761    /// pre-computed factor arrays — the document's `const`-op variables plus
762    /// any caller-supplied factor arrays — whose gathers resolve an
763    /// out-of-range index by a declared boundary policy (default
764    /// `E_TREEWALK_CONSTARRAY_OOB`) rather than by the state gather's zero
765    /// ghost. `Rc` so the per-call scratch can hold the same registry without
766    /// cloning the name set every step. Empty for a model with no const
767    /// factors, which reads byte-identically to the pre-§5.5.5 evaluator.
768    const_scope: std::sync::Arc<ConstArrayScope>,
769    /// The single-model namespace (the top-level `models` map key), set by
770    /// [`Self::from_file`]. The raw single-model path keys params/states by their
771    /// BARE variable names (`R_0`, `psi[i,j]`), but the `flatten` path and the
772    /// Julia toolkit namespace them (`Model.R_0`).
773    /// So a caller's `parameters` / `initial_conditions` override key is accepted
774    /// in EITHER form: a `<namespace>.` prefix is stripped before lookup (WS3
775    /// cross-toolkit override-naming parity). `None` on the `from_flattened`
776    /// path, whose names are already fully namespaced.
777    namespace: Option<String>,
778    /// The working precision this model was COMPILED under
779    /// (`crate::precision`), captured so evaluation reproduces it even when the
780    /// caller reaches an `ArrayCompiled` directly (`compile_array`, the
781    /// `debug_*` entry points) rather than through an `EsmProblem`. Build-time
782    /// folding and run-time evaluation must round identically, and the
783    /// artifact carrying its own precision is what makes that true by
784    /// construction rather than by the caller remembering.
785    precision: crate::precision::Env,
786    /// Every name this model DECLARES — states, parameters and observeds —
787    /// captured at compile so the evaluator's fault arm can tell a name the
788    /// document never declared from one it declared but had not produced a
789    /// value for yet. See [`EvalCtx::declared`]; issue #181.
790    declared_names: HashSet<String>,
791    /// The field initial conditions construction resolved
792    /// ([`Self::field_ic_records`]), held for the first solve to take instead
793    /// of resolving them again. See [`driver::FieldIcMemo`].
794    #[cfg(feature = "solve")]
795    field_ic_memo: RefCell<Option<driver::FieldIcMemo>>,
796    /// Each shaped parameter whose declared default the build lowered into a
797    /// `const` observed, with the dense row-major `(shape, values)` it was
798    /// lowered from. The literal's JSON is the per-cell oracle's form; the tape
799    /// reads these numbers instead of parsing it back.
800    inline_param_arrays: HashMap<String, (Vec<usize>, Vec<f64>)>,
801    /// Every parameter refreshed from OUTSIDE the model (a data-source `from`
802    /// binding or a registered handler), with its declared shape resolved
803    /// against the index-set registry: the names the forcing buffer serves
804    /// ([`Self::forcing`]). `None` when the declared shape names a set the
805    /// registry cannot size. The tape compiles a read of one of these into a
806    /// forcing load (`tape::Instr::LoadForcing`) rather than declining it.
807    forcing_decls: IndexMap<String, Option<Vec<usize>>>,
808    /// The declared `default` of each [`Self::forcing_decls`] entry that has
809    /// one, as the dense row-major field it denotes over the resolved shape
810    /// (`None` while that shape does not resolve). Construction serves it from
811    /// the forcing buffer when no provider or caller array supplies the name:
812    /// the default is the value until data arrives (esm-spec §6.3).
813    #[allow(clippy::type_complexity)]
814    forcing_defaults: HashMap<String, Option<(Vec<usize>, Vec<f64>)>>,
815    /// Every SHAPED parameter the build left with no value at all: no
816    /// `default`, no inline data, no `distribution`, and no refresh from
817    /// outside the model. It holds a scalar-table slot with no value, which a
818    /// per-cell read cannot index; [`crate::problem::esm_problem`] refuses the
819    /// build naming each one (esm-spec §10.10).
820    unvalued_shaped_params: Vec<String>,
821    /// The state variables with neither a `D` equation nor an algebraic
822    /// definition, held at their initial value. Not ODE states, so the
823    /// front door's missing-initial-value gate leaves them alone, as the
824    /// other bindings do.
825    held_at_ic: HashSet<String>,
826    /// The tape programs this model has built ([`tape::TapeCache`], which says
827    /// what a kept program depends on).
828    tape_cache: tape::TapeCache,
829}
830
831/// A reuse pool of `f64` backing buffers for vectorized kernel intermediates.
832/// `take_array*`/`give_array` recycle buffers; after a warm-up call the pool
833/// holds enough output-box-sized buffers that no further allocation occurs.
834///
835/// Recycling is **shape-bucketed**: [`Pool::give_array`] parks the whole
836/// [`ArrayD`] (buffer *plus* its already-built `IxDyn` dim and strides) in the
837/// bucket for its exact shape, and [`Pool::take_array_uninit`] pops it straight
838/// back out. In a steady-state RHS the same handful of output-box shapes repeat
839/// on every call, so the common checkout is a hash lookup and a `Vec::pop` —
840/// neither the `memset` that a `Vec::resize`-based refill costs nor the
841/// `IxDyn(shape)` / `from_shape_vec` shape+stride reconstruction that a
842/// buffer-only pool has to redo per checkout (both were measurable shares of
843/// the solve: ~9.5% and ~8.3% respectively on the simpleclimate PPM core).
844///
845/// `free` is the spill list for the cases the buckets cannot serve: a
846/// first-of-its-shape checkout, and a returned array that is not in standard
847/// layout (so its buffer, but not its shape, is reusable).
848#[derive(Default)]
849struct Pool {
850    /// Whole recycled arrays, bucketed by exact shape. Keyed by [`DimU`], which
851    /// hashes and compares as a `[usize]` slice, so the hot lookup borrows the
852    /// caller's `&[usize]` with no key materialization.
853    shaped: HashMap<DimU, Vec<ArrayD<f64>>, FxBuildHasher>,
854    /// Bare buffers with no usable shape attached.
855    free: Vec<Vec<f64>>,
856}
857
858impl Pool {
859    /// Check out a zero-filled buffer of `len` elements, reusing a free buffer
860    /// whose capacity already covers `len` (no reallocation in steady state).
861    fn take(&mut self, len: usize) -> Vec<f64> {
862        if let Some(pos) = self.free.iter().position(|b| b.capacity() >= len) {
863            let mut b = self.free.swap_remove(pos);
864            b.clear();
865            b.resize(len, 0.0);
866            b
867        } else if let Some(mut b) = self.free.pop() {
868            // A free buffer exists but is too small; grow it (warm-up only).
869            b.clear();
870            b.resize(len, 0.0);
871            b
872        } else {
873            vec![0.0; len]
874        }
875    }
876
877    /// Check out an owned `ArrayD` of the given row-major `shape` whose element
878    /// contents are **unspecified** — whatever the previous holder of the
879    /// recycled buffer left there.
880    ///
881    /// For kernels that write every element of the box before anything reads it
882    /// (an elementwise `Zip`, a `fill`, a coordinate ramp, an `assign` of a
883    /// same-shape source). A caller that writes only PART of the box — a
884    /// Dirichlet-ghost gather, a reduction accumulator seeded with the identity,
885    /// a region-tiled `makearray` — must use [`Pool::take_array`] instead, or it
886    /// will read the previous kernel's data.
887    fn take_array_uninit(&mut self, shape: &[usize]) -> ArrayD<f64> {
888        if let Some(bucket) = self.shaped.get_mut(shape) {
889            if let Some(arr) = bucket.pop() {
890                // Recycled whole: right shape, right strides, right length —
891                // nothing to rebuild and nothing to refill.
892                return arr;
893            }
894        }
895        let len = shape.iter().copied().product::<usize>().max(1);
896        let buf = self.take(len);
897        ArrayD::from_shape_vec(IxDyn(shape), buf).expect("pool buffer length matches shape")
898    }
899
900    /// Check out a zero-filled owned `ArrayD` of the given row-major `shape`,
901    /// backed by a pooled buffer.
902    fn take_array(&mut self, shape: &[usize]) -> ArrayD<f64> {
903        let mut arr = self.take_array_uninit(shape);
904        arr.fill(0.0);
905        arr
906    }
907
908    /// Return an owned `ArrayD` to the pool, preserving its buffer capacity and
909    /// — when it is in standard (contiguous, row-major) layout, as every buffer
910    /// this module hands out is and the in-place kernels keep it — its built
911    /// shape and strides too, so the next same-shape checkout rebuilds nothing.
912    fn give_array(&mut self, arr: ArrayD<f64>) {
913        if arr.is_standard_layout() {
914            if let Some(bucket) = self.shaped.get_mut(arr.shape()) {
915                bucket.push(arr);
916                return;
917            }
918            let shape: DimU = arr.shape().iter().copied().collect();
919            self.shaped.insert(shape, vec![arr]);
920            return;
921        }
922        let (buf, _offset) = arr.into_raw_vec_and_offset();
923        self.free.push(buf);
924    }
925}
926
927/// The rule-invariant slice of an evaluation environment: everything an
928/// [`EvalCtx`] carries EXCEPT the observed-array map (whose borrow must be
929/// re-taken after each observed rule writes its output) and the per-loop
930/// index binds. Built once per RHS call / materialization pass from the
931/// natural source objects; the `EvalCtx` literal itself is then written
932/// exactly once — in [`EvalEnv::ctx`] — instead of being hand-assembled at
933/// every evaluation site. Field semantics are documented on [`EvalCtx`].
934#[derive(Clone, Copy)]
935struct EvalEnv<'a> {
936    state_arrays: &'a ArrMap,
937    params: &'a [f64],
938    param_names: &'a [String],
939    t: f64,
940    /// See [`EvalCtx::derived_rings`].
941    derived_rings: &'a RefCell<HashMap<String, ArrayD<f64>>>,
942    /// See [`EvalCtx::derived_extents`]. The compiled-RHS paths pass
943    /// [`empty_derived_extents`] (their derived sets were densified to
944    /// intervals at build time); the standalone expression entry point
945    /// carries real extents.
946    derived_extents: &'a HashMap<String, i64>,
947    /// See [`EvalCtx::forcing`].
948    forcing: &'a RefCell<HashMap<String, ArrayD<f64>>>,
949    /// See [`EvalCtx::cse`].
950    cse: Option<&'a CseRt>,
951    /// See [`EvalCtx::const_lits`].
952    const_lits: Option<&'a ConstLitMemo>,
953    /// See [`EvalCtx::const_arrays`].
954    const_arrays: &'a ConstArrayScope,
955    /// See [`EvalCtx::declared`].
956    declared: &'a HashSet<String>,
957}
958
959impl<'a> EvalEnv<'a> {
960    /// A fresh evaluation context (empty loop binds) over this environment,
961    /// reading observed values from `observed_arrays`.
962    fn ctx(&self, observed_arrays: &'a ArrMap) -> EvalCtx<'a> {
963        EvalCtx {
964            state_arrays: self.state_arrays,
965            observed_arrays,
966            params: self.params,
967            param_names: self.param_names,
968            loop_binds: IdxMap::default(),
969            t: self.t,
970            derived_rings: self.derived_rings,
971            derived_extents: self.derived_extents,
972            forcing: self.forcing,
973            cse: self.cse,
974            const_lits: self.const_lits,
975            const_arrays: self.const_arrays,
976            recur: None,
977            declared: self.declared,
978        }
979    }
980
981    /// Like [`Self::ctx`] but with no CSE memo — the per-cell oracle loops
982    /// evaluate without one.
983    fn oracle_ctx(&self, observed_arrays: &'a ArrMap) -> EvalCtx<'a> {
984        EvalCtx {
985            cse: None,
986            ..self.ctx(observed_arrays)
987        }
988    }
989}
990
991struct EvalCtx<'a> {
992    state_arrays: &'a ArrMap,
993    observed_arrays: &'a ArrMap,
994    params: &'a [f64],
995    param_names: &'a [String],
996    loop_binds: IdxMap,
997    t: f64,
998    /// Runtime registry of FAQ-materialized derived rings (RFC §8.1): an
999    /// `intersect_polygon` clip self-registers its closed overlap ring here
1000    /// under its node `id`, so a downstream `faq` over a `kind:"derived"`
1001    /// index set (`from_faq: <id>`) resolves its extent (the distinct-vertex
1002    /// count) via [`derived_extent`]. Interior-mutable so the producer can
1003    /// register while the same borrow chain reads it; empty for models with no
1004    /// derived sets (byte-identical to the pre-geometry path).
1005    derived_rings: &'a RefCell<HashMap<String, ArrayD<f64>>>,
1006    /// Build-time value-invention derived-index-set extents (RFC §6.1 / §5.5),
1007    /// keyed by the PRODUCING aggregate's `id` — the same key a
1008    /// `kind:"derived"` index set names in its `from_faq`, and the same keying
1009    /// the Python `EvalContext.derived_extents` and the Julia
1010    /// `_declared_shape_extents` use.
1011    ///
1012    /// This is the relational counterpart of `derived_rings`. A clip ring is
1013    /// materialized *during* evaluation (the producer node runs, then the
1014    /// consumer reads its ring); the value-invention engine instead runs ONCE at
1015    /// setup, off the per-step hot path, and hands its distinct-set cardinality
1016    /// here. A derived range consults this first ([`derived_extent`]) and falls
1017    /// back to the ring, so the two producers coexist in one model.
1018    ///
1019    /// NOT interior-mutable, unlike `derived_rings`: nothing produces into it
1020    /// during evaluation. Empty on the compiled-RHS path, where
1021    /// [`crate::value_invention::rewrite_derived_index_sets`] has already
1022    /// densified every materialized derived set to an `interval` *before* ranges
1023    /// were resolved — those axes therefore arrive as plain static bounds and
1024    /// this map is never consulted. It carries the weight on the standalone
1025    /// [`eval_expression_with_extents`] entry point, where a runner evaluates an
1026    /// expression that still holds a [`RangeSpec::DerivedDyn`] bound.
1027    derived_extents: &'a HashMap<String, i64>,
1028    /// External refreshable forcing-array channel (PR-1, ess-14f.7). Unlike
1029    /// `derived_rings` (rebuilt fresh every RHS call), this borrows the
1030    /// model-lifetime [`ArrayCompiled::forcing`] buffer a driver refreshes
1031    /// between cadence segments; a forcing-fed variable name resolves to its
1032    /// entry here (see [`lookup_variable`]). Empty for models with no loader
1033    /// forcing, so the scalar-`p` path reads identically.
1034    forcing: &'a RefCell<HashMap<String, ArrayD<f64>>>,
1035    /// Common-subexpression memo for the vectorized overlay (ess-cse). `None`
1036    /// on the entry points that evaluate a one-off expression (`eval_expression`,
1037    /// the build-time observed materialization), where there is nothing to amortize
1038    /// a structural analysis over; the overlay then behaves exactly as it did
1039    /// before. See [`cse`] for the scoping rule that makes sharing sound.
1040    cse: Option<&'a CseRt>,
1041    /// Per-model memo for INLINE array-valued `const` literals (esm-spec §4.3.3),
1042    /// keyed by the literal node's address and the active working precision.
1043    ///
1044    /// A `const` node carries its payload as raw JSON, so materializing it is a
1045    /// recursive `serde_json::Value` walk plus a fresh `ArrayD` allocation.
1046    /// Nothing about that depends on the state, the time or the loop binds, yet
1047    /// `index(const([...]), i)` — the shape every transcribed lookup table takes
1048    /// (an RRTM k-distribution band is ~14k numbers) — needs the whole walk
1049    /// once per cell of the enclosing aggregate without it, which is the
1050    /// dominant cost of a right-hand side that gathers such a table.
1051    ///
1052    /// The memo makes it once per model instead. Entries are `Rc`-shared so a
1053    /// hit costs a refcount bump, never a copy, and [`eval_index`] reads the
1054    /// element straight out of the shared array rather than cloning it into a
1055    /// [`Value::Array`].
1056    ///
1057    /// `None` on every path that carries no memo to hang the entry on: the
1058    /// standalone `eval_expression*` entry points, the recurrence sweep, the
1059    /// build-time observed snapshots, the output-node observed pass and the
1060    /// test-only reference tape executor. Those walk the payload directly and
1061    /// answer exactly the same values; the output-node pass repeats the same
1062    /// bodies per node and would benefit from a memo of its own, which it does
1063    /// not yet have. The memo itself lives on [`RhsScratch`], whose lifetime is
1064    /// the rule set's, and is retargeted on the same key [`CseRt::retarget`]
1065    /// uses.
1066    const_lits: Option<&'a ConstLitMemo>,
1067    /// The CONST-ARRAY provenance of this evaluation (CONFORMANCE_SPEC §5.5.5).
1068    ///
1069    /// A gather whose target is named here is a **const-array gather** — a
1070    /// pre-computed factor (Fornberg weights, mesh connectivity, a per-cell
1071    /// metric / geometry array) — and its out-of-range resolution follows the
1072    /// factor's declared per-dimension boundary policy, defaulting to
1073    /// `E_TREEWALK_CONSTARRAY_OOB`. Every other gather is a state / observed
1074    /// gather and keeps the homogeneous-Dirichlet zero-ghost convention, which
1075    /// §5.5.5 says is **never** applied to a const-array gather.
1076    ///
1077    /// [`ConstArrayScope::empty`] on every path with no const-array registry,
1078    /// which reads byte-identically to the pre-§5.5.5 evaluator.
1079    const_arrays: &'a ConstArrayScope,
1080    /// The partially materialized array of the [`AlgebraicRule::Recurrence`]
1081    /// currently sweeping, if any (esm-spec §4.3.1.1). An `index` gather whose
1082    /// operand names [`RecurScope::name`] resolves HERE — before state,
1083    /// observed, parameters and forcing — because during the sweep the array
1084    /// exists nowhere else; and it resolves fail-closed, unlike every other
1085    /// gather channel.
1086    ///
1087    /// `None` on every path in the runtime except the recurrence sweep itself,
1088    /// which is what makes the feature cost nothing for a document that does
1089    /// not use it: one `Option` test on the `index` fast path and one on each
1090    /// vectorized-overlay entry gate.
1091    recur: Option<&'a RecurScope<'a>>,
1092    /// Every name the compiled model DECLARES — states, parameters and
1093    /// observeds — regardless of whether a value for it exists in any of the
1094    /// maps above *at this moment*.
1095    ///
1096    /// Read on exactly one path: the final arm of [`lookup_variable`], where
1097    /// every resolution scope has been tried and none matched. Without it that
1098    /// arm can only report `E_TREEWALK_UNBOUND_NAME` — "bound by NOTHING in
1099    /// scope" — which is a claim about the DOCUMENT that the evaluator is not
1100    /// in a position to make. It was routinely false: a declared observed
1101    /// whose rule had not run yet (a dependency cycle stalls the sweep,
1102    /// esm-spec §4.9.6) resolved through the same arm and was reported as
1103    /// undeclared, naming a variable that is defined and referenced perfectly
1104    /// well while saying nothing about the actual defect (issue #181).
1105    ///
1106    /// With it the arm splits the two cases: a name in this set is
1107    /// `E_TREEWALK_UNRESOLVED_ORDER` (declared, no value HERE), a name outside
1108    /// it keeps `E_TREEWALK_UNBOUND_NAME` (declared nowhere, CONFORMANCE_SPEC
1109    /// §5.23). [`empty_declared_names`] on the standalone expression entry
1110    /// points, which evaluate against a scope the caller supplies rather than a
1111    /// compiled model — an empty set never claims a name IS declared, so those
1112    /// paths keep exactly the diagnosis they had.
1113    declared: &'a HashSet<String>,
1114}
1115
1116/// The shared empty declared-name set (see [`EvalCtx::declared`]).
1117///
1118/// Handed to every evaluation seam with no compiled model behind it. Empty
1119/// means "this evaluation cannot vouch for any name", which is the
1120/// conservative reading: the fault arm then reports the general
1121/// `E_TREEWALK_UNBOUND_NAME` exactly as it did before this field existed.
1122fn empty_declared_names() -> &'static HashSet<String> {
1123    static EMPTY: std::sync::OnceLock<HashSet<String>> = std::sync::OnceLock::new();
1124    EMPTY.get_or_init(HashSet::new)
1125}
1126
1127/// The name `E_TREEWALK_CONSTARRAY_OOB` reports for a `const` literal written
1128/// inline as an `index` base, which has no variable name of its own.
1129pub(crate) const INLINE_CONST_NAME: &str = "inline const";
1130
1131/// Per-model memo for inline array-valued `const` literals — see
1132/// [`EvalCtx::const_lits`] for why it exists.
1133///
1134/// Keyed by `(node address, working precision)`. The address is the payload
1135/// address of the `Arc<ExpressionNode>` the rule set holds, so it is stable and
1136/// unique for exactly as long as that rule set is alive — the same identity
1137/// argument the CSE class table rests on, and [`Self::retarget`] discards the
1138/// table on the same key when a scratch is handed a different rule set.
1139/// Precision is part of the key because a `const` literal rounds on ingress
1140/// (esm-spec §11.3.1), so the binary32 materialization is a different array.
1141#[derive(Default)]
1142pub(super) struct ConstLitMemo {
1143    inner: RefCell<ConstLitInner>,
1144}
1145
1146#[derive(Default)]
1147struct ConstLitInner {
1148    /// The rule-set key this table was built against (see [`CseRt::retarget`]).
1149    key: u64,
1150    /// `None` marks a payload with no array materialization — a ragged or
1151    /// non-numeric literal. Recorded rather than left absent so a repeat read
1152    /// is still a single probe; the caller then falls back to the plain walk,
1153    /// which is what produces that literal's NaN sentinel.
1154    map: HashMap<(usize, bool), Option<std::sync::Arc<ArrayD<f64>>>>,
1155}
1156
1157/// Walk an inline `const` node's JSON payload into a shared array. `None` for
1158/// any payload [`eval::json_to_value`] does not read as a numeric array — a
1159/// scalar, or a ragged/non-numeric one.
1160fn walk_const_lit(node: &ExpressionNode) -> Option<std::sync::Arc<ArrayD<f64>>> {
1161    match node.value.as_ref().map(eval::json_to_value) {
1162        Some(Some(Value::Array(a))) => Some(std::sync::Arc::new(*a)),
1163        _ => None,
1164    }
1165}
1166
1167impl ConstLitMemo {
1168    /// Bind this memo to a rule set. A different key discards every entry: the
1169    /// addresses it holds belonged to the previous rule set and a new node may
1170    /// land on one of them.
1171    pub(super) fn retarget(&self, key: u64) {
1172        let mut inner = self.inner.borrow_mut();
1173        if inner.key != key {
1174            inner.key = key;
1175            inner.map.clear();
1176        }
1177    }
1178
1179    /// The materialized array for the inline `const` literal `node`, computing
1180    /// (and remembering) it on the first read. `None` for a scalar or malformed
1181    /// payload, which the caller handles exactly as an un-memoized walk would.
1182    pub(super) fn get(&self, node: &ExpressionNode) -> Option<std::sync::Arc<ArrayD<f64>>> {
1183        let k = (
1184            node as *const ExpressionNode as usize,
1185            crate::precision::active().is_f32(),
1186        );
1187        if let Some(hit) = self.inner.borrow().map.get(&k) {
1188            return hit.clone();
1189        }
1190        let built = walk_const_lit(node);
1191        self.inner.borrow_mut().map.insert(k, built.clone());
1192        built
1193    }
1194}
1195
1196/// The materialized array for an inline `const` literal, through `memo` when
1197/// one is in scope and by a direct walk otherwise. The two answers are equal by
1198/// construction: the memo stores exactly what the walk produces.
1199pub(super) fn const_lit_array(
1200    node: &ExpressionNode,
1201    memo: Option<&ConstLitMemo>,
1202) -> Option<std::sync::Arc<ArrayD<f64>>> {
1203    match memo {
1204        Some(m) => m.get(node),
1205        None => walk_const_lit(node),
1206    }
1207}
1208
1209/// Which arrays in an evaluation are CONST-ARRAY factors, and each one's
1210/// declared per-dimension out-of-range boundary policy
1211/// (CONFORMANCE_SPEC §5.5.5 / ess-gj4).
1212///
1213/// This is the dense evaluator's counterpart of the Julia reference's
1214/// `const_arrays` registry (whose `BoundedConstArray` wrapper carries the same
1215/// per-dimension policy tuple) and of [`crate::value_invention`]'s
1216/// `const_array_boundaries` input — the two engines must agree byte-for-byte on
1217/// a resolved gather, so they resolve it by the same table.
1218#[derive(Debug, Clone, Default)]
1219pub struct ConstArrayScope {
1220    /// Names that are const-array factors.
1221    names: HashSet<String>,
1222    /// Per-name, per-dimension policy. A name absent here (or a dimension
1223    /// beyond the declared vec) is [`BoundaryKind::Error`] — the default that
1224    /// keeps genuine connectivity / stencil-weight bugs caught.
1225    boundaries: HashMap<String, Vec<BoundaryKind>>,
1226}
1227
1228impl ConstArrayScope {
1229    /// The shared empty scope: nothing is a const array, so every gather keeps
1230    /// the zero-ghost convention exactly as before.
1231    pub fn empty() -> &'static ConstArrayScope {
1232        static EMPTY: std::sync::OnceLock<ConstArrayScope> = std::sync::OnceLock::new();
1233        EMPTY.get_or_init(ConstArrayScope::default)
1234    }
1235
1236    /// A scope over `names`, with no declared boundary policies (so every one
1237    /// of them resolves an out-of-range gather as `E_TREEWALK_CONSTARRAY_OOB`).
1238    pub fn from_names<I: IntoIterator<Item = String>>(names: I) -> ConstArrayScope {
1239        ConstArrayScope {
1240            names: names.into_iter().collect(),
1241            boundaries: HashMap::new(),
1242        }
1243    }
1244
1245    /// Declare `name`'s per-dimension boundary policy (and mark it const).
1246    pub fn with_boundary(mut self, name: impl Into<String>, dims: Vec<BoundaryKind>) -> Self {
1247        let name = name.into();
1248        self.names.insert(name.clone());
1249        self.boundaries.insert(name, dims);
1250        self
1251    }
1252
1253    /// Is this scope empty (nothing is a const array)?
1254    pub fn is_empty(&self) -> bool {
1255        self.names.is_empty()
1256    }
1257
1258    /// Is `name` a const-array factor here?
1259    pub fn is_const(&self, name: &str) -> bool {
1260        self.names.contains(name)
1261    }
1262
1263    /// Is `base` a `const` literal written inline as an `index` base? It is a
1264    /// const array in its own right (esm-spec §4.3.3), with no declared policy.
1265    pub fn is_inline_const(base: &Expr) -> bool {
1266        matches!(base, Expr::Operator(n) if n.op == "const")
1267    }
1268
1269    /// Is an `index` over `base` a const-array gather: a named factor in this
1270    /// scope, or a `const` literal written inline?
1271    pub fn is_const_base(&self, base: &Expr) -> bool {
1272        match base {
1273            Expr::Variable(v) => self.is_const(v),
1274            other => Self::is_inline_const(other),
1275        }
1276    }
1277
1278    /// `name`'s declared policy for dimension `d`, defaulting to
1279    /// [`BoundaryKind::Error`].
1280    pub fn boundary(&self, name: &str, d: usize) -> BoundaryKind {
1281        self.boundaries
1282            .get(name)
1283            .and_then(|dims| dims.get(d))
1284            .copied()
1285            .unwrap_or(BoundaryKind::Error)
1286    }
1287}