use std::collections::HashMap;
use num_traits::{One, ToPrimitive};
use rustc_hash::FxHashMap;
use super::dim::ConstDim;
use crate::base::arena::Arena;
use crate::base::node::{ExprId, ExprNode};
use crate::base::numeric::Q;
use crate::base::walk;
use crate::prelude::Ex;
#[derive(Clone, Debug, Default)]
pub struct DimMap {
map: HashMap<String, ConstDim>,
}
impl DimMap {
pub fn new() -> Self {
Self::default()
}
pub fn with(mut self, name: &str, dim: ConstDim) -> Self {
self.map.insert(name.to_string(), dim);
self
}
pub fn with_var(mut self, var: &Ex, dim: ConstDim) -> Self {
let name = format!("{}", var);
self.map.insert(name, dim);
self
}
pub fn get(&self, name: &str) -> Option<&ConstDim> {
self.map.get(name)
}
pub fn insert(&mut self, name: &str, dim: ConstDim) {
self.map.insert(name.to_string(), dim);
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
}
impl core::fmt::Display for ConstDim {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
let n = self.name();
if n != "Unknown" {
return write!(f, "{}", n);
}
write!(
f,
"L^{} M^{} T^{} I^{} Θ^{} N^{} J^{}",
self.l, self.m, self.t, self.i, self.th, self.n, self.j
)
}
}
pub fn infer_dimension(expr: &Ex, dims: &DimMap) -> Result<ConstDim, String> {
let inner = expr.inner.read();
infer_in_arena(&inner.arena, expr.raw_id(), dims)
}
fn dim_pow(d: ConstDim, r: &Q) -> Result<ConstDim, String> {
let scale = |e: i8| -> Result<i8, String> {
let v = Q::from_integer(e.into()) * r;
if !v.is_integer() {
return Err(format!(
"Non-integer power of dimensioned quantity: ({d})^{r} is not a whole dimension"
));
}
v.to_integer()
.to_i8()
.ok_or_else(|| format!("Dimension exponent overflow in ({d})^{r}"))
};
Ok(ConstDim {
l: scale(d.l)?,
m: scale(d.m)?,
t: scale(d.t)?,
i: scale(d.i)?,
th: scale(d.th)?,
n: scale(d.n)?,
j: scale(d.j)?,
})
}
fn dim_mul(a: ConstDim, b: ConstDim) -> Result<ConstDim, String> {
let add = |x: i8, y: i8| {
x.checked_add(y)
.ok_or_else(|| format!("Dimension exponent overflow in ({a})·({b})"))
};
Ok(ConstDim {
l: add(a.l, b.l)?,
m: add(a.m, b.m)?,
t: add(a.t, b.t)?,
i: add(a.i, b.i)?,
th: add(a.th, b.th)?,
n: add(a.n, b.n)?,
j: add(a.j, b.j)?,
})
}
fn dim_div(a: ConstDim, b: ConstDim) -> Result<ConstDim, String> {
dim_mul(a, dim_pow(b, &-Q::one())?)
}
fn exponent_value(arena: &Arena, e: ExprId) -> Option<Q> {
if let Some(r) = arena.as_num(e) {
return Some(r.clone());
}
if let ExprNode::Neg(inner) = arena.node(e)
&& let Some(r) = arena.as_num(*inner)
{
return Some(-r.clone());
}
let v = crate::transforms::evalf::evalf_f64(arena, e).ok()?;
let n = v.round();
((v - n).abs() < 1e-10 && n.abs() < 1e6).then(|| Q::from_integer((n as i64).into()))
}
fn infer_in_arena(arena: &Arena, root: ExprId, dims: &DimMap) -> Result<ConstDim, String> {
let dimensionless = ConstDim::DIMENSIONLESS;
let mut dim_of: FxHashMap<ExprId, ConstDim> = FxHashMap::default();
for id in walk::post_order_ids(arena, root) {
let get = |c: &ExprId| -> Result<ConstDim, String> {
dim_of
.get(c)
.copied()
.ok_or_else(|| "internal error: operand without a dimension".to_string())
};
let polymorphic = |c: ExprId| {
arena.is_zero_structural(c)
|| matches!(
arena.node(c),
ExprNode::Infinity | ExprNode::NegInfinity | ExprNode::ComplexInfinity
)
};
let same = |ops: &[ExprId], what: &str| -> Result<ConstDim, String> {
let mut first: Option<(usize, ConstDim)> = None;
for (i, &c) in ops.iter().enumerate() {
if polymorphic(c) {
continue;
}
let d = get(&c)?;
match first {
None => first = Some((i, d)),
Some((i0, d0)) if !d0.eq(d) => {
return Err(format!(
"Dimension mismatch in {what}: term {i0} has dimension {d0} \
but term {i} has dimension {d}"
));
}
Some(_) => {}
}
}
Ok(first.map_or(dimensionless, |(_, d)| d))
};
let node = arena.node(id);
let d = match node {
ExprNode::Num(_) => dimensionless,
ExprNode::Symbol(sid) => {
let name = arena.symbol_name(*sid);
*dims
.get(name)
.ok_or_else(|| format!("Unknown variable '{name}' — not in dimension map"))?
}
ExprNode::PhysicalConstant(..) => dims
.get(&arena.display(id).to_string())
.copied()
.unwrap_or(dimensionless),
ExprNode::Add(ops) => same(ops, "addition")?,
ExprNode::Min(ops) | ExprNode::Max(ops) => same(ops, "min/max")?,
ExprNode::Mul(ops) => {
let mut acc = dimensionless;
for c in ops.iter() {
acc = dim_mul(acc, get(c)?)?;
}
acc
}
ExprNode::Pow(b, e) => {
let exp_dim = get(e)?;
if !exp_dim.eq(dimensionless) {
return Err(format!("Exponent must be dimensionless, got {exp_dim}"));
}
let base_dim = get(b)?;
if base_dim.eq(dimensionless) {
dimensionless
} else {
let r = exponent_value(arena, *e).ok_or_else(|| {
format!(
"Non-constant power of a dimensioned quantity (base dimension: {base_dim})"
)
})?;
dim_pow(base_dim, &r)?
}
}
ExprNode::Neg(a)
| ExprNode::Abs(a)
| ExprNode::Re(a)
| ExprNode::Im(a)
| ExprNode::Conjugate(a) => get(a)?,
ExprNode::Sign(a) | ExprNode::Heaviside(a) | ExprNode::Arg(a) => {
get(a)?;
dimensionless
}
ExprNode::DiracDelta(a) => dim_div(dimensionless, get(a)?)?,
ExprNode::Atan2(y, x) => {
same(&[*y, *x], "atan2")?;
dimensionless
}
ExprNode::Gt(a, b) | ExprNode::Ge(a, b) | ExprNode::Eq_(a, b) | ExprNode::Ne(a, b) => {
same(&[*a, *b], "a relation")?;
dimensionless
}
ExprNode::And(_)
| ExprNode::Or(_)
| ExprNode::Not(_)
| ExprNode::BoolTrue
| ExprNode::BoolFalse => dimensionless,
ExprNode::Piecewise(pieces) => {
let values: Vec<ExprId> = pieces.iter().map(|&(v, _)| v).collect();
same(&values, "a piecewise expression")?
}
ExprNode::Derivative(f, x) => dim_div(get(f)?, get(x)?)?,
ExprNode::Integral(f, x) => dim_mul(get(f)?, get(x)?)?,
ExprNode::DefiniteIntegral(f, x, lo, hi) => {
same(&[*x, *lo, *hi], "the bounds of an integral")?;
dim_mul(get(f)?, get(x)?)?
}
ExprNode::Sum(f, k, lo, hi) => {
for c in [k, lo, hi] {
let d = get(c)?;
if !d.eq(dimensionless) {
return Err(format!(
"Summation index and bounds must be dimensionless, got {d}"
));
}
}
get(f)?
}
ExprNode::Product_(f, k, lo, hi) => {
for c in [k, lo, hi] {
let d = get(c)?;
if !d.eq(dimensionless) {
return Err(format!(
"Product index and bounds must be dimensionless, got {d}"
));
}
}
let df = get(f)?;
if df.eq(dimensionless) {
dimensionless
} else {
let count = exponent_value(arena, *hi)
.zip(exponent_value(arena, *lo))
.map(|(h, l)| h - l + Q::one())
.filter(|n| n.is_integer() && *n >= Q::from_integer(0.into()))
.ok_or_else(|| {
format!("Product of a dimensioned factor ({df}) over a symbolic range")
})?;
dim_pow(df, &count)?
}
}
ExprNode::Limit(f, x, point) => {
same(&[*x, *point], "a limit point")?;
get(f)?
}
ExprNode::Series(f, x, point, order) => {
same(&[*x, *point], "a series expansion point")?;
let d = get(order)?;
if !d.eq(dimensionless) {
return Err(format!("Series order must be dimensionless, got {d}"));
}
get(f)?
}
ExprNode::Residue(f, z, point) => {
same(&[*z, *point], "a residue point")?;
dim_mul(get(f)?, get(z)?)?
}
ExprNode::LaplaceTransform(f, t, s) | ExprNode::InverseLaplaceTransform(f, s, t) => {
let (dt, ds) = (get(t)?, get(s)?);
if !dim_mul(dt, ds)?.eq(dimensionless) {
return Err(format!(
"Laplace variables must have reciprocal dimensions, got {dt} and {ds}"
));
}
let wrt = if matches!(node, ExprNode::LaplaceTransform(..)) {
dt
} else {
ds
};
dim_mul(get(f)?, wrt)?
}
ExprNode::Subs(f, x, value) => {
same(&[*x, *value], "a substitution")?;
get(f)?
}
ExprNode::Apply(..) if dims.get(&arena.display(id).to_string()).is_some() => *dims
.get(&arena.display(id).to_string())
.unwrap_or(&dimensionless),
ExprNode::Pi
| ExprNode::E
| ExprNode::ImaginaryUnit
| ExprNode::EulerGamma
| ExprNode::Catalan
| ExprNode::GoldenRatio
| ExprNode::Infinity
| ExprNode::NegInfinity
| ExprNode::ComplexInfinity
| ExprNode::NaN
| ExprNode::EmptySet
| ExprNode::UniversalSet
| ExprNode::Interval(..)
| ExprNode::FiniteSet(_)
| ExprNode::SetUnion(_)
| ExprNode::SetIntersection(_)
| ExprNode::SetComplement(..)
| ExprNode::ConditionSet(..)
| ExprNode::ImageSet(..) => dimensionless,
other => {
let mut bad: Option<(usize, ConstDim)> = None;
let mut k = 0usize;
other.for_each_child(|c| {
let d = dim_of.get(&c).copied().unwrap_or(dimensionless);
if bad.is_none() && !d.eq(dimensionless) {
bad = Some((k, d));
}
k += 1;
});
if let Some((i, d)) = bad {
return Err(if matches!(other, ExprNode::Apply(..)) {
format!("Applied function argument {i} must be dimensionless, got {d}")
} else {
format!(
"Function argument {i} must be dimensionless, got {d} \
(in expression {})",
arena.display(id)
)
});
}
dimensionless
}
};
dim_of.insert(id, d);
}
dim_of
.get(&root)
.copied()
.ok_or_else(|| "internal error: no dimension for the expression".to_string())
}
pub fn check_dimensions(expr: &Ex, dims: &DimMap) -> Result<(), String> {
infer_dimension(expr, dims).map(|_| ())
}
pub fn assert_dimension(expr: &Ex, dims: &DimMap, expected: ConstDim) -> Result<(), String> {
let actual = infer_dimension(expr, dims)?;
if actual.eq(expected) {
Ok(())
} else {
Err(format!(
"Expected dimension {} but expression has dimension {}",
expected, actual,
))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn dims() -> DimMap {
DimMap::new()
.with("m", ConstDim::MASS)
.with("a", ConstDim::ACCELERATION)
.with("g", ConstDim::ACCELERATION)
.with("v", ConstDim::VELOCITY)
.with("t", ConstDim::TIME)
.with("x", ConstDim::LENGTH)
.with("k", ConstDim::STIFFNESS)
.with("F", ConstDim::FORCE)
.with("R", ConstDim::RESISTANCE)
.with("I", ConstDim::CURRENT)
}
#[test]
fn infer_mass_times_accel_is_force() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, a);
let expr = &m * &a;
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::FORCE), "Expected Force, got {}", d);
}
#[test]
fn infer_half_mv_squared_is_energy() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, v);
let half = ctx.int(1) / ctx.int(2);
let expr = &half * &m * v.powi(2);
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::ENERGY), "Expected Energy, got {}", d);
}
#[test]
fn infer_add_mismatch_is_error() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, a);
let expr = &m + &a;
let result = infer_dimension(&expr, &dims());
assert!(
result.is_err(),
"Adding Mass + Acceleration should be an error"
);
}
#[test]
fn infer_voltage_is_current_times_resistance() {
let ctx = crate::api::context::Context::new();
let i_var = ctx.symbol("I");
let r_var = ctx.symbol("R");
let expr = &i_var * &r_var;
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::VOLTAGE), "Expected Voltage, got {}", d);
}
#[test]
fn infer_spring_force() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; k, x);
let expr = &k * &x;
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::FORCE), "Expected Force, got {}", d);
}
#[test]
fn infer_pure_number_is_dimensionless() {
let ctx = crate::api::context::Context::new();
let d = infer_dimension(&ctx.int(42), &dims()).unwrap();
assert!(d.eq(ConstDim::DIMENSIONLESS));
}
#[test]
fn infer_unknown_variable_is_error() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; unknown);
let result = infer_dimension(&unknown, &dims());
assert!(result.is_err(), "Unknown variable should produce an error");
}
#[test]
fn pow_length_squared_is_area() {
let d = ConstDim::LENGTH.pow(2);
assert!(d.eq(ConstDim::AREA));
}
#[test]
fn pow_length_cubed_is_volume() {
let d = ConstDim::LENGTH.pow(3);
assert!(d.eq(ConstDim::VOLUME));
}
#[test]
fn pow_zero_is_dimensionless() {
let d = ConstDim::FORCE.pow(0);
assert!(d.eq(ConstDim::DIMENSIONLESS));
}
#[test]
fn pow_one_is_identity() {
let d = ConstDim::VELOCITY.pow(1);
assert!(d.eq(ConstDim::VELOCITY));
}
#[test]
fn name_known_dimensions() {
assert_eq!(ConstDim::FORCE.name(), "Force");
assert_eq!(ConstDim::ENERGY.name(), "Energy");
assert_eq!(ConstDim::VOLTAGE.name(), "Voltage");
assert_eq!(ConstDim::DIMENSIONLESS.name(), "Dimensionless");
assert_eq!(ConstDim::MASS.name(), "Mass");
assert_eq!(ConstDim::LENGTH.name(), "Length");
assert_eq!(ConstDim::TIME.name(), "Time");
}
#[test]
fn name_unknown_dimension() {
let exotic = ConstDim::new(3, 2, -1, 0, 0, 0, 0);
assert_eq!(exotic.name(), "Unknown");
}
#[test]
fn dimmap_with_var() {
let ctx = crate::api::context::Context::new();
let x = ctx.symbol("x");
let dm = DimMap::new().with_var(&x, ConstDim::LENGTH);
assert_eq!(dm.get("x"), Some(&ConstDim::LENGTH));
}
#[test]
fn dimmap_len_and_empty() {
let dm = DimMap::new();
assert!(dm.is_empty());
assert_eq!(dm.len(), 0);
let dm = dm.with("x", ConstDim::LENGTH);
assert!(!dm.is_empty());
assert_eq!(dm.len(), 1);
}
#[test]
fn infer_add_consistent_is_ok() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, a, g);
let expr = &m * &a + &m * &g;
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::FORCE), "Expected Force, got {}", d);
}
#[test]
#[allow(non_snake_case)]
fn infer_negation_preserves_dimension() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; F);
let expr = -&F;
let d = infer_dimension(&expr, &dims()).unwrap();
assert!(d.eq(ConstDim::FORCE), "Expected Force, got {}", d);
}
#[test]
fn assert_dimension_ok() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, a);
let expr = &m * &a;
assert!(assert_dimension(&expr, &dims(), ConstDim::FORCE).is_ok());
}
#[test]
fn assert_dimension_mismatch() {
let ctx = crate::api::context::Context::new();
crate::syms!(ctx; m, a);
let expr = &m * &a;
assert!(assert_dimension(&expr, &dims(), ConstDim::ENERGY).is_err());
}
}