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}