fidget_core/var/
mod.rs

1//! Input variables to math expressions
2//!
3//! A [`Var`] maintains a persistent identity from
4//! [`Tree`](crate::context::Tree) to [`Context`](crate::context::Node) (where
5//! it is wrapped in a [`Op::Input`](crate::context::Op::Input)) to evaluation
6//! (where [`Tape::vars`](crate::eval::Tape::vars) maps from `Var` to index in
7//! the argument list).
8use crate::context::{BadNode, Context, IntoNode, Node};
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11
12/// The [`Var`] type is an input to a math expression
13///
14/// We pre-define common variables (e.g. `X`, `Y`, `Z`) but also allow for fully
15/// customized values (using [`Var::V`]).
16///
17/// Variables are "global", in that every instance of `Var::X` represents the
18/// same thing.  To generate a "local" variable, [`Var::new`] picks a random
19/// 64-bit value, which is very unlikely to collide with anything else.
20#[derive(
21    Copy,
22    Clone,
23    Debug,
24    Hash,
25    Eq,
26    PartialEq,
27    Ord,
28    PartialOrd,
29    Serialize,
30    Deserialize,
31)]
32pub enum Var {
33    /// Variable representing the X axis for 2D / 3D shapes
34    X,
35    /// Variable representing the Y axis for 2D / 3D shapes
36    Y,
37    /// Variable representing the Z axis for 3D shapes
38    Z,
39    /// Generic variable
40    V(VarIndex),
41}
42
43/// Type for a variable index (implemented as a `u64`), used in [`Var::V`]
44#[derive(
45    Copy,
46    Clone,
47    Debug,
48    Hash,
49    Eq,
50    PartialEq,
51    Ord,
52    PartialOrd,
53    Serialize,
54    Deserialize,
55)]
56#[serde(transparent)]
57pub struct VarIndex(u64);
58
59impl Var {
60    /// Returns a new variable, with a random 64-bit index
61    ///
62    /// The odds of collision with any previous variable are infintesimally
63    /// small; if you are generating billions of random variables, something
64    /// else in the system is likely to break before collisions become an issue.
65    #[allow(clippy::new_without_default)]
66    pub fn new() -> Self {
67        let v: u64 = rand::random();
68        Var::V(VarIndex(v))
69    }
70
71    /// Returns the [`VarIndex`] from a [`Var::V`] instance, or `None`
72    pub fn index(&self) -> Option<VarIndex> {
73        if let Var::V(i) = *self { Some(i) } else { None }
74    }
75}
76
77impl std::fmt::Display for Var {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        match self {
80            Var::X => write!(f, "X"),
81            Var::Y => write!(f, "Y"),
82            Var::Z => write!(f, "Z"),
83            Var::V(VarIndex(v)) if *v < 256 => write!(f, "v_{v}"),
84            Var::V(VarIndex(v)) => write!(f, "V({v:x})"),
85        }
86    }
87}
88
89impl IntoNode for Var {
90    fn into_node(self, ctx: &mut Context) -> Result<Node, BadNode> {
91        Ok(ctx.var(self))
92    }
93}
94
95/// Map from [`Var`] to a particular index
96///
97/// Variable indexes are automatically assigned the first time
98/// [`VarMap::insert`] is called on that variable.
99///
100/// Indexes are guaranteed to be tightly packed, i.e. a map `vars` will contains
101/// values from `0..vars.len()`.
102///
103/// For efficiency, this type does not allocate heap memory for `Var::X/Y/Z`.
104#[derive(Default, Serialize, Deserialize)]
105pub struct VarMap {
106    x: Option<usize>,
107    y: Option<usize>,
108    z: Option<usize>,
109    v: HashMap<VarIndex, usize>,
110}
111
112#[allow(missing_docs)]
113impl VarMap {
114    pub fn new() -> Self {
115        Self::default()
116    }
117    pub fn len(&self) -> usize {
118        self.x.is_some() as usize
119            + self.y.is_some() as usize
120            + self.z.is_some() as usize
121            + self.v.len()
122    }
123    pub fn is_empty(&self) -> bool {
124        self.x.is_none()
125            && self.y.is_none()
126            && self.z.is_none()
127            && self.v.is_empty()
128    }
129    pub fn get(&self, v: &Var) -> Option<usize> {
130        match v {
131            Var::X => self.x,
132            Var::Y => self.y,
133            Var::Z => self.z,
134            Var::V(v) => self.v.get(v).cloned(),
135        }
136    }
137    /// Inserts a variable if not already present in the map
138    ///
139    /// The index is automatically assigned.
140    pub fn insert(&mut self, v: Var) {
141        let next = self.len();
142        match v {
143            Var::X => self.x.get_or_insert(next),
144            Var::Y => self.y.get_or_insert(next),
145            Var::Z => self.z.get_or_insert(next),
146            Var::V(v) => self.v.entry(v).or_insert(next),
147        };
148    }
149
150    /// Checks whether tracing arguments are valid
151    pub fn check_tracing_arguments<T>(
152        &self,
153        vars: &[T],
154    ) -> Result<(), TracingArgError> {
155        if vars.len() < self.len() {
156            // It's okay to be passed extra vars, because expressions may have
157            // been simplified.
158            Err(TracingArgError::BadVarSlice(BadVarSlice {
159                actual: vars.len(),
160                expected: self.len(),
161            }))
162        } else {
163            Ok(())
164        }
165    }
166
167    /// Check whether bulk arguments are valid
168    pub fn check_bulk_arguments<T, V: std::ops::Deref<Target = [T]>>(
169        &self,
170        vars: &[V],
171    ) -> Result<(), BulkArgError> {
172        // It's fine if the caller has given us extra variables (e.g. due to
173        // tape simplification), but it must have given us enough.
174        if vars.len() < self.len() {
175            Err(BulkArgError::BadVarSlice(BadVarSlice {
176                actual: vars.len(),
177                expected: self.len(),
178            }))
179        } else {
180            let Some(n) = vars.first().map(|v| v.len()) else {
181                return Ok(());
182            };
183            if let Some((i, v)) =
184                vars.iter().enumerate().find(|(_i, v)| v.len() != n)
185            {
186                Err(MismatchedSlices {
187                    index_a: 0,
188                    length_a: n,
189                    index_b: i,
190                    length_b: v.len(),
191                }
192                .into())
193            } else {
194                Ok(())
195            }
196        }
197    }
198
199    /// Iterates over `(variable, index)` tuples
200    pub fn iter(&self) -> impl Iterator<Item = (Var, usize)> {
201        [
202            self.x.map(|x| (Var::X, x)),
203            self.y.map(|y| (Var::Y, y)),
204            self.z.map(|z| (Var::Z, z)),
205        ]
206        .into_iter()
207        .flatten()
208        .chain(self.v.iter().map(|(v, k)| (Var::V(*v), *k)))
209    }
210}
211
212/// Error type for when variable slice length does not match expected count
213#[derive(thiserror::Error, Debug)]
214#[error(
215    "variable slice length ({actual}) does not match \
216     expected count ({expected})"
217)]
218pub struct BadVarSlice {
219    /// Slice length provided by the user
220    pub actual: usize,
221    /// Expected slice length
222    pub expected: usize,
223}
224
225/// Error type for mismatched slice lengths in bulk evaluation
226#[derive(thiserror::Error, Debug)]
227#[error(
228    "slice lengths are mismatched: \
229    slice {index_a} has {length_a} items; \
230    slice {index_b} has {length_b} items;"
231)]
232pub struct MismatchedSlices {
233    index_a: usize,
234    length_a: usize,
235    index_b: usize,
236    length_b: usize,
237}
238
239/// Error type for checking bulk arguments for evaluation
240#[derive(thiserror::Error, Debug)]
241pub enum BulkArgError {
242    /// Variable slice length does not match expected count
243    #[error(transparent)]
244    BadVarSlice(#[from] BadVarSlice),
245
246    /// Variable slice lengths are mismatched
247    #[error(transparent)]
248    MismatchedSlices(#[from] MismatchedSlices),
249}
250
251/// Error type for checking tracing arguments for evaluation
252#[derive(thiserror::Error, Debug)]
253pub enum TracingArgError {
254    /// Variable slice length does not match expected count
255    #[error(transparent)]
256    BadVarSlice(BadVarSlice),
257}
258
259impl std::ops::Index<&Var> for VarMap {
260    type Output = usize;
261    fn index(&self, v: &Var) -> &Self::Output {
262        match v {
263            Var::X => self.x.as_ref().unwrap(),
264            Var::Y => self.y.as_ref().unwrap(),
265            Var::Z => self.z.as_ref().unwrap(),
266            Var::V(v) => &self.v[v],
267        }
268    }
269}
270
271#[cfg(test)]
272mod test {
273    use super::*;
274
275    #[test]
276    fn var_identity() {
277        let v1 = Var::new();
278        let v2 = Var::new();
279        assert_ne!(v1, v2);
280    }
281
282    #[test]
283    fn var_map() {
284        let v = Var::new();
285        let mut m = VarMap::new();
286        assert!(m.get(&v).is_none());
287        m.insert(v);
288        assert_eq!(m.get(&v), Some(0));
289        m.insert(v);
290        assert_eq!(m.get(&v), Some(0));
291
292        let u = Var::new();
293        assert!(m.get(&u).is_none());
294        m.insert(u);
295        assert_eq!(m.get(&u), Some(1));
296    }
297}
fidget_core/var/mod.rs

fidget_core/var/
mod.rs
