earthsci_ast/simulate_array/compile.rs
1//! Compile path: model → [`ArrayCompiled`]. Array-op / spatial-model file
2//! detection, subsystem mounting (esm-spec §4.6), ragged keyed-factor scope
3//! resolution (RFC §5.4), the staged [`ArrayCompiled::from_model`] build, the
4//! build-time field evaluators, and the shape-inference / LHS-parsing
5//! lowering helpers.
6
7use super::*;
8use crate::faq::{effective_reduce_kind, is_faq_op, validate_oplus_spellings};
9use crate::flatten::FlattenedSystem;
10use crate::op_registry::{OpError, is_builtin_function_name};
11use crate::simulate::{CompileError, SimulateError};
12use crate::types::{
13 EsmFile, ExpressionNode, JoinClause, Model, ModelVariable, OverlapClause, VariableType,
14};
15use crate::value_invention::{
16 ValueInventionResult, materialize_value_invention, rewrite_derived_index_sets,
17};
18use indexmap::IndexMap;
19use serde_json::Value as JsonValue;
20use std::collections::{BTreeMap, BTreeSet, HashSet};
21
22// ============================================================================
23// Detection: does the file contain array-op expressions anywhere?
24// ============================================================================
25
26/// Names of the array-op sidecar operators introduced in gt-t5c. `faq`
27/// and `makearray` are the composition primitives; the rest are shape /
28/// extraction helpers that are only meaningful when operating on array
29/// intermediates.
30pub(super) const ARRAY_OP_NAMES: &[&str] = &[
31 "faq", // unified Functional Aggregate Query op (RFC semiring-faq-unified-ir §5.6)
32 "makearray",
33 "reshape",
34 "transpose",
35 "concat",
36 "broadcast",
37];
38
39/// Return true if any expression in the file uses a gt-t5c array op.
40pub fn file_has_array_ops(file: &EsmFile) -> bool {
41 let Some(models) = &file.models else {
42 return false;
43 };
44 models
45 .values()
46 .any(|model| model_tree_any(model, &model_has_array_ops))
47}
48
49/// Whether `pred` holds for `model` or for any component mounted beneath it
50/// through `subsystems` (esm-spec §4.6). A mounted component's equations and
51/// variables are part of the model [`mount_subsystems`] builds, so a routing
52/// question about the model is a question about its whole mount tree: a
53/// document whose only array or derivative lives in a subsystem is otherwise
54/// routed as if it had none. An entry [`parse_subsystem_model`] cannot read
55/// contributes nothing here; the build reports it.
56pub(crate) fn model_tree_any(model: &Model, pred: &dyn Fn(&Model) -> bool) -> bool {
57 if pred(model) {
58 return true;
59 }
60 model.subsystems.as_ref().is_some_and(|subs| {
61 subs.iter().any(|(name, value)| {
62 parse_subsystem_model(name, value).is_ok_and(|(sub, _)| model_tree_any(&sub, pred))
63 })
64 })
65}
66
67/// Return true if the file has spatial structure: any model with array-shaped
68/// state variables (`shape` field non-empty).
69///
70/// Used by [`crate::compile_array`] to accept discretized-PDE files even when
71/// the equations do not yet contain explicit `faq`/`index` nodes (e.g. a
72/// spatial model whose equations were rewritten using indexed-scalar
73/// `D(u[i]) = ...` form rather than the `faq` wrapper).
74pub fn file_has_spatial_model(file: &EsmFile) -> bool {
75 let Some(models) = &file.models else {
76 return false;
77 };
78 let has_shaped_variable = |model: &Model| {
79 model
80 .variables
81 .values()
82 .any(|var| var.shape.as_ref().is_some_and(|shape| !shape.is_empty()))
83 };
84 models
85 .values()
86 .any(|model| model_tree_any(model, &has_shaped_variable))
87}
88
89pub(super) fn model_has_array_ops(model: &Model) -> bool {
90 for eq in &model.equations {
91 if expr_has_array_op(&eq.lhs) || expr_has_array_op(&eq.rhs) {
92 return true;
93 }
94 }
95 // Also detect by the presence of bracketed initial conditions in the
96 // variable definitions — not strictly an AST signal but a strong hint.
97 for name in model.variables.keys() {
98 if name.contains('[') {
99 return true;
100 }
101 }
102 false
103}
104
105pub(super) fn expr_has_array_op(expr: &Expr) -> bool {
106 match expr {
107 Expr::Number(_) | Expr::Integer(_) | Expr::Variable(_) => false,
108 Expr::Operator(node) => {
109 if ARRAY_OP_NAMES.contains(&node.op.as_str()) {
110 return true;
111 }
112 if node.op == "index" {
113 // `index` is only meaningful when there is an array to index
114 // into — always recognise it as an array-op signal.
115 return true;
116 }
117 node.any_child(&mut expr_has_array_op)
118 }
119 }
120}
121
122/// Walk an expression and reject every operator node that may not reach an
123/// evaluator (esm-spec §4.2 / §9.6.8).
124///
125/// This is the crate's single compile-time operator gate. It used to check only
126/// for unlowered *spatial* ops — which meant a malformed-but-schema-valid node
127/// (`atan2` with one argument, `min` with one, a ragged `makearray`, a typo'd
128/// `"expp"`) sailed straight into the evaluators. There it either **panicked**
129/// on an out-of-bounds `args[1]`, or, more insidiously, was quietly assigned two
130/// *different* values by the per-cell oracle and the vectorized overlay
131/// depending on whether the enclosing body happened to vectorize.
132///
133/// Delegating to [`crate::op_registry`] closes all of that at once: past this
134/// gate, every surviving node is an evaluable-core op with a legal arity, so the
135/// evaluators only ever have to agree on nodes that are *legal* — and for those
136/// they agree by construction.
137///
138/// The walk uses `for_each_child`, so it descends into sidecar expression fields
139/// (`aggregate.expr`, `makearray.values`, `filter`, `key`, …), not just `args`.
140///
141/// # Errors
142///
143/// [`CompileError::UnloweredOperatorError`] for a rewrite-target op (sugar, a
144/// spatial `D`, a user op, or a misspelling); [`CompileError::InvalidOperatorArity`]
145/// for a core op with the wrong argument count;
146/// [`CompileError::MakearrayRegionInvalid`] for a ragged or inverted `makearray`;
147/// [`CompileError::InvalidBroadcastFn`] for a `broadcast` whose `fn` is absent
148/// or names no scalar operator.
149pub(super) fn check_no_spatial_ops(expr: &Expr) -> Result<(), CompileError> {
150 crate::op_registry::check_expr(expr).map_err(|e| match e {
151 OpError::Unlowered { op } => CompileError::UnloweredOperatorError { op },
152 OpError::Arity { op, got, expected } => {
153 CompileError::InvalidOperatorArity { op, got, expected }
154 }
155 OpError::MakearrayRegion { reason } => CompileError::MakearrayRegionInvalid { reason },
156 OpError::BroadcastFn { reason, .. } => CompileError::InvalidBroadcastFn { reason },
157 })
158}
159
160// ============================================================================
161// Subsystem mounting (esm-spec §4.6 dot notation).
162// ============================================================================
163
164/// Coerce one resolved `subsystems` entry into a typed [`Model`] plus the
165/// document `index_sets` registry it ships (empty for a bare model fragment).
166/// The loader ([`crate::ref_loading::resolve_subsystem_refs_raw`]) inlines each
167/// `{ "ref": … }` as the referenced file's full JSON, so the common shape is a
168/// whole ESM document carrying exactly one model (the MPAS mesh contract:
169/// `grids/mpas/mesh/level0.esm`); a bare `{ "variables": …, "equations": … }`
170/// fragment is also accepted. An unresolved `{ "ref": … }` — a document built
171/// programmatically without the loader — is a hard error, never a silent drop.
172pub(crate) fn parse_subsystem_model(
173 sub_name: &str,
174 value: &serde_json::Value,
175) -> Result<(Model, HashMap<String, IndexSet>), CompileError> {
176 let obj = value.as_object().ok_or_else(|| {
177 CompileError::build_err(format!("subsystem '{sub_name}' is not a JSON object"))
178 })?;
179 if obj.contains_key("models") {
180 let file: EsmFile = serde_json::from_value(value.clone()).map_err(|e| {
181 CompileError::build_err(format!(
182 "subsystem '{sub_name}' does not parse as an ESM file: {e}"
183 ))
184 })?;
185 let models = file.models.unwrap_or_default();
186 if models.len() != 1 {
187 return Err(CompileError::build_err(format!(
188 "subsystem '{sub_name}' resolves to a file with {} models; exactly one is \
189 required to mount it",
190 models.len()
191 )));
192 }
193 let model = models.into_values().next().expect("len checked above");
194 Ok((
195 model,
196 file.index_sets.unwrap_or_default().into_iter().collect(),
197 ))
198 } else if obj.contains_key("variables") || obj.contains_key("equations") {
199 let model: Model = serde_json::from_value(value.clone()).map_err(|e| {
200 CompileError::build_err(format!(
201 "subsystem '{sub_name}' does not parse as a model: {e}"
202 ))
203 })?;
204 Ok((model, HashMap::new()))
205 } else if obj.contains_key("ref") {
206 Err(CompileError::build_err(format!(
207 "subsystem '{sub_name}' is an unresolved {{\"ref\": …}}; load the document \
208 through the official loader (crate::parse::load_path) so \
209 resolve_subsystem_refs_raw inlines it first"
210 )))
211 } else {
212 Err(CompileError::build_err(format!(
213 "subsystem '{sub_name}' has neither 'models' nor 'variables'/'equations'"
214 )))
215 }
216}
217
218/// Mount every subsystem of `model` into the model's own registries under
219/// dot-prefixed names (esm-spec §4.6): each subsystem variable `x` becomes
220/// `"<sub>.x"` (with sibling references inside its expression renamed to the
221/// mounted names), each subsystem equation is appended with its references
222/// renamed the same way, and the subsystem file's `index_sets` merge into the
223/// document registry (the parent's declaration wins on a name collision).
224/// Recursive, so a nested subsystem mounts as `"<sub>.<subsub>.x"`. This is the
225/// array-runtime analogue of the Julia flatten's subsystem namespacing — it is
226/// what makes the MPAS keyed-factor wiring contract (`nEdgesOnCell :=
227/// mesh.nEdgesOnCell`, a bare-name observed alias of a mounted const factor)
228/// resolvable. A model without subsystems is untouched (byte-identical build).
229pub(crate) fn mount_subsystems(
230 model: &mut Model,
231 index_sets: &mut HashMap<String, IndexSet>,
232) -> Result<(), CompileError> {
233 let Some(subs) = model.subsystems.take() else {
234 return Ok(());
235 };
236 let mut names: Vec<String> = subs.keys().cloned().collect();
237 names.sort();
238 for sub_name in names {
239 let (mut sub_model, mut sub_sets) = parse_subsystem_model(&sub_name, &subs[&sub_name])?;
240 // Grandchildren first, so their variables are already dot-prefixed
241 // within `sub_model` when this level's prefix is applied.
242 mount_subsystems(&mut sub_model, &mut sub_sets)?;
243 for (k, v) in sub_sets {
244 index_sets.entry(k).or_insert(v);
245 }
246 let siblings: Vec<String> = {
247 let mut s: Vec<String> = sub_model.variables.keys().cloned().collect();
248 s.sort();
249 s
250 };
251 let rename_all = |expr: &Expr| -> Expr {
252 let mut out = expr.clone();
253 for s in &siblings {
254 out = rename_free_symbol(&out, s, &format!("{sub_name}.{s}"));
255 }
256 out
257 };
258 for (vname, var) in &sub_model.variables {
259 let mounted = format!("{sub_name}.{vname}");
260 if model.variables.contains_key(&mounted) {
261 return Err(CompileError::build_err(format!(
262 "mounting subsystem '{sub_name}' would overwrite existing variable \
263 '{mounted}'"
264 )));
265 }
266 let mut var = var.clone();
267 // A parameter `update`'s Expressions are the only ones a variable
268 // still carries (esm 1.0.0); an unknown's definition is an equation,
269 // renamed with the rest below.
270 var.for_each_expression_mut(&mut |expr| *expr = rename_all(expr));
271 model.variables.insert(mounted, var);
272 }
273 for eq in &sub_model.equations {
274 model.equations.push(crate::types::Equation {
275 lhs: rename_all(&eq.lhs),
276 rhs: rename_all(&eq.rhs),
277 comment: eq.comment.clone(),
278 });
279 }
280 }
281 Ok(())
282}
283
284// ============================================================================
285// Ragged keyed-factor scope resolution (esm-spec §4.3.1 / RFC §5.4).
286// ============================================================================
287
288/// Resolve each ragged index set's `offsets`/`values` keyed factors against the
289/// model scope, rewriting the registry copy in place. A keyed factor binds by
290/// BARE name (the document-scoped registry keeps the authored name), but
291/// flattening/mounting prefixes every variable with its owning component path
292/// (`nEdgesOnCell` → `Divergence.nEdgesOnCell` alias and
293/// `Divergence.mesh.nEdgesOnCell` const). Resolution rule (mirror of the Julia
294/// tree_walk `_factor_scope`): an exact-name variable wins; otherwise the
295/// dot-suffix match at the SHALLOWEST namespace depth (the model's own
296/// re-exposed alias, not the mounted subsystem's original). Multiple matches at
297/// that depth are a hard error — never a silent empty contraction. A factor
298/// with no in-scope match is left bare (it may be supplied by the caller's
299/// runtime channels), preserving existing behavior. No-op (byte-identical) for
300/// documents without ragged index sets.
301pub(super) fn apply_ragged_factor_scope(
302 index_sets: &mut HashMap<String, IndexSet>,
303 variables: &IndexMap<String, ModelVariable>,
304) -> Result<(), CompileError> {
305 let scope_one = |fname: &str, set_name: &str| -> Result<Option<String>, CompileError> {
306 if variables.contains_key(fname) {
307 return Ok(None); // exact name is in scope; keep as authored
308 }
309 let suffix = format!(".{fname}");
310 let cands: Vec<&String> = variables.keys().filter(|n| n.ends_with(&suffix)).collect();
311 if cands.is_empty() {
312 return Ok(None); // leave bare; a genuinely unbound read surfaces later
313 }
314 let mindepth = cands
315 .iter()
316 .map(|c| c.matches('.').count())
317 .min()
318 .expect("non-empty");
319 let best: Vec<&&String> = cands
320 .iter()
321 .filter(|c| c.matches('.').count() == mindepth)
322 .collect();
323 if best.len() > 1 {
324 let mut names: Vec<String> = best.iter().map(|s| (**s).clone()).collect();
325 names.sort();
326 return Err(CompileError::build_err(format!(
327 "ragged index set '{set_name}': keyed factor '{fname}' is ambiguous in the \
328 model scope — {} candidates at namespace depth {mindepth}: {}",
329 names.len(),
330 names.join(", ")
331 )));
332 }
333 Ok(Some((*best[0]).clone()))
334 };
335 let set_names: Vec<String> = index_sets
336 .iter()
337 .filter(|(_, s)| s.kind == "ragged")
338 .map(|(n, _)| n.clone())
339 .collect();
340 for name in set_names {
341 let (off, vals) = {
342 let s = &index_sets[&name];
343 (s.offsets.clone(), s.values.clone())
344 };
345 if let Some(f) = off
346 && let Some(scoped) = scope_one(&f, &name)?
347 {
348 index_sets.get_mut(&name).expect("present").offsets = Some(scoped);
349 }
350 if let Some(f) = vals
351 && let Some(scoped) = scope_one(&f, &name)?
352 {
353 index_sets.get_mut(&name).expect("present").values = Some(scoped);
354 }
355 }
356 Ok(())
357}
358
359/// The rewrites `flatten` makes to a model's equations that a route reading
360/// the AUTHORED model — the array runtime's single-model route, and the build
361/// pipeline (`crate::prepare`) — must make itself, so every route answers one
362/// document with one number.
363///
364/// * esm-spec §4.2, the two halves of the right-hand-side `D` rule (`flatten`'s
365/// phase 5b′): each RHS structural `D` resolves into the tendency the model
366/// defines for it, and one that resolves to nothing is refused rather than
367/// answered with `0` (or, through a reference evaluator, `NaN`). The same two
368/// functions, in the same order, as [`ArrayCompiled::from_flattened`].
369/// * esm-spec §4.8.3 "Angles are the ONE exception" (phase 5g): see
370/// [`normalize_model_angle_arguments`]. Without it `sin(theta)` with `theta`
371/// declared `deg` evaluated `sin(90 radians)` = 0.894 where the scalar and
372/// coupled routes, and the other bindings, returned 1.
373pub(crate) fn apply_flatten_rewrites(model: &mut Model) -> Result<(), CompileError> {
374 let time_invariant: std::collections::HashSet<String> = model
375 .variables
376 .iter()
377 .filter(|(_, v)| v.var_type == crate::types::VariableType::Parameter)
378 .map(|(name, _)| name.clone())
379 .collect();
380 let carried =
381 crate::flatten::resolve_rhs_time_derivatives(&mut model.equations, &time_invariant);
382 if carried
383 && crate::flatten::first_unresolved_rhs_time_derivative_in(&model.equations).is_some()
384 {
385 return Err(CompileError::UnloweredOperatorError {
386 op: "D".to_string(),
387 });
388 }
389 normalize_model_angle_arguments(model);
390 Ok(())
391}
392
393/// Fold a declared angle's scale into every `sin`/`cos`/`tan` argument of
394/// `model`, so the argument reaches the evaluator in RADIANS (esm-spec §4.8.3,
395/// issue #409).
396///
397/// The array runtime's SINGLE-MODEL route deliberately never flattens
398/// (`ArrayCompiled::from_file`), and neither does the build pipeline
399/// (`crate::prepare`), so `flatten`'s phase 5g
400/// (`crate::flatten::normalize_angle_arguments`) never runs for either. This is
401/// that phase, against the authored model's own declarations, so every route
402/// answers one document with one number.
403fn normalize_model_angle_arguments(model: &mut Model) {
404 // A document that declares no angle at a scale other than 1 cannot be
405 // rewritten, so decide that from the DECLARATIONS (each distinct unit
406 // string once) and walk nothing.
407 let mut seen: HashSet<&str> = HashSet::new();
408 let declares_angle = model
409 .variables
410 .values()
411 .filter_map(|v| v.units.as_deref())
412 .filter(|u| seen.insert(u))
413 .any(|u| {
414 crate::units::parse_unit(u)
415 .is_ok_and(|unit| crate::units::angle_normalization_factor(&unit).is_some())
416 });
417 if !declares_angle {
418 return;
419 }
420 let (env, _) = crate::units::build_unit_env(&model.variables);
421 for eq in &mut model.equations {
422 if let Some(next) = crate::units::normalize_angle_arguments(&eq.rhs, &env) {
423 eq.rhs = next;
424 }
425 }
426}
427
428// ============================================================================
429// Compile path: model → ArrayCompiled.
430// ============================================================================
431
432impl ArrayCompiled {
433 pub fn from_file(file: &EsmFile) -> Result<Self, CompileError> {
434 Self::from_file_with(file, false)
435 }
436
437 /// [`Self::from_file`], with `runtime_array_parameters` as
438 /// [`Self::from_flattened_with`] takes it.
439 pub(crate) fn from_file_with(
440 file: &EsmFile,
441 runtime_array_parameters: bool,
442 ) -> Result<Self, CompileError> {
443 let Some(models) = &file.models else {
444 return Err(CompileError::build_err("File has no models to simulate"));
445 };
446 if models.len() != 1 {
447 return Err(CompileError::build_err(
448 "Array-op path currently only supports a single model file (no coupling)",
449 ));
450 }
451 let (model_name, model) = models.iter().next().unwrap();
452 // v0.8.0: `index_sets` is document-scoped (one registry shared by all
453 // models), so source it from the file rather than the model.
454 let index_sets: HashMap<String, IndexSet> = file
455 .index_sets
456 .clone()
457 .unwrap_or_default()
458 .into_iter()
459 .collect();
460 // `from_model` clones the model internally anyway, so taking the clone
461 // here costs nothing extra — same shape as `from_file_owned` below,
462 // which owns its model outright. The model's name rides along so the
463 // pipeline can resolve self-qualified references once the subsystems
464 // are mounted.
465 let mut compiled = Self::compile_pipeline(
466 model.clone(),
467 &index_sets,
468 None,
469 Some(model_name),
470 runtime_array_parameters,
471 )?;
472 // Record the model's namespace so overrides may be keyed `Model.param`
473 // (the scalar/flatten/Julia convention) as well as the raw `param` this
474 // single-model path builds (WS3 override-naming parity).
475 compiled.namespace = Some(model_name.clone());
476 compiled.qualify_data_fed(model_name);
477 Ok(compiled)
478 }
479
480 /// [`Self::from_file`], consuming the file. The peak-memory-lean build for
481 /// a large expanded discretization: the single model is MOVED out of the
482 /// file (nothing is cloned), the rest of the document is dropped here, and
483 /// [`Self::compile_pipeline`] then moves each observed body into its
484 /// compiled rule instead of deep-copying it. For `simpleclimate.esm` at
485 /// the production grid the borrowed `from_file` holds three ~1 GiB copies
486 /// live at once (the caller's file, the compile's private model clone, the
487 /// compiled rules); this path holds ~one. Same build stages, same result.
488 pub fn from_file_owned(file: EsmFile) -> Result<Self, CompileError> {
489 let Some(models) = file.models else {
490 return Err(CompileError::build_err("File has no models to simulate"));
491 };
492 if models.len() != 1 {
493 return Err(CompileError::build_err(
494 "Array-op path currently only supports a single model file (no coupling)",
495 ));
496 }
497 let index_sets: HashMap<String, IndexSet> =
498 file.index_sets.unwrap_or_default().into_iter().collect();
499 let (model_name, model) = models.into_iter().next().unwrap();
500 let mut compiled =
501 Self::compile_pipeline(model, &index_sets, None, Some(&model_name), false)?;
502 compiled.qualify_data_fed(&model_name);
503 compiled.namespace = Some(model_name);
504 Ok(compiled)
505 }
506
507 /// Build from a [`FlattenedSystem`].
508 ///
509 /// [`crate::flatten::flatten`] already merges a coupled, multi-component
510 /// file into a single dot-namespaced system (coupling rules applied, every
511 /// variable reference namespaced). The array path historically only had
512 /// [`Self::from_file`], which rejects `models.len() != 1` outright because
513 /// it operates on a raw [`Model`] and has no coupling machinery of its own.
514 /// This constructor closes that seam: it consumes the already-coupled
515 /// flatten output directly, so a discretized **coupled** spatial model
516 /// compiles + evaluates through the vectorized array runtime, reusing
517 /// `flatten.rs`'s coupling verbatim (no new coupling logic here). The raw
518 /// single-model `from_file` guard is intentionally left intact — the real
519 /// pipeline flattens first and reaches the array runtime through here
520 /// (ess-14f.8).
521 ///
522 /// The flattened system splits variables into typed maps; [`Self::from_model`]
523 /// expects a single registry discriminated by [`ModelVariable::var_type`].
524 /// We merge them back into one synthetic [`Model`] (each variable already
525 /// carries its `var_type`) and delegate, so every downstream stage — shape
526 /// inference, faq lowering, the diffsol RHS build — is shared bit-for-bit
527 /// with the single-model path.
528 pub fn from_flattened(flat: &FlattenedSystem) -> Result<Self, CompileError> {
529 Self::from_flattened_with(flat, false)
530 }
531
532 /// [`Self::from_flattened`], with every shaped parameter that has a value
533 /// and no `update` left as a run-time input read from the forcing
534 /// buffer (`runtime_array_parameters`, API_SPEC §5.8) rather than lowered
535 /// into build-time data.
536 pub(crate) fn from_flattened_with(
537 flat: &FlattenedSystem,
538 runtime_array_parameters: bool,
539 ) -> Result<Self, CompileError> {
540 // A coupling can mix precisions the per-model inference never saw
541 // (esm-spec §4.7 "Working precision crosses a mount", §11.3.1).
542 crate::precision_infer::check_flattened(flat)?;
543 // Reject hybrid dimensionality and model events. The data-loader
544 // refresh path that drives
545 // this seam is event-free by design (a driver-level segmented solve,
546 // not an in-solver event), so rejecting here loses no in-scope
547 // capability while preventing a model that *does* declare events from
548 // compiling with its events silently dropped.
549 //
550 // esm-spec §4.2: a right-hand-side structural `D` that `flatten`'s
551 // tendency resolution left standing is a rewrite-target reaching
552 // evaluation, and no implementation may invent a value for it — in
553 // particular not the `0` this runtime's `eval`/tape lowering return.
554 // Outside the dimensionality guard below, which only admits a system
555 // with a SPATIAL independent variable: this defect lives in 0-D
556 // documents. Same check, same code, as the scalar
557 // `reject_unsupported_features`.
558 if crate::flatten::first_unresolved_rhs_time_derivative(flat).is_some() {
559 return Err(CompileError::UnloweredOperatorError {
560 op: "D".to_string(),
561 });
562 }
563 if flat.independent_variables != ["t"] {
564 // A spatial independent variable means a rewrite-target operator was
565 // never discretized. Report THAT, with the uniform
566 // `unlowered_operator` code esm-spec §4.2 / §9.6.8 specifies for
567 // an op reaching evaluation unlowered; the dimensionality error
568 // is the fallback for a spatial axis with no such op behind it.
569 if let Some(op) = crate::flatten::first_unlowered_operator(flat) {
570 return Err(CompileError::UnloweredOperatorError { op });
571 }
572 return Err(CompileError::UnsupportedDimensionalityError {
573 independent_variables: flat.independent_variables.clone(),
574 });
575 }
576 if let Some(event) = flat.continuous_events.first() {
577 return Err(crate::compile_error::continuous_event_refusal(
578 crate::compile_error::ARRAY_EVALUATOR,
579 event.name.as_deref(),
580 ));
581 }
582 if let Some(event) = flat.discrete_events.first() {
583 return Err(crate::compile_error::discrete_event_refusal(
584 crate::compile_error::ARRAY_EVALUATOR,
585 event.name.as_deref(),
586 ));
587 }
588
589 // Re-merge the typed variable maps into one registry. The maps are
590 // disjoint by construction (a variable has exactly one `var_type`), so
591 // no key collides. `parameters` is the WHOLE parameter set of every
592 // cadence (esm-libraries-spec §4.7.5 step 4: the wiener / discrete
593 // subsets partition it rather than sitting beside it), so a Brownian
594 // parameter still reaches `from_model` and gets its explicit "no SDE"
595 // rejection, and a `data`-kind discrete parameter still reaches the
596 // provider forcing seam `classify_variables` routes.
597 let mut variables: IndexMap<String, ModelVariable> = IndexMap::new();
598 for (name, var) in &flat.state_variables {
599 variables.insert(name.clone(), var.clone());
600 }
601 for (name, var) in &flat.parameters {
602 variables.insert(name.clone(), var.clone());
603 }
604 for (name, var) in &flat.observed_variables {
605 variables.insert(name.clone(), var.clone());
606 }
607
608 // `index_sets` is not carried through flatten today, so coupled models
609 // that address `faq` ranges via `{ "from": <set> }`
610 // are not yet resolvable on this path (tracked as follow-up). Dense
611 // `[lo, hi]` ranges — what discretized stencils emit — need no
612 // registry and work here.
613 let model = Model {
614 variables,
615 equations: flat.equations.clone(),
616 ..Default::default()
617 };
618 // The document `index_sets` registry is carried through flatten
619 // (`FlattenedSystem::index_sets`), so a coupled array system can resolve
620 // `faq` `ranges` `{ "from": <set> }`, `join.on` gates, and
621 // derived-set references exactly as the single-model `from_file` path
622 // does against `file.index_sets`. Empty for a file with no index sets, so
623 // dense `[lo, hi]`-range discretized stencils are unaffected.
624 let index_sets: HashMap<String, IndexSet> = flat
625 .index_sets
626 .iter()
627 .map(|(k, v)| (k.clone(), v.clone()))
628 .collect();
629 let mut compiled =
630 Self::compile_pipeline(model, &index_sets, None, None, runtime_array_parameters)?;
631 // `flatten` routed every `ic` equation out of `flat.equations`, so the
632 // synthetic model's stage-(0) walk never saw an initial-condition
633 // right-hand side. Walk them here, before anything evaluates one
634 // (esm-spec §9.6.3 constraint 6), so an unlowered op in an initial
635 // condition is refused at build exactly as it is in an equation.
636 for (_, rhs) in &flat.field_ics {
637 check_evaluable(rhs)?;
638 }
639 // Carry the classified scoped-reference `ic` equations through so `u0` is
640 // folded from the provider-served loaded initial fields at build time.
641 compiled.field_ics = flat.field_ics.clone();
642 // …and the merge's rename map, so an override key naming a state an
643 // `operator_compose` renaming match DELETED resolves to the survivor
644 // rather than designating nothing (issue #230). The synthetic `Model`
645 // above cannot carry it: it is a property of the COUPLING, not of any
646 // one component.
647 compiled.merged_renames = flat
648 .metadata
649 .merged_variable_renames
650 .iter()
651 .map(|(k, v)| (k.clone(), v.clone()))
652 .collect();
653 Ok(compiled)
654 }
655
656 /// Build from a single [`Model`] and the document-scoped `index_sets`
657 /// registry (RFC semiring-faq-unified-ir §5.2, v0.8.0). The registry lives
658 /// on the top-level document and is shared by all models, so it is passed in
659 /// explicitly; pass an empty map for a model with no `{ "from": <set> }`
660 /// range references.
661 pub fn from_model(
662 model: &Model,
663 index_sets: &HashMap<String, IndexSet>,
664 ) -> Result<Self, CompileError> {
665 Self::from_model_with_arrays(model, index_sets, None)
666 }
667
668 /// [`from_model`](Self::from_model) with a caller-supplied factor-array
669 /// channel for the build-time value-invention engine.
670 ///
671 /// `vi_arrays` is overlaid on the `const`-literal factors the document
672 /// carries (caller wins on a name collision — see `vi_factor_arrays`) and
673 /// handed to the relational engine. Supply it whenever a producer's factors
674 /// arrive from OUTSIDE the document: ISRM's overlap gate reads its envelopes
675 /// from `type: "parameter"` variables that data loaders fill through
676 /// `coupling` `param_to_var` edges, which the `const` scan cannot see, so
677 /// without this channel the gate has no envelopes to build and the producer
678 /// invents nothing.
679 ///
680 /// `None` is exactly [`from_model`](Self::from_model) — every existing
681 /// caller is unaffected.
682 pub fn from_model_with_arrays(
683 model: &Model,
684 index_sets: &HashMap<String, IndexSet>,
685 vi_arrays: Option<&HashMap<String, ArrayD<f64>>>,
686 ) -> Result<Self, CompileError> {
687 // The rewrite passes below need an owned model; clone so the caller's
688 // model — and its serialized form — is untouched. A caller that can
689 // give up its model avoids this copy via [`Self::compile_pipeline`]
690 // (reached through [`Self::from_file_owned`] / `compile_array`).
691 Self::from_model_owned_with_arrays(model.clone(), index_sets, vi_arrays)
692 }
693
694 /// [`Self::from_model_with_arrays`], consuming the model: the compile
695 /// pipeline's rewrite passes mutate it in place (no private clone), and
696 /// the observed bodies — the dominant allocation of a large expanded
697 /// discretization — are MOVED into the compiled rules rather than
698 /// deep-copied (`build_observed_rules` takes each `var.expression`).
699 /// Behaviourally identical to `from_model`: every stage runs in the same
700 /// order on the same values.
701 fn from_model_owned_with_arrays(
702 model_owned: Model,
703 index_sets: &HashMap<String, IndexSet>,
704 vi_arrays: Option<&HashMap<String, ArrayD<f64>>>,
705 ) -> Result<Self, CompileError> {
706 Self::compile_pipeline(model_owned, index_sets, vi_arrays, None, false)
707 }
708
709 /// The compile pipeline itself; every public entry point above delegates
710 /// here. `model_name` names the document model being compiled on the
711 /// single-model route. With a name, every self-qualified reference (`M.x`,
712 /// `M.sub.x`) is resolved to its local spelling right after the subsystems
713 /// are mounted, so one authored inside a subsystem body resolves as well
714 /// as one in the parent's own equations.
715 fn compile_pipeline(
716 mut model_owned: Model,
717 index_sets: &HashMap<String, IndexSet>,
718 vi_arrays: Option<&HashMap<String, ArrayD<f64>>>,
719 model_name: Option<&str>,
720 runtime_array_parameters: bool,
721 ) -> Result<Self, CompileError> {
722 // Resolve `{ "from": <index set> }` range references (RFC
723 // semiring-faq-unified-ir §5.2) into concrete `[lo, hi]` intervals
724 // before any shape inference or rule building; every downstream
725 // consumer then sees only dense interval ranges.
726 //
727 // Mount subsystems under dot-prefixed names (esm-spec §4.6) and resolve
728 // each ragged index set's keyed factors against the resulting model
729 // scope (RFC §5.4; the Julia `_factor_scope` mirror). Both are no-ops —
730 // and the registry copy is byte-identical — for models without
731 // subsystems / ragged sets.
732 // An event, continuous or discrete, is refused before anything is
733 // built. This is the SINGLE-MODEL route's check: `from_flattened` checks
734 // the flattened event lists itself, because the synthetic model it hands
735 // down carries no events. Subsystems are searched too, since mounting
736 // keeps only their variables and equations.
737 if let Some((construct, name)) = crate::compile_error::first_event(&model_owned) {
738 return Err(crate::compile_error::event_refusal(
739 construct,
740 crate::compile_error::ARRAY_EVALUATOR,
741 name.as_deref(),
742 ));
743 }
744 let mut index_sets_owned = index_sets.clone();
745 mount_subsystems(&mut model_owned, &mut index_sets_owned)?;
746 if let Some(name) = model_name {
747 resolve_self_references(&mut model_owned, name);
748 }
749 // `flatten`'s rewrites, which this SINGLE-MODEL route never gets
750 // (see `from_file_owned`). Runs after mounting so a subsystem's
751 // equations and declarations are in scope under their mounted names.
752 apply_flatten_rewrites(&mut model_owned)?;
753 // Lower every SHAPED parameter whose value the document supplies —
754 // inline array data, or one scalar broadcast over the grid (esm-spec
755 // §6.3 / §6.6.2) — into the `const`-observed channel this runtime
756 // already reads. Runs after mounting so a subsystem's shaped parameter
757 // is lowered under its namespaced name too; `vi_arrays` is passed so a
758 // caller-supplied factor array stays authoritative over a declared
759 // default.
760 let (inline_param_arrays, runtime_params) = lower_inline_array_parameters(
761 &mut model_owned,
762 &index_sets_owned,
763 vi_arrays,
764 runtime_array_parameters,
765 )?;
766 apply_ragged_factor_scope(&mut index_sets_owned, &model_owned.variables)?;
767 // Under `element_type: "Float32"`, reject an index set whose subscripts
768 // binary32 cannot address exactly. Index expressions share the value
769 // kernels — there is no integer type in the expression language — so
770 // `i + 1` above 2^24 would round, and a gather would silently read the
771 // wrong cell. A no-op under Float64, and for every realistic extent.
772 if crate::precision::is_f32() {
773 let mut names: Vec<&String> = index_sets_owned.keys().collect();
774 names.sort();
775 for name in names {
776 if let Some(size) = index_sets_owned[name].size {
777 crate::precision::check_index_set_extent(name, size)?;
778 }
779 }
780 }
781 // Materialize genuine relational OUTPUTS — the arg-witness reducer
782 // (`argmin`/`argmax`, RFC §5.7 rule 6) and the grouped/derived SCVT chain
783 // (`group_aggregate`) — to CONSTANT DATA at build setup, then rewrite each
784 // output's defining equation to a `const` literal the per-cell oracle
785 // already evaluates. This runs the byte-conformant [`crate::value_invention`]
786 // engine (the previously-unwired front door) and mirrors the Julia
787 // reference's "materialize to data" and the live Python interpreter, so
788 // `argmin` / `group_aggregate` now SIMULATE end-to-end instead of raising
789 // [`CompileError::UnevaluableOperatorError`]. Runs BEFORE
790 // [`strip_value_invention`] so a bin-skolem `join` feeding an argmin is still
791 // intact when the buffer is computed. A NO-OP (byte-identical) for every
792 // model without an arg-witness op; `skolem`/`distinct` producers are
793 // handled by the two passes below.
794 // One walk says which value-invention producers the model has at all;
795 // the passes below look only for those. A pass that rewrites the
796 // equations is followed by a fresh look.
797 let mut vi = ViCensus::of(&model_owned);
798 if vi.arg_witness {
799 materialize_vi_outputs_to_data(&mut model_owned, &mut index_sets_owned, vi_arrays)?;
800 vi = ViCensus::of(&model_owned);
801 }
802 // Build-time value invention for every non-geometry derived index set
803 // (`skolem`/`distinct`/`rank`, RFC §6.1): the producer's member count sizes
804 // a range over the set. Runs before `strip_value_invention` drops the
805 // producer. A producer that cannot run is recorded rather than raised, and
806 // is refused below only if a surviving expression ranges over its set.
807 let derived = materialize_derived_extents(&model_owned, &index_sets_owned, vi_arrays, &vi);
808 // A skolem map buffer (a broad-phase bin key per cell) becomes constant
809 // data, so a `join.on` gate comparing two of them stays a gate — a
810 // value-equality filter on data columns — rather than being dropped.
811 let mut map_names: Vec<&String> = derived.map_codes.keys().collect();
812 map_names.sort();
813 if !map_names.is_empty() {
814 for name in map_names {
815 rewrite_equation_to_const(&mut model_owned, name, &derived.map_codes[name]);
816 }
817 vi = ViCensus::of(&model_owned);
818 }
819 let index_sets = &index_sets_owned;
820 // Drop value-invention (relational) scaffolding — skolem-id bin maps and
821 // membership sets over `kind: "derived"` index sets — plus the broad-phase
822 // `join.on` gates keyed on them, BEFORE join/range resolution. The extents
823 // counted above are all that survive of such a set: its members and
824 // skolem maps are not materialized, and the elided gate is numerically
825 // inert only where a dense narrow phase follows (see
826 // `strip_value_invention`). A no-op unless a `skolem` op or a
827 // derived-set-shaped variable is present.
828 // A producer that could not run leaves its derived set with no
829 // members: a stage that then trips over the stripped producer (a
830 // surviving read of one of its outputs, an output shaped on its set)
831 // names that set (`derived_index_set_unmaterialized`, esm-spec §4.2)
832 // rather than the operator or the axis it met.
833 strip_value_invention(&mut model_owned, index_sets, &vi)
834 .or_else(|e| derived.refuse_on_failure(index_sets).and(Err(e)))?;
835 // The model's CONST-ARRAY registry (CONFORMANCE_SPEC §5.5.5): the
836 // `const`-literal factor variables — Fornberg weights, mesh
837 // connectivity, a geometry table — that [`collect_const_factor_arrays`]
838 // already identifies as the Rust analogue of the Julia reference's
839 // `const_arrays`. Captured HERE, before stage (6) MOVES each observed
840 // body out of the variable registry. A gather on one of these resolves
841 // an out-of-range index by its declared boundary policy instead of the
842 // state gather's zero ghost, which §5.5.5 says is never a const array's.
843 let const_scope = std::sync::Arc::new(ConstArrayScope::from_names(
844 observed_bodies(&model_owned)
845 .into_iter()
846 .filter(|(_, body)| matches!(body, Expr::Operator(n) if n.op == "const"))
847 .map(|(n, _)| n),
848 ));
849 // Resolve `join.on` value-equality clauses (RFC §5.3) FIRST, while each
850 // aggregate range still carries its `{ "from": <index set> }` linkage so
851 // the join key columns' member values can be read. A join whose key
852 // columns resolve to the same loop symbol is the degenerate positional
853 // no-op (byte-identical to the no-join form); a join over two distinct
854 // loop symbols is the data-derived value-equality case and is lowered
855 // into a member-equality `filter` over the contraction; a join over a
856 // genuine (non-loop) data column is rejected rather than mis-combined.
857 crate::join::resolve_aggregate_joins(&mut model_owned, index_sets)?;
858 // Then rewrite every `{ "from": <index set> }` range reference (§5.2)
859 // into a concrete `[lo, hi]` interval before shape inference / rule
860 // building, so every downstream consumer sees only dense intervals.
861 // The next three passes look only at nodes with `ranges` and at `faq`
862 // nodes, and shape inference only at `index` nodes; no pass from here
863 // to there adds any of them, so one walk says whether there are any.
864 let census = OpCensus::of(&model_owned);
865 if census.ranges {
866 crate::faq::resolve_aggregate_ranges_with_extents(
867 &mut model_owned,
868 index_sets,
869 &derived.extents,
870 )
871 .or_else(|e| derived.refuse_on_failure(index_sets).and(Err(e)))?;
872 // A range over a non-geometry derived set that value invention did
873 // not size would contract as empty and read 0 (esm-spec §9.6.6):
874 // refuse it.
875 refuse_unmaterialized_derived_ranges(&model_owned, index_sets, &derived)?;
876 }
877 // Reject any aggregate whose ⊕ is spelled outside the schema's closed
878 // `reduce` / `semiring` enums. The gate lives here, at the one funnel
879 // every array-runtime build passes through, because the seams that
880 // actually resolve ⊕ (`extract_derivative_faq`, `faq_spec`)
881 // return `Option` and so could only decline silently.
882 if census.faq {
883 validate_oplus_spellings(&model_owned)?;
884 }
885
886 // Stages (0)-(5) read the model immutably; only OWNED products leave
887 // this block, so stage (6) below can borrow the model mutably (it
888 // moves each observed body out of the variable registry).
889 let (
890 observed_names,
891 forcing_decls,
892 eliminated,
893 held_at_ic,
894 slots,
895 param_names,
896 param_index,
897 param_defaults,
898 unvalued_shaped_params,
899 forcing_defaults,
900 data_fed,
901 ) = {
902 let model = &model_owned;
903
904 // (0) Reject spatial differential operators anywhere in the model's
905 // equations or observed-variable expressions (esm-i7b).
906 reject_unlowered_spatial_ops(model)?;
907
908 // (0a) Reject an implicit equation: this runtime has no algebraic
909 // solve, and every stage below would skip the equation, leaving the
910 // unknown at its initial value.
911 if let Some(eq) = crate::compile_error::first_implicit_equation(&model.equations) {
912 return Err(crate::compile_error::implicit_equation_refusal(
913 crate::compile_error::ARRAY_EVALUATOR,
914 eq,
915 ));
916 }
917
918 // (0a′) Reject an unknown carrying BOTH a derivative equation and a
919 // bare-LHS one. esm-spec §4.9.4 counts both, so the document has one
920 // more equation than it has unknowns and `validate` reports
921 // `equation_count_mismatch`; the build says the same rather than
922 // integrating a system free of the constraint the file declares.
923 // After mounting, so a pair inside a subsystem is seen under its
924 // mounted names, and ahead of `partition_states`, which would
925 // otherwise pick one of the two without saying so.
926 if let Some((name, diff, alg)) =
927 crate::compile_error::first_doubly_defined_unknown(&model.equations)
928 {
929 return Err(crate::compile_error::doubly_defined_unknown_refusal(
930 &name, diff, alg,
931 ));
932 }
933
934 // (0b) Reject a reference to a variable bound in NONE of the model's
935 // binding categories — the same function `crate::prepare` calls on
936 // the build-pipeline route
937 // (CONFORMANCE_SPEC §5.23; `extra_bound` is empty here because this
938 // path's evaluation scope is the model's own declarations). Without
939 // it a typo'd/undeclared bare name reaches `lookup_variable`'s final
940 // arm, which now fails closed but only names the fault at the drain;
941 // catching it HERE names the variable before anything is built.
942 check_free_variables(model, index_sets, &[])?;
943
944 // (1) Collect state / parameter / observed variables.
945 let (state_vars, param_vars, observed_vars, data_fed) =
946 classify_variables(model, &runtime_params)?;
947
948 // (2)+(2b) Infer state shapes from every equation usage, seeding
949 // declared array shapes where the index-usage inference left an
950 // array state scalar.
951 let shape_map = infer_state_shapes(model, &state_vars, index_sets, census.index)?;
952
953 // (3) Partition state variables into integrated / eliminated /
954 // held-at-ic.
955 let (final_states, eliminated, held_at_ic) = partition_states(model, &state_vars);
956
957 // (4) Build flat offsets and scalar-slot names per state variable.
958 let slots = build_slot_tables(model, &final_states, &shape_map)?;
959
960 // (5) Build the param tables.
961 let (param_names, param_index, param_defaults) = build_param_tables(model, ¶m_vars);
962 // A SHAPED parameter that reached the scalar table with no value —
963 // no `default`, no `distribution`, no inline data, not refreshed
964 // from outside. The front door refuses these by name
965 // (`crate::problem::esm_problem`, esm-spec §10.10); a lower-level
966 // caller reaches them as a solve-time missing parameter.
967 let unvalued_shaped_params: Vec<String> = param_vars
968 .iter()
969 .filter(|name| {
970 model.variables.get(**name).is_some_and(|v| {
971 v.shape.as_ref().is_some_and(|s| !s.is_empty())
972 && v.default.is_none()
973 && v.distribution.is_none()
974 })
975 })
976 .map(|name| (*name).clone())
977 .collect();
978
979 // Stage (6) consumes the observed bodies by NAME (sorted order,
980 // exactly the order `classify_variables` produced them in).
981 let observed_names: Vec<String> =
982 observed_vars.iter().map(|(n, _)| (*n).clone()).collect();
983 // The externally refreshed parameters ride in `observed_vars` with
984 // no defining rule (see `classify_variables`); these are the
985 // forcing buffer's names.
986 let mut forcing_decls: IndexMap<String, Option<Vec<usize>>> = observed_vars
987 .iter()
988 .filter(|(_, var)| var.var_type == VariableType::Parameter)
989 .map(|(name, var)| {
990 let shape = match var.shape.as_deref() {
991 None => Some(Vec::new()),
992 Some(decl) => resolve_declared_shape(decl, index_sets),
993 };
994 ((*name).clone(), shape)
995 })
996 .collect();
997 infer_unsized_forcing_shapes(&mut forcing_decls, &model.equations);
998 let forcing_defaults = forcing_default_values(&observed_vars, &forcing_decls)?;
999 (
1000 observed_names,
1001 forcing_decls,
1002 eliminated,
1003 held_at_ic,
1004 slots,
1005 param_names,
1006 param_index,
1007 param_defaults,
1008 unvalued_shaped_params,
1009 forcing_defaults,
1010 data_fed,
1011 )
1012 };
1013
1014 // (5c) The build-time array scope a field `ic` may read (esm-spec
1015 // §6.6.5) — captured before stage (6) moves the bodies out.
1016 let ic_scope = capture_ic_scope(
1017 &model_owned,
1018 &observed_names,
1019 &forcing_decls,
1020 &slots.var_shapes,
1021 &eliminated,
1022 );
1023
1024 // (6)+(6b) Build the dependency-ordered observed algebraic rules,
1025 // MOVING each declared observed's body expression out of the model.
1026 let observed_rules =
1027 build_observed_rules(&mut model_owned, &observed_names, &eliminated, index_sets)?;
1028
1029 let model = &model_owned;
1030
1031 // Classify scoped-reference / array `ic` equations (esm-spec §11.4.1)
1032 // out of the rule builder into `field_ics` (see [`classify_field_ics`]).
1033 let field_ics = classify_field_ics(model);
1034 let init_faqs = classify_initialization_faqs(model);
1035
1036 // (7)+(7b)+(8) Build the RHS rules, cover held-at-ic slots, and
1037 // validate that every state slot has a defining equation.
1038 let rhs_rules = build_rhs_rules(model, &slots, &held_at_ic)?;
1039 let held_at_ic_names = held_at_ic.clone();
1040
1041 let SlotTables {
1042 var_shapes,
1043 state_defaults,
1044 n_states,
1045 } = slots;
1046 // Every name this model declares (issue #181): the state slots, the
1047 // parameters, and each observed rule's target. Read only by the
1048 // evaluator's fault arm, to tell a name declared NOWHERE from one
1049 // declared here but not yet materialized — see [`EvalCtx::declared`].
1050 let declared_names: HashSet<String> = var_shapes
1051 .keys()
1052 .cloned()
1053 .chain(param_names.iter().cloned())
1054 .chain(observed_rules.iter().map(|r| observed_rule_var(r).clone()))
1055 .collect();
1056
1057 Ok(ArrayCompiled {
1058 runtime_mode: crate::simulate_array::RuntimeMode::default(),
1059 #[cfg(feature = "xla")]
1060 xla_rhs: std::cell::OnceCell::new(),
1061 var_shapes,
1062 state_names: std::cell::OnceCell::new(),
1063 qualified_state_names: std::cell::OnceCell::new(),
1064 state_defaults,
1065 param_names,
1066 param_index,
1067 param_defaults,
1068 observed_rules,
1069 rhs_rules,
1070 rule_set: next_rule_set_id(),
1071 n_states,
1072 declared_names,
1073 forcing: Rc::new(RefCell::new(HashMap::new())),
1074 data_fed,
1075 forcing_generation: std::cell::Cell::new(0),
1076 field_ics,
1077 init_faqs,
1078 ic_scope,
1079 index_sets: index_sets.clone(),
1080 namespace: None,
1081 const_scope,
1082 precision: crate::precision::Env::capture(),
1083 merged_renames: HashMap::new(),
1084 #[cfg(feature = "solve")]
1085 field_ic_memo: RefCell::new(None),
1086 inline_param_arrays,
1087 forcing_decls,
1088 forcing_defaults,
1089 unvalued_shaped_params,
1090 held_at_ic: held_at_ic_names,
1091 tape_cache: tape::TapeCache::new(),
1092 })
1093 }
1094}
1095
1096// ============================================================================
1097// `from_model` build stages. Each function is one numbered stage of the
1098// compile pipeline (its number matches the stage comment in
1099// [`ArrayCompiled::from_model`], which composes them in order); the bodies are
1100// extracted verbatim from the former inline implementation.
1101// ============================================================================
1102
1103/// (0) Reject, at BUILD, every operator this runtime cannot evaluate — the
1104/// open rewrite-target tier (`grad`/`div`/`laplacian`, a spatial `D`, a user op,
1105/// a typo), which the canonical pipeline requires ESD discretization to have
1106/// rewritten before the simulator sees it (esm-i7b), AND the evaluable-core ops
1107/// that have no rule in this evaluator because an earlier stage was supposed to
1108/// eliminate them (`skolem`, `rank`, `distinct`, `argmin`, `argmax`, `ic`,
1109/// `enum`, `table_lookup`, `apply_expression_template`).
1110///
1111/// The second half is [`check_evaluable`]'s `is_evaluable_op` gate, and it is
1112/// here because [`eval_op`]'s backstop for an ungated op is `unreachable!` — a
1113/// PANIC, which is what an author got (exit 101, no diagnostic) from a document
1114/// that `esm validate` had just passed. The public `eval_expression` gated;
1115/// this path did not, and `hoist_static_observeds` walked straight into the
1116/// backstop. Gating the whole model once, at the single funnel every
1117/// array-runtime build passes through, closes the class rather than the one op
1118/// that was reported: any of the nine now raises `unevaluable_operator` naming
1119/// itself.
1120///
1121/// Ordering matters and is already right: `materialize_vi_outputs_to_data`,
1122/// `materialize_derived_extents` and `strip_value_invention` run BEFORE the
1123/// staged build, so by the time this sees the model an arg-witness output has
1124/// become `const` data and a skolem or derived-set producer has been dropped —
1125/// what remains is genuinely unevaluable.
1126fn reject_unlowered_spatial_ops(model: &Model) -> Result<(), CompileError> {
1127 for eq in &model.equations {
1128 // The `ic` exemption is for the LHS and only the LHS (esm-spec
1129 // §11.4.1 puts `ic` there). Applying it to the RHS as well let an
1130 // `ic`-headed RHS — the one shape that is NOT an initial-condition
1131 // statement — walk past the gate and reach the evaluator, where it is
1132 // an ordinary unevaluable-core op with no rule.
1133 check_evaluable_side(&eq.lhs)?;
1134 check_evaluable(&eq.rhs)?;
1135 }
1136 for var in model.variables.values() {
1137 let mut failure = None;
1138 var.for_each_expression(&mut |expr| {
1139 if failure.is_none()
1140 && let Err(e) = check_evaluable(expr)
1141 {
1142 failure = Some(e);
1143 }
1144 });
1145 if let Some(e) = failure {
1146 return Err(e);
1147 }
1148 }
1149 Ok(())
1150}
1151
1152/// [`check_evaluable`] over the LEFT side of an equation, with the ONE
1153/// structural wrapper the array path consumes rather than evaluates unwrapped
1154/// first: an `ic` LHS (esm-spec §11.4). `ic` is evaluable-core and has no `eval_op` arm —
1155/// correctly, since initial-condition assembly reads the equation and the
1156/// evaluator never sees the node — so gating it as an ordinary expression would
1157/// reject every document that states an initial condition. Its OPERAND is
1158/// checked, so nothing inside it escapes the gate.
1159fn check_evaluable_side(expr: &Expr) -> Result<(), CompileError> {
1160 if let Expr::Operator(node) = expr
1161 && node.op == "ic"
1162 {
1163 for arg in &node.args {
1164 check_evaluable(arg)?;
1165 }
1166 return Ok(());
1167 }
1168 check_evaluable(expr)
1169}
1170
1171/// A set of names, hashed with Fx, as the free-variable check credits them
1172/// from the expressions and probes them at every variable reference.
1173type NameSet = HashSet<String, rustc_hash::FxBuildHasher>;
1174
1175/// (0b) Reject a reference to a variable that is bound in NONE of the model's
1176/// binding categories. Without it a typo'd or undeclared bare name falls
1177/// through [`lookup_variable`]'s final arm, which fails closed only at the
1178/// drain, instead of failing loudly at build time. The error is
1179/// `InterpreterBuildError` / `Unknown variable '{name}' referenced in
1180/// expression`.
1181///
1182/// **ONE gate, and EVERY route calls it** (CONFORMANCE_SPEC §5.23). This is not
1183/// a compile-path-local check: [`crate::prepare::run_prepare`] — the build
1184/// pipeline any ingesting document and every `esm simulate` of a relational
1185/// document takes — calls it too, before it evaluates a single observed. A
1186/// second copy is exactly how one route came to reach the check and the other
1187/// to skip it (finding F24, then F25), so this is a shared function rather than
1188/// a pattern repeated per caller. `extra_bound` carries the names the CALLER's
1189/// evaluation scope binds outside the model's own declarations.
1190///
1191/// The bound set MIRRORS [`lookup_variable`]'s runtime resolution scope so a
1192/// legitimately runtime-bound name is never rejected — a false positive here
1193/// would reject a valid model, which is strictly worse than the silent-NaN it
1194/// closes:
1195///
1196/// * `t` — the independent variable (the array path only supports `t`);
1197/// * every declared model variable — state / parameter / observed, the keys of
1198/// `model.variables` (a discrete variable is rejected earlier, but its key is
1199/// still credited);
1200/// * every equation LHS-defined target (a name defined by an equation even if
1201/// not carried in `variables`);
1202/// * `_var` (§6.4 operator placeholder) and the document `index_sets` axis
1203/// names;
1204/// * spatial-coordinate symbols — the free symbols of every `ic` RHS (§11.4
1205/// defines these to BE coordinate expressions) and every spatial-op `dim`;
1206/// * loop / index binders introduced anywhere in the equation — `faq` /
1207/// `makearray` `output_idx` & `ranges`, bare `index(array, i…)` subscript
1208/// positions, an `integral` `int_var`, an argmin/argmax `arg`, and
1209/// `apply_expression_template` `bindings` keys. Collected over the WHOLE
1210/// equation (both sides unioned) so a stencil offset `index(u, i+1)` sees the
1211/// `i` bound as a bare position elsewhere in the same equation, and so a
1212/// symbol bound on the LHS is in scope on the RHS.
1213///
1214/// DELIBERATELY CONSERVATIVE SKIPS — a construct here is treated as BOUND, never
1215/// rejected, because a genuine typo in it is not provably distinguishable at
1216/// build time from a legitimate runtime-bound name:
1217///
1218/// * a DOTTED name (`A.b`) — a qualified cross-namespace reference or, on the
1219/// coupled/flatten path, an external *forcing* channel name (`M.src`,
1220/// `Box.scale`) that is UNDECLARED by design and resolved at runtime through
1221/// the forcing buffer (see [`lookup_variable`]'s final `forcing` arm and the
1222/// `segmented_refresh_solve` / `refresh_conformance` fixtures, which strip the
1223/// loader-fed `discrete` declarations precisely so these resolve as bare —
1224/// but post-flatten DOTTED — forcing names). A dotted typo is genuinely
1225/// indistinguishable from a dotted forcing name, so it is skipped;
1226/// * any name used as the HEAD of an `index(name, …)` op ANYWHERE in the model
1227/// — an array-valued leaf. On a single-model (`from_file`) path a loader-fed
1228/// forcing FIELD stays BARE and undeclared (the `observed_cadence_tier`
1229/// fixture's `f`), and a forcing field is always array-valued and read via
1230/// `index` (a SCALAR forcing goes through `params`/`set_params`, not the
1231/// buffer). Crediting every `index` head therefore keeps a bare forcing field
1232/// in scope; the residual conservative cost is that a typo appearing *only* as
1233/// an `index` head is not caught (reported as a deliberate skip);
1234/// * an `ic` equation — its RHS is a coordinate expression resolved at `u0`
1235/// build time by the field evaluator against grid geometry, not by the
1236/// per-cell RHS oracle; its free symbols are already credited as coordinates
1237/// above;
1238/// * a builtin function name spelled as a bare leaf (`exp`, `min`, …).
1239pub(crate) fn check_free_variables(
1240 model: &Model,
1241 index_sets: &HashMap<String, IndexSet>,
1242 extra_bound: &[String],
1243) -> Result<(), CompileError> {
1244 // ---- Build the bound set. ------------------------------------------------
1245 // The names the model and its caller declare, borrowed; `bound` holds what
1246 // the expressions themselves credit. A name is bound when either has it.
1247 let mut declared: HashSet<&str, rustc_hash::FxBuildHasher> = HashSet::default();
1248 declared.insert("t");
1249 declared.insert("_var");
1250 declared.extend(model.variables.keys().map(String::as_str));
1251 declared.extend(index_sets.keys().map(String::as_str));
1252 // Names the CALLER's evaluation scope binds that this model does not declare
1253 // — a build-pipeline caller's const arrays / provider slabs / member-factor
1254 // columns. Widening the bound set can only ever prevent a false positive,
1255 // which is this gate's cardinal requirement; the compiled path passes none.
1256 declared.extend(extra_bound.iter().map(String::as_str));
1257 let mut bound: NameSet = NameSet::default();
1258 // Each checked equation's binders, collected in this pass over the
1259 // equations rather than in a second one.
1260 let mut eq_binders: Vec<NameSet> = Vec::with_capacity(model.equations.len());
1261 for eq in &model.equations {
1262 if let Some(v) = equation_defined_var(&eq.lhs) {
1263 bound.insert(v);
1264 }
1265 // §11.4: an `ic` RHS's free symbols name spatial coordinates that are
1266 // implicitly in scope (e.g. the ignition front `psi(x)` over the bare
1267 // coordinate `x`). Spatial-op `dim`s name the same axes.
1268 // And array-valued leaves (potential bare forcing FIELDS) — see the
1269 // doc note on `collect_dims_and_index_heads`.
1270 let mut binders = NameSet::default();
1271 if is_ic_lhs(&eq.lhs) {
1272 collect_free_bare_symbols(&eq.rhs, &mut bound);
1273 collect_dims_and_index_heads(&eq.lhs, &mut bound);
1274 collect_dims_and_index_heads(&eq.rhs, &mut bound);
1275 } else {
1276 collect_dims_heads_and_binders(&eq.lhs, &mut bound, &mut binders, None);
1277 collect_dims_heads_and_binders(&eq.rhs, &mut bound, &mut binders, Some(false));
1278 }
1279 eq_binders.push(binders);
1280 }
1281 for var in model.variables.values() {
1282 var.for_each_expression(&mut |expr| collect_dims_and_index_heads(expr, &mut bound));
1283 }
1284
1285 // ---- Check every equation (skipping `ic`) and observed expression. -------
1286 // Each check sees the shared bound set plus its own expression's binders,
1287 // added for the check and taken back out after it (see `with_binders`).
1288 let mut binders: NameSet = NameSet::default();
1289 for (eq, mut eq_binders) in model.equations.iter().zip(eq_binders) {
1290 if is_ic_lhs(&eq.lhs) {
1291 continue;
1292 }
1293 with_binders(&mut bound, &mut eq_binders, |scope| {
1294 let is_bound = |name: &str| declared.contains(name) || scope.contains(name);
1295 check_expr_free_vars(&eq.lhs, &is_bound)?;
1296 check_expr_free_vars(&eq.rhs, &is_bound)
1297 })?;
1298 }
1299 for var in model.variables.values() {
1300 let mut failure = None;
1301 var.for_each_expression(&mut |expr| {
1302 if failure.is_some() {
1303 return;
1304 }
1305 collect_binders(expr, &mut binders);
1306 if let Err(e) = with_binders(&mut bound, &mut binders, |scope| {
1307 check_expr_free_vars(expr, &|name: &str| {
1308 declared.contains(name) || scope.contains(name)
1309 })
1310 }) {
1311 failure = Some(e);
1312 }
1313 });
1314 if let Some(e) = failure {
1315 return Err(e);
1316 }
1317 }
1318 Ok(())
1319}
1320
1321/// Run `check` against `bound` widened by `binders`, then restore `bound`:
1322/// only the binders it did not already hold are inserted, and exactly those
1323/// are removed again. `binders` is left empty for the next expression. The
1324/// cost is the expression's own binders, not the size of the bound set, so
1325/// the whole check stays linear in the model.
1326fn with_binders<T>(
1327 bound: &mut NameSet,
1328 binders: &mut NameSet,
1329 check: impl FnOnce(&NameSet) -> T,
1330) -> T {
1331 let mut added: Vec<String> = Vec::new();
1332 for name in binders.drain() {
1333 if !bound.contains(&name) {
1334 bound.insert(name.clone());
1335 added.push(name);
1336 }
1337 }
1338 let out = check(bound);
1339 for name in &added {
1340 bound.remove(name);
1341 }
1342 out
1343}
1344
1345/// Is this LHS an initial-condition marker (`{"op": "ic", …}`)?
1346fn is_ic_lhs(lhs: &Expr) -> bool {
1347 matches!(lhs, Expr::Operator(op) if op.op == "ic")
1348}
1349
1350/// The index / integration symbols a single node BINDS for its body (mirrors
1351/// `structural.rs` `bound_index_symbols`): `output_idx` / `ranges` keys, an
1352/// `integral` `int_var`, an argmin/argmax `arg`, the BARE subscript positions of
1353/// an `index(array, i…)` node, and `apply_expression_template` `bindings` keys.
1354fn node_binders(node: &ExpressionNode, out: &mut NameSet) {
1355 node_loop_binders(node, out);
1356 if node.op == "index" {
1357 // Only a BARE position (`index(u, i)`) is a binder; an index EXPRESSION
1358 // (`index(u, i+1)`) is a USE of a symbol bound elsewhere and is checked.
1359 for arg in node.args.iter().skip(1) {
1360 if let Expr::Variable(name) = arg {
1361 out.insert(name.clone());
1362 }
1363 }
1364 }
1365}
1366
1367/// [`node_binders`] less the bare `index` subscripts: the names a node binds
1368/// for its body and nothing else.
1369fn node_loop_binders(node: &ExpressionNode, out: &mut NameSet) {
1370 if let Some(idx) = &node.output_idx {
1371 out.extend(idx.iter().cloned());
1372 }
1373 if let Some(ranges) = &node.ranges {
1374 out.extend(ranges.keys().cloned());
1375 }
1376 if let Some(v) = &node.int_var {
1377 out.insert(v.clone());
1378 }
1379 if let Some(a) = &node.arg {
1380 out.insert(a.clone());
1381 }
1382 if let Some(bindings) = &node.bindings {
1383 out.extend(bindings.keys().cloned());
1384 }
1385}
1386
1387/// Union every binder introduced anywhere in the subtree (whole-tree, like
1388/// `structural.rs` `collect_bound_symbols`). Widening the bound set only ever
1389/// prevents a false positive, which is the cardinal requirement here.
1390fn collect_binders(expr: &Expr, out: &mut NameSet) {
1391 if let Expr::Operator(node) = expr {
1392 node_binders(node, out);
1393 node.for_each_child(&mut |child| collect_binders(child, out));
1394 }
1395}
1396
1397/// Collect every free BARE (non-dotted, non-builtin) symbol in the subtree —
1398/// used to credit an `ic` RHS's coordinate symbols into the bound set.
1399fn collect_free_bare_symbols(expr: &Expr, out: &mut NameSet) {
1400 match expr {
1401 Expr::Variable(name) if !name.contains('.') && !is_builtin_function_name(name) => {
1402 out.insert(name.clone());
1403 }
1404 Expr::Operator(node) => {
1405 node.for_each_child(&mut |child| collect_free_bare_symbols(child, out));
1406 }
1407 _ => {}
1408 }
1409}
1410
1411/// Collect, in one walk of the subtree:
1412///
1413/// - the `dim` axis of every node that carries one, regardless of `op`
1414/// (esm-spec §4.9.1 (ii), as revised): a coordinate axis is resolved
1415/// STRUCTURALLY by the presence of a `dim` field, not by a hardcoded
1416/// spatial-operator name list (the sugar ops carry no privileged status).
1417/// Defensive: the array path rejects unlowered spatial ops earlier, but a
1418/// coordinate an `ic` RHS shares with a node's `dim` stays creditable;
1419/// - the bare name at the HEAD (first arg) of every `index(name, …)` op — an
1420/// array-valued leaf (a declared state/observed, or a bare, undeclared,
1421/// loader-fed forcing FIELD read at runtime through the forcing buffer).
1422/// Crediting these keeps a legitimate bare forcing field in scope.
1423fn collect_dims_and_index_heads(expr: &Expr, out: &mut NameSet) {
1424 if let Expr::Operator(node) = expr {
1425 if let Some(dim) = &node.dim {
1426 out.insert(dim.clone());
1427 }
1428 if node.op == "index"
1429 && let Some(Expr::Variable(name)) = node.args.first()
1430 && !name.contains('.')
1431 {
1432 out.insert(name.clone());
1433 }
1434 node.for_each_child(&mut |child| collect_dims_and_index_heads(child, out));
1435 }
1436}
1437
1438/// [`collect_dims_and_index_heads`] into `credited` and, in the same walk, the
1439/// binders into `binders`: as [`collect_binders`] for a left-hand side
1440/// (`rhs: None`); for a right-hand side (`rhs: Some(enclosed)`), the same less
1441/// one case: a bare `index(array, i)` subscript that no enclosing node binds
1442/// and the left-hand side does not bind either. There it is a READ of `i` with
1443/// nothing in scope to give it a value — `d ~ index(faq{i}(…), i)` reads the
1444/// `i` of no loop, since the `faq`'s own `i` is bound only inside it — so it is
1445/// checked rather than credited. Under any node that carries binders (`faq`,
1446/// `makearray`, `integral`, a template call…) every bare subscript still
1447/// counts as a binder.
1448fn collect_dims_heads_and_binders(
1449 expr: &Expr,
1450 credited: &mut NameSet,
1451 binders: &mut NameSet,
1452 rhs: Option<bool>,
1453) {
1454 let Expr::Operator(node) = expr else {
1455 return;
1456 };
1457 if let Some(dim) = &node.dim {
1458 credited.insert(dim.clone());
1459 }
1460 if node.op == "index"
1461 && let Some(Expr::Variable(name)) = node.args.first()
1462 && !name.contains('.')
1463 {
1464 credited.insert(name.clone());
1465 }
1466 let child_rhs = match rhs {
1467 None => {
1468 node_binders(node, binders);
1469 None
1470 }
1471 Some(enclosed) => {
1472 let carries = node.output_idx.is_some()
1473 || node.ranges.is_some()
1474 || node.int_var.is_some()
1475 || node.arg.is_some()
1476 || node.bindings.is_some();
1477 if enclosed || node.op != "index" {
1478 node_binders(node, binders);
1479 }
1480 Some(enclosed || carries)
1481 }
1482 };
1483 node.for_each_child(&mut |child| {
1484 collect_dims_heads_and_binders(child, credited, binders, child_rhs)
1485 });
1486}
1487
1488/// Reject the first bare (non-dotted) variable reference bound in none of the
1489/// categories in `scope`, as `Unknown variable '{name}' referenced in
1490/// expression`. The full expression-bearing
1491/// child set is descended via [`ExpressionNode::for_each_child`] (args plus the
1492/// sidecar fields), so a reference hidden in an aggregate body, filter, integral
1493/// bound, table axis, aggregate key, or template binding is not missed. A `fn`
1494/// op's callee lives in `node.name` (not a child), so it is never mistaken for a
1495/// variable.
1496fn check_expr_free_vars(expr: &Expr, scope: &dyn Fn(&str) -> bool) -> Result<(), CompileError> {
1497 match expr {
1498 Expr::Variable(name) => {
1499 // A dotted name is a qualified / forcing reference resolved at
1500 // runtime; builtins and derivative markers are always valid.
1501 if name.contains('.') || is_builtin_function_name(name) || name.starts_with("d(") {
1502 return Ok(());
1503 }
1504 if scope(name) {
1505 return Ok(());
1506 }
1507 Err(CompileError::build_err(format!(
1508 "Unknown variable '{name}' referenced in expression"
1509 )))
1510 }
1511 Expr::Operator(node) => {
1512 let mut first_err: Option<CompileError> = None;
1513 node.for_each_child(&mut |child| {
1514 if first_err.is_none()
1515 && let Err(e) = check_expr_free_vars(child, scope)
1516 {
1517 first_err = Some(e);
1518 }
1519 });
1520 match first_err {
1521 Some(e) => Err(e),
1522 None => Ok(()),
1523 }
1524 }
1525 _ => Ok(()),
1526 }
1527}
1528
1529/// One state-free observed in the build-time scope a field `ic` reads.
1530#[derive(Clone, Debug)]
1531pub(super) struct IcScopeDef {
1532 pub(super) name: String,
1533 pub(super) body: Expr,
1534 /// The other scope definitions the body reads. Each sorts before this one
1535 /// in [`IcScope::defs`].
1536 pub(super) deps: Vec<String>,
1537 /// Every name the body reads that must be bound before it is evaluated:
1538 /// the observeds it reads, in the scope or not, and the provider-served
1539 /// fields. An absent one is never handed to the evaluator, which could
1540 /// only walk the body per cell (and a strict compiler refuse it) before
1541 /// failing on it; the definition is left out of the scope instead.
1542 pub(super) needs: Vec<String>,
1543}
1544
1545/// The STATE-FREE observeds a field `ic` right-hand side reads (esm-spec
1546/// §6.6.5 "Build-time evaluation scope"), captured by [`capture_ic_scope`].
1547#[derive(Clone, Debug, Default)]
1548pub(super) struct IcScope {
1549 /// The definitions in the dependency cone of the `ic` right-hand sides,
1550 /// in dependency order: one forward pass evaluates each after everything
1551 /// it reads.
1552 pub(super) defs: Vec<IcScopeDef>,
1553 /// The scope definitions an `ic` right-hand side names directly — where
1554 /// the cone a provider-served field may cut short starts.
1555 pub(super) roots: Vec<String>,
1556 /// The observeds in that cone left out of the scope at compile time, with
1557 /// the reason, for the diagnostic of an `ic` that reads one.
1558 pub(super) excluded: Vec<(String, String)>,
1559}
1560
1561/// The names `expr` reads FREE: every bare variable reference not bound by an
1562/// enclosing node (a `faq`'s `output_idx` / `ranges` keys, an `integral`'s
1563/// variable, …). A loop index is no read of the observed or state that
1564/// happens to share its name.
1565fn collect_free_names(expr: &Expr, bound: &mut Vec<String>, out: &mut HashSet<String>) {
1566 match expr {
1567 Expr::Variable(name) if !bound.contains(name) => {
1568 out.insert(name.clone());
1569 }
1570 Expr::Operator(node) => {
1571 let mut own = NameSet::default();
1572 node_loop_binders(node, &mut own);
1573 let depth = bound.len();
1574 bound.extend(own);
1575 node.for_each_child(&mut |child| collect_free_names(child, bound, out));
1576 bound.truncate(depth);
1577 }
1578 _ => {}
1579 }
1580}
1581
1582pub(super) fn free_names(expr: &Expr) -> HashSet<String> {
1583 let mut out = HashSet::new();
1584 collect_free_names(expr, &mut Vec::new(), &mut out);
1585 out
1586}
1587
1588/// Memoized walk over the observed definitions marking each one that reads
1589/// state or `t`, directly or through another observed — linear in the
1590/// definitions' reads.
1591struct StateTaint<'a> {
1592 defs: &'a BTreeMap<String, (Expr, HashSet<String>)>,
1593 is_state: &'a dyn Fn(&str) -> bool,
1594 /// `None` while a definition is being visited (a cycle reads it as clean:
1595 /// the cycle itself keeps it out of the scope), then why it is tainted, if
1596 /// it is.
1597 marks: HashMap<&'a str, Option<Option<String>>>,
1598}
1599
1600impl<'a> StateTaint<'a> {
1601 fn cause(&mut self, name: &'a str) -> Option<String> {
1602 match self.marks.get(name) {
1603 Some(Some(cause)) => return cause.clone(),
1604 Some(None) => return None,
1605 None => {}
1606 }
1607 self.marks.insert(name, None);
1608 let defs = self.defs;
1609 let reads = &defs[name].1;
1610 let mut cause = reads
1611 .iter()
1612 .filter(|n| (self.is_state)(n))
1613 .min()
1614 .map(|n| format!("reads '{n}'"));
1615 if cause.is_none() {
1616 let mut deps: Vec<&'a String> = reads
1617 .iter()
1618 .filter(|n| n.as_str() != name && defs.contains_key(n.as_str()))
1619 .collect();
1620 deps.sort();
1621 for dep in deps {
1622 if let Some(inner) = self.cause(dep) {
1623 cause = Some(format!("reads '{dep}', which {inner}"));
1624 break;
1625 }
1626 }
1627 }
1628 self.marks.insert(name, Some(cause.clone()));
1629 cause
1630 }
1631}
1632
1633/// The STATE-FREE observed definitions the `ic` right-hand sides read
1634/// (esm-spec §6.6.5 "Build-time evaluation scope").
1635///
1636/// A state-free observed — one whose defining expression closes over
1637/// parameters, inline `const` data, provider-served fields and other
1638/// state-free observeds, with no state and no `t` — is resolvable BEFORE the
1639/// simulation runs; the build already materializes exactly this class of field
1640/// as a setup array for `BuildInspection`. Admitting it as an `ic` RHS
1641/// (esm-spec §11.4.1) is the same evaluator reached from one more place, and it
1642/// is what lets a column test seed `u` from a `const` gather instead of minting
1643/// a per-regime rewrite-rule library.
1644///
1645/// State-freedom is TRANSITIVE: an observed that reads another observed
1646/// reading state (`Mz_conv` ← `div_h` ← `conv_x` ← `u`) is not state-free
1647/// either, and a state is any unknown that is not observed — an eliminated
1648/// algebraic one too. Only the cone the `ic` right-hand sides actually read is
1649/// kept, as the Python binding's `_buildtime_observed_arrays` keeps it, so an
1650/// unrelated observed is never evaluated; it is ordered so one pass evaluates
1651/// each definition after everything it reads. A definition on a dependency
1652/// cycle has no such place, and is left out.
1653///
1654/// Captured HERE, before stage (6) MOVES each observed's body out of the model,
1655/// and evaluated lazily by [`ArrayCompiled::resolve_field_ics`] — only when the
1656/// document actually has a field `ic` to fold.
1657fn capture_ic_scope(
1658 model: &Model,
1659 observed_names: &[String],
1660 forcing_names: &IndexMap<String, Option<Vec<usize>>>,
1661 state_names: &IndexMap<String, VarShape>,
1662 eliminated: &HashSet<String>,
1663) -> IcScope {
1664 let ic_reads: HashSet<String> = model
1665 .equations
1666 .iter()
1667 .filter(|eq| is_ic_lhs(&eq.lhs))
1668 .flat_map(|eq| free_names(&eq.rhs))
1669 .collect();
1670 if ic_reads.is_empty() {
1671 return IcScope::default();
1672 }
1673 let observed: HashSet<&str> = observed_names.iter().map(String::as_str).collect();
1674 let defs: BTreeMap<String, (Expr, HashSet<String>)> = observed_bodies(model)
1675 .into_iter()
1676 .filter(|(name, _)| observed.contains(name.as_str()))
1677 .map(|(name, body)| {
1678 let reads = free_names(&body);
1679 (name, (body, reads))
1680 })
1681 .collect();
1682 let is_state = |n: &str| n == "t" || state_names.contains_key(n) || eliminated.contains(n);
1683 let mut taint = StateTaint {
1684 defs: &defs,
1685 is_state: &is_state,
1686 marks: HashMap::new(),
1687 };
1688
1689 // The state-free cone of the `ic` reads.
1690 let mut excluded: Vec<(String, String)> = Vec::new();
1691 let mut roots: Vec<String> = ic_reads
1692 .iter()
1693 .filter(|n| defs.contains_key(n.as_str()))
1694 .cloned()
1695 .collect();
1696 roots.sort();
1697 let mut wanted: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
1698 let mut stack: Vec<&str> = roots.iter().map(String::as_str).collect();
1699 while let Some(name) = stack.pop() {
1700 if wanted.contains_key(name) || excluded.iter().any(|(n, _)| n == name) {
1701 continue;
1702 }
1703 if let Some(cause) = taint.cause(name) {
1704 excluded.push((name.to_string(), format!("is not state-free: it {cause}")));
1705 continue;
1706 }
1707 let deps: Vec<&str> = defs[name]
1708 .1
1709 .iter()
1710 .map(String::as_str)
1711 .filter(|n| *n != name && defs.contains_key(*n))
1712 .collect();
1713 stack.extend(&deps);
1714 wanted.insert(name, deps);
1715 }
1716 // A clean definition can still read a tainted one through a cycle the
1717 // memoized walk saw only half of; keep only reads of the cone.
1718 for deps in wanted.values_mut() {
1719 deps.retain(|d| !excluded.iter().any(|(n, _)| n == d));
1720 }
1721
1722 // Dependency order (Kahn), ties by name so the order is deterministic.
1723 let mut pending: BTreeMap<&str, usize> = BTreeMap::new();
1724 let mut readers: HashMap<&str, Vec<&str>> = HashMap::new();
1725 for (name, deps) in &wanted {
1726 pending.insert(name, deps.len());
1727 for dep in deps {
1728 readers.entry(dep).or_default().push(name);
1729 }
1730 }
1731 let mut ready: std::collections::BTreeSet<&str> = pending
1732 .iter()
1733 .filter(|(_, n)| **n == 0)
1734 .map(|(name, _)| *name)
1735 .collect();
1736 let mut order: Vec<&str> = Vec::with_capacity(wanted.len());
1737 while let Some(name) = ready.pop_first() {
1738 pending.remove(name);
1739 order.push(name);
1740 for reader in readers.get(name).into_iter().flatten() {
1741 if let Some(n) = pending.get_mut(reader) {
1742 *n -= 1;
1743 if *n == 0 {
1744 ready.insert(reader);
1745 }
1746 }
1747 }
1748 }
1749 if !pending.is_empty() {
1750 let cycle: Vec<&str> = pending.keys().copied().collect();
1751 let reason = format!(
1752 "is on, or reads, a dependency cycle among the observeds {}, so no evaluation \
1753 order exists",
1754 cycle.join(", ")
1755 );
1756 excluded.extend(cycle.iter().map(|n| (n.to_string(), reason.clone())));
1757 }
1758
1759 let defs = order
1760 .into_iter()
1761 .map(|name| {
1762 let (body, reads) = &defs[name];
1763 let mut deps: Vec<String> = wanted[name].iter().map(|d| d.to_string()).collect();
1764 deps.sort();
1765 let mut needs: Vec<String> = reads
1766 .iter()
1767 .filter(|n| {
1768 n.as_str() != name
1769 && (defs.contains_key(n.as_str()) || forcing_names.contains_key(n.as_str()))
1770 })
1771 .cloned()
1772 .collect();
1773 needs.sort();
1774 IcScopeDef {
1775 name: name.to_string(),
1776 body: body.clone(),
1777 deps,
1778 needs,
1779 }
1780 })
1781 .collect();
1782 excluded.sort();
1783 IcScope {
1784 defs,
1785 roots,
1786 excluded,
1787 }
1788}
1789
1790/// Lower every SHAPED parameter whose value the document itself supplies —
1791/// INLINE ARRAY DATA (a row-major nested JSON array on its `default`, esm-spec
1792/// §6.3, and the `parameter_overrides` a test resolves into it, §6.6.2) or a
1793/// SCALAR broadcast over the declared `shape` — into the `const` observed
1794/// channel this runtime already evaluates.
1795///
1796/// A parameter's value reaches the RHS through the positional scalar `params`
1797/// vector, which has exactly one f64 per parameter: a whole column has nowhere
1798/// to live there. It does have somewhere to live as a `const`-op observed —
1799/// the established self-contained build-time array channel ([`eval_const`],
1800/// [`collect_const_factor_arrays`], the `ConstArrayScope` gather policy) — so
1801/// the parameter is re-declared as an observed defined by its own data. That is
1802/// exactly the lowering a column component had to author by hand as a per-regime
1803/// rewrite-rule library; doing it here is what lets one shared component carry
1804/// its profiles inline.
1805///
1806/// A SCALAR on a shaped parameter takes the SAME channel, because esm-spec §6.3
1807/// gives it the same meaning — "the one value applies to every element" — and a
1808/// scalar `params` slot cannot express that: every per-cell read of the name
1809/// (`index(ramp, k)`, and the gather `index_array_leaves_by_loops` inserts for a
1810/// bare whole-array RHS) indexes a scalar, which [`eval_index`] resolves to
1811/// `NaN`. The solver then fails at `t = 0` naming no parameter at all
1812/// (EarthSciML/EarthSciAST#219). Broadcasting the scalar here is the shaped
1813/// UNKNOWN's rule — [`build_slot_tables`] already writes one scalar `default`
1814/// into every per-cell slot — applied to the parameter side.
1815///
1816/// A 0-D parameter is untouched, and so is one whose value comes from OUTSIDE
1817/// the document — an `update` binding (the loader / handler forcing seam) or a
1818/// caller-supplied `vi_arrays` entry — because that channel owns the buffer and
1819/// its data beats a declared default (the precedence the Julia reference's
1820/// `_register_inline_array_parameters` takes). So this is a no-op for every
1821/// document that does not declare a shaped parameter's value inline.
1822///
1823/// # Errors
1824///
1825/// [`CompileError::InterpreterBuildError`] when the authored array is ragged or
1826/// its extents do not match the declared `shape` resolved against the index-set
1827/// registry — the load-time shape check esm-spec §6.6.2 requires.
1828#[allow(clippy::type_complexity)]
1829fn lower_inline_array_parameters(
1830 model: &mut Model,
1831 index_sets: &HashMap<String, IndexSet>,
1832 external: Option<&HashMap<String, ArrayD<f64>>>,
1833 keep_at_runtime: bool,
1834) -> Result<(HashMap<String, (Vec<usize>, Vec<f64>)>, HashSet<String>), CompileError> {
1835 let mut lowered: Vec<(String, JsonValue)> = Vec::new();
1836 let mut dense: HashMap<String, (Vec<usize>, Vec<f64>)> = HashMap::new();
1837 let mut kept: HashSet<String> = HashSet::new();
1838 for (name, var) in &mut model.variables {
1839 if var.var_type != VariableType::Parameter {
1840 continue;
1841 }
1842 let Some(default) = var.default.as_ref() else {
1843 continue;
1844 };
1845 // `runtime_array_parameters`: a shaped parameter with a value and no
1846 // `update` stays a parameter, read at run time from the forcing buffer
1847 // (its value seeds the buffer), instead of becoming build-time data.
1848 // A caller's array counts: `esm_problem` binds it onto `default`.
1849 if keep_at_runtime
1850 && var.shape.as_deref().is_some_and(|s| !s.is_empty())
1851 && var.update.is_none()
1852 {
1853 kept.insert(name.clone());
1854 continue;
1855 }
1856 if !default.is_array() {
1857 // esm-spec §6.3 broadcast: one scalar over the whole declared grid.
1858 // Only for a parameter this document is the sole source of — a 0-D
1859 // one has no grid to broadcast over, an externally refreshed one is
1860 // filled by the forcing buffer, and a caller `vi_arrays` entry is
1861 // authoritative over any declared default.
1862 let declared = var.shape.as_deref().unwrap_or(&[]);
1863 if declared.is_empty()
1864 || var.update.is_some()
1865 || external.is_some_and(|a| a.contains_key(name))
1866 {
1867 continue;
1868 }
1869 // A shape that does not resolve (an unmaterialized derived set) has
1870 // no extent to broadcast over; leave the parameter alone so the
1871 // build's own extent checks report any real disagreement.
1872 let Some(want) = resolve_declared_shape(declared, index_sets) else {
1873 continue;
1874 };
1875 let Some(scalar) = default.as_scalar() else {
1876 continue;
1877 };
1878 let values = vec![scalar; want.iter().copied().product::<usize>().max(1)];
1879 var.var_type = VariableType::Unknown;
1880 var.default = None;
1881 lowered.push((name.clone(), dense_to_json(&want, &values)));
1882 dense.insert(name.clone(), (want, values));
1883 continue;
1884 }
1885 let (shape, values) = default.to_dense().map_err(|e| {
1886 CompileError::build_err(format!("parameter '{name}': {e} (esm-spec §6.3)"))
1887 })?;
1888 let declared = var.shape.as_deref().unwrap_or(&[]);
1889 if declared.is_empty() {
1890 return Err(CompileError::build_err(format!(
1891 "parameter '{name}' carries inline array data but declares no `shape`; only a \
1892 SHAPED variable takes a nested array (esm-spec §6.3)"
1893 )));
1894 }
1895 if let Some(want) = resolve_declared_shape(declared, index_sets)
1896 && want != shape
1897 {
1898 return Err(CompileError::build_err(format!(
1899 "parameter '{name}': inline array data has shape {shape:?}, which does not match \
1900 the declared shape {want:?} (esm-spec §6.6.2 — the array MUST match the \
1901 variable's declared shape after metaparameter folding)"
1902 )));
1903 }
1904 var.var_type = VariableType::Unknown;
1905 var.default = None;
1906 lowered.push((name.clone(), dense_to_json(&shape, &values)));
1907 dense.insert(name.clone(), (shape, values));
1908 }
1909 for (name, value) in lowered {
1910 // A bare-variable LHS makes it an OBSERVED (esm-spec §6.3.1), and a
1911 // `const` RHS makes it build-time data — the same pair
1912 // `rewrite_equation_to_const` produces for a materialized relational
1913 // output.
1914 model.equations.push(crate::types::Equation {
1915 lhs: Expr::Variable(name),
1916 rhs: Expr::operator(ExpressionNode {
1917 op: "const".to_string(),
1918 value: Some(value),
1919 ..Default::default()
1920 }),
1921 comment: None,
1922 });
1923 }
1924 Ok((dense, kept))
1925}
1926
1927/// The values [`lower_inline_array_parameters`] gives each shaped parameter
1928/// the document supplies — inline array data, or one scalar broadcast over
1929/// its declared grid — as dense arrays, without rewriting the model. The
1930/// build pipeline binds these, as the compile path does, so an observed it
1931/// evaluates can read them. A parameter whose data does not resolve to its
1932/// declared shape is left out: the compile path reports it.
1933#[cfg(not(target_arch = "wasm32"))]
1934pub(crate) fn inline_parameter_arrays(
1935 model: &Model,
1936 index_sets: &HashMap<String, IndexSet>,
1937) -> Vec<(String, ArrayD<f64>)> {
1938 let mut out = Vec::new();
1939 for (name, var) in &model.variables {
1940 if var.var_type != VariableType::Parameter {
1941 continue;
1942 }
1943 let Some(default) = var.default.as_ref() else {
1944 continue;
1945 };
1946 let declared = var.shape.as_deref().unwrap_or(&[]);
1947 if declared.is_empty() {
1948 continue;
1949 }
1950 let want = resolve_declared_shape(declared, index_sets);
1951 let (shape, values) = if default.is_array() {
1952 match default.to_dense() {
1953 Ok((shape, values)) if want.as_ref().is_none_or(|w| *w == shape) => (shape, values),
1954 _ => continue,
1955 }
1956 } else {
1957 if var.update.is_some() {
1958 continue;
1959 }
1960 let (Some(want), Some(scalar)) = (want, default.as_scalar()) else {
1961 continue;
1962 };
1963 let n = want.iter().copied().product::<usize>().max(1);
1964 (want, vec![scalar; n])
1965 };
1966 let prec = crate::precision::active();
1967 let values: Vec<f64> = values.into_iter().map(|v| prec.round(v)).collect();
1968 if let Ok(a) = ArrayD::from_shape_vec(IxDyn(&shape), values) {
1969 out.push((name.clone(), a));
1970 }
1971 }
1972 out
1973}
1974
1975/// Re-nest a row-major dense buffer into the nested JSON array a `const` node
1976/// carries. The inverse of [`crate::types::InlineValue::to_dense`].
1977fn dense_to_json(shape: &[usize], values: &[f64]) -> JsonValue {
1978 fn build(shape: &[usize], values: &[f64]) -> JsonValue {
1979 let Some((&axis, rest)) = shape.split_first() else {
1980 return serde_json::Number::from_f64(values[0])
1981 .map(JsonValue::Number)
1982 .unwrap_or(JsonValue::Null);
1983 };
1984 let stride = rest.iter().product::<usize>().max(1);
1985 JsonValue::Array(
1986 (0..axis)
1987 .map(|i| build(rest, &values[i * stride..(i + 1) * stride]))
1988 .collect(),
1989 )
1990 }
1991 build(shape, values)
1992}
1993
1994/// (1) Collect state / parameter / observed variables (sorted by name for a
1995/// deterministic build).
1996///
1997/// Every category is DERIVED (esm-spec §6.3.1), never read off a declared type:
1998/// an unknown is an ODE state or an observed according to the equation that
1999/// defines it, and a parameter is Brownian or discrete according to its
2000/// `update`. A Brownian parameter is refused (`unsupported_construct`), never
2001/// a silent drop, and so is a discrete one this backend cannot refresh.
2002#[allow(clippy::type_complexity)]
2003fn classify_variables<'m>(
2004 model: &'m Model,
2005 runtime_params: &HashSet<String>,
2006) -> Result<
2007 (
2008 Vec<&'m String>,
2009 Vec<&'m String>,
2010 Vec<(&'m String, &'m ModelVariable)>,
2011 Vec<(String, String)>,
2012 ),
2013 CompileError,
2014> {
2015 let mut state_vars: Vec<&'m String> = Vec::new();
2016 let mut param_vars: Vec<&'m String> = Vec::new();
2017 let mut observed_vars: Vec<(&'m String, &'m ModelVariable)> = Vec::new();
2018 // The DATA-FED parameters this classification routes to the forcing
2019 // channel, each with the `data_sources` key its `update` names
2020 // (esm-spec §9.6.6 `data_source_unbound`). Recorded here because this is
2021 // the point that decides a parameter's fate, so the list carries exactly
2022 // the names the runtime will look up in the forcing buffer — no second
2023 // flatten, and no re-derivation of the document's namespacing.
2024 let mut data_fed: Vec<(String, String)> = Vec::new();
2025
2026 let class = crate::classification::VariableRoles::of(model);
2027
2028 let mut var_keys: Vec<&String> = model.variables.keys().collect();
2029 var_keys.sort_unstable();
2030 for name in var_keys {
2031 let var = &model.variables[name];
2032 match var.var_type {
2033 // An ALGEBRAIC unknown joins the states: it is solved for rather
2034 // than eliminated, exactly as it was before 1.0.0 when it was
2035 // declared `state` and pinned by an expression-LHS equation.
2036 VariableType::Unknown => {
2037 if class.is_observed(name) {
2038 observed_vars.push((name, var));
2039 } else {
2040 state_vars.push(name);
2041 }
2042 }
2043 VariableType::Parameter => {
2044 // A shaped parameter left as a run-time input
2045 // (`runtime_array_parameters`) is read from the forcing
2046 // buffer, as an externally refreshed one is.
2047 if runtime_params.contains(name) {
2048 observed_vars.push((name, var));
2049 continue;
2050 }
2051 if class.is_brownian(name) {
2052 return Err(CompileError::UnsupportedConstruct {
2053 construct: crate::compile_error::WIENER_NOISE,
2054 evaluator: crate::compile_error::ARRAY_EVALUATOR,
2055 detail: format!("parameter '{name}'"),
2056 });
2057 }
2058 if class.is_discrete_parameter(name) {
2059 // A parameter refreshed from OUTSIDE the model — a `from`
2060 // binding to a data source, or a registered `handler` — IS
2061 // the forcing seam: its value is written into the forcing
2062 // buffer between segments (by `crate::provider`'s refresh
2063 // executor, or by the host through `forcing_handle`) and
2064 // read back at the RHS. That is exactly the slot an
2065 // observed with no defining rule occupies, so it goes
2066 // there: `build_observed_rules` lowers only observeds that
2067 // HAVE a definition, and `lookup_variable`'s forcing arm
2068 // resolves the name at evaluation time.
2069 if externally_refreshed(var) {
2070 if let Some(source) = var.update.as_ref().and_then(data_feed_source) {
2071 data_fed.push((name.clone(), source.to_string()));
2072 }
2073 observed_vars.push((name, var));
2074 continue;
2075 }
2076 // What is left is a parameter that recomputes itself from a
2077 // symbolic `expression` at each refresh, which needs event
2078 // machinery this backend does not have: esm-spec §5.4 makes
2079 // `schedule` / `condition` / `crossing` updates the 1.0.0
2080 // spelling of the discrete and continuous events that
2081 // wrote parameters. Binning it as a state (integrated) or a
2082 // plain parameter (frozen) would both be WRONG — and
2083 // silently so. Fail loudly instead, naming the rule kinds;
2084 // the document still VALIDATES, it just cannot be simulated
2085 // by this backend yet.
2086 let kinds: Vec<&str> = var
2087 .update
2088 .iter()
2089 .flat_map(|spec| spec.rules())
2090 .filter(|rule| {
2091 rule.value()
2092 .is_some_and(|v| v.from.is_none() && v.handler.is_none())
2093 })
2094 .map(|rule| rule.kind())
2095 .collect();
2096 return Err(CompileError::UnsupportedFeatureError {
2097 feature: "discrete".to_string(),
2098 message: format!(
2099 "Rust array simulation backend does not yet support a discrete parameter that recomputes itself symbolically; parameter '{name}' carries an `expression` update of kind {} (esm-spec §5.4: the event-driven parameter update)",
2100 kinds
2101 .iter()
2102 .map(|k| format!("`{k}`"))
2103 .collect::<Vec<_>>()
2104 .join(", ")
2105 ),
2106 });
2107 }
2108 param_vars.push(name);
2109 }
2110 }
2111 }
2112 Ok((state_vars, param_vars, observed_vars, data_fed))
2113}
2114
2115/// (2) Infer shapes for state variables from all equation usages, then (2b)
2116/// overwrite with each state's declared array `shape` where one is declared
2117/// and resolves. The declared shape (index-set names resolved to sizes via
2118/// the document registry) is AUTHORITATIVE over usage inference (esm-spec
2119/// §11), matching the Python binding: a whole-array `D(state)` never
2120/// index-uses the state, so inference alone collapses it to a scalar; and an
2121/// observed's halo gather (`index(q, clamp(i±k))` inside an extended-grid
2122/// aggregate, e.g. the duo grid's `duo_extend`) index-uses the state at
2123/// offsets past its true extent, so inference alone WIDENS it past the grid.
2124/// Usage inference remains the fallback for states with no (resolvable)
2125/// declared shape.
2126/// `any_index`: whether the model has an `index` node at all; without one no
2127/// equation indexes a state, and every state's inferred shape is empty.
2128fn infer_state_shapes(
2129 model: &Model,
2130 state_vars: &[&String],
2131 index_sets: &HashMap<String, IndexSet>,
2132 any_index: bool,
2133) -> Result<HashMap<String, Vec<usize>>, CompileError> {
2134 let mut shape_map = if any_index {
2135 infer_shapes(state_vars, &model.equations)?
2136 } else {
2137 state_vars
2138 .iter()
2139 .map(|name| ((*name).clone(), Vec::new()))
2140 .collect()
2141 };
2142
2143 // (2a) A declared shape naming an index set the registry does not hold, on a
2144 // state no equation indexes, has no extent from either source. Laying it out
2145 // as one scalar slot would integrate a field the document says is arrayed, so
2146 // refuse it as Julia's tree-walk does (`E_TREEWALK_UNDECLARED_INDEX_SET`). A
2147 // name the registry holds but cannot size yet (an unmaterialized derived set)
2148 // is not this case, and neither is a state its equations index.
2149 for name in state_vars {
2150 let Some(decl) = model.variables.get(*name).and_then(|v| v.shape.as_ref()) else {
2151 continue;
2152 };
2153 let undeclared: Vec<&str> = decl
2154 .iter()
2155 .filter(|axis| !index_sets.contains_key(axis.as_str()))
2156 .map(String::as_str)
2157 .collect();
2158 if undeclared.is_empty() || shape_map.get(*name).is_some_and(|s| !s.is_empty()) {
2159 continue;
2160 }
2161 return Err(CompileError::build_err(format!(
2162 "state '{name}' declares shape {decl:?}, but index set(s) {undeclared:?} are not \
2163 declared in the document `index_sets` registry and no equation indexes the state, \
2164 so it has no extent; declare them, or inject the grid that does (esm-spec §6.3, \
2165 §9.7.10: a name still unresolved after injection is an error at the build)"
2166 )));
2167 }
2168
2169 // (2b) Declared shapes are authoritative wherever they resolve.
2170 for name in state_vars {
2171 if let Some(decl) = model.variables.get(*name).and_then(|v| v.shape.as_ref()) {
2172 if !decl.is_empty() {
2173 if let Some(resolved) = resolve_declared_shape(decl, index_sets) {
2174 shape_map.insert((*name).clone(), resolved);
2175 }
2176 }
2177 }
2178 }
2179 Ok(shape_map)
2180}
2181
2182/// (3) Partition state variables into `(final_states, eliminated, held_at_ic)`.
2183/// A state with a `D` equation is integrated; one defined by an algebraic
2184/// equation (but no `D`) is eliminated to an observed; one with neither (an
2185/// `ic`-only field) is carried at its ic with zero derivative — kept as a
2186/// state so its cells are enumerated and held constant.
2187fn partition_states(
2188 model: &Model,
2189 state_vars: &[&String],
2190) -> (Vec<String>, HashSet<String>, HashSet<String>) {
2191 // The scalar derivative targets borrow their names (one per state in a
2192 // scalar document); the rarer `faq` targets are owned.
2193 let mut scalar_targets: rustc_hash::FxHashSet<&str> = rustc_hash::FxHashSet::default();
2194 let mut faq_targets: HashSet<String> = HashSet::new();
2195 for eq in &model.equations {
2196 if let Some(name) = derivative_scalar_target(&eq.lhs) {
2197 scalar_targets.insert(name);
2198 }
2199 if let Some(DerivArrayop { var: name, .. }) = extract_derivative_faq(&eq.lhs, &eq.rhs) {
2200 faq_targets.insert(name);
2201 }
2202 }
2203 let is_derivative_target =
2204 |name: &str| scalar_targets.contains(name) || faq_targets.contains(name);
2205 let algebraic_defined = collect_algebraic_defined(&model.equations);
2206
2207 let mut final_states: Vec<String> = Vec::new();
2208 let mut eliminated: HashSet<String> = HashSet::new();
2209 let mut held_at_ic: HashSet<String> = HashSet::new();
2210 for name in state_vars {
2211 if is_derivative_target(name) {
2212 final_states.push((*name).clone());
2213 } else if algebraic_defined.contains(*name) {
2214 // No D equation, but an algebraic equation defines it.
2215 eliminated.insert((*name).clone());
2216 } else {
2217 // No D and no algebraic definition: hold at ic (zero derivative).
2218 final_states.push((*name).clone());
2219 held_at_ic.insert((*name).clone());
2220 }
2221 }
2222 (final_states, eliminated, held_at_ic)
2223}
2224
2225/// Flat state-vector tables built by [`build_slot_tables`] (stage 4),
2226/// mirroring the corresponding [`ArrayCompiled`] fields: per-variable
2227/// shape/offset descriptions and per-variable defaults. Per-slot names are
2228/// not built: they follow from the layout (see `layout::slot_name`).
2229struct SlotTables {
2230 var_shapes: IndexMap<String, VarShape>,
2231 state_defaults: Vec<StateDefault>,
2232 n_states: usize,
2233}
2234
2235/// (4) Build the flat offset and extents per state variable (column-major
2236/// slot enumeration within each), with each variable's default.
2237///
2238/// # Errors
2239///
2240/// [`CompileError::InterpreterBuildError`] when a state's inline array `default`
2241/// (esm-spec §6.3) is ragged or does not match its resolved grid shape.
2242fn build_slot_tables(
2243 model: &Model,
2244 final_states: &[String],
2245 shape_map: &HashMap<String, Vec<usize>>,
2246) -> Result<SlotTables, CompileError> {
2247 let mut var_shapes: IndexMap<String, VarShape> = IndexMap::new();
2248 let mut state_defaults: Vec<StateDefault> = Vec::with_capacity(final_states.len());
2249 let mut flat_offset: usize = 0;
2250
2251 for name in final_states {
2252 let shape = shape_map.get(name).cloned().unwrap_or_default();
2253 let origin: Vec<i64> = if shape.is_empty() {
2254 Vec::new()
2255 } else {
2256 vec![1i64; shape.len()]
2257 };
2258 let var = model.variables.get(name);
2259 // A SHAPED unknown may declare its whole initial profile as INLINE
2260 // ARRAY DATA (esm-spec §6.3) instead of one broadcast scalar. It is
2261 // kept ROW-major (the authored nesting's order) and gathered into the
2262 // column-major slots when the initial state is built. A ragged array
2263 // or a shape that disagrees with the slots is a build error, never a
2264 // silently mis-seeded state vector.
2265 let default = match var.and_then(|v| v.default_array()) {
2266 Some(Ok((dshape, values))) => {
2267 if dshape != shape {
2268 return Err(CompileError::build_err(format!(
2269 "state '{name}': inline array `default` has shape {dshape:?}, which does \
2270 not match the resolved grid shape {shape:?} (esm-spec §6.6.2)"
2271 )));
2272 }
2273 if shape.is_empty() {
2274 StateDefault::Scalar(var.and_then(|v| v.default_scalar()))
2275 } else {
2276 StateDefault::Field(values)
2277 }
2278 }
2279 Some(Err(e)) => {
2280 return Err(CompileError::build_err(format!(
2281 "state '{name}': {e} (esm-spec §6.3)"
2282 )));
2283 }
2284 _ => StateDefault::Scalar(var.and_then(|v| v.default_scalar())),
2285 };
2286 let total = shape.iter().copied().product::<usize>().max(1);
2287 state_defaults.push(default);
2288 var_shapes.insert(
2289 name.clone(),
2290 VarShape {
2291 shape,
2292 origin,
2293 flat_offset,
2294 },
2295 );
2296 flat_offset += total;
2297 }
2298
2299 Ok(SlotTables {
2300 var_shapes,
2301 state_defaults,
2302 n_states: flat_offset,
2303 })
2304}
2305
2306/// (5) Build the param tables: positional names, name → position index, and
2307/// per-position defaults.
2308fn build_param_tables(
2309 model: &Model,
2310 param_vars: &[&String],
2311) -> (Vec<String>, HashMap<String, usize>, Vec<Option<f64>>) {
2312 let param_names: Vec<String> = param_vars.iter().map(|s| (*s).clone()).collect();
2313 let param_index: HashMap<String, usize> = param_names
2314 .iter()
2315 .enumerate()
2316 .map(|(i, n)| (n.clone(), i))
2317 .collect();
2318 let param_defaults: Vec<Option<f64>> = param_vars
2319 .iter()
2320 .map(|n| model.variables.get(*n).and_then(|v| v.default_scalar()))
2321 .collect();
2322 (param_names, param_index, param_defaults)
2323}
2324
2325// ============================================================================
2326// Causal self-reference (recurrence) along one index axis — esm-spec §4.3.1.1
2327// ============================================================================
2328//
2329// An equation defining an array-shaped unknown `V` whose RHS reads
2330// `index(V, …)` at a strictly earlier position along one of the defining
2331// aggregate's output axes is a RECURRENCE DEFINITION of `V`. Detection is
2332// structural and costs a document that has no such read one walk of each
2333// algebraic RHS at build time; everything below is skipped when
2334// `collect_self_reads` finds nothing, which is the overwhelmingly common case.
2335//
2336// Rationale, alternatives, and the reason the lag is DERIVED rather than
2337// declared: `docs/content/rfcs/causal-self-reference-recurrence.md`.
2338
2339/// One causal self-read found in a variable's own defining RHS.
2340struct SelfRead {
2341 /// The read's index arguments (`args[1..]` of the `index` node).
2342 args: Vec<Expr>,
2343 /// Bounds of every index symbol bound by an enclosing `faq` at the
2344 /// point the read was found — the scope a symbol-valued lag is proved in.
2345 /// Innermost binding wins, matching the evaluator's `loop_binds` shadowing.
2346 env: HashMap<String, (i64, i64)>,
2347}
2348
2349/// A recurrence definition, lowered to the frame + per-cell body the sweep
2350/// runs (see [`AlgebraicRule::Recurrence`]).
2351pub(super) struct RecurLowering {
2352 pub(super) idx_names: Vec<String>,
2353 pub(super) ranges: Vec<(i64, i64)>,
2354 pub(super) body: Expr,
2355 pub(super) axis: usize,
2356 pub(super) max_lag: i64,
2357 pub(super) lag_proven: bool,
2358}
2359
2360fn recur_err(code: &str, detail: String) -> CompileError {
2361 CompileError::InterpreterBuildError {
2362 details: format!("{code}: {detail}"),
2363 }
2364}
2365
2366/// The affine form of an index expression with respect to the frame symbol
2367/// `sym`: the coefficient of `sym`, plus the integer bounds of the symbol-free
2368/// part — `None` when those bounds cannot be proved.
2369///
2370/// The two halves carry different weight. The **coefficient** must be provable:
2371/// unless the expression carries `sym` exactly once with coefficient 1 it does
2372/// not name a position relative to the cell being written, and there is nothing
2373/// to interpret. The **constant part** need not be: an unprovable offset is a
2374/// lag whose sign is unknown, which esm-spec §4.3.1.1 admits and leaves to the
2375/// runtime's fail-closed read — a cell the sweep has not published cannot be
2376/// read, so an unprovable lag cannot produce a wrong number, only a fault.
2377///
2378/// Treating an unresolvable symbol as fatal instead would make the VALIDATOR
2379/// reject documents the evaluator accepts (it resolves ranges against the
2380/// registry first and so proves more), which is the one disagreement between
2381/// the two that is never defensible.
2382struct Affine {
2383 coef: i64,
2384 konst: Option<(i64, i64)>,
2385}
2386
2387fn affine_in_sym(e: &Expr, sym: &str, env: &HashMap<String, (i64, i64)>) -> Option<Affine> {
2388 let konst = |lo: i64, hi: i64| {
2389 Some(Affine {
2390 coef: 0,
2391 konst: Some((lo, hi)),
2392 })
2393 };
2394 match e {
2395 Expr::Integer(n) => konst(*n, *n),
2396 Expr::Number(f) if f.fract() == 0.0 && f.is_finite() => {
2397 let n = *f as i64;
2398 konst(n, n)
2399 }
2400 Expr::Variable(v) if v == sym => Some(Affine {
2401 coef: 1,
2402 konst: Some((0, 0)),
2403 }),
2404 Expr::Variable(v) => Some(Affine {
2405 coef: 0,
2406 konst: env.get(v).copied(),
2407 }),
2408 Expr::Operator(node) if node.args.len() == 2 => {
2409 let a = affine_in_sym(&node.args[0], sym, env)?;
2410 let b = affine_in_sym(&node.args[1], sym, env)?;
2411 let both = a.konst.zip(b.konst);
2412 match node.op.as_str() {
2413 "+" => Some(Affine {
2414 coef: a.coef + b.coef,
2415 konst: both.map(|((la, ha), (lb, hb))| (la + lb, ha + hb)),
2416 }),
2417 "-" => Some(Affine {
2418 coef: a.coef - b.coef,
2419 konst: both.map(|((la, ha), (lb, hb))| (la - hb, ha - lb)),
2420 }),
2421 // Scaling only by a symbol-free integer whose value is PINNED
2422 // (`lo == hi`): otherwise the product is not affine with a known
2423 // coefficient, and the coefficient is the half that must be
2424 // provable.
2425 "*" => {
2426 let (k, other) = match (a.coef, a.konst, b.coef, b.konst) {
2427 (0, Some((lo, hi)), _, _) if lo == hi => (lo, &b),
2428 (_, _, 0, Some((lo, hi))) if lo == hi => (lo, &a),
2429 _ => return None,
2430 };
2431 Some(Affine {
2432 coef: other.coef * k,
2433 konst: other.konst.map(|(lo, hi)| {
2434 let (p, q) = (lo * k, hi * k);
2435 (p.min(q), p.max(q))
2436 }),
2437 })
2438 }
2439 _ => None,
2440 }
2441 }
2442 _ => None,
2443 }
2444}
2445
2446/// How one index argument of a causal self-read relates to the frame symbol at
2447/// that axis. Written in terms of the LAG `sym - <index expression>`, so a
2448/// positive lag is an earlier position.
2449enum SelfIdx {
2450 /// The lag is exactly zero: the read stays on this axis's own cell, so this
2451 /// axis is not the one the recurrence folds along.
2452 Identity,
2453 /// A lag that is not provably zero-or-forward, so this axis IS the
2454 /// recurrence axis. `proven` records whether every value of the lag was
2455 /// shown `>= 1` statically; when it is false the read may name the cell
2456 /// being written or a later one for some values of a contracted symbol, and
2457 /// the runtime's fail-closed read is what rules those out (esm-spec
2458 /// §4.3.1.1 point 5) — a cell that has not been published cannot be read,
2459 /// whether the sweep will publish it later or never.
2460 Offset { proven: bool, max_lag: i64 },
2461 /// A lag that is provably `<= 0` for every value: a same-cell or forward
2462 /// read, which no sweep order can satisfy.
2463 Forward,
2464 /// Not affine in the frame symbol at all.
2465 Bad,
2466}
2467
2468fn classify_self_index(arg: &Expr, sym: &str, env: &HashMap<String, (i64, i64)>) -> SelfIdx {
2469 // lag = sym - arg. `arg` must carry `sym` with coefficient exactly 1, or
2470 // "the previous position along this axis" is not what it names.
2471 let Some(Affine { coef, konst }) = affine_in_sym(arg, sym, env) else {
2472 return SelfIdx::Bad;
2473 };
2474 if coef != 1 {
2475 return SelfIdx::Bad;
2476 }
2477 let Some((clo, chi)) = konst else {
2478 // A lag whose sign could not be proved. Admitted as the recurrence
2479 // axis: it is not provably wrong, and the cells where it would be
2480 // ill-founded cannot be read (they are not published), so the runtime
2481 // faults there instead of returning a number.
2482 return SelfIdx::Offset {
2483 proven: false,
2484 max_lag: 0,
2485 };
2486 };
2487 let (lag_lo, lag_hi) = (-chi, -clo);
2488 if lag_lo == 0 && lag_hi == 0 {
2489 return SelfIdx::Identity;
2490 }
2491 if lag_hi <= 0 {
2492 return SelfIdx::Forward;
2493 }
2494 SelfIdx::Offset {
2495 proven: lag_lo >= 1,
2496 max_lag: lag_hi,
2497 }
2498}
2499
2500/// Bounds every `faq` range in `node` contributes to the symbol scope.
2501fn aggregate_range_env(node: &ExpressionNode) -> Vec<(String, (i64, i64))> {
2502 node.ranges
2503 .as_ref()
2504 .map(|m| {
2505 m.iter()
2506 .filter_map(|(k, spec)| spec.bounds().map(|b| (k.clone(), (b[0], b[1]))))
2507 .collect()
2508 })
2509 .unwrap_or_default()
2510}
2511
2512/// Walk `e` collecting every `index(var, …)` read, and note whether `var` is
2513/// ever read BARE (which is never a causal read — esm-spec §4.3.1.1 rejection
2514/// 4 — because the whole array does not exist during the sweep).
2515fn collect_self_reads(
2516 e: &Expr,
2517 var: &str,
2518 env: &mut Vec<(String, (i64, i64))>,
2519 out: &mut Vec<SelfRead>,
2520 bare: &mut bool,
2521) {
2522 match e {
2523 Expr::Integer(_) | Expr::Number(_) => {}
2524 Expr::Variable(v) => {
2525 if v == var {
2526 *bare = true;
2527 }
2528 }
2529 Expr::Operator(node) => {
2530 let pushed = if crate::faq::is_faq_op(&node.op) {
2531 let add = aggregate_range_env(node);
2532 let n = add.len();
2533 env.extend(add);
2534 n
2535 } else {
2536 0
2537 };
2538 let is_self_index = node.op == "index"
2539 && matches!(node.args.first(), Some(Expr::Variable(v)) if v == var);
2540 if is_self_index {
2541 // Innermost binding wins: later entries overwrite earlier ones.
2542 let mut snapshot: HashMap<String, (i64, i64)> = HashMap::new();
2543 for (k, v) in env.iter() {
2544 snapshot.insert(k.clone(), *v);
2545 }
2546 out.push(SelfRead {
2547 args: node.args[1..].to_vec(),
2548 env: snapshot,
2549 });
2550 }
2551 // `args[0]` of a self-read is the name itself, not a bare read of
2552 // the array; every other operand is walked normally (an index
2553 // expression may itself contain a self-read).
2554 let skip = usize::from(is_self_index);
2555 for a in node.args.iter().skip(skip) {
2556 collect_self_reads(a, var, env, out, bare);
2557 }
2558 for side in [
2559 node.expr.as_deref(),
2560 node.filter.as_deref(),
2561 node.key.as_deref(),
2562 node.lower.as_deref(),
2563 node.upper.as_deref(),
2564 ]
2565 .into_iter()
2566 .flatten()
2567 {
2568 collect_self_reads(side, var, env, out, bare);
2569 }
2570 if let Some(vs) = node.values.as_ref() {
2571 for v in vs {
2572 collect_self_reads(v, var, env, out, bare);
2573 }
2574 }
2575 env.truncate(env.len() - pushed);
2576 }
2577 }
2578}
2579
2580/// Is `e` an `index(var, sym…)` whose indices are exactly `idx_names` in order
2581/// — the §4.3 indexed-`faq` LHS form?
2582fn lhs_identity_gather(e: &Expr, var: &str, idx_names: &[String]) -> bool {
2583 let Expr::Operator(node) = e else {
2584 return false;
2585 };
2586 if node.op != "index" || node.args.len() != idx_names.len() + 1 {
2587 return false;
2588 }
2589 if !matches!(node.args.first(), Some(Expr::Variable(v)) if v == var) {
2590 return false;
2591 }
2592 node.args[1..]
2593 .iter()
2594 .zip(idx_names)
2595 .all(|(a, want)| matches!(a, Expr::Variable(v) if v == want))
2596}
2597
2598/// Cell-restrict a `faq` that produces the whole frame: move its output
2599/// indices out to the enclosing sweep, keeping its contraction, `filter`,
2600/// `reduce`, `join` and `key` intact so the body evaluates at one cell exactly
2601/// as §4.3.1 specifies for a non-recurrent aggregate.
2602///
2603/// When there is nothing left to contract or gate, the restriction IS the body,
2604/// so the wrapper is dropped — the common shape (`s[k] = f(s[k-1])`) then walks
2605/// one expression per cell rather than re-deriving an empty contraction.
2606fn cell_restrict_aggregate(node: &ExpressionNode, idx_names: &[String]) -> Expr {
2607 let body = node
2608 .expr
2609 .as_deref()
2610 .cloned()
2611 .unwrap_or(Expr::Number(f64::NAN));
2612 let remaining: HashMap<String, crate::types::RangeSpec> = node
2613 .ranges
2614 .as_ref()
2615 .map(|m| {
2616 m.iter()
2617 .filter(|(k, _)| !idx_names.iter().any(|n| n == *k))
2618 .map(|(k, v)| (k.clone(), v.clone()))
2619 .collect()
2620 })
2621 .unwrap_or_default();
2622 let gated = node.filter.is_some()
2623 || node.join.is_some()
2624 || node.key.is_some()
2625 || node.distinct == Some(true);
2626 if remaining.is_empty() && !gated {
2627 return body;
2628 }
2629 let mut restricted = node.clone();
2630 restricted.output_idx = Some(Vec::new());
2631 restricted.ranges = Some(remaining);
2632 Expr::operator(restricted)
2633}
2634
2635/// Recognize and lower a recurrence definition of `var`.
2636///
2637/// `Ok(None)` — no causal self-read, so this equation is compiled exactly as it
2638/// was before this feature existed. `Err(…)` — a self-read that is not a
2639/// well-founded causal read, reported with the esm-spec §4.3.1.1 code. Never a
2640/// silent fallback: the pre-feature behaviour for every rejected shape here was
2641/// a plausible wrong number.
2642pub(super) fn lower_recurrence(
2643 var: &str,
2644 lhs: &Expr,
2645 rhs: &Expr,
2646) -> Result<Option<RecurLowering>, CompileError> {
2647 let mut env: Vec<(String, (i64, i64))> = Vec::new();
2648 let mut reads: Vec<SelfRead> = Vec::new();
2649 let mut bare = false;
2650 collect_self_reads(rhs, var, &mut env, &mut reads, &mut bare);
2651 if reads.is_empty() {
2652 return Ok(None);
2653 }
2654 if bare {
2655 return Err(recur_err(
2656 "recurrence_not_wellfounded",
2657 format!(
2658 "'{var}' is read BARE in its own defining equation as well as through `index`. \
2659 A bare read names the whole array, which does not exist while the recurrence \
2660 sweeps it; read every self-reference through `index` at a strictly earlier \
2661 position (esm-spec §4.3.1.1)."
2662 ),
2663 ));
2664 }
2665
2666 // ---- the cell frame ----------------------------------------------------
2667 // Either the indexed-`faq` LHS form (`faq{expr: V[k…]} ~ …`) or a
2668 // bare LHS whose RHS is an aggregate over V's axes. Anything else has no
2669 // frame to sweep.
2670 let frame_node: &ExpressionNode = match (lhs, rhs) {
2671 (Expr::Operator(l), _) if crate::faq::is_faq_op(&l.op) => l.as_ref(),
2672 (Expr::Variable(v), Expr::Operator(r)) if v == var && crate::faq::is_faq_op(&r.op) => {
2673 r.as_ref()
2674 }
2675 _ => {
2676 return Err(recur_err(
2677 "recurrence_unsupported_form",
2678 format!(
2679 "the definition of '{var}' reads '{var}' at an earlier position, but its \
2680 shape gives the runtime no cell frame to sweep. Write the recurrence as one \
2681 `faq` over the variable's axes — either `{var} ~ aggregate{{…}}` or \
2682 `aggregate{{expr: index({var}, k…)}} ~ …` — with the base case as an \
2683 `ifelse` guard in the body (esm-spec §4.3.1.1)."
2684 ),
2685 ));
2686 }
2687 };
2688 let idx_names: Vec<String> = frame_node.output_idx.clone().unwrap_or_default();
2689 if idx_names.is_empty() || idx_names.iter().any(|n| n.parse::<i64>().is_ok()) {
2690 return Err(recur_err(
2691 "recurrence_unsupported_form",
2692 format!(
2693 "the recurrence definition of '{var}' has no symbolic output index to fold \
2694 along (`output_idx` is {idx_names:?}). A literal singleton dimension cannot be \
2695 a recurrence axis (esm-spec §4.3.1.1)."
2696 ),
2697 ));
2698 }
2699 let ranges_map = frame_node.ranges.clone().unwrap_or_default();
2700 let mut ranges: Vec<(i64, i64)> = Vec::with_capacity(idx_names.len());
2701 for n in &idx_names {
2702 match ranges_map.get(n) {
2703 // `Interval` specifically: a `Strided` axis has no unambiguous
2704 // "previous position", and a ragged / derived / unresolved axis has
2705 // no static total order to fold along at all.
2706 Some(crate::types::RangeSpec::Interval([lo, hi])) => ranges.push((*lo, *hi)),
2707 other => {
2708 return Err(recur_err(
2709 "recurrence_not_wellfounded",
2710 format!(
2711 "axis '{n}' of the recurrence definition of '{var}' is not a static \
2712 unit-step ascending interval (range: {other:?}). A ragged, derived, \
2713 strided or unresolved axis carries no total order to fold along \
2714 (esm-spec §4.3.1.1)."
2715 ),
2716 ));
2717 }
2718 }
2719 }
2720
2721 // ---- the axis and the derived lag bound --------------------------------
2722 let frame_env: HashMap<String, (i64, i64)> = idx_names
2723 .iter()
2724 .cloned()
2725 .zip(ranges.iter().copied())
2726 .collect();
2727 let mut axis: Option<usize> = None;
2728 let mut max_lag: i64 = 0;
2729 let mut lag_proven = true;
2730 for read in &reads {
2731 if read.args.len() != idx_names.len() {
2732 return Err(recur_err(
2733 "recurrence_not_wellfounded",
2734 format!(
2735 "a self-read of '{var}' supplies {} indices but its frame has {} axes \
2736 ({idx_names:?}). Every causal self-read indexes every axis (esm-spec \
2737 §4.3.1.1).",
2738 read.args.len(),
2739 idx_names.len()
2740 ),
2741 ));
2742 }
2743 let mut env = frame_env.clone();
2744 for (k, v) in &read.env {
2745 env.insert(k.clone(), *v);
2746 }
2747 let mut lagged: Option<(usize, i64)> = None;
2748 for (d, arg) in read.args.iter().enumerate() {
2749 match classify_self_index(arg, &idx_names[d], &env) {
2750 SelfIdx::Identity => {}
2751 SelfIdx::Forward => {
2752 return Err(recur_err(
2753 "recurrence_not_wellfounded",
2754 format!(
2755 "index {d} of a self-read of '{var}' names the cell being written, \
2756 or a LATER one, on axis '{}'. A causal self-reference reads \
2757 strictly EARLIER positions; no sweep order can satisfy a forward \
2758 or same-cell read (esm-spec §4.3.1.1).",
2759 idx_names[d]
2760 ),
2761 ));
2762 }
2763 SelfIdx::Offset {
2764 proven,
2765 max_lag: hi,
2766 } => {
2767 lag_proven &= proven;
2768 if lagged.is_some() {
2769 return Err(recur_err(
2770 "recurrence_not_wellfounded",
2771 format!(
2772 "a self-read of '{var}' is offset on more than one axis. A \
2773 causal self-reference folds along exactly ONE axis; every \
2774 other index must be the bare frame symbol (esm-spec §4.3.1.1)."
2775 ),
2776 ));
2777 }
2778 lagged = Some((d, hi));
2779 }
2780 SelfIdx::Bad => {
2781 return Err(recur_err(
2782 "recurrence_not_wellfounded",
2783 format!(
2784 "index {d} of a self-read of '{var}' is not an offset of the frame \
2785 symbol '{}'. A causal self-read names a position RELATIVE to the \
2786 cell being written — `{} - c`, `{} - a`, `{} - a - c` — so that \
2787 which axis the recurrence folds along, and in which direction, is \
2788 decidable. An index that does not carry '{}' with coefficient 1 \
2789 (a bare constant, `2*{}`, another axis's symbol) is rejected \
2790 rather than guessed at (esm-spec §4.3.1.1). {}",
2791 idx_names[d],
2792 idx_names[d],
2793 idx_names[d],
2794 idx_names[d],
2795 idx_names[d],
2796 idx_names[d],
2797 crate::structural::data_lag_guidance(var, &idx_names[d])
2798 ),
2799 ));
2800 }
2801 }
2802 }
2803 let Some((d, hi)) = lagged else {
2804 return Err(recur_err(
2805 "recurrence_not_wellfounded",
2806 format!(
2807 "a self-read of '{var}' is at the SAME cell on every axis, so it defines \
2808 '{var}' in terms of itself rather than of an earlier position. A causal \
2809 self-reference must be strictly earlier along one axis (esm-spec §4.3.1.1)."
2810 ),
2811 ));
2812 };
2813 match axis {
2814 None => axis = Some(d),
2815 Some(prev) if prev == d => {}
2816 Some(prev) => {
2817 return Err(recur_err(
2818 "recurrence_not_wellfounded",
2819 format!(
2820 "the self-reads of '{var}' disagree on the recurrence axis: one folds \
2821 along '{}' and another along '{}'. A definition folds along exactly one \
2822 axis (esm-spec §4.3.1.1).",
2823 idx_names[prev], idx_names[d]
2824 ),
2825 ));
2826 }
2827 }
2828 max_lag = max_lag.max(hi);
2829 }
2830 let axis = axis.expect("at least one self-read, each of which set the axis");
2831
2832 // ---- the per-cell body -------------------------------------------------
2833 let body = match rhs {
2834 Expr::Operator(r)
2835 if crate::faq::is_faq_op(&r.op)
2836 && r.output_idx.as_deref() == Some(idx_names.as_slice()) =>
2837 {
2838 cell_restrict_aggregate(r, &idx_names)
2839 }
2840 // The indexed-`faq` LHS form whose RHS is already a cell body, or an
2841 // RHS aggregate over a DIFFERENT frame (which cannot be restricted onto
2842 // this one).
2843 other => {
2844 if lhs_identity_gather(
2845 frame_node.expr.as_deref().unwrap_or(&Expr::Integer(0)),
2846 var,
2847 &idx_names,
2848 ) {
2849 other.clone()
2850 } else {
2851 return Err(recur_err(
2852 "recurrence_unsupported_form",
2853 format!(
2854 "the recurrence definition of '{var}' cannot be restricted to one cell: \
2855 its RHS is not a `faq` over the frame {idx_names:?}, and its LHS \
2856 is not the identity gather `index({var}, {})`. A self-read reached only \
2857 through a `makearray` region, a `reshape`/`transpose`/`concat` operand, \
2858 or an aggregate over a different frame cannot be sequenced cell by cell \
2859 — the region order of a `makearray` fixes which write WINS, not the \
2860 order cells are EVALUATED in (esm-spec §4.3.1.1, §4.3.2).",
2861 idx_names.join(", ")
2862 ),
2863 ));
2864 }
2865 }
2866 };
2867
2868 Ok(Some(RecurLowering {
2869 idx_names,
2870 ranges,
2871 body,
2872 axis,
2873 max_lag,
2874 lag_proven,
2875 }))
2876}
2877
2878/// (6) Build observed algebraic rules from eliminated state variables AND
2879/// from declared observed variables that define an expression, then (6b)
2880/// dependency-order them so each rule is evaluated only after the observeds it
2881/// reads (RFC §8.1): the geometry chain `const` polygons → `clip =
2882/// intersect_polygon` → `area = FAQ(clip)` must materialize the ring before
2883/// the FAQ over it. The rules are collected in sorted/equation order, which is
2884/// NOT dependency order; the stable Kahn sweep ([`dependency_order_observed`])
2885/// preserves declaration order among independent observeds (mirrors Python
2886/// `simulation._order_observed_equations`).
2887fn build_observed_rules(
2888 model: &mut Model,
2889 observed_names: &[String],
2890 eliminated: &HashSet<String>,
2891 index_sets: &HashMap<String, IndexSet>,
2892) -> Result<Vec<AlgebraicRule>, CompileError> {
2893 let mut observed_rules: Vec<AlgebraicRule> = Vec::new();
2894 let array_axes = declared_axis_names(model);
2895 // Read only for a bare-index definition.
2896 let declared: std::cell::OnceCell<HashSet<String>> = std::cell::OnceCell::new();
2897
2898 // Declared observed variables with an `expression` field. An array-shaped
2899 // observed — a discretization-agnostic PDE leaf's `psi_x`, `grad_mag`,
2900 // `U_n`, `S_n`, a `const`-op field, a keyed-factor alias — is evaluated
2901 // WHOLESALE here: `eval` looks each array-valued observed reference up in
2902 // the observed-array map and broadcasts the elementwise ops over it, so a
2903 // readable intermediate decomposition (WS4) already runs as authored. The
2904 // rules are dependency-ordered below, so `grad_mag` materializes before
2905 // `U_n`/`S_n` read it.
2906 //
2907 // The definition is the EQUATION whose LHS names the unknown — esm 1.0.0
2908 // states an unknown's behaviour in `equations` and nowhere else, so this
2909 // reads the equation rather than the removed `variables[..].expression`
2910 // field. `observed_names` preserves the sorted order `classify_variables`
2911 // produced, so the rule order — and the stable dependency ordering below —
2912 // is unchanged.
2913 //
2914 // WHICH rule form it lowers to is decided by the LHS, and both forms matter:
2915 //
2916 // * a `faq` LHS (`faq{expr: w[i]} ~ aggregate{…}`) is a PER-CELL
2917 // definition and lowers to [`AlgebraicRule::ArrayLoop`], the form the
2918 // whole-array overlay vectorizes. Lowering it as a wholesale body instead
2919 // silently drops a varying array observed onto the per-cell oracle — the
2920 // dominant per-step cost for a coupled behaviour stack;
2921 // * a bare-variable LHS lowers WHOLESALE through [`lower_algebraic_body`]:
2922 // `eval` materializes the body's arrays and broadcasts the elementwise ops
2923 // over them, so a readable intermediate decomposition runs as authored.
2924 let mut def_eq: HashMap<&str, &crate::types::Equation> = HashMap::new();
2925 for eq in &model.equations {
2926 if let crate::classification::LhsFormRef::Bare(name) =
2927 crate::classification::lhs_form_ref(&eq.lhs)
2928 {
2929 def_eq.entry(name).or_insert(eq);
2930 }
2931 }
2932 for name in observed_names {
2933 let Some(eq) = def_eq.get(name.as_str()) else {
2934 continue;
2935 };
2936 // A bare-index LHS (`index(V, k…) ~ rhs`, esm-spec §6.3.1) binds none of
2937 // its subscripts, so it runs only in the form whose range the RHS
2938 // supplies; see [`check_bare_index_definition`].
2939 if let Expr::Operator(lhs) = &eq.lhs
2940 && lhs.op == "index"
2941 {
2942 let declared = declared.get_or_init(|| model.variables.keys().cloned().collect());
2943 check_bare_index_definition(name, lhs, &eq.rhs, &array_axes, declared)?;
2944 }
2945 // A CAUSAL SELF-REFERENCE (esm-spec §4.3.1.1) is recognized before
2946 // either ordinary lowering, because both of them would compile the
2947 // self-read as a gather on a variable that is not bound anywhere — the
2948 // pre-feature behaviour, which validated and then produced nothing.
2949 // `None` here means the RHS contains no self-read at all, so every
2950 // existing document takes exactly the paths below, unchanged.
2951 if let Some(r) = lower_recurrence(name, &eq.lhs, &eq.rhs)? {
2952 observed_rules.push(AlgebraicRule::Recurrence {
2953 var: name.clone(),
2954 output_idx_names: r.idx_names,
2955 output_ranges: r.ranges,
2956 body: std::sync::Arc::new(r.body),
2957 axis: r.axis,
2958 max_lag: r.max_lag,
2959 lag_proven: r.lag_proven,
2960 });
2961 continue;
2962 }
2963 if let Some(a) = extract_algebraic_faq(&eq.lhs, &eq.rhs) {
2964 observed_rules.push(AlgebraicRule::ArrayLoop {
2965 var: a.var,
2966 output_idx_names: a.idx_names,
2967 output_ranges: a.ranges,
2968 body: std::sync::Arc::new(a.body),
2969 });
2970 } else {
2971 observed_rules.push(lower_algebraic_body(
2972 name,
2973 eq.rhs.clone(),
2974 &array_axes,
2975 index_sets,
2976 )?);
2977 }
2978 }
2979
2980 // Algebraic faq equations for ELIMINATED state variables — the same two
2981 // forms, for a name the DAE pass removed from the state vector rather than
2982 // one the classification calls observed. An observed's own defining equation
2983 // was lowered above, so it is skipped here rather than emitted twice.
2984 let is_observed: HashSet<&str> = if eliminated.is_empty() {
2985 HashSet::new()
2986 } else {
2987 observed_names.iter().map(String::as_str).collect()
2988 };
2989 for eq in &model.equations {
2990 if let Some(a) = extract_algebraic_faq(&eq.lhs, &eq.rhs) {
2991 if eliminated.contains(&a.var) && !is_observed.contains(a.var.as_str()) {
2992 if let Some(r) = lower_recurrence(&a.var, &eq.lhs, &eq.rhs)? {
2993 observed_rules.push(AlgebraicRule::Recurrence {
2994 var: a.var.clone(),
2995 output_idx_names: r.idx_names,
2996 output_ranges: r.ranges,
2997 body: std::sync::Arc::new(r.body),
2998 axis: r.axis,
2999 max_lag: r.max_lag,
3000 lag_proven: r.lag_proven,
3001 });
3002 continue;
3003 }
3004 observed_rules.push(AlgebraicRule::ArrayLoop {
3005 var: a.var,
3006 output_idx_names: a.idx_names,
3007 output_ranges: a.ranges,
3008 body: std::sync::Arc::new(a.body),
3009 });
3010 }
3011 continue;
3012 }
3013 if let Expr::Variable(name) = &eq.lhs
3014 && eliminated.contains(name)
3015 && !is_observed.contains(name.as_str())
3016 {
3017 observed_rules.push(lower_algebraic_body(
3018 name,
3019 eq.rhs.clone(),
3020 &array_axes,
3021 index_sets,
3022 )?);
3023 }
3024 }
3025 dependency_order_observed(observed_rules)
3026}
3027
3028/// Refuse a bare-index observed definition that is not the runnable form.
3029///
3030/// `index(V, k…) ~ rhs` binds none of `k…`, so the range has to come from the
3031/// right-hand side: the definition runs exactly when `rhs` is a `faq` whose
3032/// `output_idx` names the subscripts, in order (esm-spec §6.3.1,
3033/// CONFORMANCE_SPEC §5.36.2). That RHS is already the whole array, which is what
3034/// the wholesale lowering evaluates. Any other spelling — a scalar RHS, an offset
3035/// or permuted subscript, a subscript count that disagrees with `V`'s declared
3036/// rank — would be lowered as if its RHS were the whole array, which writes the
3037/// wrong cells for an offset and cannot integrate for a scalar, so it is refused.
3038///
3039/// The gather must also be the DIRECT one, `index(V, k…)`: the base name is read
3040/// through nested `index` / `broadcast` wrappers, but `index(index(V, j), k)`
3041/// addresses a cell of a cell, not the whole of `V`, so it is refused rather
3042/// than lowered as though its RHS were all of `V`.
3043///
3044/// Flatten namespaces a free subscript (`k` becomes `Model.k`) but not a `faq`
3045/// binder, so a subscript matches its binder in either spelling.
3046///
3047/// Checked first: a plain-symbol subscript at any level of the gather that the
3048/// right-hand `faq` does not bind and that names no declared variable is an
3049/// index symbol nothing binds — `unbound_index_symbol` (esm-spec §6.3.1), the
3050/// finding `validate` reports for the same left-hand side.
3051fn check_bare_index_definition(
3052 name: &str,
3053 lhs: &ExpressionNode,
3054 rhs: &Expr,
3055 array_axes: &HashMap<String, Vec<String>>,
3056 declared: &HashSet<String>,
3057) -> Result<(), CompileError> {
3058 let head_is_the_variable = matches!(lhs.args.first(), Some(Expr::Variable(_)));
3059 let subs = lhs.args.get(1..).unwrap_or_default();
3060 let prefix = name.rfind('.').map_or("", |p| &name[..=p]);
3061 let names_binder = |sub: &Expr, binder: &String| match sub {
3062 Expr::Variable(v) => v == binder || v.strip_prefix(prefix) == Some(binder.as_str()),
3063 _ => false,
3064 };
3065 let frame: &[String] = match rhs {
3066 Expr::Operator(node) if is_faq_op(&node.op) => node.output_idx.as_deref().unwrap_or(&[]),
3067 _ => &[],
3068 };
3069 let mut gather = Some(lhs);
3070 while let Some(g) = gather {
3071 for sub in g.args.iter().skip(1) {
3072 if let Expr::Variable(v) = sub
3073 && !declared.contains(v)
3074 && !frame.iter().any(|b| names_binder(sub, b))
3075 {
3076 let sym = v.strip_prefix(prefix).unwrap_or(v);
3077 return Err(CompileError::InterpreterBuildError {
3078 details: format!(
3079 "{}: equation defining '{name}' subscripts its left-hand side with \
3080 '{sym}', which no faq binds (esm-spec §6.3.1)",
3081 crate::diagnostic::codes::UNBOUND_INDEX_SYMBOL
3082 ),
3083 });
3084 }
3085 }
3086 gather = match g.args.first() {
3087 Some(Expr::Operator(inner)) if inner.op == "index" => Some(inner),
3088 _ => None,
3089 };
3090 }
3091 let binds_subscripts = match rhs {
3092 Expr::Operator(node) if is_faq_op(&node.op) => match node.output_idx.as_deref() {
3093 Some(frame) => {
3094 !subs.is_empty()
3095 && frame.len() == subs.len()
3096 && subs.iter().zip(frame).all(|(s, b)| names_binder(s, b))
3097 }
3098 None => false,
3099 },
3100 _ => false,
3101 };
3102 let rank_agrees = match array_axes.get(name) {
3103 Some(axes) => axes.len() == subs.len(),
3104 None => true,
3105 };
3106 if head_is_the_variable && binds_subscripts && rank_agrees {
3107 return Ok(());
3108 }
3109 Err(CompileError::InterpreterBuildError {
3110 details: format!(
3111 "indexed_definition_unsupported_form: '{name}' is defined by a bare-index LHS that \
3112 is not runnable; the RHS must be a `faq` whose `output_idx` names the LHS \
3113 subscripts in order (esm-spec §6.3.1)"
3114 ),
3115 })
3116}
3117
3118/// Wrap one algebraic body — a declared observed's `expression`, or the RHS of
3119/// a bare-`Variable`-LHS equation — in the rule form that evaluates it
3120/// correctly.
3121///
3122/// The default, and what every already-correct model keeps byte for byte, is
3123/// the WHOLESALE [`AlgebraicRule::Scalar`]: `eval` materializes the body's
3124/// arrays and broadcasts the elementwise ops over them. That broadcast is
3125/// POSITIONAL, so it is only right when the operands already sit on the
3126/// result's axes in the result's order.
3127///
3128/// When the target declares a `shape` and an operand does not carry all of it
3129/// (`p3: [lon,lat,lev] = w2 * z1` with `w2: [lon,lat]` and `z1: [lev]`),
3130/// positional broadcasting cannot express what the index-set names mean — each
3131/// operand must replicate along the axes it does NOT carry (esm-spec §4.3.4).
3132/// Such a body is lowered to the per-cell [`AlgebraicRule::ArrayLoop`] form
3133/// instead, which gathers each operand at its own axes exactly as the
3134/// equivalent `faq` spelling would. An operand carrying an index set the
3135/// result does not have is rejected by [`build_gather_plan`].
3136fn lower_algebraic_body(
3137 name: &str,
3138 body: Expr,
3139 array_axes: &HashMap<String, Vec<String>>,
3140 index_sets: &HashMap<String, IndexSet>,
3141) -> Result<AlgebraicRule, CompileError> {
3142 let declared_shape = array_axes
3143 .get(name)
3144 .and_then(|axes| resolve_declared_shape(axes, index_sets));
3145 let scalar_rule = |body: Expr| AlgebraicRule::Scalar {
3146 var: name.to_string(),
3147 body: std::sync::Arc::new(body),
3148 declared_shape: declared_shape.clone(),
3149 };
3150 let Some(target_axes) = array_axes.get(name) else {
3151 return Ok(scalar_rule(body));
3152 };
3153 let plan = build_gather_plan(&body, array_axes, name, Some(target_axes.as_slice()), false)?;
3154 if plan_is_identity(&plan, target_axes.len()) {
3155 return Ok(scalar_rule(body));
3156 }
3157 // Name alignment IS needed, but the declared axes cannot be densely sized
3158 // (a derived / ragged index set): keep the wholesale rule rather than
3159 // fabricate loop bounds.
3160 let Some(extents) = resolve_declared_shape(target_axes, index_sets) else {
3161 return Ok(scalar_rule(body));
3162 };
3163 let loops: Vec<String> = (0..extents.len())
3164 .map(|d| format!("_lp{d}_{name}"))
3165 .collect();
3166 let output_ranges: Vec<(i64, i64)> = extents.iter().map(|n| (1i64, *n as i64)).collect();
3167 let looped = index_array_leaves_by_loops(&body, array_axes, Some(&plan), &loops);
3168 Ok(AlgebraicRule::ArrayLoop {
3169 var: name.to_string(),
3170 output_idx_names: loops,
3171 output_ranges,
3172 body: std::sync::Arc::new(looped),
3173 })
3174}
3175
3176/// Classify scoped-reference / array `ic` equations (esm-spec §11.4.1) out of
3177/// the rule builder into `field_ics`, mirroring the flatten-path
3178/// classification — a single-model (`from_file`) build must fold
3179/// coordinate-expression / broadcast-constant ics into `u0` exactly as the
3180/// coupled path does. The RHS collected here has already been range-resolved
3181/// against the document registry by [`ArrayCompiled::from_model`].
3182fn classify_field_ics(model: &Model) -> Vec<(String, Expr)> {
3183 let mut field_ics: Vec<(String, Expr)> = Vec::new();
3184 for eq in &model.equations {
3185 if let Some(target) = crate::flatten::extract_ic_target(&eq.lhs) {
3186 field_ics.push((target, eq.rhs.clone()));
3187 }
3188 }
3189 field_ics
3190}
3191
3192/// The `faq`-valued initialization equations (esm-spec §6.2, "equations that
3193/// hold only at t=0"): `u ~ faq(…)` with a bare-variable left-hand side and an
3194/// array-valued right-hand side, in document order. Each assigns the cells of
3195/// its ranges at `u0` build time ([`ArrayCompiled::seed_initialization_faqs`]);
3196/// the other spellings name no cells to assign and are not evaluated, as in
3197/// the Julia reference's `_seed_faq_init_u0!`.
3198fn classify_initialization_faqs(model: &Model) -> Vec<(String, Expr)> {
3199 model
3200 .initialization_equations
3201 .iter()
3202 .flatten()
3203 .filter_map(|eq| match (&eq.lhs, &eq.rhs) {
3204 (Expr::Variable(target), Expr::Operator(node))
3205 if is_faq_op(&node.op)
3206 && node.output_idx.as_ref().is_some_and(|o| !o.is_empty()) =>
3207 {
3208 Some((target.clone(), eq.rhs.clone()))
3209 }
3210 _ => None,
3211 })
3212 .collect()
3213}
3214
3215/// (7) Build the RHS rules. Each equation with a derivative LHS produces
3216/// either a scalar slot write, an indexed scalar slot write, or an array
3217/// loop — one lowering function per form below. Then (7b) held-at-ic states
3218/// (no `D`, no algebraic definition) carry every cell at its ic with zero
3219/// derivative — their slots are marked covered without emitting a rule (the
3220/// RHS zero-initializes `dy` each call and never writes them) — and (8) every
3221/// state slot must end up with a defining equation.
3222fn build_rhs_rules(
3223 model: &Model,
3224 slots: &SlotTables,
3225 held_at_ic: &HashSet<String>,
3226) -> Result<Vec<RhsRule>, CompileError> {
3227 let var_shapes = &slots.var_shapes;
3228 let mut rhs_rules: Vec<RhsRule> = Vec::new();
3229 let mut covered_slots = SlotCoverage::new(slots.n_states);
3230
3231 // Declared index-set axis NAMES of every array-shaped variable (state /
3232 // parameter / observed), used to lower a whole-array `D(state)` RHS into
3233 // per-cell gathers that align each operand BY NAME (esm-spec §4.3.4).
3234 let array_axes = declared_axis_names(model);
3235
3236 for eq in &model.equations {
3237 if let Some(d) = extract_derivative_faq(&eq.lhs, &eq.rhs) {
3238 lower_faq_derivative(d, var_shapes, &mut covered_slots, &mut rhs_rules)?;
3239 continue;
3240 }
3241 // Scalar D(var, t) = rhs.
3242 if let Some((var, idx_opt)) = extract_derivative_scalar(&eq.lhs) {
3243 match idx_opt {
3244 Some(indices) => lower_indexed_derivative(
3245 var,
3246 &indices,
3247 &eq.rhs,
3248 var_shapes,
3249 &mut covered_slots,
3250 &mut rhs_rules,
3251 )?,
3252 None => lower_bare_derivative(
3253 var,
3254 &eq.rhs,
3255 model,
3256 &array_axes,
3257 var_shapes,
3258 &mut covered_slots,
3259 &mut rhs_rules,
3260 )?,
3261 }
3262 continue;
3263 }
3264 // Otherwise: algebraic equation (or something we don't support).
3265 // If the LHS is algebraic for an eliminated variable it was
3266 // already consumed above; ignore here.
3267 }
3268
3269 cover_held_at_ic_slots(held_at_ic, var_shapes, &mut covered_slots);
3270 check_state_slots_covered(slots, &covered_slots)?;
3271
3272 Ok(rhs_rules)
3273}
3274
3275/// Lower an array-op derivative over `(idx_names, ranges)` to one
3276/// [`RhsRule::ArrayLoop`], marking the covered slots.
3277fn lower_faq_derivative(
3278 d: DerivArrayop,
3279 var_shapes: &IndexMap<String, VarShape>,
3280 covered_slots: &mut SlotCoverage,
3281 rhs_rules: &mut Vec<RhsRule>,
3282) -> Result<(), CompileError> {
3283 let DerivArrayop {
3284 var,
3285 idx_names,
3286 ranges,
3287 lhs_idx_exprs,
3288 body,
3289 contract_names,
3290 contract_dims,
3291 reduce,
3292 filter,
3293 } = d;
3294 if !var_shapes.contains_key(&var) {
3295 return Err(CompileError::build_err(format!(
3296 "Array-op derivative targets unknown state variable '{var}'"
3297 )));
3298 }
3299 // Mark the covered slots.
3300 mark_faq_lhs_coverage(
3301 &var_shapes[&var],
3302 &idx_names,
3303 &ranges,
3304 &lhs_idx_exprs,
3305 covered_slots,
3306 );
3307 rhs_rules.push(RhsRule::ArrayLoop {
3308 var_name: var,
3309 output_idx_names: idx_names,
3310 output_ranges: ranges,
3311 lhs_idx_exprs,
3312 body: Box::new(body),
3313 contract_names,
3314 contract_dims,
3315 reduce,
3316 filter,
3317 });
3318 Ok(())
3319}
3320
3321/// Lower an indexed scalar derivative `D(var[i,…], t) = rhs` to one
3322/// [`RhsRule::IndexedScalar`] at the resolved slot.
3323fn lower_indexed_derivative(
3324 var: String,
3325 indices: &[i64],
3326 rhs: &Expr,
3327 var_shapes: &IndexMap<String, VarShape>,
3328 covered_slots: &mut SlotCoverage,
3329 rhs_rules: &mut Vec<RhsRule>,
3330) -> Result<(), CompileError> {
3331 // Indexed: find slot.
3332 let shape = var_shapes.get(&var).ok_or_else(|| {
3333 CompileError::build_err(format!(
3334 "Scalar derivative targets unknown state variable '{var}'"
3335 ))
3336 })?;
3337 let flat = multi_to_flat_col_major(indices, &shape.shape, &shape.origin);
3338 let slot = shape.flat_offset + flat;
3339 covered_slots.insert(slot);
3340 rhs_rules.push(RhsRule::IndexedScalar {
3341 slot,
3342 body: Box::new(rhs.clone()),
3343 });
3344 Ok(())
3345}
3346
3347/// Lower a bare-variable derivative `D(var, t) = rhs`: a plain scalar slot
3348/// write when the state is 0-D, otherwise one of the two whole-array lifts
3349/// ([`lower_wholearray_producer_lift`] / [`lower_wholearray_percell`]).
3350fn lower_bare_derivative(
3351 var: String,
3352 rhs: &Expr,
3353 model: &Model,
3354 array_axes: &HashMap<String, Vec<String>>,
3355 var_shapes: &IndexMap<String, VarShape>,
3356 covered_slots: &mut SlotCoverage,
3357 rhs_rules: &mut Vec<RhsRule>,
3358) -> Result<(), CompileError> {
3359 let shape = var_shapes
3360 .get(&var)
3361 .ok_or_else(|| {
3362 CompileError::build_err(format!(
3363 "Scalar derivative targets unknown state variable '{var}'"
3364 ))
3365 })?
3366 .clone();
3367 // The result's declared index-set axis names, when they line up
3368 // with the shape actually compiled for it. This is what makes
3369 // the bare RHS's operands alignable BY NAME below; a state
3370 // whose shape was inferred from index usage rather than
3371 // declared (or declared at a different rank) has no usable
3372 // names and keeps the positional lowering.
3373 let target_axes: Option<&[String]> = model
3374 .variables
3375 .get(&var)
3376 .and_then(|v| v.shape.as_deref())
3377 .filter(|d| d.len() == shape.shape.len());
3378 if shape.shape.is_empty() {
3379 // Plain scalar D(var, t) = rhs.
3380 let slot = shape.flat_offset;
3381 covered_slots.insert(slot);
3382 rhs_rules.push(RhsRule::Scalar {
3383 slot,
3384 body: Box::new(rhs.clone()),
3385 });
3386 } else if rhs_has_array_producer(rhs) {
3387 lower_wholearray_producer_lift(
3388 var,
3389 rhs,
3390 &shape,
3391 target_axes,
3392 array_axes,
3393 covered_slots,
3394 rhs_rules,
3395 )?;
3396 } else {
3397 lower_wholearray_percell(
3398 var,
3399 rhs,
3400 &shape,
3401 target_axes,
3402 array_axes,
3403 var_shapes,
3404 covered_slots,
3405 rhs_rules,
3406 )?;
3407 }
3408 Ok(())
3409}
3410
3411/// Whole-array `D(var) = <rhs containing a lowered stencil>`
3412/// (an array-PRODUCING `makearray`/`faq` in elementwise
3413/// position, the form a §9.6.3 discretization rewrite emits):
3414/// lift to the per-cell `faq` (ArrayLoop) form the
3415/// derivative partition consumes — output loops over the full
3416/// declared shape, each array leaf and array producer gathered
3417/// per cell via `index(node, loops…)`. This is the loop-form
3418/// analog of the Julia `_lift_wholearray_deriv_equations`
3419/// (shape_promotion.jl) and keeps the rule eligible for the
3420/// vectorized whole-array fast path (ess-bdm).
3421fn lower_wholearray_producer_lift(
3422 var: String,
3423 rhs: &Expr,
3424 shape: &VarShape,
3425 target_axes: Option<&[String]>,
3426 array_axes: &HashMap<String, Vec<String>>,
3427 covered_slots: &mut SlotCoverage,
3428 rhs_rules: &mut Vec<RhsRule>,
3429) -> Result<(), CompileError> {
3430 let ndim = shape.shape.len();
3431 let loops: Vec<String> = (0..ndim).map(|d| format!("_lp{d}_{var}")).collect();
3432 let output_ranges: Vec<(i64, i64)> = shape
3433 .shape
3434 .iter()
3435 .zip(shape.origin.iter())
3436 .map(|(sz, o)| (*o, *o + *sz as i64 - 1))
3437 .collect();
3438 let lhs_idx_exprs: Vec<Expr> = loops.iter().map(|l| Expr::Variable(l.clone())).collect();
3439 let plan = build_gather_plan(rhs, array_axes, &var, target_axes, false)?;
3440 let body = index_array_leaves_by_loops(rhs, array_axes, Some(&plan), &loops);
3441 let total = shape.shape.iter().copied().product::<usize>().max(1);
3442 covered_slots.insert_range(shape.flat_offset, total);
3443 rhs_rules.push(RhsRule::ArrayLoop {
3444 var_name: var.clone(),
3445 output_idx_names: loops,
3446 output_ranges,
3447 lhs_idx_exprs,
3448 body: Box::new(body),
3449 contract_names: Vec::new(),
3450 contract_dims: Vec::new(),
3451 // No contracted index, so ⊕ never folds anything; carry the schema's
3452 // stated default. (This used to be spelled `effective_reduce_kind(None,
3453 // None)`, which only obscured that there is no node to read it off.)
3454 reduce: ReduceKind::Sum,
3455 filter: None,
3456 });
3457 Ok(())
3458}
3459
3460/// Whole-array `D(var) = <array-valued rhs>` over a declared
3461/// array shape, with elementwise semantics: each array-shaped
3462/// RHS leaf is indexed by the output cell. This is the
3463/// array-runtime analog of the Julia
3464/// `_lift_wholearray_deriv_equations` lift.
3465///
3466/// The rule is ONE [`RhsRule::ArrayLoop`] over the whole shape, its body
3467/// gathering every leaf by the loop symbols ([`CellCoords::Loops`]) — the
3468/// same rewrite the per-cell form applies with integer coordinates, so each
3469/// cell evaluates the identical scalar expression. Only a body
3470/// [`wholearray_body_loops_as_percell`] rejects keeps the per-cell form, one
3471/// [`RhsRule::IndexedScalar`] per cell.
3472#[allow(clippy::too_many_arguments)]
3473fn lower_wholearray_percell(
3474 var: String,
3475 rhs: &Expr,
3476 shape: &VarShape,
3477 target_axes: Option<&[String]>,
3478 array_axes: &HashMap<String, Vec<String>>,
3479 var_shapes: &IndexMap<String, VarShape>,
3480 covered_slots: &mut SlotCoverage,
3481 rhs_rules: &mut Vec<RhsRule>,
3482) -> Result<(), CompileError> {
3483 let plan = build_gather_plan(rhs, array_axes, &var, target_axes, true)?;
3484 let total = shape.shape.iter().copied().product::<usize>().max(1);
3485 if wholearray_body_loops_as_percell(rhs, array_axes, var_shapes) {
3486 let ndim = shape.shape.len();
3487 let loops: Vec<String> = (0..ndim).map(|d| format!("_lp{d}_{var}")).collect();
3488 let output_ranges: Vec<(i64, i64)> = shape
3489 .shape
3490 .iter()
3491 .zip(shape.origin.iter())
3492 .map(|(sz, o)| (*o, *o + *sz as i64 - 1))
3493 .collect();
3494 let lhs_idx_exprs: Vec<Expr> = loops.iter().map(|l| Expr::Variable(l.clone())).collect();
3495 let body = index_array_leaves(rhs, array_axes, Some(&plan), CellCoords::Loops(&loops));
3496 covered_slots.insert_range(shape.flat_offset, total);
3497 rhs_rules.push(RhsRule::ArrayLoop {
3498 var_name: var,
3499 output_idx_names: loops,
3500 output_ranges,
3501 lhs_idx_exprs,
3502 body: Box::new(body),
3503 contract_names: Vec::new(),
3504 contract_dims: Vec::new(),
3505 reduce: ReduceKind::Sum,
3506 filter: None,
3507 });
3508 return Ok(());
3509 }
3510 for flat in 0..total {
3511 let multi0 = flat_to_multi_col_major(flat, &shape.shape);
3512 let cell: Vec<i64> = multi0
3513 .iter()
3514 .zip(shape.origin.iter())
3515 .map(|(m, o)| *m as i64 + *o)
3516 .collect();
3517 let body = index_array_leaves(rhs, array_axes, Some(&plan), CellCoords::Ints(&cell));
3518 let slot = shape.flat_offset + flat;
3519 covered_slots.insert(slot);
3520 rhs_rules.push(RhsRule::IndexedScalar {
3521 slot,
3522 body: Box::new(body),
3523 });
3524 }
3525 Ok(())
3526}
3527
3528/// (7b) Held-at-ic states (no `D`, no algebraic definition) carry every
3529/// cell at its ic with zero derivative: mark their slots covered
3530/// without emitting a rule. The RHS zero-initializes `dy` each call
3531/// and never writes these slots, so they stay constant (a state that
3532/// feeds an observed — e.g. `phi` into `heat_release` — must not
3533/// drift).
3534fn cover_held_at_ic_slots(
3535 held_at_ic: &HashSet<String>,
3536 var_shapes: &IndexMap<String, VarShape>,
3537 covered_slots: &mut SlotCoverage,
3538) {
3539 for name in held_at_ic {
3540 if let Some(vs) = var_shapes.get(name) {
3541 let total = vs.shape.iter().copied().product::<usize>().max(1);
3542 covered_slots.insert_range(vs.flat_offset, total);
3543 }
3544 }
3545}
3546
3547/// (8) Every state slot must have a defining equation.
3548fn check_state_slots_covered(
3549 slots: &SlotTables,
3550 covered_slots: &SlotCoverage,
3551) -> Result<(), CompileError> {
3552 match covered_slots.first_uncovered() {
3553 Some(slot) => {
3554 let name = slot_name(&slots.var_shapes, slot).unwrap_or_default();
3555 Err(CompileError::build_err(format!(
3556 "State slot '{name}' has no defining derivative equation."
3557 )))
3558 }
3559 None => Ok(()),
3560 }
3561}
3562
3563/// Evaluate a state-free build-time expression (grid geometry, §11.4.1
3564/// coordinate-expression `ic` RHSs, §6.6.5 analytic `reference`s) through the
3565/// official array evaluator.
3566///
3567/// **The reference evaluator, under every compiler.** No compiler tier has a
3568/// form for initial-state assembly, so a field `ic` is evaluated here whatever
3569/// the Problem names. What keeps that honest is that construction EXERCISES
3570/// every such evaluation ([`ArrayCompiled::field_ic_records`]): one the
3571/// whole-array overlay serves is reported as vectorized, and one that walks a
3572/// `faq` per cell is refused by a strict compiler, naming the target
3573/// (esm-libraries-spec §2.5.10). An inline test's analytic `reference` also
3574/// comes here, and stays here under every compiler: it is the test's oracle,
3575/// not the model (`crate::inline_tests::evaluate_cellwise`).
3576/// Array-producing `faq`/`makearray` nodes
3577/// yield arrays; elementwise ops broadcast over them. Any `{ "from": <set> }`
3578/// range references are resolved against `index_sets` first, so a raw
3579/// (pre-compile) expression evaluates exactly as an equation expression does
3580/// after [`crate::faq::resolve_aggregate_ranges`].
3581///
3582/// STATE references are not in scope — the context carries no states. Model
3583/// PARAMETERS (load-time constants) ARE in scope when supplied via `params`
3584/// (name → value): a parameter-dependent coordinate expression / reference then
3585/// resolves (esm-spec §6.6.5). Mirrors the Python `_eval_buildtime_field` /
3586/// Julia `_eval_cellwise` machinery.
3587pub(crate) fn eval_buildtime_field(
3588 expr: &Expr,
3589 index_sets: &HashMap<String, IndexSet>,
3590 params: &HashMap<String, f64>,
3591) -> Result<Value, CompileError> {
3592 eval_buildtime_field_in_scope(expr, index_sets, params, &HashMap::new())
3593}
3594
3595/// [`eval_buildtime_field`] with a build-time ARRAY scope bound by name — a
3596/// provider-served forcing field, a shaped parameter's inline column, and the
3597/// STATE-FREE array observeds an `ic` RHS may read (esm-spec §6.6.5). The arrays
3598/// bind exactly where state arrays would, so a bare read yields the whole field
3599/// and `index(name, k)` gathers a cell.
3600pub(crate) fn eval_buildtime_field_in_scope(
3601 expr: &Expr,
3602 index_sets: &HashMap<String, IndexSet>,
3603 params: &HashMap<String, f64>,
3604 scope: &HashMap<String, ArrayD<f64>>,
3605) -> Result<Value, CompileError> {
3606 let mut resolved = expr.clone();
3607 crate::faq::resolve_expr_ranges(&mut resolved, index_sets)?;
3608 let param_names: Vec<String> = params.keys().cloned().collect();
3609 let param_vec: Vec<f64> = param_names.iter().map(|n| params[n]).collect();
3610 eval_expression(&resolved, scope, ¶m_vec, ¶m_names, 0.0)
3611}
3612
3613/// Resolve one grid cell's initial value for a scoped-reference / array `ic`
3614/// equation (esm-spec §11.4.1). `cell` is the 0-based multi-index of the element
3615/// within the target's grid shape. Supported RHS forms, in order:
3616///
3617/// 1. A LOADED FIELD — a bare reference to a provider-served forcing entry that
3618/// supplies the initial field over the lifted grid. The cell is read directly
3619/// when the field's rank matches the target grid; a single-element field is
3620/// broadcast.
3621/// 2. A BROADCAST CONSTANT — an RHS that const-folds to a finite scalar.
3622/// 3. A COORDINATE EXPRESSION — an elementwise expression over array-producing
3623/// `faq`/`makearray` nodes (e.g. `cos(pi * x_coord)` where `x_coord`
3624/// is a grid-geometry aggregate expanded from a §9.7 template import),
3625/// evaluated through the official array evaluator ([`eval_buildtime_field`])
3626/// in a state-free context and indexed at this cell.
3627///
3628/// Anything else is a hard error, so a scoped-reference ic that cannot be resolved
3629/// is never silently dropped. Mirrors tree_walk.jl `_resolve_field_ic` and the
3630/// Python `_resolve_field_ic`.
3631pub(super) fn resolve_field_ic_cell(
3632 target: &str,
3633 rhs: &Expr,
3634 cell: &[usize],
3635 forcing: &HashMap<String, ArrayD<f64>>,
3636 index_sets: &HashMap<String, IndexSet>,
3637 params: &HashMap<String, f64>,
3638 // Per-target memo of the case-(2) constant and the case-(3) whole-field
3639 // evaluation (both cell-independent), so the expression is evaluated once
3640 // per target rather than once per cell. `None` on entry for the first cell;
3641 // filled on first use.
3642 cached_field: &mut Option<Value>,
3643) -> Result<f64, SimulateError> {
3644 // (1) Loaded field served through the provider forcing buffer.
3645 if let Expr::Variable(name) = rhs
3646 && let Some(arr) = forcing.get(name)
3647 {
3648 if arr.ndim() == cell.len() {
3649 return Ok(arr[IxDyn(cell)]);
3650 } else if arr.len() == 1 {
3651 return Ok(arr.iter().copied().next().unwrap());
3652 }
3653 return Err(SimulateError::InvalidFieldInitialCondition {
3654 name: target.to_string(),
3655 details: format!(
3656 "loaded field '{name}' has ndim={} which does not match the {}-D lifted target grid",
3657 arr.ndim(),
3658 cell.len()
3659 ),
3660 });
3661 }
3662 // (2) Broadcast constant. Finite-only, and `Ok`-only: an op with no scalar
3663 // value (a `faq` grid-geometry node) must fall through to the
3664 // coordinate-expression path below — never silently seed the state vector.
3665 // `evaluate` refuses one by name, which the `if let Ok(..)` catches; the
3666 // `is_finite()` guard rejects a genuine `1.0/0.0`.
3667 //
3668 // Evaluated ONCE per target and memoized with the case-(3) field: the
3669 // value does not depend on the cell, so walking the expression again for
3670 // each cell would be a per-cell tree walk producing one number.
3671 if let Some(Value::Scalar(c)) = cached_field.as_ref()
3672 && c.is_finite()
3673 {
3674 return Ok(*c);
3675 }
3676 if cached_field.is_none()
3677 && let Ok(c) = crate::expression::evaluate(rhs, params)
3678 && c.is_finite()
3679 {
3680 *cached_field = Some(Value::Scalar(c));
3681 return Ok(c);
3682 }
3683 // (3) Coordinate expression over grid-geometry aggregates (model
3684 // parameters — e.g. a free-name geometry `x0`/`dx` — bind via `params`).
3685 // The whole-field evaluation is memoized in `cached_field` (see caller): it
3686 // is cell-independent, so it runs once per target instead of once per cell.
3687 if let Expr::Operator(_) = rhs {
3688 if cached_field.is_none() {
3689 // On a `CompileError` the memo stays empty and we fall through to
3690 // the case-(4) hard error below (byte-identical to the old
3691 // `match … { _ => {} }` arm, which likewise dropped the error).
3692 if let Ok(v) = eval_buildtime_field(rhs, index_sets, params) {
3693 *cached_field = Some(v);
3694 }
3695 }
3696 match cached_field.as_ref() {
3697 Some(Value::Scalar(s)) if s.is_finite() => return Ok(*s),
3698 Some(Value::Array(arr)) => {
3699 if arr.ndim() != cell.len() {
3700 return Err(SimulateError::InvalidFieldInitialCondition {
3701 name: target.to_string(),
3702 details: format!(
3703 "coordinate expression evaluates to ndim={}, which does not match the {}-D lifted target grid",
3704 arr.ndim(),
3705 cell.len()
3706 ),
3707 });
3708 }
3709 let v = arr[IxDyn(cell)];
3710 if v.is_finite() {
3711 return Ok(v);
3712 }
3713 }
3714 _ => {}
3715 }
3716 }
3717 // (4) Unsupported RHS — a clear error, never a silent drop.
3718 let hint = match rhs {
3719 Expr::Variable(name) => format!(" (no provider field named '{name}')"),
3720 _ => String::new(),
3721 };
3722 Err(SimulateError::InvalidFieldInitialCondition {
3723 name: target.to_string(),
3724 details: format!(
3725 "RHS is neither a provider-served loaded field, a constant, nor a per-cell coordinate expression{hint}"
3726 ),
3727 })
3728}
3729
3730// ============================================================================
3731// Shape inference + LHS parsing helpers.
3732// ============================================================================
3733
3734/// The variable a top-level equation defines, if any: `v = …`, `index(v, …) = …`,
3735/// `D(v) = …` / `D(index(v, …)) = …`, `ic(v) = …`, or a `faq`
3736/// whose body is `D(index(v, …))` / `index(v, …)`. Used to prune value-invention
3737/// equations and to classify algebraic definitions.
3738pub(super) fn equation_defined_var(lhs: &Expr) -> Option<String> {
3739 match lhs {
3740 Expr::Variable(v) => Some(v.clone()),
3741 Expr::Operator(node) => match node.op.as_str() {
3742 "index" => match node.args.first() {
3743 Some(Expr::Variable(v)) => Some(v.clone()),
3744 _ => None,
3745 },
3746 "D" | "ic" => match node.args.first() {
3747 Some(Expr::Variable(v)) => Some(v.clone()),
3748 Some(inner) => equation_defined_var(inner),
3749 None => None,
3750 },
3751 op if is_faq_op(op) => node.expr.as_ref().and_then(|b| equation_defined_var(b)),
3752 _ => None,
3753 },
3754 _ => None,
3755 }
3756}
3757
3758/// The state variable an *algebraic* (non-`D`, non-`ic`) equation defines, if any.
3759/// A state so-defined is eliminated to an observed rather than integrated.
3760pub(super) fn algebraic_defined_var(lhs: &Expr) -> Option<String> {
3761 match lhs {
3762 Expr::Variable(v) => Some(v.clone()),
3763 Expr::Operator(node) => match node.op.as_str() {
3764 "index" => match node.args.first() {
3765 Some(Expr::Variable(v)) => Some(v.clone()),
3766 _ => None,
3767 },
3768 op if is_faq_op(op) => {
3769 // `faq(expr = index(v, …))` — but NOT `expr = D(index(v, …))`,
3770 // which is a derivative, not an algebraic definition.
3771 let body = node.expr.as_ref()?;
3772 if let Expr::Operator(b) = body.as_ref() {
3773 if b.op == "D" {
3774 return None;
3775 }
3776 }
3777 equation_defined_var(body)
3778 }
3779 _ => None,
3780 },
3781 _ => None,
3782 }
3783}
3784
3785/// Every state variable defined by an algebraic equation (see
3786/// [`algebraic_defined_var`]).
3787pub(super) fn collect_algebraic_defined(equations: &[crate::types::Equation]) -> HashSet<String> {
3788 let mut out = HashSet::new();
3789 for eq in equations {
3790 if let Some(v) = algebraic_defined_var(&eq.lhs) {
3791 out.insert(v);
3792 }
3793 }
3794 out
3795}
3796
3797/// True if `expr` (or any subexpression) uses a `skolem` op — the marker of a
3798/// value-invention (relational) producer whose integer-id buffer the dense
3799/// array evaluator does not materialize.
3800pub(super) fn expr_contains_skolem(expr: &Expr) -> bool {
3801 match expr {
3802 Expr::Number(_) | Expr::Integer(_) | Expr::Variable(_) => false,
3803 Expr::Operator(node) => node.op == "skolem" || node.any_child(&mut expr_contains_skolem),
3804 }
3805}
3806
3807/// Whether `expr` contains a `rank` over a `kind: "derived"` index set — the
3808/// dense ids of an invented set, a build-time relational output.
3809fn ranks_a_derived_set(expr: &Expr, index_sets: &HashMap<String, IndexSet>) -> bool {
3810 match expr {
3811 Expr::Operator(node) => {
3812 (node.op == "rank"
3813 && node.args.iter().any(|a| {
3814 matches!(a, Expr::Variable(v)
3815 if index_sets.get(v).is_some_and(|is| is.kind == "derived"))
3816 }))
3817 || node.any_child(&mut |c| ranks_a_derived_set(c, index_sets))
3818 }
3819 _ => false,
3820 }
3821}
3822
3823/// Collect the `id`s of every geometry ring producer (`intersect_polygon` /
3824/// `polygon_intersection_area`) reachable from `expr`. A `kind: "derived"` index
3825/// set whose `from_faq` names one of these is the materialized clip ring (kept),
3826/// distinguishing it from a relational membership set (dropped).
3827pub(super) fn collect_geometry_producer_ids(expr: &Expr, out: &mut HashSet<String>) {
3828 let Expr::Operator(node) = expr else {
3829 return;
3830 };
3831 if matches!(
3832 node.op.as_str(),
3833 "intersect_polygon" | "polygon_intersection_area"
3834 ) {
3835 if let Some(id) = &node.id {
3836 out.insert(id.clone());
3837 }
3838 }
3839 node.for_each_child(&mut |child| collect_geometry_producer_ids(child, out));
3840}
3841
3842/// Strip every `join.on` key-pair that references a dropped value-invention
3843/// variable, dropping a clause whose pairs all vanish and clearing an empty
3844/// `join`. A bin-skolem broad-phase gate keyed on such a column cannot be
3845/// evaluated by the dense array runtime (the integer-id buffer is not
3846/// materialized); eliding it degrades the aggregate to the dense contraction
3847/// over all index combinations — the pruned combinations contribute the
3848/// additive identity, which for a geometric narrow phase (`polygon_intersection_area`,
3849/// zero on non-overlapping pairs) they already do, so the result is unchanged.
3850pub(super) fn strip_vi_joins(expr: &mut Expr, vi_cols: &HashSet<String>) {
3851 // Sharing-aware gate: after load-time interning (`crate::intern`)
3852 // operator payloads are shared `Arc`s, and a mutable descent
3853 // copy-on-write splits every node it touches. Only branches actually
3854 // carrying a `join` clause are descended; a join-free subtree (every
3855 // subtree of a §9.7-expanded discretization) is left fully shared.
3856 fn contains_join(e: &Expr) -> bool {
3857 match e {
3858 Expr::Operator(node) => node.join.is_some() || node.any_child(&mut contains_join),
3859 _ => false,
3860 }
3861 }
3862 if !contains_join(expr) {
3863 return;
3864 }
3865 let Some(node) = expr.node_mut() else {
3866 return;
3867 };
3868 if let Some(joins) = &mut node.join {
3869 for clause in joins.iter_mut() {
3870 clause
3871 .on
3872 .retain(|pair| !pair.iter().any(|c| vi_cols.contains(c)));
3873 }
3874 // An OVERLAP clause carries no `on` pairs at all and is kept: it names
3875 // const-array envelope FACTORS, not value-invention id columns, and it
3876 // DRIVES the dense enumeration (CONFORMANCE_SPEC §5.5.6). Only a
3877 // bin-equality clause whose every pair referenced a dropped column is
3878 // now empty and inert.
3879 joins.retain(|clause| !clause.on.is_empty() || clause.overlap.is_some());
3880 if joins.is_empty() {
3881 node.join = None;
3882 }
3883 // A `faq`'s `args` is the declarative operand list (esm-spec §4.3.1);
3884 // the key columns the gate read leave it with the gate.
3885 if is_faq_op(&node.op) {
3886 node.args
3887 .retain(|a| !matches!(a, Expr::Variable(v) if vi_cols.contains(v)));
3888 }
3889 }
3890 node.for_each_child_mut(&mut |child| strip_vi_joins(child, vi_cols));
3891}
3892
3893/// Drop value-invention (relational) variables and their defining equations, and
3894/// strip broad-phase `join.on` gates that reference them, BEFORE join / range
3895/// resolution and shape inference.
3896///
3897/// The dense Rust array runtime evaluates FAQ aggregates and the fused geometry
3898/// leaf, but does NOT materialize value-invention buffers — skolem-id maps
3899/// (`skolem`/`rank`) or a membership set over a `kind: "derived"` (FAQ-produced)
3900/// index set. [`materialize_derived_extents`] has already counted each such
3901/// set's members, so a range over it is sized; its per-element values are not
3902/// available. A variable that is one of these, and the `join.on` gate keyed on
3903/// it, are relational scaffolding around a densely-evaluable narrow phase. For a
3904/// conservative regrid the narrow phase is `polygon_intersection_area`, which is
3905/// zero on exactly the pairs the bin-skolem gate would prune, so the dense
3906/// contraction is numerically identical (see [`strip_vi_joins`]). Genuine
3907/// (loop-symbol) joins and non-VI models stay byte-identical: the pass is a
3908/// no-op unless a `skolem` op or a derived-set-shaped variable is present.
3909pub(super) fn strip_value_invention(
3910 model: &mut Model,
3911 index_sets: &HashMap<String, IndexSet>,
3912 vi: &ViCensus,
3913) -> Result<(), CompileError> {
3914 let mut vi_vars: HashSet<String> = HashSet::new();
3915 // Ids of geometry ring producers (`intersect_polygon` / `polygon_intersection_area`).
3916 // A `kind: "derived"` index set whose `from_faq` names one of these IS
3917 // materialized by the dense runtime (the clipped overlap ring), so a variable
3918 // shaped over it — e.g. a geometry `clip` — must be KEPT.
3919 let any_derived = index_sets.values().any(|is| is.kind == "derived");
3920 // Only read for a shape over a derived set, so not collected without one.
3921 let mut geom_ids: HashSet<String> = HashSet::new();
3922 if any_derived && vi.geometry {
3923 for eq in &model.equations {
3924 collect_geometry_producer_ids(&eq.rhs, &mut geom_ids);
3925 }
3926 // An observed unknown's defining body is one of those equation RHSs from
3927 // esm 1.0.0, so the loop above already covers it; what is left on a
3928 // variable is a parameter `update`'s expressions.
3929 for var in model.variables.values() {
3930 var.for_each_expression(&mut |expr| collect_geometry_producer_ids(expr, &mut geom_ids));
3931 }
3932 }
3933 // (a) A variable shaped over a `kind: "derived"` index set whose FAQ producer
3934 // is NOT a geometry ring producer — a relational membership / candidate
3935 // set the dense runtime does not enumerate. Build-time DATA over such a
3936 // set (a parameter, or an unknown defined by a `const` node, e.g. a
3937 // per-edge length over invented edges) is an ordinary array sized by the
3938 // set's extent, and is kept.
3939 let const_defined: HashSet<String> = model
3940 .equations
3941 .iter()
3942 .filter(|eq| matches!(&eq.rhs, Expr::Operator(n) if n.op == "const"))
3943 .filter_map(|eq| equation_defined_var(&eq.lhs))
3944 .collect();
3945 for (name, var) in &model.variables {
3946 if var.var_type == VariableType::Parameter || const_defined.contains(name) {
3947 continue;
3948 }
3949 if let Some(shape) = &var.shape {
3950 if shape.iter().any(|s| {
3951 index_sets
3952 .get(s)
3953 .filter(|is| is.kind == "derived")
3954 .map(|is| {
3955 !is.from_faq
3956 .as_deref()
3957 .map(|f| geom_ids.contains(f))
3958 .unwrap_or(false)
3959 })
3960 .unwrap_or(false)
3961 }) {
3962 vi_vars.insert(name.clone());
3963 }
3964 }
3965 }
3966 // (b) A variable defined by an equation whose RHS produces a skolem id, or
3967 // a `rank` dense id over an invented set: build-time relational
3968 // outputs, dropped as the other bindings drop them.
3969 for eq in &model.equations {
3970 if (vi.skolem && expr_contains_skolem(&eq.rhs))
3971 || (any_derived && ranks_a_derived_set(&eq.rhs, index_sets))
3972 {
3973 if let Some(v) = equation_defined_var(&eq.lhs) {
3974 vi_vars.insert(v);
3975 }
3976 }
3977 }
3978 if vi_vars.is_empty() {
3979 return Ok(());
3980 }
3981 // Kept for the diagnostic at the end: which op defined each stripped
3982 // variable.
3983 let stripped_bodies: HashMap<String, Expr> = model
3984 .equations
3985 .iter()
3986 .filter_map(|eq| {
3987 let v = equation_defined_var(&eq.lhs)?;
3988 vi_vars.contains(&v).then(|| (v, eq.rhs.clone()))
3989 })
3990 .collect();
3991 model.variables.retain(|k, _| !vi_vars.contains(k));
3992 model.equations.retain(|eq| {
3993 equation_defined_var(&eq.lhs)
3994 .map(|v| !vi_vars.contains(&v))
3995 .unwrap_or(true)
3996 });
3997 // `join.on` columns are NOT namespaced by flatten (they are bare strings, not
3998 // expressions), while the dropped variable keys ARE model-prefixed. Match a
3999 // join column against both the qualified VI name and its unqualified suffix so
4000 // a coupled `join.on [[rg_src_bin, rg_tgt_bin]]` gate is stripped even though
4001 // its columns stayed unqualified.
4002 let mut vi_cols: HashSet<String> = vi_vars.clone();
4003 for v in &vi_vars {
4004 if let Some(pos) = v.rfind('.') {
4005 vi_cols.insert(v[pos + 1..].to_string());
4006 }
4007 }
4008 for eq in &mut model.equations {
4009 strip_vi_joins(&mut eq.lhs, &vi_cols);
4010 strip_vi_joins(&mut eq.rhs, &vi_cols);
4011 }
4012 for var in model.variables.values_mut() {
4013 var.for_each_expression_mut(&mut |expr| strip_vi_joins(expr, &vi_cols));
4014 }
4015 // A stripped producer that something SURVIVING still reads was never a
4016 // legitimate strip: its only definition was an evaluable-core op with no
4017 // evaluation rule, and an earlier stage was supposed to have eliminated
4018 // it. Report THAT — `unevaluable_operator`, naming the op — rather than
4019 // letting the free-variable gate below report the dangling name, which
4020 // sends an author looking for a typo in a variable they declared.
4021 //
4022 // Both outcomes are refusals, so nothing that built before stops building;
4023 // only the wording changes, and only for a document that was already dead.
4024 for eq in &model.equations {
4025 let mut hit: Option<String> = None;
4026 let mut find = |expr: &Expr| {
4027 if hit.is_none() {
4028 hit = first_variable_reference_in(expr, &vi_vars);
4029 }
4030 };
4031 find(&eq.lhs);
4032 find(&eq.rhs);
4033 if let Some(name) = hit {
4034 let op = stripped_producer_op(&name, &stripped_bodies)
4035 .unwrap_or_else(|| "skolem".to_string());
4036 return Err(CompileError::UnevaluableOperatorError { op });
4037 }
4038 }
4039 Ok(())
4040}
4041
4042/// The first name in `expr` that `stripped` contains, if any.
4043fn first_variable_reference_in(expr: &Expr, stripped: &HashSet<String>) -> Option<String> {
4044 match expr {
4045 Expr::Variable(v) if stripped.contains(v) => Some(v.clone()),
4046 Expr::Operator(node) => {
4047 let mut out = None;
4048 node.any_child(&mut |child| {
4049 out = first_variable_reference_in(child, stripped);
4050 out.is_some()
4051 });
4052 out
4053 }
4054 _ => None,
4055 }
4056}
4057
4058/// The evaluable-core op at the head of the body that defined the stripped
4059/// variable `name`, for the diagnostic above.
4060fn stripped_producer_op(name: &str, bodies: &HashMap<String, Expr>) -> Option<String> {
4061 fn head(expr: &Expr) -> Option<String> {
4062 match expr {
4063 Expr::Operator(node) if !crate::simulate_array::is_evaluable_op(&node.op) => {
4064 Some(node.op.clone())
4065 }
4066 Expr::Operator(node) => {
4067 let mut out = None;
4068 node.any_child(&mut |child| {
4069 out = head(child);
4070 out.is_some()
4071 });
4072 out
4073 }
4074 _ => None,
4075 }
4076 }
4077 head(bodies.get(name)?)
4078}
4079
4080// ===========================================================================
4081// Value-invention OUTPUT front-door — materialize the arg-witness / grouped
4082// relational outputs to constant data (the previously-unwired engine).
4083// ===========================================================================
4084
4085/// True iff any equation or observed expression in `model` contains an
4086/// arg-witness reducer (`argmin` / `argmax`). Mirrors [`expr_contains_skolem`]:
4087/// its presence is the marker of a genuine relational OUTPUT — a per-element
4088/// nearest-witness INDEX buffer (RFC §5.7 rule 6) — that the dense evaluator
4089/// cannot run and [`strip_value_invention`] does not remove (it is neither a
4090/// derived-set-shaped var nor a skolem producer). This gates the build-time
4091/// materialize pass so every model WITHOUT one stays byte-identical.
4092fn expr_contains_arg_witness(expr: &Expr) -> bool {
4093 match expr {
4094 Expr::Number(_) | Expr::Integer(_) | Expr::Variable(_) => false,
4095 Expr::Operator(node) => {
4096 node.op == "argmin"
4097 || node.op == "argmax"
4098 || node.any_child(&mut expr_contains_arg_witness)
4099 }
4100 }
4101}
4102
4103/// Which value-invention producers a model's expressions (the equations, both
4104/// sides, and the variables' own expressions) contain anywhere: an
4105/// `argmin`/`argmax` witness, a `skolem`, a geometry ring producer.
4106pub(super) struct ViCensus {
4107 arg_witness: bool,
4108 skolem: bool,
4109 geometry: bool,
4110}
4111
4112impl ViCensus {
4113 fn of(model: &Model) -> Self {
4114 fn walk(expr: &Expr, c: &mut ViCensus) {
4115 let Expr::Operator(node) = expr else {
4116 return;
4117 };
4118 match node.op.as_str() {
4119 "argmin" | "argmax" => c.arg_witness = true,
4120 "skolem" => c.skolem = true,
4121 "intersect_polygon" | "polygon_intersection_area" => c.geometry = true,
4122 _ => {}
4123 }
4124 node.for_each_child(&mut |child| walk(child, c));
4125 }
4126 let mut c = ViCensus {
4127 arg_witness: false,
4128 skolem: false,
4129 geometry: false,
4130 };
4131 for eq in &model.equations {
4132 walk(&eq.lhs, &mut c);
4133 walk(&eq.rhs, &mut c);
4134 }
4135 for var in model.variables.values() {
4136 var.for_each_expression(&mut |expr| walk(expr, &mut c));
4137 }
4138 c
4139 }
4140}
4141
4142fn model_contains_arg_witness(model: &Model) -> bool {
4143 model
4144 .equations
4145 .iter()
4146 .any(|eq| expr_contains_arg_witness(&eq.lhs) || expr_contains_arg_witness(&eq.rhs))
4147 || model.variables.values().any(|v| {
4148 let mut found = false;
4149 v.for_each_expression(&mut |e| found |= expr_contains_arg_witness(e));
4150 found
4151 })
4152}
4153
4154/// Is this parameter refreshed from OUTSIDE the model — every update rule
4155/// reading either a data source (`from`) or a registered handler?
4156///
4157/// Both are the FORCING seam: the value arrives from the runtime between
4158/// segments rather than being computed by the model, which is what lets this
4159/// backend serve it out of the forcing buffer with no event machinery of its
4160/// own (CONFORMANCE_SPEC §5.10.1, §5.13.2). A rule with an `expression` value
4161/// form is the opposite case — the model computes it, and something has to run
4162/// that computation on each refresh.
4163/// The declared `default` of every parameter the forcing buffer serves that
4164/// declares one, as the dense field it denotes over the parameter's resolved
4165/// shape: a scalar broadcast over the whole grid (esm-spec §6.3), or its inline
4166/// array data. `None` for a parameter whose shape does not resolve yet (an
4167/// unmaterialized derived set), which has no field to fill.
4168///
4169/// This is the parameter's value when the caller supplies no data for it:
4170/// construction writes it into the buffer ([`crate::problem::esm_problem`]).
4171///
4172/// # Errors
4173///
4174/// [`CompileError::InterpreterBuildError`] when inline array data does not
4175/// match the declared shape (esm-spec §6.6.2).
4176#[allow(clippy::type_complexity)]
4177fn forcing_default_values(
4178 observed_vars: &[(&String, &ModelVariable)],
4179 forcing_decls: &IndexMap<String, Option<Vec<usize>>>,
4180) -> Result<HashMap<String, Option<(Vec<usize>, Vec<f64>)>>, CompileError> {
4181 let mut out = HashMap::new();
4182 for (name, var) in observed_vars {
4183 let (Some(default), Some(decl)) = (var.default.as_ref(), forcing_decls.get(name.as_str()))
4184 else {
4185 continue;
4186 };
4187 let Some(want) = decl else {
4188 out.insert((*name).clone(), None);
4189 continue;
4190 };
4191 let values = if let Some(scalar) = default.as_scalar() {
4192 vec![scalar; want.iter().product::<usize>()]
4193 } else {
4194 let (shape, values) = default.to_dense().map_err(|e| {
4195 CompileError::build_err(format!("parameter '{name}': {e} (esm-spec §6.3)"))
4196 })?;
4197 if &shape != want {
4198 return Err(CompileError::build_err(format!(
4199 "parameter '{name}': inline array data has shape {shape:?}, which does not \
4200 match the declared shape {want:?} (esm-spec §6.6.2)"
4201 )));
4202 }
4203 values
4204 };
4205 out.insert((*name).clone(), Some((want.clone(), values)));
4206 }
4207 Ok(out)
4208}
4209
4210/// The `data_sources` key of the first data feed in `spec` — a rule of
4211/// `kind: "data"` carrying a `from` binding (esm-spec §5.4). Lives here rather
4212/// than in `data_fed` because that module is not built for wasm32, while the
4213/// classifier is.
4214pub(crate) fn data_feed_source(spec: &crate::types::ParameterUpdateSpec) -> Option<&str> {
4215 spec.rules().iter().find_map(|r| {
4216 r.data_source()
4217 .filter(|_| r.value().is_some_and(|v| v.from.is_some()))
4218 })
4219}
4220
4221fn externally_refreshed(var: &ModelVariable) -> bool {
4222 let Some(spec) = &var.update else {
4223 return false;
4224 };
4225 spec.rules().iter().all(|rule| {
4226 rule.value()
4227 .is_some_and(|v| v.from.is_some() || v.handler.is_some())
4228 })
4229}
4230
4231/// name → defining RHS for every OBSERVED unknown of `model`
4232/// (esm-spec §6.3.1).
4233///
4234/// The one place this pipeline asks "what defines this observed?" — from esm
4235/// 1.0.0 the answer is the equation whose LHS is the bare variable, never a
4236/// `variables[v].expression` field. Sorted by name (a `BTreeMap`), so every
4237/// consumer iterates deterministically.
4238fn observed_bodies(model: &Model) -> std::collections::BTreeMap<String, Expr> {
4239 crate::classification::observed_definitions_of_parts(&model.variables, &model.equations)
4240}
4241
4242/// Gather the build-time-CONSTANT factor arrays the value-invention engine reads
4243/// (`index(gx, g)` etc.): every variable whose `expression` is a `const` op (the
4244/// established self-contained build-time array channel — see the geometry
4245/// `src_poly`/`tgt_poly` fixtures). Each is evaluated once, with no state /
4246/// params / `t` (a `const` literal needs none), into its dense `ArrayD`. This is
4247/// the Rust analogue of the Julia reference's `const_arrays` registry and the
4248/// Python interpreter's join-free const-observed pre-materialization.
4249fn collect_const_factor_arrays(model: &Model) -> HashMap<String, ArrayD<f64>> {
4250 let mut out: HashMap<String, ArrayD<f64>> = HashMap::new();
4251 for (name, body) in observed_bodies(model) {
4252 let expr = &body;
4253 let Expr::Operator(node) = expr else { continue };
4254 if node.op != "const" {
4255 continue;
4256 }
4257 match eval_expression(expr, &HashMap::new(), &[], &[], 0.0) {
4258 Ok(Value::Array(a)) => {
4259 out.insert(name.clone(), *a);
4260 }
4261 Ok(Value::Scalar(s)) => {
4262 out.insert(name.clone(), ArrayD::from_elem(IxDyn(&[]), s));
4263 }
4264 Err(_) => {}
4265 }
4266 }
4267 out
4268}
4269
4270/// The value-invention engine's factor arrays for `model`: the `const`-literal
4271/// variables [`collect_const_factor_arrays`] finds, OVERLAID with the caller's
4272/// arrays (caller wins on a name collision).
4273///
4274/// The overlay exists because a real model's factors do not live in the
4275/// document. In ISRM the overlap gate's envelope factors — `X`/`Y` (emission
4276/// points) and `W`/`S`/`E`/`N` (cell rectangles) — are declared
4277/// `type: "parameter"` and filled by data loaders through `coupling`
4278/// `param_to_var` edges. They carry no `const` expression, so the scan finds
4279/// nothing for them, the gate cannot build an envelope, and the producer
4280/// invents no members. Only the caller holds those arrays, and before this
4281/// there was no way to hand them over on the in-tree build path.
4282///
4283/// Caller-wins (rather than const-wins) is the deliberate direction: a caller
4284/// that supplies an array is asserting the value it loaded, and a document
4285/// `const` — typically a small placeholder or a default — must not silently
4286/// override real data. It also makes the channel usable for OVERRIDES, matching
4287/// the Julia reference's `const_arrays=` kwarg.
4288fn vi_factor_arrays<S: std::hash::BuildHasher>(
4289 model: &Model,
4290 caller_arrays: Option<&HashMap<String, ArrayD<f64>, S>>,
4291) -> HashMap<String, ArrayD<f64>> {
4292 let mut arrays = collect_const_factor_arrays(model);
4293 if let Some(extra) = caller_arrays {
4294 for (name, arr) in extra {
4295 arrays.insert(name.clone(), arr.clone());
4296 }
4297 }
4298 arrays
4299}
4300
4301/// The value-invention engine's scalar `params` map (e.g. the bin width of a
4302/// broad-phase skolem quantization): each 0-D parameter's caller override —
4303/// keyed by its name or its model-local tail — else its `default`. An array
4304/// parameter carries no scalar and is supplied (if at all) through
4305/// [`collect_const_factor_arrays`].
4306fn collect_scalar_params(model: &Model, overrides: &HashMap<String, f64>) -> HashMap<String, f64> {
4307 let mut out: HashMap<String, f64> = HashMap::new();
4308 for (name, var) in &model.variables {
4309 if var.var_type != VariableType::Parameter
4310 || !var.shape.as_ref().map(|s| s.is_empty()).unwrap_or(true)
4311 {
4312 continue;
4313 }
4314 let tail = name.rsplit('.').next().unwrap_or(name);
4315 let value = overrides
4316 .get(name)
4317 .or_else(|| overrides.get(tail))
4318 .copied()
4319 .or_else(|| var.default_scalar());
4320 if let Some(v) = value {
4321 out.insert(name.clone(), v);
4322 }
4323 }
4324 out
4325}
4326
4327/// Rewrite the equation (or observed `expression`) that DEFINES `name` into a
4328/// whole-array `const` literal carrying the materialized dense `buf` — the
4329/// "materialize to data" step. The relational op (`argmin` / `group_aggregate`)
4330/// that produced `name` is thereby replaced by data the existing oracle
4331/// evaluates, so it never reaches the run path as an
4332/// [`CompileError::UnevaluableOperatorError`]. The LHS collapses to the bare
4333/// variable (a whole-array assignment); the eliminated-state machinery then
4334/// materializes it as an ordinary constant observed.
4335fn rewrite_equation_to_const(model: &mut Model, name: &str, buf: &[f64]) {
4336 let value = JsonValue::Array(
4337 buf.iter()
4338 .map(|&v| {
4339 serde_json::Number::from_f64(v)
4340 .map(JsonValue::Number)
4341 .unwrap_or(JsonValue::Null)
4342 })
4343 .collect(),
4344 );
4345 let const_node = Expr::operator(ExpressionNode {
4346 op: "const".to_string(),
4347 value: Some(value),
4348 ..Default::default()
4349 });
4350 let mut replaced = false;
4351 for eq in &mut model.equations {
4352 if equation_defined_var(&eq.lhs).as_deref() == Some(name) {
4353 eq.lhs = Expr::Variable(name.to_string());
4354 eq.rhs = const_node.clone();
4355 replaced = true;
4356 }
4357 }
4358 // No equation defines `name` yet: esm 1.0.0 states an unknown's behaviour
4359 // in `equations`, so the folded const becomes a new bare-LHS equation
4360 // rather than a `variables[name].expression` field.
4361 if !replaced && model.variables.contains_key(name) {
4362 model.equations.push(crate::types::Equation {
4363 lhs: Expr::Variable(name.to_string()),
4364 rhs: const_node,
4365 comment: None,
4366 });
4367 }
4368}
4369
4370/// Run the build-time value-invention engine over a TYPED [`Model`] plus the
4371/// document-scoped `index_sets` registry, with a caller-supplied factor-array
4372/// channel — the front door a Rust runner needs, and the one the in-tree build
4373/// path (`materialize_vi_outputs_to_data`) uses internally.
4374///
4375/// [`materialize_value_invention`] itself takes a raw `serde_json::Value` and a
4376/// finished array map, which leaves a caller holding a typed model with two
4377/// chores it has no library support for: assembling the JSON view the engine
4378/// expects (model + registry merged as a sibling), and gathering the factor
4379/// arrays. This does both. Concretely it is what turns an ISRM document plus
4380/// loaded `X`/`Y`/`W`/`S`/`E`/`N` into
4381/// [`ValueInventionResult::extents`](crate::ValueInventionResult::extents) — the
4382/// map that then sizes every `emis_src_cells` axis through
4383/// [`crate::faq::resolve_aggregate_ranges_with_extents`] and
4384/// [`crate::simulate_array::eval_expression_with_extents`].
4385///
4386/// `caller_arrays` overlays the document's `const` factors (caller wins; see
4387/// `vi_factor_arrays`). Scalar `params` come from the model's own 0-D
4388/// parameter defaults.
4389///
4390/// A model with no `skolem`/`distinct`/`rank` producer yields an empty result
4391/// rather than an error — the engine's own no-op contract.
4392///
4393/// # Errors
4394///
4395/// [`CompileError::InterpreterBuildError`] if the model or registry cannot be
4396/// serialized to the engine's JSON view, or if the engine rejects the document
4397/// (e.g. a producer that classifies CONTINUOUS, §5.7 guard 2).
4398pub fn run_value_invention<S: std::hash::BuildHasher>(
4399 model: &Model,
4400 index_sets: &HashMap<String, IndexSet>,
4401 caller_arrays: Option<&HashMap<String, ArrayD<f64>, S>>,
4402) -> Result<ValueInventionResult, CompileError> {
4403 run_value_invention_with_params(model, index_sets, caller_arrays, &HashMap::new())
4404}
4405
4406/// [`run_value_invention`] with the caller's scalar parameter overrides (the
4407/// SciML `p`), which a producer reads in place of the declared defaults: a
4408/// bin width is build-time data, so a `p` entry for it sizes the invented set.
4409pub fn run_value_invention_with_params<S: std::hash::BuildHasher>(
4410 model: &Model,
4411 index_sets: &HashMap<String, IndexSet>,
4412 caller_arrays: Option<&HashMap<String, ArrayD<f64>, S>>,
4413 overrides: &HashMap<String, f64>,
4414) -> Result<ValueInventionResult, CompileError> {
4415 let const_arrays = vi_factor_arrays(model, caller_arrays);
4416 let params = collect_scalar_params(model, overrides);
4417
4418 // The engine walks the RAW `serde_json::Value` document (it preserves the
4419 // aggregate `key`/`distinct`/`arg` fields), with the document-scoped
4420 // `index_sets` registry merged down as a sibling — mirroring the engine's own
4421 // `model_json` fixture helper and `crate::cadence`.
4422 let mut model_json = serde_json::to_value(model).map_err(|e| {
4423 CompileError::build_err(format!("value-invention: could not serialize model: {e}"))
4424 })?;
4425 if let JsonValue::Object(m) = &mut model_json {
4426 let is_json = serde_json::to_value(index_sets).map_err(|e| {
4427 CompileError::build_err(format!(
4428 "value-invention: could not serialize index_sets: {e}"
4429 ))
4430 })?;
4431 m.insert("index_sets".to_string(), is_json);
4432 }
4433
4434 materialize_value_invention(&model_json, &const_arrays, ¶ms, &HashMap::new()).map_err(|e| {
4435 CompileError::build_err(format!("value-invention materialize failed: {}", e.0))
4436 })
4437}
4438
4439/// The build-time extents of a model's non-geometry derived index sets, and why
4440/// value invention failed when it did.
4441pub(super) struct DerivedMaterialization {
4442 /// `from_faq` producer id → member count, as [`run_value_invention`] reports.
4443 pub(super) extents: HashMap<String, i64>,
4444 /// Ids of the geometry ring producers, whose sets the runtime sizes itself.
4445 geometry_ids: HashSet<String>,
4446 /// The engine's error, if it could not run.
4447 failure: Option<String>,
4448 /// The skolem map buffers as integer codes
4449 /// ([`ValueInventionResult::map_codes`]).
4450 map_codes: HashMap<String, Vec<f64>>,
4451}
4452
4453/// Run value invention when the model has a derived index set whose producer is
4454/// not a geometry ring. A model without one pays only the scan.
4455fn materialize_derived_extents(
4456 model: &Model,
4457 index_sets: &HashMap<String, IndexSet>,
4458 caller_arrays: Option<&HashMap<String, ArrayD<f64>>>,
4459 vi: &ViCensus,
4460) -> DerivedMaterialization {
4461 let mut geometry_ids = HashSet::new();
4462 if vi.geometry {
4463 for eq in &model.equations {
4464 collect_geometry_producer_ids(&eq.rhs, &mut geometry_ids);
4465 }
4466 for var in model.variables.values() {
4467 var.for_each_expression(&mut |expr| {
4468 collect_geometry_producer_ids(expr, &mut geometry_ids)
4469 });
4470 }
4471 }
4472 let mut out = DerivedMaterialization {
4473 extents: HashMap::new(),
4474 geometry_ids,
4475 failure: None,
4476 map_codes: HashMap::new(),
4477 };
4478 // A skolem bin map that a `join.on` gate compares needs the engine too,
4479 // derived set or not: its codes keep the gate (`map_codes` below), which
4480 // is what makes the gated contraction cost its admitted pairs rather than
4481 // the full product on every evaluation.
4482 let needs_value_invention = index_sets.values().any(|is| {
4483 is.kind == "derived"
4484 && is
4485 .from_faq
4486 .as_deref()
4487 .is_some_and(|f| !out.geometry_ids.contains(f))
4488 }) || (vi.skolem && skolem_map_keys_a_join(model));
4489 if needs_value_invention {
4490 let factors = value_invention_factors(model, index_sets, caller_arrays);
4491 let lean = without_factor_literals(model, &factors);
4492 match run_value_invention(&lean, index_sets, Some(&factors)) {
4493 Ok(result) => {
4494 out.extents = result.extents;
4495 out.map_codes = result.map_codes;
4496 }
4497 Err(e) => out.failure = Some(e.to_string()),
4498 }
4499 }
4500 out
4501}
4502
4503/// Whether a `join.on` key pair anywhere in `model` names a variable defined by
4504/// a `skolem` (a broad-phase bin map). A column matches by its qualified name
4505/// or its unqualified suffix, as [`strip_value_invention`] matches it.
4506fn skolem_map_keys_a_join(model: &Model) -> bool {
4507 let mut maps: HashSet<String> = HashSet::new();
4508 for eq in &model.equations {
4509 if expr_contains_skolem(&eq.rhs)
4510 && let Some(v) = equation_defined_var(&eq.lhs)
4511 {
4512 if let Some(pos) = v.rfind('.') {
4513 maps.insert(v[pos + 1..].to_string());
4514 }
4515 maps.insert(v);
4516 }
4517 }
4518 if maps.is_empty() {
4519 return false;
4520 }
4521 fn keys_on(e: &Expr, maps: &HashSet<String>) -> bool {
4522 let Expr::Operator(node) = e else {
4523 return false;
4524 };
4525 node.join.as_ref().is_some_and(|j| {
4526 j.iter()
4527 .any(|c| c.on.iter().flatten().any(|col| maps.contains(col)))
4528 }) || node.any_child(&mut |c| keys_on(c, maps))
4529 }
4530 model.equations.iter().any(|eq| keys_on(&eq.rhs, &maps))
4531}
4532
4533/// The factor arrays value invention reads: the `const`-literal variables and
4534/// the caller's arrays ([`vi_factor_arrays`]), plus the build-time coordinates
4535/// a skolem map gathers from (`src_lon` in `skolem("bin", floor(index(src_lon,
4536/// i) / dx), ...)`) when they are observeds of that data alone, so the engine
4537/// can read them as factors. The analogue of Julia's `_derive_binning_coords`.
4538/// A coordinate that reads a parameter, a state or anything else stays out
4539/// (the engine then reports it as before).
4540fn value_invention_factors(
4541 model: &Model,
4542 index_sets: &HashMap<String, IndexSet>,
4543 caller_arrays: Option<&HashMap<String, ArrayD<f64>>>,
4544) -> HashMap<String, ArrayD<f64>> {
4545 fn index_targets(e: &Expr, out: &mut BTreeSet<String>) {
4546 let Expr::Operator(node) = e else {
4547 return;
4548 };
4549 if node.op == "index"
4550 && let Some(Expr::Variable(v)) = node.args.first()
4551 {
4552 out.insert(v.clone());
4553 }
4554 node.for_each_child(&mut |c| index_targets(c, out));
4555 }
4556 /// The names `e` reads that it does not bind itself.
4557 fn free_names(e: &Expr, bound: &mut Vec<String>, out: &mut BTreeSet<String>) {
4558 match e {
4559 Expr::Variable(v) if !bound.contains(v) => {
4560 out.insert(v.clone());
4561 }
4562 Expr::Operator(node) => {
4563 let mark = bound.len();
4564 bound.extend(node.output_idx.iter().flatten().cloned());
4565 bound.extend(node.ranges.iter().flat_map(|r| r.keys().cloned()));
4566 node.for_each_child(&mut |c| free_names(c, bound, out));
4567 bound.truncate(mark);
4568 }
4569 _ => {}
4570 }
4571 }
4572 struct Cx<'a> {
4573 defs: BTreeMap<String, Expr>,
4574 index_sets: &'a HashMap<String, IndexSet>,
4575 scope: HashMap<String, ArrayD<f64>>,
4576 }
4577 fn derive(cx: &mut Cx<'_>, name: &str, depth: usize) -> bool {
4578 if cx.scope.contains_key(name) {
4579 return true;
4580 }
4581 if depth > 16 {
4582 return false;
4583 }
4584 let Some(body) = cx.defs.get(name).cloned() else {
4585 return false;
4586 };
4587 let mut free = BTreeSet::new();
4588 free_names(&body, &mut Vec::new(), &mut free);
4589 for v in &free {
4590 if !cx.index_sets.contains_key(v) && !derive(cx, v, depth + 1) {
4591 return false;
4592 }
4593 }
4594 match eval_buildtime_field_in_scope(&body, cx.index_sets, &HashMap::new(), &cx.scope) {
4595 Ok(Value::Array(a)) => {
4596 cx.scope.insert(name.to_string(), *a);
4597 true
4598 }
4599 _ => false,
4600 }
4601 }
4602
4603 let mut targets = BTreeSet::new();
4604 for eq in &model.equations {
4605 if expr_contains_skolem(&eq.rhs) {
4606 index_targets(&eq.rhs, &mut targets);
4607 }
4608 }
4609 let mut cx = Cx {
4610 defs: observed_bodies(model),
4611 index_sets,
4612 scope: vi_factor_arrays(model, caller_arrays),
4613 };
4614 for t in &targets {
4615 derive(&mut cx, t, 0);
4616 }
4617 cx.scope
4618}
4619
4620/// `model` with the literal of every `const` equation whose variable is
4621/// already one of `factors` replaced by a placeholder. Value invention reads
4622/// factors from its array map, never from the document's literals, and the
4623/// literals are most of what serializing the model for it would cost.
4624fn without_factor_literals(model: &Model, factors: &HashMap<String, ArrayD<f64>>) -> Model {
4625 let mut lean = model.clone();
4626 for eq in &mut lean.equations {
4627 let is_factor = matches!(&eq.rhs, Expr::Operator(n) if n.op == "const")
4628 && equation_defined_var(&eq.lhs).is_some_and(|v| factors.contains_key(&v));
4629 if is_factor {
4630 eq.rhs = Expr::operator(ExpressionNode {
4631 op: "const".into(),
4632 value: Some(serde_json::Value::from(0)),
4633 ..Default::default()
4634 });
4635 }
4636 }
4637 lean
4638}
4639
4640/// The `from_faq` of the first range in `expr` over a derived set that neither
4641/// `extents` sizes nor a geometry producer in `geometry_ids` materializes.
4642pub(super) fn first_unmaterialized_derived_range(
4643 expr: &Expr,
4644 extents: &HashMap<String, i64>,
4645 geometry_ids: &HashSet<String>,
4646) -> Option<String> {
4647 let Expr::Operator(node) = expr else {
4648 return None;
4649 };
4650 if let Some(ranges) = &node.ranges {
4651 let mut unmaterialized: Vec<&str> = ranges
4652 .values()
4653 .filter_map(|spec| spec.derived())
4654 .filter(|f| !extents.contains_key(*f) && !geometry_ids.contains(*f))
4655 .collect();
4656 unmaterialized.sort_unstable();
4657 if let Some(f) = unmaterialized.first() {
4658 return Some((*f).to_string());
4659 }
4660 }
4661 let mut found = None;
4662 node.any_child(&mut |child| {
4663 found = first_unmaterialized_derived_range(child, extents, geometry_ids);
4664 found.is_some()
4665 });
4666 found
4667}
4668
4669/// The refusal for a range over the derived set produced by `from_faq`. A
4670/// producer refused by the cadence guard keeps that guard's code.
4671pub(super) fn unmaterialized_derived_error(
4672 from_faq: &str,
4673 index_sets: &HashMap<String, IndexSet>,
4674 failure: Option<&str>,
4675) -> CompileError {
4676 let mut names: Vec<&str> = index_sets
4677 .iter()
4678 .filter(|(_, is)| is.from_faq.as_deref() == Some(from_faq))
4679 .map(|(name, _)| name.as_str())
4680 .collect();
4681 names.sort_unstable();
4682 // The standalone evaluator entry has no index-set registry to name the set from.
4683 let subject = match names.first() {
4684 Some(set) => format!("derived index set '{set}' (from_faq '{from_faq}')"),
4685 None => format!("the derived index set produced by '{from_faq}'"),
4686 };
4687 match failure {
4688 Some(reason)
4689 if reason.contains(crate::diagnostic::codes::RELATIONAL_NODE_IN_CONTINUOUS) =>
4690 {
4691 CompileError::ValueInventionRefused {
4692 code: crate::diagnostic::codes::RELATIONAL_NODE_IN_CONTINUOUS,
4693 reason: format!("{subject}: {reason}"),
4694 }
4695 }
4696 _ => CompileError::ValueInventionRefused {
4697 code: crate::diagnostic::codes::DERIVED_INDEX_SET_UNMATERIALIZED,
4698 reason: format!(
4699 "{subject} is not materialized: {}",
4700 failure.unwrap_or("no geometry or value-invention producer supplied its extent")
4701 ),
4702 },
4703 }
4704}
4705
4706impl DerivedMaterialization {
4707 /// The refusal for the first (by set name) non-geometry derived set, when
4708 /// value invention failed.
4709 fn refuse_on_failure(
4710 &self,
4711 index_sets: &HashMap<String, IndexSet>,
4712 ) -> Result<(), CompileError> {
4713 let Some(failure) = self.failure.as_deref() else {
4714 return Ok(());
4715 };
4716 let mut sets: Vec<(&String, &str)> = index_sets
4717 .iter()
4718 .filter(|(_, is)| is.kind == "derived")
4719 .filter_map(|(n, is)| Some((n, is.from_faq.as_deref()?)))
4720 .filter(|(_, f)| !self.geometry_ids.contains(*f))
4721 .collect();
4722 sets.sort();
4723 match sets.first() {
4724 Some((_, from_faq)) => Err(unmaterialized_derived_error(
4725 from_faq,
4726 index_sets,
4727 Some(failure),
4728 )),
4729 None => Ok(()),
4730 }
4731 }
4732}
4733
4734/// Whether any expression of a model has a node carrying `ranges`, a `faq`
4735/// node, or an `index` node: the equations (both sides), the initialization
4736/// equations, and the variables' own expressions.
4737struct OpCensus {
4738 ranges: bool,
4739 faq: bool,
4740 index: bool,
4741}
4742
4743impl OpCensus {
4744 fn of(model: &Model) -> Self {
4745 fn walk(expr: &Expr, c: &mut OpCensus) {
4746 let Expr::Operator(node) = expr else {
4747 return;
4748 };
4749 c.ranges |= node.ranges.is_some();
4750 c.faq |= crate::faq::is_faq_op(&node.op);
4751 c.index |= node.op == "index";
4752 if !(c.ranges && c.faq && c.index) {
4753 node.for_each_child(&mut |child| walk(child, c));
4754 }
4755 }
4756 let mut c = OpCensus {
4757 ranges: false,
4758 faq: false,
4759 index: false,
4760 };
4761 for eq in model
4762 .equations
4763 .iter()
4764 .chain(model.initialization_equations.iter().flatten())
4765 {
4766 walk(&eq.lhs, &mut c);
4767 walk(&eq.rhs, &mut c);
4768 }
4769 for var in model.variables.values() {
4770 var.for_each_expression(&mut |expr| walk(expr, &mut c));
4771 }
4772 c
4773 }
4774}
4775
4776fn refuse_unmaterialized_derived_ranges(
4777 model: &Model,
4778 index_sets: &HashMap<String, IndexSet>,
4779 derived: &DerivedMaterialization,
4780) -> Result<(), CompileError> {
4781 let mut found: Option<String> = None;
4782 let mut visit = |expr: &Expr| {
4783 if found.is_none() {
4784 found =
4785 first_unmaterialized_derived_range(expr, &derived.extents, &derived.geometry_ids);
4786 }
4787 };
4788 for eq in &model.equations {
4789 visit(&eq.lhs);
4790 visit(&eq.rhs);
4791 }
4792 for eq in model.initialization_equations.iter().flatten() {
4793 visit(&eq.lhs);
4794 visit(&eq.rhs);
4795 }
4796 for var in model.variables.values() {
4797 var.for_each_expression(&mut |expr| visit(expr));
4798 }
4799 match found {
4800 Some(from_faq) => Err(unmaterialized_derived_error(
4801 &from_faq,
4802 index_sets,
4803 derived.failure.as_deref(),
4804 )),
4805 None => Ok(()),
4806 }
4807}
4808
4809/// Wire the value-invention front door into the array run path: run the
4810/// byte-conformant [`materialize_value_invention`] engine over the raw-JSON model
4811/// and rewrite each materialized relational OUTPUT to constant data
4812/// ([`rewrite_equation_to_const`]), so `argmin` / `argmax` / `group_aggregate`
4813/// simulate end-to-end. Derived index sets named by a materialized producer are
4814/// densified to intervals via [`rewrite_derived_index_sets`] (the same handoff
4815/// [`apply_value_invention`] performs). A NO-OP — and byte-identical — for any
4816/// model without an arg-witness op (gated by [`model_contains_arg_witness`]). A
4817/// `skolem`/`distinct` producer in such a model is sized by
4818/// [`materialize_derived_extents`] and then dropped by
4819/// [`strip_value_invention`].
4820///
4821/// `caller_arrays` is the caller-supplied factor-array channel (see
4822/// [`vi_factor_arrays`]): loader-fed envelope/connectivity factors that are not
4823/// `const` literals in the document. `None` reproduces the previous behaviour
4824/// exactly.
4825fn materialize_vi_outputs_to_data(
4826 model: &mut Model,
4827 index_sets: &mut HashMap<String, IndexSet>,
4828 caller_arrays: Option<&HashMap<String, ArrayD<f64>>>,
4829) -> Result<(), CompileError> {
4830 if !model_contains_arg_witness(model) {
4831 return Ok(());
4832 }
4833 let result = run_value_invention(model, index_sets, caller_arrays)?;
4834
4835 // Densify any derived index set named by a materialized producer (§8.1 handoff).
4836 rewrite_derived_index_sets(index_sets, &result.extents);
4837
4838 // Materialize-to-data: the arg-witness assignment (integer nearest-witness
4839 // index) and the grouped/derived SCVT chain (num / den / centroid) become
4840 // constant observeds.
4841 for (name, buf) in &result.assignments {
4842 let as_f64: Vec<f64> = buf.iter().map(|&i| i as f64).collect();
4843 rewrite_equation_to_const(model, name, &as_f64);
4844 }
4845 for (name, buf) in &result.groups {
4846 rewrite_equation_to_const(model, name, buf);
4847 }
4848 Ok(())
4849}
4850
4851/// Resolve a declared `shape` (index-set names) to concrete dense sizes against
4852/// the document registry: an `interval` set contributes its `size`, a
4853/// `categorical` set its member count. Returns `None` if any entry is a set the
4854/// registry cannot densely size (derived / ragged / unknown).
4855pub(super) fn resolve_declared_shape(
4856 decl: &[String],
4857 index_sets: &HashMap<String, IndexSet>,
4858) -> Option<Vec<usize>> {
4859 let mut out = Vec::with_capacity(decl.len());
4860 for s in decl {
4861 let is = index_sets.get(s)?;
4862 let sz = match is.kind.as_str() {
4863 "interval" => is.size? as usize,
4864 "categorical" => is.members.as_ref()?.len(),
4865 _ => return None,
4866 };
4867 out.push(sz);
4868 }
4869 Some(out)
4870}
4871
4872/// True iff `expr` is an array-PRODUCING node: a `makearray`, or an
4873/// `faq` with a non-empty `output_idx` (a scalar reduction has
4874/// an empty `output_idx` and produces a scalar). Mirrors the Julia
4875/// `_is_array_producer` (shape_promotion.jl).
4876pub(super) fn is_array_producer(node: &ExpressionNode) -> bool {
4877 if node.op == "makearray" {
4878 return true;
4879 }
4880 is_faq_op(&node.op)
4881 && node
4882 .output_idx
4883 .as_ref()
4884 .map(|v| !v.is_empty())
4885 .unwrap_or(false)
4886}
4887
4888/// True iff the RHS of a whole-array `D(state)` equation contains an
4889/// array-producing node in elementwise position — the signature of a
4890/// discretization rule's lowered stencil (`alpha * makearray(…)`). Follows the
4891/// same descent rules as the Julia `_index_array_leaves`: `index` gathers and
4892/// aggregate-family nodes are already scalar/self-indexed and are not entered.
4893pub(super) fn rhs_has_array_producer(expr: &Expr) -> bool {
4894 match expr {
4895 Expr::Operator(node) => {
4896 if is_array_producer(node) {
4897 return true;
4898 }
4899 if node.op == "index" || is_faq_op(&node.op) {
4900 return false;
4901 }
4902 node.args.iter().any(rhs_has_array_producer)
4903 }
4904 _ => false,
4905 }
4906}
4907
4908/// Resolve every self-qualified reference in `model` (`M.a` written inside
4909/// model `M`) to its local name, as the single-model route does before it
4910/// compiles — for the build pipeline, which reads the same authored model.
4911/// Run it after [`mount_subsystems`], so a reference authored inside a
4912/// subsystem body is resolved too.
4913#[cfg(not(target_arch = "wasm32"))]
4914pub(crate) fn resolve_model_self_references(model: &mut Model, model_name: &str) {
4915 resolve_self_references(model, model_name);
4916}
4917
4918fn resolve_self_references(model: &mut Model, model_name: &str) {
4919 let hits = self_qualified_references(model, model_name);
4920 if !hits.is_empty() {
4921 resolve_self_qualified_references(model, &hits);
4922 }
4923}
4924
4925/// esm-spec §4.6: inside model `M`, `M.x` — and `M.sub.x` for a mounted
4926/// subsystem `sub` — is the FULLY QUALIFIED spelling of the local `x`
4927/// (`sub.x`): the same variable the bare / subsystem-relative form names. The
4928/// single-model array build keeps the model's own unqualified names (the
4929/// flattened multi-model build and the scalar path qualify everything, which is
4930/// why the same reference resolves there), so a self-qualified reference has to
4931/// be brought back to the local spelling before the build. This collects the
4932/// `(qualified, local)` pairs to rewrite: every variable reference in the
4933/// model's expressions that starts with `<model_name>.` and whose remainder is
4934/// a declared variable or is rooted at a declared subsystem. Anything else is
4935/// left alone for the usual unbound-name diagnostics. Empty — the common case,
4936/// costing one read-only walk — means the model is used untouched.
4937///
4938/// Scope: every expression of the model as it stands. The callers run this
4939/// after `mount_subsystems`, so a subsystem's equations are already the
4940/// model's own under their mounted names, and a self-qualified reference
4941/// authored inside a subsystem body (`M.Ocean.Surface.sst` in
4942/// `M.Atmosphere`'s equation) is collected like one in the parent's.
4943fn self_qualified_references(model: &Model, model_name: &str) -> Vec<(String, String)> {
4944 let prefix = format!("{model_name}.");
4945 let subsystems: HashSet<&str> = model
4946 .subsystems
4947 .as_ref()
4948 .map(|s| s.keys().map(String::as_str).collect())
4949 .unwrap_or_default();
4950 // A `BTreeMap` both DEDUPES and orders, where the first cut rescanned a
4951 // growing `Vec` per name; and the walk looks only at variable LEAVES rather
4952 // than materializing `free_variables`' whole name set for every node.
4953 let mut hits: BTreeMap<String, String> = BTreeMap::new();
4954 fn collect(
4955 expr: &Expr,
4956 prefix: &str,
4957 model: &Model,
4958 subsystems: &HashSet<&str>,
4959 hits: &mut BTreeMap<String, String>,
4960 ) {
4961 match expr {
4962 Expr::Variable(name) => {
4963 let Some(rest) = name.strip_prefix(prefix) else {
4964 return;
4965 };
4966 let local = model.variables.contains_key(rest)
4967 || rest
4968 .split_once('.')
4969 .is_some_and(|(head, _)| subsystems.contains(head));
4970 if local && !hits.contains_key(name) {
4971 hits.insert(name.clone(), rest.to_string());
4972 }
4973 }
4974 Expr::Operator(node) => {
4975 node.for_each_child(&mut |c| collect(c, prefix, model, subsystems, hits));
4976 }
4977 Expr::Number(_) | Expr::Integer(_) => {}
4978 }
4979 }
4980 {
4981 // Scoped so the closure's `&mut hits` borrow ends before the map is
4982 // consumed below.
4983 let mut visit = |expr: &Expr| collect(expr, prefix.as_str(), model, &subsystems, &mut hits);
4984 // The expression-bearing fields the array build compiles: every equation
4985 // (both sides) and every variable's own expressions (a parameter
4986 // `update`'s `when` / `expression`, esm 1.0.0).
4987 for eq in &model.equations {
4988 visit(&eq.lhs);
4989 visit(&eq.rhs);
4990 }
4991 for var in model.variables.values() {
4992 var.for_each_expression(&mut visit);
4993 }
4994 }
4995 hits.into_iter().collect()
4996}
4997
4998/// Apply the `(qualified, local)` rewrites of [`self_qualified_references`] to
4999/// every expression of `model`, in ONE capture-aware pass.
5000///
5001/// Capture-aware in BOTH directions, which is why this is not a loop over
5002/// [`rename_free_symbol`]. That helper shadows on the name it renames FROM,
5003/// which for a `<model>.`-qualified source is never a loop symbol; the direction
5004/// that matters here is the TARGET. A node that binds the LOCAL name — model
5005/// variable `i` read as `M.i` inside `aggregate(i from x; …)` — would otherwise
5006/// capture the rewritten reference and silently read the loop index instead of
5007/// the variable. The rewrite is therefore dropped for that subtree and the
5008/// qualified spelling survives into the build's ordinary unbound-name
5009/// diagnostic: a named error, not a wrong number. One pass rather than one per
5010/// hit also stops an H-hit model from being rebuilt H times.
5011fn resolve_self_qualified_references(model: &mut Model, hits: &[(String, String)]) {
5012 let map: HashMap<&str, &str> = hits.iter().map(|(q, l)| (q.as_str(), l.as_str())).collect();
5013 let mut rewrite = |expr: &mut Expr| {
5014 *expr = rewrite_self_qualified(expr, &map);
5015 };
5016 for eq in &mut model.equations {
5017 rewrite(&mut eq.lhs);
5018 rewrite(&mut eq.rhs);
5019 }
5020 for var in model.variables.values_mut() {
5021 var.for_each_expression_mut(&mut rewrite);
5022 }
5023}
5024
5025/// One capture-aware substitution pass for [`resolve_self_qualified_references`]:
5026/// every free `Variable(q)` with `hits[q] == l` becomes `Variable(l)`, and a
5027/// node that BINDS either spelling — as an `output_idx` / `ranges` loop symbol
5028/// or an `integral`'s `var` — drops that hit for its whole subtree.
5029fn rewrite_self_qualified(expr: &Expr, hits: &HashMap<&str, &str>) -> Expr {
5030 if hits.is_empty() {
5031 return expr.clone();
5032 }
5033 match expr {
5034 Expr::Variable(v) => hits
5035 .get(v.as_str())
5036 .map_or_else(|| expr.clone(), |l| Expr::Variable((*l).to_string())),
5037 Expr::Operator(node) => {
5038 let binds = |s: &str| {
5039 node.output_idx
5040 .as_ref()
5041 .is_some_and(|ix| ix.iter().any(|x| x == s))
5042 || node.ranges.as_ref().is_some_and(|r| r.contains_key(s))
5043 || node.int_var.as_deref() == Some(s)
5044 };
5045 // Only reallocate the map on the rare node that shadows something.
5046 let narrowed: Option<HashMap<&str, &str>> =
5047 hits.iter().any(|(q, l)| binds(q) || binds(l)).then(|| {
5048 hits.iter()
5049 .filter(|&(q, l)| !binds(q) && !binds(l))
5050 .map(|(q, l)| (*q, *l))
5051 .collect()
5052 });
5053 let sub = narrowed.as_ref().unwrap_or(hits);
5054 let mut out = node.map_children(&mut |a| rewrite_self_qualified(a, sub));
5055 // `map_children` preserves `join` verbatim, so the node's
5056 // string-encoded variable references catch up with its child ones
5057 // here — see [`rename_join_names`].
5058 if let Some(join) = &out.join {
5059 let mut renamed = join.clone();
5060 for (q, l) in sub {
5061 renamed = rename_join_names(&renamed, q, l);
5062 }
5063 out.join = Some(renamed);
5064 }
5065 Expr::operator(out)
5066 }
5067 Expr::Number(_) | Expr::Integer(_) => expr.clone(),
5068 }
5069}
5070
5071/// Capture-aware rename of a free loop symbol: every free `Variable(from)`
5072/// becomes `Variable(to)`; a node that BINDS `from` itself (its `output_idx`
5073/// or `ranges` declare it) shadows the outer symbol and is left untouched.
5074///
5075/// Renames the plain-string names a `join` clause carries as well — see
5076/// [`rename_join_names`]. That is not an extra: `map_children` is an
5077/// `Expr`-CHILD walker, and a `join.on` key column or an `overlap` envelope
5078/// factor is a variable reference that happens to be encoded as a string on the
5079/// node (CONFORMANCE_SPEC §5.5.6), so a walker that skipped them would rename
5080/// half of what one node says about the same variable.
5081pub(super) fn rename_free_symbol(expr: &Expr, from: &str, to: &str) -> Expr {
5082 match expr {
5083 Expr::Variable(v) if v == from => Expr::Variable(to.to_string()),
5084 Expr::Operator(node) => {
5085 let binds = node
5086 .output_idx
5087 .as_ref()
5088 .map(|ix| ix.iter().any(|s| s == from))
5089 .unwrap_or(false)
5090 || node
5091 .ranges
5092 .as_ref()
5093 .map(|r| r.contains_key(from))
5094 .unwrap_or(false);
5095 if binds {
5096 return expr.clone();
5097 }
5098 let mut out = node.map_children(&mut |a| rename_free_symbol(a, from, to));
5099 // `map_children` preserves `join` verbatim, so this is where the
5100 // node's string-encoded references catch up with its child ones.
5101 // The binder test above has already returned for a node that binds
5102 // `from`, which is the shadowing rule an `on` key column needs: a
5103 // column naming one of THIS node's own loop symbols resolves against
5104 // its `ranges` (`join.rs::resolve_side`), never against the variable
5105 // registry, so renaming it would make it resolve to nothing.
5106 if let Some(join) = &out.join {
5107 out.join = Some(rename_join_names(join, from, to));
5108 }
5109 Expr::operator(out)
5110 }
5111 _ => expr.clone(),
5112 }
5113}
5114
5115/// Rename `from` → `to` in every plain-string VARIABLE REFERENCE a node's
5116/// `join` clauses carry (CONFORMANCE_SPEC §5.5.6): each `on` key column, each
5117/// `overlap` envelope factor, and the data-column names of an already-resolved
5118/// `on_gate`. The array-runtime twin of `flatten.rs::rename_join_names`, which
5119/// does the same job for a `variable_map` that removes a parameter.
5120///
5121/// A clause's LOOP SYMBOLS are left alone — `OverlapClause::sym_src` /
5122/// `sym_tgt` and the `on_gate`'s two gated symbols are binders of the node, not
5123/// references into the variable registry. (`map_column_names` renames the gate's
5124/// COLUMNS only, for exactly that reason.)
5125fn rename_join_names(join: &[JoinClause], from: &str, to: &str) -> Vec<JoinClause> {
5126 let ren = |n: &String| -> String { if n == from { to.to_string() } else { n.clone() } };
5127 join.iter()
5128 .map(|c| JoinClause {
5129 on: c.on.iter().map(|[l, r]| [ren(l), ren(r)]).collect(),
5130 // Binders of the node, not references — renaming leaves them alone.
5131 syms: c.syms.clone(),
5132 overlap: c.overlap.as_ref().map(|ov| OverlapClause {
5133 src_env: ov.src_env.iter().map(&ren).collect(),
5134 tgt_env: ov.tgt_env.iter().map(&ren).collect(),
5135 eps: ov.eps,
5136 sym_src: ov.sym_src.clone(),
5137 sym_tgt: ov.sym_tgt.clone(),
5138 }),
5139 on_gate: c.on_gate.as_ref().map(|g| g.map_column_names(ren)),
5140 })
5141 .collect()
5142}
5143
5144/// Inline a `makearray`'s ARRAY-VALUED aggregate region values into the
5145/// enclosing loop symbols: a region value that is a pointwise
5146/// `faq` whose output ranges equal the region bounds exactly
5147/// (no contraction, no filter) is replaced by its body with each output
5148/// symbol renamed to the enclosing loop symbol. This turns the discretized
5149/// `makearray([interior], [aggregate_i(stencil)])` form a §9.6.3 rewrite rule
5150/// emits into the scalar-region-value form the vectorized whole-array kernel
5151/// consumes directly (the build-time `index(faq, …)` collapse the Julia
5152/// reference performs). Values that do not match are left untouched.
5153pub(super) fn inline_region_aggregates(node: &ExpressionNode, loops: &[String]) -> ExpressionNode {
5154 let (Some(regions), Some(values)) = (&node.regions, &node.values) else {
5155 return node.clone();
5156 };
5157 if regions.len() != values.len() {
5158 return node.clone();
5159 }
5160 let mut out = node.clone();
5161 out.values = Some(
5162 regions
5163 .iter()
5164 .zip(values.iter())
5165 .map(|(region, value)| {
5166 let Expr::Operator(v) = value else {
5167 return value.clone();
5168 };
5169 if !is_faq_op(&v.op) || v.filter.is_some() {
5170 return value.clone();
5171 }
5172 let (Some(idx), Some(ranges), Some(body)) = (&v.output_idx, &v.ranges, &v.expr)
5173 else {
5174 return value.clone();
5175 };
5176 if idx.len() != region.len()
5177 || ranges.len() != idx.len()
5178 || loops.len() != idx.len()
5179 {
5180 return value.clone();
5181 }
5182 for (d, sym) in idx.iter().enumerate() {
5183 let want = crate::types::region_bounds(®ion[d]);
5184 match (ranges.get(sym).and_then(|r| r.bounds()), want) {
5185 (Some(b), Some(w)) if b == w => {}
5186 _ => return value.clone(),
5187 }
5188 }
5189 let mut inlined = body.as_ref().clone();
5190 for (sym, loop_name) in idx.iter().zip(loops.iter()) {
5191 if sym != loop_name {
5192 inlined = rename_free_symbol(&inlined, sym, loop_name);
5193 }
5194 }
5195 inlined
5196 })
5197 .collect(),
5198 );
5199 out
5200}
5201
5202/// Declared index-set axis NAMES of every array-shaped variable — the ordered
5203/// `shape` list (esm-spec §6.3), which names the index set each axis ranges
5204/// over. [`resolve_declared_shape`] throws these names away in favour of
5205/// extents; the whole-array lowering below needs them, because a BARE
5206/// array-level expression aligns its operands by index-set NAME, not by
5207/// position (esm-spec §4.3.4).
5208pub(super) fn declared_axis_names(model: &Model) -> HashMap<String, Vec<String>> {
5209 model
5210 .variables
5211 .iter()
5212 .filter_map(|(k, v)| {
5213 v.shape
5214 .as_ref()
5215 .filter(|s| !s.is_empty())
5216 .map(|s| (k.clone(), s.clone()))
5217 })
5218 .collect()
5219}
5220
5221/// Where each array-shaped operand's axes sit in the RESULT's axis list.
5222///
5223/// A bare array-level expression is lowered by enumerating the result's cells
5224/// and gathering every operand at that cell. WHICH of the cell's coordinates an
5225/// operand is gathered by is decided by index-set NAME (esm-spec §4.3.4): axis
5226/// `d` of an operand declared over `["lat"]` is the result's `lat` axis
5227/// wherever that sits, and every result axis the operand does not declare is
5228/// one it BROADCASTS along. A name-keyed plan makes `D(dp) = w1` (with
5229/// `dp: [lon,lat,lev]`, `w1: [lat]`) compute exactly what the `faq`
5230/// spelling `sum_{i,j,k} w1[j]` computes, and makes axis ORDER immaterial — a
5231/// `[lat,lon]` operand transposes rather than being reinterpreted.
5232///
5233/// An entry is present only for a leaf that aligns by name. A leaf ABSENT from
5234/// the plan keeps the legacy POSITIONAL lowering (leading axes) — the fallback
5235/// for an operand, or a result, that carries no declared index-set names.
5236type GatherPlan = HashMap<String, Vec<usize>>;
5237
5238/// Collect the array-shaped `Variable` leaves that a whole-array lowering will
5239/// wrap in a NAME-ALIGNED `index(…)` gather: the ones standing in genuinely
5240/// ELEMENTWISE position (see [`crate::op_registry::is_elementwise_node`]),
5241/// which is the only place element alignment is defined. That includes a
5242/// `broadcast` node, whose `fn` IS the scalar operator, so the `broadcast` and
5243/// bare spellings of one expression align identically. A leaf reached through
5244/// an op that consumes its operands whole — a `faq`, a `makearray`, an
5245/// `index` target, a shape op, a relational or geometry kernel — is left to
5246/// that op's own operand contract and keeps the legacy positional lowering.
5247///
5248/// The descent otherwise follows the lowering it serves, so the plan covers
5249/// exactly the leaves that get rewritten: [`index_array_leaves`] rewrites
5250/// every child (`into_binders = true`), while [`index_array_leaves_by_loops`]
5251/// stops at an `index` gather, an aggregate node, or an array producer
5252/// (`into_binders = false`).
5253fn collect_wrapped_array_leaves(
5254 expr: &Expr,
5255 array_axes: &HashMap<String, Vec<String>>,
5256 into_binders: bool,
5257 out: &mut Vec<String>,
5258) {
5259 match expr {
5260 Expr::Variable(v) if array_axes.contains_key(v) => {
5261 out.push(v.clone());
5262 }
5263 Expr::Operator(node) => {
5264 if !crate::op_registry::is_elementwise_node(node) {
5265 return;
5266 }
5267 if !into_binders
5268 && (is_array_producer(node) || node.op == "index" || is_faq_op(&node.op))
5269 {
5270 return;
5271 }
5272 node.for_each_child(&mut |c| {
5273 collect_wrapped_array_leaves(c, array_axes, into_binders, out)
5274 });
5275 }
5276 _ => {}
5277 }
5278}
5279
5280/// Build the [`GatherPlan`] for lowering `rhs` onto a result declared over
5281/// `target_axes`, and REJECT an operand that cannot be aligned.
5282///
5283/// An operand whose declared index sets are a SUBSET of the result's aligns by
5284/// name and broadcasts along the axes it does not carry. An operand carrying an
5285/// index set the result does NOT have has no axis to align to and no defensible
5286/// value to take, so it is a hard build error rather than a positional
5287/// reinterpretation (esm-spec §4.3.4; issue #100). The shapes are fully known
5288/// statically, so `validate()` reports the same defect as an
5289/// `array_shape_mismatch` structural error before a build is ever attempted.
5290///
5291/// `target_axes` is `None` — and the plan consequently empty, keeping the
5292/// legacy positional lowering everywhere — when the result carries no usable
5293/// declared names.
5294fn build_gather_plan(
5295 rhs: &Expr,
5296 array_axes: &HashMap<String, Vec<String>>,
5297 target: &str,
5298 target_axes: Option<&[String]>,
5299 into_binders: bool,
5300) -> Result<GatherPlan, CompileError> {
5301 let Some(target_axes) = target_axes else {
5302 return Ok(GatherPlan::new());
5303 };
5304 // A repeated axis name offers no unambiguous position to align to; keep the
5305 // positional lowering rather than guess which occurrence was meant.
5306 if (1..target_axes.len()).any(|d| target_axes[..d].contains(&target_axes[d])) {
5307 return Ok(GatherPlan::new());
5308 }
5309 let mut leaves: Vec<String> = Vec::new();
5310 collect_wrapped_array_leaves(rhs, array_axes, into_binders, &mut leaves);
5311 let mut plan = GatherPlan::new();
5312 for leaf in leaves {
5313 let Some(leaf_axes) = array_axes.get(&leaf) else {
5314 continue;
5315 };
5316 let mut positions = Vec::with_capacity(leaf_axes.len());
5317 for axis in leaf_axes {
5318 match target_axes.iter().position(|t| t == axis) {
5319 Some(p) => positions.push(p),
5320 None => {
5321 return Err(CompileError::UnalignableArrayShape {
5322 operand: leaf.clone(),
5323 axis: axis.clone(),
5324 result: target.to_string(),
5325 result_axes: target_axes.to_vec(),
5326 });
5327 }
5328 }
5329 }
5330 plan.insert(leaf, positions);
5331 }
5332 Ok(plan)
5333}
5334
5335/// True iff every planned operand already sits on exactly the result's axes in
5336/// the result's order, so the name-aligned lowering and the legacy positional
5337/// one agree cell for cell and no rewrite is needed.
5338fn plan_is_identity(plan: &GatherPlan, target_rank: usize) -> bool {
5339 plan.values()
5340 .all(|pos| pos.len() == target_rank && pos.iter().enumerate().all(|(d, &p)| d == p))
5341}
5342
5343/// Rewrite a whole-array `D(state)` RHS into its per-cell body over the given
5344/// LOOP SYMBOLS (the loop-name dual of [`index_array_leaves`], mirroring the
5345/// Julia `_index_array_leaves` in shape_promotion.jl): each bare array-shaped
5346/// `Variable` leaf and each array-PRODUCING node (a `makearray` — whose
5347/// aggregate region values are first inlined via [`inline_region_aggregates`]
5348/// — or a `faq` with output axes) is wrapped in
5349/// `index(node, loops…)`; `index` gathers and scalar reductions stay
5350/// untouched; other operators recurse elementwise.
5351///
5352/// A declared leaf listed in `plan` is gathered by the loop symbols of ITS OWN
5353/// axes (esm-spec §4.3.4), so a rank-1 `[lat]` operand under a rank-3
5354/// `[lon,lat,lev]` result reads `index(w1, <lat loop>)` and replicates along
5355/// the other two. Handing it all three loops — what this did before — indexed
5356/// axes it does not have, which `index_into` reads as out of bounds and
5357/// resolves to the ghost value 0.0 everywhere. An ANONYMOUS operand (no
5358/// declared names) and every array PRODUCER keep the positional lowering.
5359pub(super) fn index_array_leaves_by_loops(
5360 expr: &Expr,
5361 array_axes: &HashMap<String, Vec<String>>,
5362 plan: Option<&GatherPlan>,
5363 loops: &[String],
5364) -> Expr {
5365 let wrap = |target: Expr, ix: Vec<&String>| {
5366 let mut args = vec![target];
5367 for l in ix {
5368 args.push(Expr::Variable(l.clone()));
5369 }
5370 Expr::operator(ExpressionNode {
5371 op: "index".to_string(),
5372 args,
5373 ..Default::default()
5374 })
5375 };
5376 match expr {
5377 Expr::Variable(v) if array_axes.contains_key(v) => {
5378 let ix: Vec<&String> = match plan.and_then(|p| p.get(v)) {
5379 // `positions` is built against this same result, so every entry
5380 // indexes `loops`.
5381 Some(positions) => positions.iter().map(|&p| &loops[p]).collect(),
5382 None => {
5383 let rank = array_axes[v].len().min(loops.len());
5384 loops[..rank].iter().collect()
5385 }
5386 };
5387 wrap(expr.clone(), ix)
5388 }
5389 Expr::Operator(node) => {
5390 if is_array_producer(node) {
5391 let target = if node.op == "makearray" {
5392 Expr::operator(inline_region_aggregates(node, loops))
5393 } else {
5394 expr.clone()
5395 };
5396 return wrap(target, loops.iter().collect());
5397 }
5398 if node.op == "index" || is_faq_op(&node.op) {
5399 return expr.clone();
5400 }
5401 // Element alignment is defined only under elementwise nodes (a
5402 // `broadcast` included — its `fn` IS the scalar op); below anything
5403 // else the operand contract is that op's own.
5404 let child_plan = plan.filter(|_| crate::op_registry::is_elementwise_node(node));
5405 let mut out = ExpressionNode::clone(node);
5406 out.args = node
5407 .args
5408 .iter()
5409 .map(|a| index_array_leaves_by_loops(a, array_axes, child_plan, loops))
5410 .collect();
5411 Expr::operator(out)
5412 }
5413 _ => expr.clone(),
5414 }
5415}
5416
5417/// The coordinates [`index_array_leaves`] gathers each leaf at: one concrete
5418/// 1-based cell, or the loop symbols of an [`RhsRule::ArrayLoop`] over the
5419/// whole shape (which bind, cell by cell, to exactly those integers).
5420#[derive(Clone, Copy)]
5421pub(super) enum CellCoords<'a> {
5422 Ints(&'a [i64]),
5423 Loops(&'a [String]),
5424}
5425
5426impl CellCoords<'_> {
5427 fn len(&self) -> usize {
5428 match self {
5429 CellCoords::Ints(c) => c.len(),
5430 CellCoords::Loops(l) => l.len(),
5431 }
5432 }
5433
5434 fn coord(&self, d: usize) -> Expr {
5435 match self {
5436 CellCoords::Ints(c) => Expr::Integer(c[d]),
5437 CellCoords::Loops(l) => Expr::Variable(l[d].clone()),
5438 }
5439 }
5440}
5441
5442/// Rewrite each bare array-shaped `Variable` leaf of a whole-array `D(state)` RHS
5443/// into an `index(var, cell…)` gather at the given cell, so the elementwise
5444/// array equation compiles to a scalar body per cell. The array target of an
5445/// existing `index` node is left untouched (it is already a gather).
5446///
5447/// A declared leaf listed in `plan` is gathered at the cell coordinates of ITS
5448/// OWN axes (esm-spec §4.3.4) — a `[lat]` operand under a `[lon,lat,lev]`
5449/// result takes the cell's `lat` coordinate and replicates along `lon`/`lev`,
5450/// and a `[lat,lon]` operand transposes. Taking the LEADING coordinates — what
5451/// this did before — laid the operand's elements along the wrong axes and
5452/// zero-filled the overhang. An ANONYMOUS operand (no declared names, or a
5453/// result with none) keeps the leading-axes lowering.
5454pub(super) fn index_array_leaves(
5455 expr: &Expr,
5456 array_axes: &HashMap<String, Vec<String>>,
5457 plan: Option<&GatherPlan>,
5458 cell: CellCoords<'_>,
5459) -> Expr {
5460 match expr {
5461 Expr::Variable(v) => {
5462 if let Some(axes) = array_axes.get(v) {
5463 let mut args = vec![Expr::Variable(v.clone())];
5464 match plan.and_then(|p| p.get(v)) {
5465 // `positions` is built against this same result, so every
5466 // entry indexes `cell`.
5467 Some(positions) => {
5468 args.extend(positions.iter().map(|&p| cell.coord(p)));
5469 }
5470 None => {
5471 let n = axes.len().min(cell.len());
5472 args.extend((0..n).map(|d| cell.coord(d)));
5473 }
5474 }
5475 Expr::operator(ExpressionNode {
5476 op: "index".to_string(),
5477 args,
5478 ..Default::default()
5479 })
5480 } else {
5481 expr.clone()
5482 }
5483 }
5484 Expr::Operator(node) => {
5485 // Element alignment is defined only under elementwise nodes (a
5486 // `broadcast` included — its `fn` IS the scalar op); below anything
5487 // else the operand contract is that op's own.
5488 let child_plan = plan.filter(|_| crate::op_registry::is_elementwise_node(node));
5489 let mut out =
5490 node.map_children(&mut |a| index_array_leaves(a, array_axes, child_plan, cell));
5491 if node.op == "index"
5492 && let Some(first) = node.args.first()
5493 {
5494 // Keep the (already array-valued) target; only the index
5495 // argument expressions are rewritten.
5496 out.args[0] = first.clone();
5497 }
5498 Expr::operator(out)
5499 }
5500 other => other.clone(),
5501 }
5502}
5503
5504/// Whether a whole-array `D(state)` RHS may be lowered as one loop over the
5505/// state's shape ([`CellCoords::Loops`]) instead of one rule per cell.
5506///
5507/// The two forms evaluate the same scalar expression at every cell, so the
5508/// oracle cannot tell them apart. The loop form is refused where the TAPE
5509/// lowers the two differently: a leaf the rewrite would gather inside an
5510/// `index` argument (a data-dependent index), any aggregate (a per-cell body
5511/// folds a scalar reduction wholesale, while one nested in a loop's box stays
5512/// per-cell), and a state leaf whose shape was inferred rather than declared,
5513/// which the rewrite leaves whole.
5514fn wholearray_body_loops_as_percell(
5515 expr: &Expr,
5516 array_axes: &HashMap<String, Vec<String>>,
5517 var_shapes: &IndexMap<String, VarShape>,
5518) -> bool {
5519 fn rewrites_a_leaf(e: &Expr, array_axes: &HashMap<String, Vec<String>>) -> bool {
5520 match e {
5521 Expr::Variable(v) => array_axes.contains_key(v),
5522 Expr::Operator(node) if node.op == "index" => node
5523 .args
5524 .iter()
5525 .skip(1)
5526 .any(|a| rewrites_a_leaf(a, array_axes)),
5527 Expr::Operator(node) => node.any_child(&mut |c| rewrites_a_leaf(c, array_axes)),
5528 _ => false,
5529 }
5530 }
5531 match expr {
5532 Expr::Variable(v) => {
5533 array_axes.contains_key(v) || var_shapes.get(v).is_none_or(|vs| vs.shape.is_empty())
5534 }
5535 Expr::Operator(node) => {
5536 if node.op == "index" {
5537 return !node
5538 .args
5539 .iter()
5540 .skip(1)
5541 .any(|a| rewrites_a_leaf(a, array_axes));
5542 }
5543 if is_faq_op(&node.op) {
5544 return false;
5545 }
5546 !node.any_child(&mut |c| !wholearray_body_loops_as_percell(c, array_axes, var_shapes))
5547 }
5548 _ => true,
5549 }
5550}
5551
5552/// If `lhs` is `D(var, t)` or `D(index(var, i1, ...), t)`, return
5553/// `(var_name, Some(indices))` for the indexed form (with all concrete
5554/// integer indices), `(var_name, None)` for the plain form. `None` result
5555/// means this LHS is neither.
5556/// The variable name [`extract_derivative_scalar`] returns for `lhs`, borrowed.
5557fn derivative_scalar_target(lhs: &Expr) -> Option<&str> {
5558 let Expr::Operator(node) = lhs else {
5559 return None;
5560 };
5561 if node.op != "D" || node.args.len() != 1 {
5562 return None;
5563 }
5564 match &node.args[0] {
5565 Expr::Variable(name) => Some(name),
5566 Expr::Operator(inner) if inner.op == "index" => {
5567 let Expr::Variable(name) = inner.args.first()? else {
5568 return None;
5569 };
5570 inner
5571 .args
5572 .iter()
5573 .skip(1)
5574 .all(|a| matches!(a, Expr::Number(_) | Expr::Integer(_)))
5575 .then_some(name.as_str())
5576 }
5577 _ => None,
5578 }
5579}
5580
5581pub(super) fn extract_derivative_scalar(lhs: &Expr) -> Option<(String, Option<Vec<i64>>)> {
5582 let Expr::Operator(node) = lhs else {
5583 return None;
5584 };
5585 if node.op != "D" {
5586 return None;
5587 }
5588 if node.args.len() != 1 {
5589 return None;
5590 }
5591 match &node.args[0] {
5592 Expr::Variable(name) => Some((name.clone(), None)),
5593 Expr::Operator(inner) if inner.op == "index" => {
5594 let name = match inner.args.first()? {
5595 Expr::Variable(v) => v.clone(),
5596 _ => return None,
5597 };
5598 let indices: Vec<i64> = inner
5599 .args
5600 .iter()
5601 .skip(1)
5602 .map(|a| match a {
5603 Expr::Number(n) => Some(*n as i64),
5604 Expr::Integer(n) => Some(*n),
5605 _ => None,
5606 })
5607 .collect::<Option<Vec<_>>>()?;
5608 Some((name, Some(indices)))
5609 }
5610 _ => None,
5611 }
5612}
5613
5614/// If `lhs` is `faq(expr=D(index(var, idx...)), ...)`, extract
5615/// `(var_name, output_idx_names, output_ranges, lhs_idx_exprs, rhs_body,
5616/// contract_names, contract_ranges, reduce)`.
5617/// `contract_names`/`contract_ranges` are indices present in the RHS ranges
5618/// but absent from `output_idx` (generalized-einsum contracted indices).
5619/// `reduce` is the semiring ⊕ resolved from the RHS node's `semiring`/`reduce`
5620/// (defaulting to `Sum` per the ESM spec).
5621/// The parsed pieces of an `aggregate(expr=D(index(var, …))) = aggregate(…)`
5622/// derivative equation, as extracted by [`extract_derivative_faq`]. The
5623/// fields mirror [`RhsRule::ArrayLoop`]'s.
5624pub(super) struct DerivArrayop {
5625 /// Target state variable name.
5626 var: String,
5627 /// Output loop index names (LHS aggregate `output_idx`).
5628 idx_names: Vec<String>,
5629 /// Concrete `(lo, hi)` bounds per output index, in `idx_names` order.
5630 ranges: Vec<(i64, i64)>,
5631 /// LHS `index(var, …)` argument expressions (may offset the loop symbols).
5632 lhs_idx_exprs: Vec<Expr>,
5633 /// Scalar RHS body evaluated per output tuple.
5634 body: Expr,
5635 /// Contracted (reduction) index names, sorted.
5636 contract_names: Vec<String>,
5637 /// Bounds of the contracted indices, parallel to `contract_names`.
5638 contract_dims: Vec<ContractDim>,
5639 /// Semiring ⊕ reducer for the contraction.
5640 reduce: ReduceKind,
5641 /// Optional §5.3 filter predicate gating the contraction.
5642 filter: Option<Box<Expr>>,
5643}
5644
5645pub(super) fn extract_derivative_faq(lhs: &Expr, rhs: &Expr) -> Option<DerivArrayop> {
5646 let Expr::Operator(node) = lhs else {
5647 return None;
5648 };
5649 if !is_faq_op(&node.op) {
5650 return None;
5651 }
5652 let body = node.expr.as_ref()?.as_ref();
5653 let idx_names = node.output_idx.clone()?;
5654 let ranges_map = node.ranges.clone()?;
5655 // Body must be D(index(var, ...)).
5656 let Expr::Operator(d_node) = body else {
5657 return None;
5658 };
5659 if d_node.op != "D" {
5660 return None;
5661 }
5662 let Expr::Operator(inner) = d_node.args.first()? else {
5663 return None;
5664 };
5665 if inner.op != "index" {
5666 return None;
5667 }
5668 let var_name = match inner.args.first()? {
5669 Expr::Variable(v) => v.clone(),
5670 _ => return None,
5671 };
5672 let lhs_idx_exprs: Vec<Expr> = inner.args.iter().skip(1).cloned().collect();
5673 // Map idx_names → ranges in order.
5674 let ranges: Vec<(i64, i64)> = idx_names
5675 .iter()
5676 .map(|n| {
5677 let r = ranges_map.get(n).and_then(|s| s.bounds()).unwrap_or([0, 0]);
5678 (r[0], r[1])
5679 })
5680 .collect();
5681 // RHS body: assume rhs is also faq with body, or pass through as
5682 // scalar-valued expr that evaluates at each tuple.
5683 // Also extract contracted (reduction) indices and the semiring ⊕ reducer.
5684 let (rhs_body, contract_names, contract_dims, reduce, filter) = match rhs {
5685 Expr::Operator(rnode) if is_faq_op(&rnode.op) => {
5686 let b = rnode.expr.as_ref().map(|b| b.as_ref().clone())?;
5687 // Declining (rather than reporting) an out-of-enum ⊕ spelling is
5688 // safe here only because this is a pattern matcher with no error
5689 // channel AND `from_model` runs `validate_oplus_spellings` over the
5690 // whole model first, so such a node cannot reach this point.
5691 let rop =
5692 effective_reduce_kind(rnode.semiring.as_deref(), rnode.reduce.as_deref()).ok()?;
5693 let mut c_names: Vec<String> = Vec::new();
5694 let mut c_dims: Vec<ContractDim> = Vec::new();
5695 if let Some(rhs_ranges) = &rnode.ranges {
5696 let mut sorted_keys: Vec<&String> = rhs_ranges.keys().collect();
5697 sorted_keys.sort();
5698 for n in sorted_keys {
5699 if !idx_names.contains(n) {
5700 // A ragged contracted index keeps its dynamic bound; all
5701 // others collapse to a static interval here.
5702 c_names.push(n.clone());
5703 c_dims.push(ContractDim::from_range(&rhs_ranges[n]));
5704 }
5705 }
5706 }
5707 // §5.3 filter rides on the RHS aggregate; carry it into the rule so
5708 // the contraction gates on it (otherwise it would be silently lost).
5709 (b, c_names, c_dims, rop, rnode.filter.clone())
5710 }
5711 other => (other.clone(), Vec::new(), Vec::new(), ReduceKind::Sum, None),
5712 };
5713 Some(DerivArrayop {
5714 var: var_name,
5715 idx_names,
5716 ranges,
5717 lhs_idx_exprs,
5718 body: rhs_body,
5719 contract_names,
5720 contract_dims,
5721 reduce,
5722 filter,
5723 })
5724}
5725
5726/// The parsed pieces of an algebraic `faq(expr=index(var, …)) = faq(…)`
5727/// definition, as extracted by [`extract_algebraic_faq`]. The counterpart
5728/// of [`DerivArrayop`] for equations that define a variable's value rather
5729/// than its derivative.
5730pub(super) struct AlgebraicArrayop {
5731 /// Defined (algebraic) variable name.
5732 pub(super) var: String,
5733 /// Output loop index names (LHS aggregate `output_idx`).
5734 pub(super) idx_names: Vec<String>,
5735 /// Concrete `(lo, hi)` bounds per output index, in `idx_names` order.
5736 pub(super) ranges: Vec<(i64, i64)>,
5737 /// Scalar RHS body evaluated per output tuple.
5738 pub(super) body: Expr,
5739}
5740
5741/// Extract an algebraic `faq(expr=index(var, idx...)) = faq(...)`
5742/// definition. Matches fixtures 02 and 04 where an algebraic variable is
5743/// defined through a faq whose body is just `index(v, i...)`.
5744pub(super) fn extract_algebraic_faq(lhs: &Expr, rhs: &Expr) -> Option<AlgebraicArrayop> {
5745 let Expr::Operator(node) = lhs else {
5746 return None;
5747 };
5748 if !is_faq_op(&node.op) {
5749 return None;
5750 }
5751 let body = node.expr.as_ref()?.as_ref();
5752 let idx_names = node.output_idx.clone()?;
5753 let ranges_map = node.ranges.clone()?;
5754 // Body must be index(var, idx...) with idx symbols matching idx_names in order.
5755 let Expr::Operator(inner) = body else {
5756 return None;
5757 };
5758 if inner.op != "index" {
5759 return None;
5760 }
5761 let var_name = match inner.args.first()? {
5762 Expr::Variable(v) => v.clone(),
5763 _ => return None,
5764 };
5765 // Indices must be exactly the output_idx names in order (v1 constraint).
5766 let idx_args: Vec<&Expr> = inner.args.iter().skip(1).collect();
5767 if idx_args.len() != idx_names.len() {
5768 return None;
5769 }
5770 for (a, want) in idx_args.iter().zip(idx_names.iter()) {
5771 match a {
5772 Expr::Variable(v) if v == want => {}
5773 _ => return None,
5774 }
5775 }
5776 let ranges: Vec<(i64, i64)> = idx_names
5777 .iter()
5778 .map(|n| {
5779 let r = ranges_map.get(n).and_then(|s| s.bounds()).unwrap_or([0, 0]);
5780 (r[0], r[1])
5781 })
5782 .collect();
5783 let rhs_body = match rhs {
5784 Expr::Operator(rnode) if is_faq_op(&rnode.op) => {
5785 // This elementwise (non-contracting) fast path does not apply a
5786 // `filter`. Bail rather than silently drop it — a filtered
5787 // definition must be compiled by a path that honors §5.3.
5788 if rnode.filter.is_some() {
5789 return None;
5790 }
5791 rnode.expr.as_ref().map(|b| b.as_ref().clone())?
5792 }
5793 other => other.clone(),
5794 };
5795 Some(AlgebraicArrayop {
5796 var: var_name,
5797 idx_names,
5798 ranges,
5799 body: rhs_body,
5800 })
5801}
5802
5803/// Shape inference: per state variable, infer its shape from every
5804/// `index(var, ...)` reference, `D(index(var, ...))` reference, and
5805/// `faq` over its elements. Returns a map var_name → shape (empty Vec
5806/// means scalar). Origins are assumed 1-based.
5807///
5808/// Two-pass design: LHS equations pin the authoritative state extent; RHS
5809/// index references (which may include stencil offsets like `i-1` or `i+1`)
5810/// are only used for variables not already shaped by the LHS. This prevents
5811/// neighbor references in PDE stencils from bloating the inferred shape.
5812pub(super) fn infer_shapes(
5813 state_vars: &[&String],
5814 equations: &[crate::types::Equation],
5815) -> Result<HashMap<String, Vec<usize>>, CompileError> {
5816 let state_set: HashSet<&str> = state_vars.iter().map(|s| s.as_str()).collect();
5817
5818 // Pass 1: LHS only — these are the authoritative (pinned) shapes.
5819 let mut per_var_min: HashMap<String, Vec<i64>> = HashMap::new();
5820 let mut per_var_max: HashMap<String, Vec<i64>> = HashMap::new();
5821 let mut seen_indexed: HashSet<String> = HashSet::new();
5822 let skip_none: HashSet<String> = HashSet::new();
5823 let no_loops: HashMap<String, (i64, i64)> = HashMap::new();
5824 {
5825 let mut walk = ShapeWalk {
5826 states: &state_set,
5827 per_var_min: &mut per_var_min,
5828 per_var_max: &mut per_var_max,
5829 seen_indexed: &mut seen_indexed,
5830 skip_shape_update: &skip_none,
5831 };
5832 for eq in equations {
5833 walk.walk(&eq.lhs, &no_loops);
5834 }
5835 }
5836
5837 // Pass 2: RHS — skip variables already pinned by LHS to prevent stencil
5838 // offsets (e.g. index(u, i-1)) from expanding the state's extent.
5839 let lhs_pinned = seen_indexed.clone();
5840 {
5841 let mut walk = ShapeWalk {
5842 states: &state_set,
5843 per_var_min: &mut per_var_min,
5844 per_var_max: &mut per_var_max,
5845 seen_indexed: &mut seen_indexed,
5846 skip_shape_update: &lhs_pinned,
5847 };
5848 for eq in equations {
5849 walk.walk(&eq.rhs, &no_loops);
5850 }
5851 }
5852
5853 let mut out: HashMap<String, Vec<usize>> = HashMap::new();
5854 for name in state_vars {
5855 let name_s = (*name).clone();
5856 if !seen_indexed.contains(&name_s) {
5857 out.insert(name_s, Vec::new());
5858 continue;
5859 }
5860 let mins = per_var_min.get(&name_s).cloned().unwrap_or_default();
5861 let maxes = per_var_max.get(&name_s).cloned().unwrap_or_default();
5862 if mins.len() != maxes.len() {
5863 return Err(CompileError::build_err(format!(
5864 "Inconsistent index rank for variable '{name_s}'"
5865 )));
5866 }
5867 let shape: Vec<usize> = mins
5868 .iter()
5869 .zip(maxes.iter())
5870 .map(|(lo, hi)| (hi - lo + 1).max(1) as usize)
5871 .collect();
5872 out.insert(name_s, shape);
5873 }
5874 Ok(out)
5875}
5876
5877/// Size, from how the equations index it, every refreshed parameter whose
5878/// declared shape names an index set the registry cannot size, as
5879/// [`infer_state_shapes`] sizes a state from its uses when no declaration
5880/// does: each axis spans `[1, hi]`, `hi` the largest subscript any `index` of
5881/// it reaches. A parameter no equation indexes, or indexes at different
5882/// ranks, stays unsized.
5883///
5884/// The box is only what a compiled read is built against. The forcing load
5885/// checks the entry the buffer actually holds against it at run time and
5886/// faults on a different shape, so a wrong guess fails loudly rather than
5887/// reading the wrong cells.
5888fn infer_unsized_forcing_shapes(
5889 decls: &mut IndexMap<String, Option<Vec<usize>>>,
5890 equations: &[crate::types::Equation],
5891) {
5892 let unsized_: HashSet<&str> = decls
5893 .iter()
5894 .filter(|(_, s)| s.is_none())
5895 .map(|(n, _)| n.as_str())
5896 .collect();
5897 if unsized_.is_empty() {
5898 return;
5899 }
5900 let mut per_var_min: HashMap<String, Vec<i64>> = HashMap::new();
5901 let mut per_var_max: HashMap<String, Vec<i64>> = HashMap::new();
5902 let mut seen_indexed: HashSet<String> = HashSet::new();
5903 let skip_none: HashSet<String> = HashSet::new();
5904 let no_loops: HashMap<String, (i64, i64)> = HashMap::new();
5905 let mut ranks: HashMap<String, HashSet<usize>> = HashMap::new();
5906 {
5907 let mut walk = ShapeWalk {
5908 states: &unsized_,
5909 per_var_min: &mut per_var_min,
5910 per_var_max: &mut per_var_max,
5911 seen_indexed: &mut seen_indexed,
5912 skip_shape_update: &skip_none,
5913 };
5914 for eq in equations {
5915 walk.walk(&eq.lhs, &no_loops);
5916 walk.walk(&eq.rhs, &no_loops);
5917 }
5918 }
5919 for eq in equations {
5920 for e in [&eq.lhs, &eq.rhs] {
5921 index_ranks(e, &unsized_, &mut ranks);
5922 }
5923 }
5924 for (name, shape) in decls.iter_mut() {
5925 if shape.is_some() || ranks.get(name).is_none_or(|r| r.len() != 1) {
5926 continue;
5927 }
5928 let Some(maxes) = per_var_max.get(name) else {
5929 continue;
5930 };
5931 if maxes.iter().all(|&hi| hi >= 1) {
5932 *shape = Some(maxes.iter().map(|&hi| hi as usize).collect());
5933 }
5934 }
5935}
5936
5937/// The subscript counts every `index` of a name in `names` uses.
5938fn index_ranks(expr: &Expr, names: &HashSet<&str>, out: &mut HashMap<String, HashSet<usize>>) {
5939 let Expr::Operator(node) = expr else {
5940 return;
5941 };
5942 if node.op == "index"
5943 && let Some(Expr::Variable(var)) = node.args.first()
5944 && names.contains(var.as_str())
5945 {
5946 out.entry(var.clone())
5947 .or_default()
5948 .insert(node.args.len() - 1);
5949 }
5950 node.for_each_child(&mut |child| index_ranks(child, names, out));
5951}
5952
5953/// Accumulator state for [`infer_shapes`]'s expression walk, so the recursion
5954/// threads one context reference instead of five parallel parameters.
5955/// `skip_shape_update` lists variables whose shapes are already pinned (by a
5956/// prior LHS pass); their bounds are not updated, though they are still
5957/// marked as seen.
5958pub(super) struct ShapeWalk<'a> {
5959 states: &'a HashSet<&'a str>,
5960 per_var_min: &'a mut HashMap<String, Vec<i64>>,
5961 per_var_max: &'a mut HashMap<String, Vec<i64>>,
5962 seen_indexed: &'a mut HashSet<String>,
5963 skip_shape_update: &'a HashSet<String>,
5964}
5965
5966impl ShapeWalk<'_> {
5967 /// Walk an expression tree collecting per-variable index bounds for shape
5968 /// inference. `loop_ranges` carries the concrete bounds of the enclosing
5969 /// aggregate loop symbols.
5970 fn walk(&mut self, expr: &Expr, loop_ranges: &HashMap<String, (i64, i64)>) {
5971 let Expr::Operator(node) = expr else {
5972 return;
5973 };
5974 if node.op == "index"
5975 && let Some(Expr::Variable(var)) = node.args.first()
5976 && self.states.contains(var.as_str())
5977 {
5978 self.seen_indexed.insert(var.clone());
5979 if !self.skip_shape_update.contains(var) {
5980 let mut dim_min: Vec<i64> = Vec::new();
5981 let mut dim_max: Vec<i64> = Vec::new();
5982 for idx_expr in node.args.iter().skip(1) {
5983 let (lo, hi) = evaluate_index_range(idx_expr, loop_ranges);
5984 dim_min.push(lo);
5985 dim_max.push(hi);
5986 }
5987 let cur_min = self.per_var_min.entry(var.clone()).or_default();
5988 let cur_max = self.per_var_max.entry(var.clone()).or_default();
5989 if cur_min.len() < dim_min.len() {
5990 cur_min.resize(dim_min.len(), i64::MAX);
5991 }
5992 if cur_max.len() < dim_max.len() {
5993 cur_max.resize(dim_max.len(), i64::MIN);
5994 }
5995 for (d, v) in dim_min.iter().enumerate() {
5996 cur_min[d] = cur_min[d].min(*v);
5997 }
5998 for (d, v) in dim_max.iter().enumerate() {
5999 cur_max[d] = cur_max[d].max(*v);
6000 }
6001 }
6002 }
6003 if is_faq_op(&node.op) {
6004 // Build loop range map from the faq's ranges. Ranges have
6005 // already been resolved to concrete intervals (RFC §5.2) by
6006 // `resolve_aggregate_ranges` at the top of `from_model`.
6007 let mut inner = loop_ranges.clone();
6008 if let Some(ranges) = &node.ranges {
6009 for (k, v) in ranges {
6010 if let Some(b) = v.bounds() {
6011 inner.insert(k.clone(), (b[0], b[1]));
6012 }
6013 }
6014 }
6015 node.for_each_child(&mut |child| self.walk(child, &inner));
6016 return;
6017 }
6018 node.for_each_child(&mut |child| self.walk(child, loop_ranges));
6019 }
6020}
6021
6022#[cfg(test)]
6023mod subsystem_ragged_and_inspection_tests {
6024 //! Subsystem mounting (esm-spec §4.6), ragged keyed-factor scope
6025 //! resolution (RFC §5.4; the Julia tree_walk `_factor_scope` mirror), and
6026 //! the [`BuildInspection`] observability surface — the Rust twins of the
6027 //! Julia `build_inspection_test.jl` cases (exact-rational overlap weights
6028 //! through the inspection surface; a 2-cell ragged CSR miniature end to
6029 //! end; build byte-identical with/without a sink).
6030 use super::*;
6031 use crate::simulate::{Alg, SolveOptions};
6032 use serde_json::json;
6033
6034 /// Typed load for inline test documents. The esm-schema pins `subsystems`
6035 /// entries to `{ "ref": … }` on disk (the official loader inlines the
6036 /// referenced file AFTER validation), so an inline-subsystem test document
6037 /// deserializes through serde directly — exactly the post-resolution shape
6038 /// the loader hands the simulator.
6039 fn typed(doc: serde_json::Value) -> EsmFile {
6040 serde_json::from_value(doc).expect("test document deserializes")
6041 }
6042
6043 /// A EsmProblem with build observability switched on at CONSTRUCTION — the
6044 /// seam that replaced threading a `&mut BuildInspection` through the run.
6045 fn inspecting_problem(file: &EsmFile) -> crate::problem::EsmProblem {
6046 crate::problem::esm_problem(
6047 file,
6048 (0.0, 1.0),
6049 crate::problem::ProblemOptions {
6050 inspect: true,
6051 rhs: crate::problem::Rhs::Always,
6052 // These documents carry a `polygon_intersection_area` regrid
6053 // and a RAGGED contraction bound, neither of which the tape
6054 // can lower, so `native` refuses them by NAME (API_SPEC
6055 // §5.8). What they pin is the ragged/overlap arithmetic and
6056 // the inspection sink, which is the reference evaluator's
6057 // answer to give.
6058 compiler: Some(crate::problem::Compiler::Interpreter),
6059 ..Default::default()
6060 },
6061 )
6062 .expect("builds")
6063 }
6064
6065 /// One `Model`, from the post-resolution JSON shape the loader hands the
6066 /// simulator (a `subsystems` entry already inlined).
6067 fn model(doc: serde_json::Value) -> Model {
6068 serde_json::from_value(doc).expect("test model deserializes")
6069 }
6070
6071 /// The aggregate RHS of the mounted equation defining `lhs`.
6072 fn agg_defining<'a>(m: &'a Model, lhs: &str) -> &'a ExpressionNode {
6073 let eq = m
6074 .equations
6075 .iter()
6076 .find(|e| matches!(&e.lhs, Expr::Variable(v) if v == lhs))
6077 .unwrap_or_else(|| panic!("no equation defines {lhs}"));
6078 match &eq.rhs {
6079 Expr::Operator(n) => n,
6080 other => panic!("{lhs} is not defined by an operator node: {other:?}"),
6081 }
6082 }
6083
6084 /// A leaf with a `join` on two DATA COLUMNS, mounted under `Leaf`.
6085 /// `join_on` is the clause's `on` list, so a test can shadow a name.
6086 fn host_mounting_a_join_leaf(join_on: serde_json::Value) -> Model {
6087 model(json!({
6088 "variables": { "doubled": { "type": "unknown" } },
6089 "equations": [
6090 { "lhs": "doubled", "rhs": { "op": "*", "args": [2.0, "Leaf.matched"] } }
6091 ],
6092 "subsystems": { "Leaf": {
6093 "variables": {
6094 "left_key": { "type": "unknown", "shape": ["leaf_left"] },
6095 "right_key": { "type": "unknown", "shape": ["leaf_right"] },
6096 "matched": { "type": "unknown" }
6097 },
6098 "equations": [
6099 { "lhs": "left_key", "rhs": { "op": "const", "args": [], "value": [7, 9, 4] } },
6100 { "lhs": "right_key", "rhs": { "op": "const", "args": [], "value": [7, 9] } },
6101 { "lhs": "matched",
6102 "rhs": { "op": "faq", "args": [], "semiring": "sum_product",
6103 "output_idx": [],
6104 "ranges": { "l": { "from": "leaf_left" },
6105 "r": { "from": "leaf_right" } },
6106 "join": [ { "on": join_on } ],
6107 "expr": 1.0 } }
6108 ]
6109 }}
6110 }))
6111 }
6112
6113 /// A `join.on` key column is a variable reference encoded as a STRING on
6114 /// the aggregate node (CONFORMANCE_SPEC §5.5.6), not an `Expr` child, so
6115 /// the mount's `map_children` walker never saw it: the leaf's variables
6116 /// became `Leaf.left_key` / `Leaf.right_key` while the join went on naming
6117 /// the bare originals, and the build died with "join key column 'left_key'
6118 /// does not resolve to a loop index of this aggregate". The names a node
6119 /// spells two ways must agree after the mount.
6120 #[test]
6121 fn mounting_carries_a_leafs_join_on_key_columns() {
6122 let mut m = host_mounting_a_join_leaf(json!([["left_key", "right_key"]]));
6123 let mut sets: HashMap<String, IndexSet> = HashMap::new();
6124 mount_subsystems(&mut m, &mut sets).expect("mounts");
6125
6126 let node = agg_defining(&m, "Leaf.matched");
6127 let on = &node.join.as_ref().expect("join survives the mount")[0].on;
6128 assert_eq!(
6129 on,
6130 &vec![["Leaf.left_key".to_string(), "Leaf.right_key".to_string()]],
6131 "both key columns follow the variables they name"
6132 );
6133 assert!(
6134 m.variables.contains_key("Leaf.left_key"),
6135 "the variables the join now names are the mounted ones"
6136 );
6137 }
6138
6139 /// The binder gate, per NAME. A leaf that declares a variable named like
6140 /// one of its own loop symbols is legal (esm-spec §4.3.1), and an `on`
6141 /// column naming that symbol resolves against the node's `ranges`
6142 /// (`join.rs::resolve_side`) — so prefixing it would make it resolve to
6143 /// nothing. The shadowed name stays bare while its SIBLING in the same
6144 /// clause is still rewritten, which is what makes this a per-name rule and
6145 /// not a per-node one.
6146 #[test]
6147 fn mounting_leaves_a_shadowed_loop_symbol_alone() {
6148 let mut m = host_mounting_a_join_leaf(json!([["l", "right_key"]]));
6149 m.subsystems.as_mut().expect("subsystems")["Leaf"]["variables"]["l"] =
6150 json!({ "type": "unknown", "shape": ["leaf_left"] });
6151 m.subsystems.as_mut().expect("subsystems")["Leaf"]["equations"]
6152 .as_array_mut()
6153 .expect("equations")
6154 .push(json!({ "lhs": "l", "rhs": { "op": "const", "args": [], "value": [1, 2, 3] } }));
6155 let mut sets: HashMap<String, IndexSet> = HashMap::new();
6156 mount_subsystems(&mut m, &mut sets).expect("mounts");
6157
6158 let node = agg_defining(&m, "Leaf.matched");
6159 let on = &node.join.as_ref().expect("join survives the mount")[0].on;
6160 assert_eq!(
6161 on,
6162 &vec![["l".to_string(), "Leaf.right_key".to_string()]],
6163 "the shadowed loop symbol stays bare; its sibling still follows the registry"
6164 );
6165 }
6166
6167 /// The same rule for the SPATIAL gate: an `overlap` clause's
6168 /// `src_env`/`tgt_env` name const-array envelope factors, which the
6169 /// broad phase resolves against the same registry the mount rewrote
6170 /// (`broad_phase::envelope_vectors`). Its `sym_src`/`sym_tgt` are range
6171 /// symbols the node binds and are left alone — they are not references.
6172 #[test]
6173 fn mounting_carries_overlap_envelope_factors() {
6174 let mut m = model(json!({
6175 "variables": { "out": { "type": "unknown" } },
6176 "equations": [ { "lhs": "out", "rhs": "Geo.hits" } ],
6177 "subsystems": { "Geo": {
6178 "variables": {
6179 "X": { "type": "unknown", "shape": ["pts"] },
6180 "W": { "type": "unknown", "shape": ["cells"] },
6181 "hits": { "type": "unknown" }
6182 },
6183 "equations": [
6184 { "lhs": "X", "rhs": { "op": "const", "args": [], "value": [1.0, 3.0] } },
6185 { "lhs": "W", "rhs": { "op": "const", "args": [], "value": [0.0, 2.0] } },
6186 { "lhs": "hits",
6187 "rhs": { "op": "faq", "args": [], "semiring": "sum_product",
6188 "output_idx": [],
6189 "ranges": { "p": { "from": "pts" }, "c": { "from": "cells" } },
6190 "join": [ { "overlap": { "src_env": ["X"], "tgt_env": ["W"] } } ],
6191 "expr": 1.0 } }
6192 ]
6193 }}
6194 }));
6195 let mut sets: HashMap<String, IndexSet> = HashMap::new();
6196 mount_subsystems(&mut m, &mut sets).expect("mounts");
6197
6198 let node = agg_defining(&m, "Geo.hits");
6199 let ov = node.join.as_ref().expect("join survives")[0]
6200 .overlap
6201 .as_ref()
6202 .expect("overlap survives");
6203 assert_eq!(ov.src_env, vec!["Geo.X".to_string()]);
6204 assert_eq!(ov.tgt_env, vec!["Geo.W".to_string()]);
6205 }
6206
6207 fn erk_opts() -> SolveOptions {
6208 SolveOptions {
6209 alg: Alg::Erk,
6210 reltol: Some(1e-10),
6211 abstol: Some(1e-12),
6212 saveat: Some(vec![1.0]),
6213 ..Default::default()
6214 }
6215 }
6216
6217 /// A 2-cell ragged CSR miniature: a `mesh` subsystem ships the const
6218 /// factors (per-cell edge counts, the padded edge-membership table, and
6219 /// per-edge weights), the parent re-exposes the offsets/values factors as
6220 /// bare-name aliases (the MPAS keyed-factor wiring contract), and an
6221 /// observed contracts over the ragged `edges_of_cell` set. Expected
6222 /// per-cell sums are exact small integers, NONZERO — so an empty ragged
6223 /// contraction (a silently unresolved offsets factor) cannot pass.
6224 fn ragged_miniature_doc() -> serde_json::Value {
6225 json!({
6226 "esm": "1.1.0",
6227 "metadata": {"name": "ragged_subsystem_miniature"},
6228 "index_sets": {
6229 "cells": {"kind": "interval", "size": 2},
6230 "edges": {"kind": "interval", "size": 3},
6231 "maxEdges": {"kind": "interval", "size": 3},
6232 "edges_of_cell": {"kind": "ragged", "of": ["cells"],
6233 "offsets": "nEdgesOnCell",
6234 "values": "edgesOnCell"}
6235 },
6236 "models": {"M": {
6237 "subsystems": {"mesh": {
6238 "esm": "1.0.0",
6239 "metadata": {"name": "mini_mesh"},
6240 "models": {"MiniMesh": {
6241 "variables": {
6242 "nEdgesOnCell": {"type": "unknown", "shape": ["cells"]},
6243 "edgesOnCell": {"type": "unknown", "shape": ["cells", "maxEdges"]},
6244 "w": {"type": "unknown", "shape": ["edges"]}
6245 },
6246 "equations": [
6247 {"lhs": "nEdgesOnCell", "rhs": {"op": "const", "value": [2, 3], "args": []}},
6248 {"lhs": "edgesOnCell", "rhs": {"op": "const", "value": [[1, 2, 0], [1, 2, 3]], "args": []}},
6249 {"lhs": "w", "rhs": {"op": "const", "value": [10.0, 20.0, 30.0], "args": []}}]
6250 }}
6251 }},
6252 "variables": {
6253 "u": {"type": "unknown", "units": "1", "shape": ["cells"]},
6254 "nEdgesOnCell": {"type": "unknown", "shape": ["cells"]},
6255 "edgesOnCell": {"type": "unknown", "shape": ["cells", "maxEdges"]},
6256 "s": {"type": "unknown", "shape": ["cells"]}
6257 },
6258 "equations": [
6259 {"lhs": "nEdgesOnCell", "rhs": "mesh.nEdgesOnCell"},
6260 {"lhs": "edgesOnCell", "rhs": "mesh.edgesOnCell"},
6261 {"lhs": "s", "rhs": {
6262 "op": "faq", "args": ["edgesOnCell", "mesh.w"],
6263 "output_idx": ["i"], "semiring": "sum_product",
6264 "ranges": {"i": {"from": "cells"},
6265 "k": {"from": "edges_of_cell", "of": ["i"]}},
6266 "expr": {"op": "index", "args": ["mesh.w",
6267 {"op": "index", "args": ["edgesOnCell", "i", "k"]}]}
6268 }},
6269 {"lhs": {"op": "ic", "args": ["u"]}, "rhs": 0.0},
6270 {"lhs": {"op": "D", "args": ["u"], "wrt": "t"}, "rhs": "s"}
6271 ]
6272 }}
6273 })
6274 }
6275
6276 /// End to end: subsystem consts mount under `mesh.*`, the bare-alias
6277 /// observeds materialize from them, the ragged offsets factor resolves in
6278 /// the model scope, and the CSR contraction yields the exact nonzero
6279 /// per-cell sums s = [10+20, 10+20+30] — both in the integrated state
6280 /// (u(1) = s from a zero ic) and, exactly, in the inspection's
6281 /// materialized setup arrays.
6282 #[test]
6283 fn ragged_csr_miniature_through_subsystem_and_aliases() {
6284 let file = typed(ragged_miniature_doc());
6285 let prob = inspecting_problem(&file);
6286 let sol = crate::problem::solve(&prob, &erk_opts()).expect("solves");
6287 let insp = prob.take_inspection();
6288 let ti = sol.time.len() - 1;
6289 let cells = crate::inline_tests::state_cells(&sol.state_variable_names, "u", "M");
6290 assert_eq!(cells.len(), 2);
6291 let u1: Vec<f64> = cells.iter().map(|(_, row)| sol.state[*row][ti]).collect();
6292 assert!((u1[0] - 30.0).abs() < 1e-8, "u[1](1) = {} != 30", u1[0]);
6293 assert!((u1[1] - 60.0).abs() < 1e-8, "u[2](1) = {} != 60", u1[1]);
6294 // The state-free rule output is captured EXACTLY at build.
6295 let s = insp.setup_arrays.get("s").expect("s captured");
6296 assert_eq!(s.shape(), [2]);
6297 assert_eq!(s[IxDyn(&[0])], 30.0);
6298 assert_eq!(s[IxDyn(&[1])], 60.0);
6299 // Mounted const factors and their aliases are captured too.
6300 for name in [
6301 "mesh.nEdgesOnCell",
6302 "mesh.edgesOnCell",
6303 "mesh.w",
6304 "nEdgesOnCell",
6305 "edgesOnCell",
6306 ] {
6307 assert!(insp.setup_arrays.contains_key(name), "missing '{name}'");
6308 }
6309 assert_eq!(
6310 insp.setup_arrays["nEdgesOnCell"],
6311 insp.setup_arrays["mesh.nEdgesOnCell"]
6312 );
6313 assert!(insp.observed_exprs.contains_key("s"));
6314 }
6315
6316 /// Filling the inspection never changes the run: the trajectory is
6317 /// bit-identical with and without a sink.
6318 #[test]
6319 fn inspection_does_not_change_the_run() {
6320 let file = typed(ragged_miniature_doc());
6321 let plain = crate::problem::esm_problem(
6322 &file,
6323 (0.0, 1.0),
6324 crate::problem::ProblemOptions {
6325 p: HashMap::new().clone(),
6326 u0: HashMap::new().clone(),
6327 rhs: crate::problem::Rhs::Always,
6328 // As `inspecting_problem`: this document's contraction bound
6329 // is RAGGED, which the tape cannot size, so `native` refuses
6330 // it by NAME (API_SPEC §5.8).
6331 compiler: Some(crate::problem::Compiler::Interpreter),
6332 ..Default::default()
6333 },
6334 )
6335 .and_then(|prob| crate::problem::solve(&prob, &erk_opts()))
6336 .expect("simulates");
6337 let prob = inspecting_problem(&file);
6338 let inspected = crate::problem::solve(&prob, &erk_opts()).expect("solves");
6339 let insp = prob.take_inspection();
6340 assert_eq!(plain.time, inspected.time);
6341 assert_eq!(plain.state, inspected.state);
6342 assert_eq!(plain.state_variable_names, inspected.state_variable_names);
6343 assert!(!insp.setup_arrays.is_empty());
6344 }
6345
6346 fn ragged_registry(offsets: &str, values: &str) -> HashMap<String, IndexSet> {
6347 HashMap::from([(
6348 "edges_of_cell".to_string(),
6349 IndexSet {
6350 kind: "ragged".to_string(),
6351 size: None,
6352 members: None,
6353 from_faq: None,
6354 member_factor: None,
6355 of: Some(vec!["cells".to_string()]),
6356 offsets: Some(offsets.to_string()),
6357 values: Some(values.to_string()),
6358 },
6359 )])
6360 }
6361
6362 fn obs_var() -> ModelVariable {
6363 serde_json::from_value(json!({"type": "unknown"})).expect("variable parses")
6364 }
6365
6366 /// `_factor_scope` semantics: an exact-name variable wins (registry
6367 /// untouched); with no exact name, the unique dot-suffix match at the
6368 /// SHALLOWEST namespace depth is substituted — for BOTH the offsets and
6369 /// values factors.
6370 #[test]
6371 fn factor_scope_exact_name_wins_and_shallowest_suffix_resolves() {
6372 // Exact name in scope: keep as authored.
6373 let mut reg = ragged_registry("nEdgesOnCell", "edgesOnCell");
6374 let vars: IndexMap<String, ModelVariable> = [
6375 ("nEdgesOnCell", obs_var()),
6376 ("mesh.nEdgesOnCell", obs_var()),
6377 ("edgesOnCell", obs_var()),
6378 ("mesh.edgesOnCell", obs_var()),
6379 ]
6380 .into_iter()
6381 .map(|(k, v)| (k.to_string(), v))
6382 .collect();
6383 apply_ragged_factor_scope(&mut reg, &vars).expect("resolves");
6384 assert_eq!(
6385 reg["edges_of_cell"].offsets.as_deref(),
6386 Some("nEdgesOnCell")
6387 );
6388 assert_eq!(reg["edges_of_cell"].values.as_deref(), Some("edgesOnCell"));
6389
6390 // No exact name: the depth-1 alias beats the depth-2 mounted const.
6391 let mut reg = ragged_registry("nEdgesOnCell", "edgesOnCell");
6392 let vars: IndexMap<String, ModelVariable> = [
6393 ("Div.nEdgesOnCell", obs_var()),
6394 ("Div.mesh.nEdgesOnCell", obs_var()),
6395 ("Div.edgesOnCell", obs_var()),
6396 ("Div.mesh.edgesOnCell", obs_var()),
6397 ]
6398 .into_iter()
6399 .map(|(k, v)| (k.to_string(), v))
6400 .collect();
6401 apply_ragged_factor_scope(&mut reg, &vars).expect("resolves");
6402 assert_eq!(
6403 reg["edges_of_cell"].offsets.as_deref(),
6404 Some("Div.nEdgesOnCell")
6405 );
6406 assert_eq!(
6407 reg["edges_of_cell"].values.as_deref(),
6408 Some("Div.edgesOnCell")
6409 );
6410
6411 // No candidate at all: left bare (the existing unbound-name behavior
6412 // surfaces downstream).
6413 let mut reg = ragged_registry("nowhere", "edgesOnCell");
6414 let vars: IndexMap<String, ModelVariable> = [("Div.edgesOnCell", obs_var())]
6415 .into_iter()
6416 .map(|(k, v)| (k.to_string(), v))
6417 .collect();
6418 apply_ragged_factor_scope(&mut reg, &vars).expect("resolves");
6419 assert_eq!(reg["edges_of_cell"].offsets.as_deref(), Some("nowhere"));
6420 }
6421
6422 /// Two dot-suffix candidates at the same (shallowest) depth are a HARD
6423 /// ERROR — never a silent pick or an empty contraction.
6424 #[test]
6425 fn factor_scope_ambiguity_is_a_hard_error() {
6426 let mut reg = ragged_registry("nEdgesOnCell", "edgesOnCell");
6427 let vars: IndexMap<String, ModelVariable> = [
6428 ("A.nEdgesOnCell", obs_var()),
6429 ("B.nEdgesOnCell", obs_var()),
6430 ("A.edgesOnCell", obs_var()),
6431 ]
6432 .into_iter()
6433 .map(|(k, v)| (k.to_string(), v))
6434 .collect();
6435 let err = apply_ragged_factor_scope(&mut reg, &vars).unwrap_err();
6436 let msg = err.to_string();
6437 assert!(msg.contains("ambiguous"), "unexpected error: {msg}");
6438 assert!(msg.contains("A.nEdgesOnCell") && msg.contains("B.nEdgesOnCell"));
6439 }
6440
6441 /// The conservative-overlap exact rationals through the inspection
6442 /// surface: two unit source squares tiling one 2x1 target rectangle. The
6443 /// per-pair overlap areas, the filtered row-sum, and the normalized
6444 /// weights are the library-shaped aggregates of the ESD fixture (narrow
6445 /// phase `polygon_intersection_area`, sliver filter `> atol`), and every
6446 /// captured value is BIT-EXACT: A_ij = [1, 1], A_j = [2], W = [1/2, 1/2].
6447 #[test]
6448 fn exact_rational_overlap_weights_through_inspection() {
6449 let doc = json!({
6450 "esm": "1.1.0",
6451 "metadata": {"name": "inspect_overlap"},
6452 "index_sets": {
6453 "src_cells": {"kind": "interval", "size": 2},
6454 "tgt_cells": {"kind": "interval", "size": 1}
6455 },
6456 "models": {"R": {
6457 "variables": {
6458 "q": {"type": "unknown", "units": "1", "shape": ["tgt_cells"],
6459 "default": 0.0},
6460 "atol": {"type": "parameter", "units": "1", "default": 1e-12},
6461 "src_poly": {"type": "unknown"},
6462 "tgt_poly": {"type": "unknown"},
6463 "A_ij": {"type": "unknown"},
6464 "A_j": {"type": "unknown"},
6465 "W_ij": {"type": "unknown"}
6466 },
6467 "equations": [
6468 {"lhs": "src_poly", "rhs": {"op": "const", "args": [], "value": [
6469 [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]],
6470 [[1.0, 0.0], [2.0, 0.0], [2.0, 1.0], [1.0, 1.0]]]}},
6471 {"lhs": "tgt_poly", "rhs": {"op": "const", "args": [], "value": [
6472 [[0.0, 0.0], [2.0, 0.0], [2.0, 1.0], [0.0, 1.0]]]}},
6473 {"lhs": "A_ij", "rhs": {
6474 "op": "faq", "args": ["src_poly", "tgt_poly"],
6475 "output_idx": ["i", "j"], "semiring": "sum_product",
6476 "ranges": {"i": {"from": "src_cells"}, "j": {"from": "tgt_cells"}},
6477 "expr": {"op": "polygon_intersection_area", "manifold": "planar",
6478 "args": [{"op": "index", "args": ["src_poly", "i"]},
6479 {"op": "index", "args": ["tgt_poly", "j"]}]}}},
6480 {"lhs": "A_j", "rhs": {
6481 "op": "faq", "args": ["A_ij"],
6482 "output_idx": ["j"], "semiring": "sum_product",
6483 "ranges": {"i": {"from": "src_cells"}, "j": {"from": "tgt_cells"}},
6484 "filter": {"op": ">", "args": [
6485 {"op": "index", "args": ["A_ij", "i", "j"]}, "atol"]},
6486 "expr": {"op": "index", "args": ["A_ij", "i", "j"]}}},
6487 {"lhs": "W_ij", "rhs": {
6488 "op": "faq", "args": ["A_ij", "A_j"],
6489 "output_idx": ["i", "j"], "semiring": "sum_product",
6490 "ranges": {"i": {"from": "src_cells"}, "j": {"from": "tgt_cells"}},
6491 "filter": {"op": ">", "args": [
6492 {"op": "index", "args": ["A_ij", "i", "j"]}, "atol"]},
6493 "expr": {"op": "/", "args": [
6494 {"op": "index", "args": ["A_ij", "i", "j"]},
6495 {"op": "index", "args": ["A_j", "j"]}]}}},
6496 {"lhs": {"op": "D", "args": ["q"], "wrt": "t"}, "rhs": 0.0}
6497 ]
6498 }}
6499 });
6500 let file = typed(doc);
6501 let prob = inspecting_problem(&file);
6502 crate::problem::solve(&prob, &erk_opts()).expect("solves");
6503 let insp = prob.take_inspection();
6504 let a_ij = insp.setup_arrays.get("A_ij").expect("A_ij captured");
6505 assert_eq!(a_ij.shape(), [2, 1]);
6506 assert_eq!(a_ij[IxDyn(&[0, 0])], 1.0);
6507 assert_eq!(a_ij[IxDyn(&[1, 0])], 1.0);
6508 let a_j = insp.setup_arrays.get("A_j").expect("A_j captured");
6509 assert_eq!(a_j.shape(), [1]);
6510 assert_eq!(a_j[IxDyn(&[0])], 2.0);
6511 let w_ij = insp.setup_arrays.get("W_ij").expect("W_ij captured");
6512 assert_eq!(w_ij.shape(), [2, 1]);
6513 assert_eq!(w_ij[IxDyn(&[0, 0])], 0.5);
6514 assert_eq!(w_ij[IxDyn(&[1, 0])], 0.5);
6515 }
6516
6517 /// esm-spec §4.6: `M.x` inside model `M` is the fully qualified spelling of
6518 /// the local `x`, and the single-model array build brings it back to the
6519 /// local one before compiling. The rewrite is capture-aware in BOTH
6520 /// directions: a node that BINDS the local name as a loop symbol keeps the
6521 /// qualified spelling, so `M.i` under `aggregate(i from x; …)` never
6522 /// silently becomes the loop index.
6523 #[test]
6524 fn self_qualified_rewrite_does_not_capture_a_bound_local_name() {
6525 let m = model(json!({
6526 "variables": {
6527 "i": {"type": "parameter", "units": "1", "default": 3.0},
6528 "u": {"type": "unknown", "units": "1", "shape": ["x"]},
6529 "z": {"type": "unknown", "units": "1"}
6530 },
6531 "equations": [
6532 // Free: rewritten to the local `i`.
6533 {"lhs": "z", "rhs": "M.i"},
6534 // Bound by the enclosing aggregate's own `i`: left alone.
6535 {"lhs": "u", "rhs": {
6536 "op": "faq", "args": [], "output_idx": ["i"],
6537 "ranges": {"i": {"from": "x"}},
6538 "expr": {"op": "*", "args": ["M.i", "i"]}}}
6539 ]
6540 }));
6541 let hits = self_qualified_references(&m, "M");
6542 assert_eq!(hits, vec![("M.i".to_string(), "i".to_string())]);
6543 let mut out = m.clone();
6544 resolve_self_qualified_references(&mut out, &hits);
6545 assert!(matches!(&out.equations[0].rhs, Expr::Variable(v) if v == "i"));
6546 let Expr::Operator(node) = &out.equations[1].rhs else {
6547 panic!("expected the aggregate");
6548 };
6549 let body = node.expr.as_deref().expect("aggregate body");
6550 let Expr::Operator(mul) = body else {
6551 panic!("expected the product");
6552 };
6553 assert!(
6554 matches!(&mul.args[0], Expr::Variable(v) if v == "M.i"),
6555 "a node binding `i` must keep the qualified spelling, got {:?}",
6556 mul.args[0]
6557 );
6558 }
6559
6560 /// esm-spec §4.6: a self-qualified reference authored INSIDE a subsystem
6561 /// body — `M.B.y` in `M.A`'s equation, `M.k` two levels down — names the
6562 /// same variable as the mounted spelling, and resolves on the single-model
6563 /// route like one in the parent's own equations. Each tendency reads a
6564 /// value no other rule supplies, so an unresolved name cannot hide.
6565 #[test]
6566 fn self_qualified_references_inside_subsystem_bodies_resolve() {
6567 let file = typed(json!({
6568 "esm": "1.0.0",
6569 "metadata": {"name": "nested_self_refs"},
6570 "models": {"M": {
6571 "variables": {"k": {"type": "parameter", "units": "1", "default": 2.0}},
6572 "equations": [],
6573 "subsystems": {
6574 "A": {
6575 "variables": {"x": {"type": "unknown", "units": "1", "default": 1.0}},
6576 "equations": [{"lhs": {"op": "D", "args": ["x"], "wrt": "t"},
6577 "rhs": {"op": "*", "args": ["M.k", "M.B.C.z"]}}],
6578 "subsystems": {"C": {
6579 "variables": {"z": {"type": "unknown", "units": "1", "default": 3.0}},
6580 "equations": [{"lhs": {"op": "D", "args": ["z"], "wrt": "t"},
6581 "rhs": {"op": "-", "args": ["M.A.x"]}}]
6582 }}
6583 },
6584 "B": {
6585 "variables": {},
6586 "equations": [],
6587 "subsystems": {"C": {
6588 "variables": {"z": {"type": "unknown", "units": "1", "default": 5.0}},
6589 "equations": [{"lhs": {"op": "D", "args": ["z"], "wrt": "t"},
6590 "rhs": {"op": "*", "args": ["M.k", "M.A.C.z"]}}]
6591 }}
6592 }
6593 }
6594 }}
6595 }));
6596 let compiled = ArrayCompiled::from_file(&file).expect("compiles");
6597 let (_, report) = compiled.build_tape(&HashSet::new());
6598 assert!(
6599 report.fallbacks.is_empty(),
6600 "every self-qualified read must resolve: {:?}",
6601 report.fallbacks
6602 );
6603 let names = compiled.state_variable_names();
6604 let at = |n: &str| names.iter().position(|s| s == n).expect(n);
6605 let mut state = vec![0.0; names.len()];
6606 state[at("A.x")] = 1.0;
6607 state[at("A.C.z")] = 3.0;
6608 state[at("B.C.z")] = 5.0;
6609 let (dy, _) = compiled.debug_eval_rhs(&state, 0.0, &HashMap::new(), true);
6610 assert_eq!(dy[at("A.x")], 10.0, "D(A.x) = k * B.C.z");
6611 assert_eq!(dy[at("A.C.z")], -1.0, "D(A.C.z) = -A.x");
6612 assert_eq!(dy[at("B.C.z")], 6.0, "D(B.C.z) = k * A.C.z");
6613 }
6614}
6615
6616#[cfg(test)]
6617mod event_driven_update_refusal_tests {
6618 use super::*;
6619
6620 /// A parameter whose `update` recomputes it from an `expression` on a
6621 /// `schedule` / `condition` / `crossing` trigger (esm-spec §5.4) is
6622 /// refused, never frozen at its default, and the refusal names the rule
6623 /// kinds it cannot run. A `data` rule on another parameter is the forcing
6624 /// seam and is not named.
6625 #[test]
6626 fn an_expression_update_is_refused_naming_its_kinds() {
6627 let file = crate::parse::load_string(
6628 r#"{
6629 "esm": "1.0.0",
6630 "metadata": { "name": "ExprUpdates", "authors": ["test"] },
6631 "models": { "M": {
6632 "variables": {
6633 "x": { "type": "unknown", "default": 0.0 },
6634 "k": { "type": "parameter", "default": 1.0, "shape": [],
6635 "update": [
6636 { "kind": "schedule", "interval": 2.0, "expression": 2.0 },
6637 { "kind": "crossing", "when": { "op": "-", "args": ["x", 1.0] }, "expression": 3.0 }
6638 ] }
6639 },
6640 "equations": [ { "lhs": { "op": "D", "args": ["x"], "wrt": "t" }, "rhs": "k" } ]
6641 } }
6642}"#,
6643 )
6644 .expect("loads");
6645 let model = &file.models.as_ref().unwrap()["M"];
6646 let err = classify_variables(model, &HashSet::new()).expect_err("refused");
6647 let text = err.to_string();
6648 assert!(
6649 text.contains("parameter 'k'") && text.contains("`schedule`, `crossing`"),
6650 "{text}"
6651 );
6652 }
6653}