Skip to main content

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(Debug, 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
138    /// Checks whether there are free (non-XYZ) variables present in the map
139    pub fn has_free_vars(&self) -> bool {
140        !self.v.is_empty()
141    }
142
143    /// Inserts a variable if not already present in the map
144    ///
145    /// The index is automatically assigned.
146    pub fn insert(&mut self, v: Var) {
147        let next = self.len();
148        match v {
149            Var::X => self.x.get_or_insert(next),
150            Var::Y => self.y.get_or_insert(next),
151            Var::Z => self.z.get_or_insert(next),
152            Var::V(v) => self.v.entry(v).or_insert(next),
153        };
154    }
155
156    /// Checks whether tracing arguments are valid
157    pub fn check_tracing_arguments<T>(
158        &self,
159        vars: &[T],
160    ) -> Result<(), TracingArgError> {
161        if vars.len() < self.len() {
162            // It's okay to be passed extra vars, because expressions may have
163            // been simplified.
164            Err(TracingArgError::BadVarSlice(BadVarSlice {
165                actual: vars.len(),
166                expected: self.len(),
167            }))
168        } else {
169            Ok(())
170        }
171    }
172
173    /// Check whether bulk arguments are valid
174    pub fn check_bulk_arguments<T, V: std::ops::Deref<Target = [T]>>(
175        &self,
176        vars: &[V],
177    ) -> Result<(), BulkArgError> {
178        // It's fine if the caller has given us extra variables (e.g. due to
179        // tape simplification), but it must have given us enough.
180        if vars.len() < self.len() {
181            Err(BulkArgError::BadVarSlice(BadVarSlice {
182                actual: vars.len(),
183                expected: self.len(),
184            }))
185        } else {
186            let Some(n) = vars.first().map(|v| v.len()) else {
187                return Ok(());
188            };
189            if let Some((i, v)) =
190                vars.iter().enumerate().find(|(_i, v)| v.len() != n)
191            {
192                Err(MismatchedSlices {
193                    index_a: 0,
194                    length_a: n,
195                    index_b: i,
196                    length_b: v.len(),
197                }
198                .into())
199            } else {
200                Ok(())
201            }
202        }
203    }
204
205    /// Iterates over `(variable, index)` tuples
206    pub fn iter(&self) -> impl Iterator<Item = (Var, usize)> {
207        [
208            self.x.map(|x| (Var::X, x)),
209            self.y.map(|y| (Var::Y, y)),
210            self.z.map(|z| (Var::Z, z)),
211        ]
212        .into_iter()
213        .flatten()
214        .chain(self.v.iter().map(|(v, k)| (Var::V(*v), *k)))
215    }
216}
217
218/// Error type for when variable slice length does not match expected count
219#[derive(thiserror::Error, Debug)]
220#[error(
221    "variable slice length ({actual}) does not match \
222     expected count ({expected})"
223)]
224pub struct BadVarSlice {
225    /// Slice length provided by the user
226    pub actual: usize,
227    /// Expected slice length
228    pub expected: usize,
229}
230
231/// Error type for mismatched slice lengths in bulk evaluation
232#[derive(thiserror::Error, Debug)]
233#[error(
234    "slice lengths are mismatched: \
235    slice {index_a} has {length_a} items; \
236    slice {index_b} has {length_b} items;"
237)]
238pub struct MismatchedSlices {
239    index_a: usize,
240    length_a: usize,
241    index_b: usize,
242    length_b: usize,
243}
244
245/// Error type for checking bulk arguments for evaluation
246#[derive(thiserror::Error, Debug)]
247pub enum BulkArgError {
248    /// Variable slice length does not match expected count
249    #[error(transparent)]
250    BadVarSlice(#[from] BadVarSlice),
251
252    /// Variable slice lengths are mismatched
253    #[error(transparent)]
254    MismatchedSlices(#[from] MismatchedSlices),
255}
256
257/// Error type for checking tracing arguments for evaluation
258#[derive(thiserror::Error, Debug)]
259pub enum TracingArgError {
260    /// Variable slice length does not match expected count
261    #[error(transparent)]
262    BadVarSlice(BadVarSlice),
263}
264
265impl std::ops::Index<&Var> for VarMap {
266    type Output = usize;
267    fn index(&self, v: &Var) -> &Self::Output {
268        match v {
269            Var::X => self.x.as_ref().unwrap(),
270            Var::Y => self.y.as_ref().unwrap(),
271            Var::Z => self.z.as_ref().unwrap(),
272            Var::V(v) => &self.v[v],
273        }
274    }
275}
276
277#[cfg(test)]
278mod test {
279    use super::*;
280
281    #[test]
282    fn var_identity() {
283        let v1 = Var::new();
284        let v2 = Var::new();
285        assert_ne!(v1, v2);
286    }
287
288    #[test]
289    fn var_map() {
290        let v = Var::new();
291        let mut m = VarMap::new();
292        assert!(m.get(&v).is_none());
293        m.insert(v);
294        assert_eq!(m.get(&v), Some(0));
295        m.insert(v);
296        assert_eq!(m.get(&v), Some(0));
297
298        let u = Var::new();
299        assert!(m.get(&u).is_none());
300        m.insert(u);
301        assert_eq!(m.get(&u), Some(1));
302    }
303}