use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::sync::Arc;
use crate::compile_error::CompileError;
use crate::types::Expr;
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash)]
pub enum Precision {
#[default]
Float64,
Float32,
}
impl Precision {
pub fn from_element_type(s: Option<&str>) -> Result<Self, CompileError> {
match s {
None | Some("Float64") => Ok(Precision::Float64),
Some("Float32") => Ok(Precision::Float32),
Some(other) => Err(CompileError::UnsupportedElementType {
element_type: other.to_string(),
}),
}
}
#[inline(always)]
#[must_use]
pub fn is_f32(self) -> bool {
matches!(self, Precision::Float32)
}
#[inline(always)]
#[must_use]
pub fn round(self, v: f64) -> f64 {
match self {
Precision::Float64 => v,
Precision::Float32 => v as f32 as f64,
}
}
#[must_use]
pub fn element_type(self) -> &'static str {
match self {
Precision::Float64 => "Float64",
Precision::Float32 => "Float32",
}
}
}
thread_local! {
static ACTIVE: Cell<Precision> = const { Cell::new(Precision::Float64) };
static DOCUMENT: Cell<Precision> = const { Cell::new(Precision::Float64) };
static VARS: RefCell<Arc<VarPrecisions>> = RefCell::new(Arc::new(VarPrecisions::default()));
}
#[derive(Debug, Default, Clone)]
pub struct VarPrecisions {
by_name: HashMap<String, Precision>,
declared: std::collections::BTreeSet<String>,
}
#[derive(Debug, Default)]
pub struct ModelDeclarations<'a> {
pub model: &'a str,
pub variables: Vec<(&'a str, Option<&'a str>)>,
pub subsystem_variables: Vec<(String, Option<String>)>,
}
impl VarPrecisions {
pub fn from_declarations(
models: &[ModelDeclarations<'_>],
document: Precision,
selected: Option<&str>,
) -> Result<Self, CompileError> {
for m in models
.iter()
.filter(|m| selected.is_none_or(|s| s == m.model))
{
for (path, et) in &m.subsystem_variables {
if et.is_some() && Precision::from_element_type(et.as_deref())? != document {
return Err(CompileError::ElementTypeUncarried {
construct: format!("subsystem variable '{}.{path}'", m.model),
reason: format!(
"it declares `element_type: \"{}\"` and the document evaluates in \
{}, but this evaluator carries a per-variable element type on \
top-level model variables only (a mounted leaf's precision reaches \
a `subsystems.<k>` mount this way, esm-spec §4.7)",
et.as_deref().unwrap_or_default(),
document.element_type()
),
});
}
}
}
let mut out = VarPrecisions::default();
let mut any = false;
let mut effective: Vec<(&str, &str, Precision)> = Vec::new();
for m in models {
for (v, et) in &m.variables {
let p = match et {
Some(_) => Precision::from_element_type(*et)?,
None => document,
};
if et.is_some() {
any = true;
out.declared.insert(format!("{}.{v}", m.model));
}
effective.push((m.model, v, p));
}
}
if !any {
return Ok(VarPrecisions::default());
}
let bare_from = selected.or(if models.len() == 1 {
Some(models[0].model)
} else {
None
});
let mut bare: HashMap<&str, Option<Precision>> = HashMap::new();
for (m, v, p) in &effective {
out.by_name.insert(format!("{m}.{v}"), *p);
match bare_from {
Some(sel) if sel == *m => {
bare.insert(v, Some(*p));
}
Some(_) => {}
None => {
let e = bare.entry(v).or_insert(Some(*p));
if *e != Some(*p) {
*e = None;
}
}
}
}
for (v, p) in bare {
if let Some(p) = p {
out.by_name.insert(v.to_string(), p);
}
}
for m in models {
for (path, _) in &m.subsystem_variables {
out.by_name.insert(format!("{}.{path}", m.model), document);
out.by_name.entry(path.clone()).or_insert(document);
}
}
Ok(out)
}
pub fn insert(&mut self, model: Option<&str>, name: &str, p: Precision) {
self.declared.insert(name.to_string());
self.by_name.insert(name.to_string(), p);
if let Some(m) = model {
self.by_name.insert(format!("{m}.{name}"), p);
}
}
#[must_use]
pub fn get(&self, name: &str) -> Option<Precision> {
if let Some(p) = self.by_name.get(name) {
return Some(*p);
}
let tail = name.rsplit('.').next()?;
if tail == name {
return None;
}
self.by_name.get(tail).copied()
}
pub fn names(&self) -> impl Iterator<Item = &str> {
self.declared.iter().map(String::as_str)
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.by_name.is_empty()
}
}
#[must_use]
pub fn subsystem_declarations(
subsystems: Option<&serde_json::Value>,
) -> Vec<(String, Option<String>)> {
fn walk(subs: &serde_json::Value, prefix: &str, out: &mut Vec<(String, Option<String>)>) {
let Some(subs) = subs.as_object() else {
return;
};
for (name, sub) in subs {
let path = if prefix.is_empty() {
name.clone()
} else {
format!("{prefix}.{name}")
};
if let Some(vars) = sub.get("variables").and_then(serde_json::Value::as_object) {
for (v, d) in vars {
out.push((
format!("{path}.{v}"),
d.get("element_type")
.and_then(serde_json::Value::as_str)
.map(str::to_string),
));
}
}
if let Some(nested) = sub.get("subsystems") {
walk(nested, &path, out);
}
}
}
let mut out = Vec::new();
if let Some(subs) = subsystems {
walk(subs, "", &mut out);
}
out
}
#[inline(always)]
#[must_use]
pub fn active() -> Precision {
ACTIVE.with(Cell::get)
}
#[inline(always)]
#[must_use]
pub fn is_f32() -> bool {
active().is_f32()
}
#[inline(always)]
#[must_use]
pub fn round(v: f64) -> f64 {
active().round(v)
}
#[inline(always)]
#[must_use]
pub fn document() -> Precision {
DOCUMENT.with(Cell::get)
}
#[inline]
#[must_use]
pub fn of_variable(name: &str) -> Precision {
VARS.with(|v| v.borrow().get(name)).unwrap_or_else(document)
}
#[inline]
#[must_use]
pub fn has_variable_overrides() -> bool {
VARS.with(|v| !v.borrow().is_empty())
}
#[must_use]
pub fn first_variable_override() -> Option<String> {
VARS.with(|v| v.borrow().names().min().map(str::to_string))
}
#[must_use]
pub fn variable_precisions() -> Arc<VarPrecisions> {
VARS.with(|v| Arc::clone(&v.borrow()))
}
#[derive(Clone, Debug, Default)]
pub struct Env {
pub document: Precision,
pub variables: Arc<VarPrecisions>,
}
impl Env {
#[must_use]
pub fn capture() -> Env {
Env {
document: document(),
variables: variable_precisions(),
}
}
pub fn enter(&self) -> EnvGuard {
enter_env(self.document, Arc::clone(&self.variables))
}
#[inline(always)]
#[must_use]
pub fn is_f32(&self) -> bool {
self.document.is_f32()
}
}
#[must_use = "the precision environment is restored as soon as the guard is dropped"]
pub struct EnvGuard {
prev_doc: Precision,
prev_active: Precision,
prev_vars: Arc<VarPrecisions>,
}
impl Drop for EnvGuard {
fn drop(&mut self) {
DOCUMENT.with(|c| c.set(self.prev_doc));
ACTIVE.with(|c| c.set(self.prev_active));
VARS.with(|v| *v.borrow_mut() = Arc::clone(&self.prev_vars));
}
}
pub fn enter_env(doc: Precision, vars: Arc<VarPrecisions>) -> EnvGuard {
let prev_doc = DOCUMENT.with(|c| c.replace(doc));
let prev_active = ACTIVE.with(|c| c.replace(doc));
let prev_vars = VARS.with(|v| std::mem::replace(&mut *v.borrow_mut(), vars));
EnvGuard {
prev_doc,
prev_active,
prev_vars,
}
}
#[must_use = "the precision is restored as soon as the guard is dropped"]
pub struct PrecisionGuard {
prev: Precision,
}
impl Drop for PrecisionGuard {
fn drop(&mut self) {
ACTIVE.with(|c| c.set(self.prev));
}
}
pub fn enter(p: Precision) -> PrecisionGuard {
let prev = ACTIVE.with(|c| c.replace(p));
PrecisionGuard { prev }
}
pub const F32_EXACT_INT_LIMIT: i64 = 1 << 24;
pub fn check_index_set_extent(name: &str, size: i64) -> Result<(), CompileError> {
if size > F32_EXACT_INT_LIMIT {
return Err(CompileError::Float32Unsupported {
construct: format!("index set `{name}` of size {size}"),
reason: format!(
"index expressions share the value kernels, so a subscript above \
2^24 ({F32_EXACT_INT_LIMIT}) would round under \
`domain.element_type: \"Float32\"`"
),
});
}
Ok(())
}
#[must_use]
pub fn f32_unsupported_reason(op: &str, fn_name: Option<&str>) -> Option<(String, String)> {
let geometry = |op: &str| {
(
format!("operator `{op}`"),
"polygon clipping and area are computed by the binary64 geometry \
kernels (`crate::geometry`), which have no binary32 form"
.to_string(),
)
};
match op {
"intersect_polygon" | "polygon_intersection_area" => Some(geometry(op)),
"fn" => match fn_name {
Some(n @ ("interp.linear" | "interp.bilinear")) => Some((
format!("closed function `{n}`"),
"table interpolation is evaluated by the binary64 closed-function \
registry, which has no binary32 form"
.to_string(),
)),
Some(n @ "datetime.julian_day") => Some((
format!("closed function `{n}`"),
"a Julian day number exceeds binary32's exact-integer range, so \
its fractional part would be destroyed by the declared precision"
.to_string(),
)),
_ => None,
},
_ => None,
}
}
pub fn check_f32_supported(expr: &Expr) -> Result<(), CompileError> {
if !is_f32() {
return Ok(());
}
check_f32_supported_tree(expr)
}
fn check_f32_supported_tree(expr: &Expr) -> Result<(), CompileError> {
let Expr::Operator(node) = expr else {
return Ok(());
};
let fn_name = node.name.as_deref();
if let Some((construct, reason)) = f32_unsupported_reason(&node.op, fn_name) {
return Err(CompileError::Float32Unsupported { construct, reason });
}
let mut first_err: Option<CompileError> = None;
node.for_each_child(&mut |child| {
if first_err.is_none()
&& let Err(e) = check_f32_supported_tree(child)
{
first_err = Some(e);
}
});
match first_err {
Some(e) => Err(e),
None => Ok(()),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_is_float64_and_round_is_identity() {
assert_eq!(active(), Precision::Float64);
let v = 0.1_f64;
assert_eq!(round(v).to_bits(), v.to_bits());
}
#[test]
fn guard_sets_and_restores() {
assert!(!is_f32());
{
let _g = enter(Precision::Float32);
assert!(is_f32());
assert_eq!(round(0.1_f64), 0.1_f32 as f64);
{
let _inner = enter(Precision::Float64);
assert!(!is_f32());
}
assert!(
is_f32(),
"the inner guard restores Float32, not the default"
);
}
assert!(!is_f32());
}
#[test]
fn element_type_parsing() {
assert_eq!(
Precision::from_element_type(None).unwrap(),
Precision::Float64
);
assert_eq!(
Precision::from_element_type(Some("Float64")).unwrap(),
Precision::Float64
);
assert_eq!(
Precision::from_element_type(Some("Float32")).unwrap(),
Precision::Float32
);
assert!(Precision::from_element_type(Some("Float16")).is_err());
assert!(Precision::from_element_type(Some("float32")).is_err());
}
#[test]
fn f32_unsupported_set() {
assert!(f32_unsupported_reason("intersect_polygon", None).is_some());
assert!(f32_unsupported_reason("polygon_intersection_area", None).is_some());
assert!(f32_unsupported_reason("fn", Some("interp.linear")).is_some());
assert!(f32_unsupported_reason("fn", Some("interp.bilinear")).is_some());
assert!(f32_unsupported_reason("fn", Some("datetime.julian_day")).is_some());
assert!(f32_unsupported_reason("fn", Some("interp.searchsorted")).is_none());
assert!(f32_unsupported_reason("fn", Some("datetime.year")).is_none());
assert!(f32_unsupported_reason("+", None).is_none());
}
#[test]
fn index_set_extent_gate() {
assert!(check_index_set_extent("rows", F32_EXACT_INT_LIMIT).is_ok());
let err = check_index_set_extent("rows", F32_EXACT_INT_LIMIT + 1).unwrap_err();
assert!(format!("{err}").contains("rows"), "{err}");
}
}