fidget_core/shape/
mod.rs

1//! Data structures for shape evaluation
2//!
3//! Types in this module are typically thin (generic) wrappers around objects
4//! that implement traits in [`fidget_core::eval`](crate::eval).  The wrapper types
5//! are specialized to operate on `x, y, z` arguments, rather than taking
6//! arbitrary numbers of variables.
7//!
8//! For example, a [`Shape`] is a wrapper which makes it easier to treat a
9//! [`Function`] as an implicit surface (with X, Y, Z axes and functions that
10//! accept a transform matrix).
11//!
12//! ```rust
13//! use fidget_core::vm::VmShape;
14//! use fidget_core::context::Context;
15//! use fidget_core::shape::EzShape;
16//!
17//! let mut ctx = Context::new();
18//! let x = ctx.x();
19//! let shape = VmShape::new(&ctx, x)?;
20//!
21//! // Let's build a single point evaluator:
22//! let mut eval = VmShape::new_point_eval();
23//! let tape = shape.ez_point_tape();
24//! let (value, _trace) = eval.eval_with_transform(
25//!     &tape, 0.25, 0.0, 0.0, &nalgebra::Matrix4::identity()
26//! )?;
27//! assert_eq!(value, 0.25);
28//! # Ok::<(), Box<dyn std::error::Error>>(())
29//! ```
30
31use crate::{
32    context::{BadNode, Context, Node, Tree},
33    eval::{
34        BulkEvalError, BulkEvaluator, Function, MathFunction, Tape,
35        TracingEvalError, TracingEvaluator,
36    },
37    types::{Grad, Interval},
38    var::{BulkArgError, TracingArgError, Var, VarIndex, VarMap},
39    vm::BadTrace,
40};
41use nalgebra::{Matrix4, Point3};
42use std::collections::HashMap;
43
44/// A shape represents an implicit surface with X/Y/Z variables
45///
46/// It is mostly agnostic to _how_ that surface is represented, wrapping a
47/// [`Function`] and a set of axes.
48///
49/// Shapes are shared between threads, so they should be cheap to clone.  In
50/// most cases, they're a thin wrapper around an `Arc<..>`.
51pub struct Shape<F> {
52    /// Wrapped function
53    f: F,
54}
55
56impl<F: Clone> Clone for Shape<F> {
57    fn clone(&self) -> Self {
58        Self { f: self.f.clone() }
59    }
60}
61
62impl<F: Function + Clone> Shape<F> {
63    /// Builds a new point evaluator
64    pub fn new_point_eval() -> ShapeTracingEval<F::PointEval> {
65        ShapeTracingEval {
66            eval: F::PointEval::default(),
67            scratch: vec![],
68        }
69    }
70
71    /// Builds a new interval evaluator
72    pub fn new_interval_eval() -> ShapeTracingEval<F::IntervalEval> {
73        ShapeTracingEval {
74            eval: F::IntervalEval::default(),
75            scratch: vec![],
76        }
77    }
78
79    /// Builds a new float slice evaluator
80    pub fn new_float_slice_eval() -> ShapeBulkEval<F::FloatSliceEval> {
81        ShapeBulkEval {
82            eval: F::FloatSliceEval::default(),
83            scratch: vec![],
84        }
85    }
86
87    /// Builds a new gradient slice evaluator
88    pub fn new_grad_slice_eval() -> ShapeBulkEval<F::GradSliceEval> {
89        ShapeBulkEval {
90            eval: F::GradSliceEval::default(),
91            scratch: vec![],
92        }
93    }
94
95    /// Returns an evaluation tape for a point evaluator
96    #[inline]
97    pub fn point_tape(
98        &self,
99        storage: F::TapeStorage,
100    ) -> ShapeTape<<F::PointEval as TracingEvaluator>::Tape> {
101        let tape = self.f.point_tape(storage);
102        ShapeTape { tape }
103    }
104
105    /// Returns an evaluation tape for a interval evaluator
106    #[inline]
107    pub fn interval_tape(
108        &self,
109        storage: F::TapeStorage,
110    ) -> ShapeTape<<F::IntervalEval as TracingEvaluator>::Tape> {
111        let tape = self.f.interval_tape(storage);
112        ShapeTape { tape }
113    }
114
115    /// Returns an evaluation tape for a float slice evaluator
116    #[inline]
117    pub fn float_slice_tape(
118        &self,
119        storage: F::TapeStorage,
120    ) -> ShapeTape<<F::FloatSliceEval as BulkEvaluator>::Tape> {
121        let tape = self.f.float_slice_tape(storage);
122        ShapeTape { tape }
123    }
124
125    /// Returns an evaluation tape for a gradient slice evaluator
126    #[inline]
127    pub fn grad_slice_tape(
128        &self,
129        storage: F::TapeStorage,
130    ) -> ShapeTape<<F::GradSliceEval as BulkEvaluator>::Tape> {
131        let tape = self.f.grad_slice_tape(storage);
132        ShapeTape { tape }
133    }
134
135    /// Computes a simplified tape using the given trace, and reusing storage
136    #[inline]
137    pub fn simplify(
138        &self,
139        trace: &F::Trace,
140        storage: F::Storage,
141        workspace: &mut F::Workspace,
142    ) -> Result<Self, BadTrace>
143    where
144        Self: Sized,
145    {
146        let f = self.f.simplify(trace, storage, workspace)?;
147        Ok(Self { f })
148    }
149
150    /// Attempt to reclaim storage from this shape
151    ///
152    /// This may fail, because shapes are `Clone` and are often implemented
153    /// using an `Arc` around a heavier data structure.
154    #[inline]
155    pub fn recycle(self) -> Option<F::Storage> {
156        self.f.recycle()
157    }
158
159    /// Returns a size associated with this shape
160    ///
161    /// This is underspecified and only used for unit testing; for tape-based
162    /// shapes, it's typically the length of the tape,
163    #[inline]
164    pub fn size(&self) -> usize {
165        self.f.size()
166    }
167
168    /// Binds a shape to a set of variable values
169    #[inline]
170    pub fn bind<'a, D>(
171        &self,
172        vars: &'a ShapeVars<D>,
173    ) -> Result<BoundShape<'a, F, D>, MissingVar> {
174        BoundShape::new(self.clone(), vars)
175    }
176}
177
178impl<F> Shape<F> {
179    /// Borrows the inner [`Function`] object
180    pub fn inner(&self) -> &F {
181        &self.f
182    }
183}
184
185/// Variables bound to values for shape evaluation
186///
187/// Note that this cannot store `X`, `Y`, `Z` variables (which are passed in as
188/// first-class arguments); it only stores [`Var::V`] values (identified by
189/// their inner [`VarIndex`]).
190pub struct ShapeVars<F>(HashMap<VarIndex, F>);
191
192impl<F> Default for ShapeVars<F> {
193    fn default() -> Self {
194        Self(HashMap::default())
195    }
196}
197
198impl<F> ShapeVars<F> {
199    /// Builds a new, empty variable set
200    pub fn new() -> Self {
201        Self(HashMap::default())
202    }
203    /// Returns the number of variables stored in the set
204    pub fn len(&self) -> usize {
205        self.0.len()
206    }
207    /// Checks whether the variable set is empty
208    pub fn is_empty(&self) -> bool {
209        self.0.is_empty()
210    }
211    /// Inserts a new variable
212    ///
213    /// Returns the previous value (if present)
214    pub fn insert(&mut self, v: VarIndex, f: F) -> Option<F> {
215        self.0.insert(v, f)
216    }
217
218    /// Iterates over values
219    pub fn values(&self) -> impl Iterator<Item = &F> {
220        self.0.values()
221    }
222
223    /// Looks up a value by [`VarIndex`]
224    pub fn get(&self, i: VarIndex) -> Option<&F> {
225        self.0.get(&i)
226    }
227
228    /// Checks whether the given [`VarIndex`] is present in the map
229    pub fn contains_key(&self, i: VarIndex) -> bool {
230        self.0.contains_key(&i)
231    }
232
233    /// Checks that this map contains all variables needed to render a shape
234    pub fn check<G: Function>(
235        &self,
236        shape: &Shape<G>,
237    ) -> Result<(), MissingVar> {
238        for (v, _) in shape.inner().vars().iter() {
239            match v {
240                Var::X | Var::Y | Var::Z => (),
241                Var::V(i) => {
242                    if !self.contains_key(i) {
243                        return Err(MissingVar { var: i });
244                    }
245                }
246            }
247        }
248        Ok(())
249    }
250}
251
252impl<'a, F> IntoIterator for &'a ShapeVars<F> {
253    type Item = (&'a VarIndex, &'a F);
254    type IntoIter = std::collections::hash_map::Iter<'a, VarIndex, F>;
255    fn into_iter(self) -> Self::IntoIter {
256        self.0.iter()
257    }
258}
259
260/// Extension trait for working with a shape without thinking much about memory
261///
262/// All of the [`Shape`] functions that use significant amounts of memory
263/// pedantically require you to pass in storage for reuse.  This trait allows
264/// you to ignore that, at the cost of performance; we require that all storage
265/// types implement [`Default`], so these functions do the boilerplate for you.
266///
267/// This trait is automatically implemented for every [`Shape`], but must be
268/// imported separately as a speed-bump to using it everywhere.
269pub trait EzShape<F: Function> {
270    /// Returns an evaluation tape for a point evaluator
271    fn ez_point_tape(
272        &self,
273    ) -> ShapeTape<<F::PointEval as TracingEvaluator>::Tape>;
274
275    /// Returns an evaluation tape for an interval evaluator
276    fn ez_interval_tape(
277        &self,
278    ) -> ShapeTape<<F::IntervalEval as TracingEvaluator>::Tape>;
279
280    /// Returns an evaluation tape for a float slice evaluator
281    fn ez_float_slice_tape(
282        &self,
283    ) -> ShapeTape<<F::FloatSliceEval as BulkEvaluator>::Tape>;
284
285    /// Returns an evaluation tape for a float slice evaluator
286    fn ez_grad_slice_tape(
287        &self,
288    ) -> ShapeTape<<F::GradSliceEval as BulkEvaluator>::Tape>;
289
290    /// Computes a simplified tape using the given trace
291    fn ez_simplify(&self, trace: &F::Trace) -> Result<Self, BadTrace>
292    where
293        Self: Sized;
294}
295
296impl<F: Function> EzShape<F> for Shape<F> {
297    fn ez_point_tape(
298        &self,
299    ) -> ShapeTape<<F::PointEval as TracingEvaluator>::Tape> {
300        self.point_tape(Default::default())
301    }
302
303    fn ez_interval_tape(
304        &self,
305    ) -> ShapeTape<<F::IntervalEval as TracingEvaluator>::Tape> {
306        self.interval_tape(Default::default())
307    }
308
309    fn ez_float_slice_tape(
310        &self,
311    ) -> ShapeTape<<F::FloatSliceEval as BulkEvaluator>::Tape> {
312        self.float_slice_tape(Default::default())
313    }
314
315    fn ez_grad_slice_tape(
316        &self,
317    ) -> ShapeTape<<F::GradSliceEval as BulkEvaluator>::Tape> {
318        self.grad_slice_tape(Default::default())
319    }
320
321    fn ez_simplify(&self, trace: &F::Trace) -> Result<Self, BadTrace> {
322        let mut workspace = Default::default();
323        self.simplify(trace, Default::default(), &mut workspace)
324    }
325}
326
327impl<F: MathFunction> Shape<F> {
328    /// Builds a new shape from a math expression with default (X, Y, Z) axes
329    pub fn new(ctx: &Context, node: Node) -> Result<Self, BadNode> {
330        let f = F::new(ctx, &[node])?;
331        Ok(Self { f })
332    }
333
334    /// Raw constructor
335    ///
336    /// # Panics
337    /// If the function `f` does not have exactly one output
338    pub fn new_raw(f: F) -> Self {
339        assert_eq!(f.output_count(), 1);
340        Self { f }
341    }
342}
343
344/// Converts a [`Tree`] to a [`Shape`] with the default axes
345impl<F: MathFunction> From<Tree> for Shape<F> {
346    fn from(t: Tree) -> Self {
347        let mut ctx = Context::new();
348        let node = ctx.import(&t);
349        Self::new(&ctx, node).unwrap()
350    }
351}
352
353/// Wrapper around a single-output function tape
354#[derive(Clone)]
355pub struct ShapeTape<T> {
356    tape: T,
357}
358
359impl<T: Tape> ShapeTape<T> {
360    /// Recycles the inner tape's storage for reuse
361    pub fn recycle(self) -> Option<T::Storage> {
362        self.tape.recycle()
363    }
364
365    /// Returns a mapping from [`Var`] to evaluation index
366    pub fn vars(&self) -> &VarMap {
367        self.tape.vars()
368    }
369}
370
371/// Wrapper around a [`TracingEvaluator`]
372///
373/// Unlike the raw tracing evaluator, a [`ShapeTracingEval`] knows about the
374/// tape's X, Y, Z axes and its evaluators take a transform matrix.
375#[derive(Debug)]
376pub struct ShapeTracingEval<E: TracingEvaluator> {
377    eval: E,
378    scratch: Vec<E::Data>,
379}
380
381impl<E: TracingEvaluator> Default for ShapeTracingEval<E> {
382    fn default() -> Self {
383        Self {
384            eval: E::default(),
385            scratch: vec![],
386        }
387    }
388}
389
390/// A [`VarIndex`] variable is missing
391#[derive(thiserror::Error, Debug, PartialEq)]
392#[error("variable {var:?} must be provided")]
393pub struct MissingVar {
394    /// Missing variable index
395    pub var: VarIndex,
396}
397
398/// Error type for shape tracing evaluation
399#[derive(thiserror::Error, Debug)]
400pub enum ShapeTracingEvalError {
401    /// Missing variable index
402    #[error(transparent)]
403    MissingVar(#[from] MissingVar),
404}
405
406/// Error type for shape bulk evaluation
407#[derive(thiserror::Error, Debug, PartialEq)]
408pub enum ShapeBulkEvalError {
409    /// Missing variable index
410    #[error(transparent)]
411    MissingVar(#[from] MissingVar),
412
413    /// Mismatched slice length
414    #[error(
415        "slice lengths are mismatched: \
416        slice for {var_a} has {length_a} items; \
417        slice for {var_b} has {length_b} items;"
418    )]
419    MismatchedVarSlices {
420        /// A specific variable
421        var_a: Var,
422        /// The length of the slice for `var_a`
423        length_a: usize,
424        /// A second variable
425        var_b: Var,
426        /// The length of the slice for `var_b`
427        length_b: usize,
428    },
429}
430
431impl<E: TracingEvaluator> ShapeTracingEval<E>
432where
433    <E as TracingEvaluator>::Data: Transformable,
434{
435    /// Tracing evaluation of the given tape with X, Y, Z input arguments
436    #[inline]
437    pub fn eval<F: Into<E::Data> + Copy>(
438        &mut self,
439        tape: &ShapeTape<E::Tape>,
440        x: F,
441        y: F,
442        z: F,
443    ) -> Result<(E::Data, Option<&E::Trace>), ShapeTracingEvalError> {
444        let h = ShapeVars::<f32>::new();
445        self.eval_raw(tape, x, y, z, None, &h)
446    }
447
448    /// Tracing evaluation with X, Y, Z arguments and a transform
449    #[inline]
450    pub fn eval_with_transform<F: Into<E::Data> + Copy>(
451        &mut self,
452        tape: &ShapeTape<E::Tape>,
453        x: F,
454        y: F,
455        z: F,
456        transform: &Matrix4<f32>,
457    ) -> Result<(E::Data, Option<&E::Trace>), ShapeTracingEvalError> {
458        let h = ShapeVars::<f32>::new();
459        self.eval_raw(tape, x, y, z, Some(transform), &h)
460    }
461
462    /// Tracing evaluation with X, Y, Z arguments, transform, and vars
463    #[inline]
464    pub fn eval_with_transform_and_vars<
465        F: Into<E::Data> + Copy,
466        V: Into<E::Data> + Copy,
467    >(
468        &mut self,
469        tape: &ShapeTape<E::Tape>,
470        x: F,
471        y: F,
472        z: F,
473        transform: &Matrix4<f32>,
474        vars: &ShapeVars<V>,
475    ) -> Result<(E::Data, Option<&E::Trace>), ShapeTracingEvalError> {
476        self.eval_raw(tape, x, y, z, Some(transform), vars)
477    }
478
479    /// Tracing evaluation with X, Y, Z arguments and vars
480    #[inline]
481    pub fn eval_with_vars<F: Into<E::Data> + Copy, V: Into<E::Data> + Copy>(
482        &mut self,
483        tape: &ShapeTape<E::Tape>,
484        x: F,
485        y: F,
486        z: F,
487        vars: &ShapeVars<V>,
488    ) -> Result<(E::Data, Option<&E::Trace>), ShapeTracingEvalError> {
489        self.eval_raw(tape, x, y, z, None, vars)
490    }
491
492    /// Innermost evaluation function
493    #[inline]
494    pub fn eval_raw<F: Into<E::Data> + Copy, V: Into<E::Data> + Copy>(
495        &mut self,
496        tape: &ShapeTape<E::Tape>,
497        x: F,
498        y: F,
499        z: F,
500        transform: Option<&Matrix4<f32>>,
501        vars: &ShapeVars<V>,
502    ) -> Result<(E::Data, Option<&E::Trace>), ShapeTracingEvalError> {
503        assert_eq!(
504            tape.tape.output_count(),
505            1,
506            "ShapeTape has multiple outputs"
507        );
508
509        let x = x.into();
510        let y = y.into();
511        let z = z.into();
512        let (x, y, z) = if let Some(t) = transform {
513            Transformable::transform(x, y, z, t)
514        } else {
515            (x, y, z)
516        };
517
518        let vs = tape.vars();
519        self.scratch.resize(vs.len(), 0f32.into());
520        for (var, index) in vs.iter() {
521            match var {
522                Var::X => self.scratch[index] = x,
523                Var::Y => self.scratch[index] = y,
524                Var::Z => self.scratch[index] = z,
525                Var::V(i) => {
526                    let Some(value) = vars.get(i) else {
527                        return Err(MissingVar { var: i }.into());
528                    };
529                    self.scratch[index] = (*value).into();
530                }
531            }
532        }
533        let (out, trace) = match self.eval.eval(&tape.tape, &self.scratch) {
534            Ok((out, trace)) => (out, trace),
535            Err(TracingEvalError(TracingArgError::BadVarSlice(..))) => {
536                unreachable!() // we resized `scratch` above
537            }
538        };
539        Ok((out[0], trace))
540    }
541}
542
543/// Wrapper around a [`BulkEvaluator`]
544///
545/// Unlike the raw bulk evaluator, a [`ShapeBulkEval`] knows about the
546/// tape's X, Y, Z axes and accepts a transform matrix.
547#[derive(Debug, Default)]
548pub struct ShapeBulkEval<E: BulkEvaluator> {
549    eval: E,
550    scratch: Vec<Vec<E::Data>>,
551}
552
553impl<E: BulkEvaluator> ShapeBulkEval<E>
554where
555    E::Data: From<f32> + Transformable,
556{
557    /// Bulk evaluation of many samples, without any variables
558    ///
559    /// If the shape includes variables other than `X`, `Y`, `Z`,
560    /// [`eval_with_vars`](Self::eval_with_vars) or
561    /// [`eval_with_var_arrays`](Self::eval_with_var_arrays) should be used
562    /// instead (and this function will return an error).
563    #[inline]
564    pub fn eval(
565        &mut self,
566        tape: &ShapeTape<E::Tape>,
567        x: &[E::Data],
568        y: &[E::Data],
569        z: &[E::Data],
570    ) -> Result<&[E::Data], ShapeBulkEvalError> {
571        self.eval_raw(tape, x, y, z, None, Self::no_vars)
572    }
573
574    /// Bulk evaluation of many samples with a transform
575    #[inline]
576    pub fn eval_with_transform(
577        &mut self,
578        tape: &ShapeTape<E::Tape>,
579        x: &[E::Data],
580        y: &[E::Data],
581        z: &[E::Data],
582        transform: &Matrix4<f32>,
583    ) -> Result<&[E::Data], ShapeBulkEvalError> {
584        self.eval_raw(tape, x, y, z, Some(transform), Self::no_vars)
585    }
586
587    /// Bulk evaluation of many samples, with slices of variables
588    ///
589    /// Each variable is a slice (or `Vec`) of values, which must be the same
590    /// length as the `x`, `y`, `z` slices.  This is in contrast with
591    /// [`eval_with_vars`](Self::eval_with_vars), where variables have a single
592    /// value used for every position in the `x`, `y,` `z` slices.
593    ///
594    /// Before evaluation, the transform matrix is applied to input coordinates.
595    #[inline]
596    pub fn eval_with_var_arrays<
597        V: std::ops::Deref<Target = [G]>,
598        G: Into<E::Data> + Copy,
599    >(
600        &mut self,
601        tape: &ShapeTape<E::Tape>,
602        x: &[E::Data],
603        y: &[E::Data],
604        z: &[E::Data],
605        vars: &ShapeVars<V>,
606    ) -> Result<&[E::Data], ShapeBulkEvalError> {
607        self.eval_raw(tape, x, y, z, None, Self::var_array(vars))
608    }
609
610    /// Bulk evaluation of many transformed samples, with slices of variables
611    #[inline]
612    pub fn eval_with_transform_and_var_arrays<
613        V: std::ops::Deref<Target = [G]>,
614        G: Into<E::Data> + Copy,
615    >(
616        &mut self,
617        tape: &ShapeTape<E::Tape>,
618        x: &[E::Data],
619        y: &[E::Data],
620        z: &[E::Data],
621        transform: &Matrix4<f32>,
622        vars: &ShapeVars<V>,
623    ) -> Result<&[E::Data], ShapeBulkEvalError> {
624        self.eval_raw(tape, x, y, z, Some(transform), Self::var_array(vars))
625    }
626
627    /// Bulk evaluation of many samples, with fixed variables
628    ///
629    /// Each variable has a single value, which is used for every position in
630    /// the `x`, `y`, `z` slices.  This is in contrast with
631    /// [`eval_with_var_arrays`](Self::eval_with_var_arrays), where variables
632    /// can be different for every position in the `x`, `y,` `z` slices.
633    #[inline]
634    pub fn eval_with_vars<G: Into<E::Data> + Copy>(
635        &mut self,
636        tape: &ShapeTape<E::Tape>,
637        x: &[E::Data],
638        y: &[E::Data],
639        z: &[E::Data],
640        vars: &ShapeVars<G>,
641    ) -> Result<&[E::Data], ShapeBulkEvalError> {
642        self.eval_raw(tape, x, y, z, None, Self::var_value(vars))
643    }
644
645    /// Bulk evaluation of many transformed samples, with fixed variables
646    #[inline]
647    pub fn eval_with_transform_and_vars<G: Into<E::Data> + Copy>(
648        &mut self,
649        tape: &ShapeTape<E::Tape>,
650        x: &[E::Data],
651        y: &[E::Data],
652        z: &[E::Data],
653        transform: &Matrix4<f32>,
654        vars: &ShapeVars<G>,
655    ) -> Result<&[E::Data], ShapeBulkEvalError> {
656        self.eval_raw(tape, x, y, z, Some(transform), Self::var_value(vars))
657    }
658
659    /// Helper function for evaluation without variables
660    #[inline]
661    fn no_vars(
662        _: &mut [E::Data],
663        var: VarIndex,
664    ) -> Result<(), ShapeBulkEvalError> {
665        Err(MissingVar { var }.into())
666    }
667
668    /// Helper to bind to a multi-value variable map
669    ///
670    /// This is a building block for working with [`eval_raw`](Self::eval_raw),
671    /// and you probably don't need it unless you're deep in the weeds.
672    #[inline]
673    pub fn var_array<
674        V: std::ops::Deref<Target = [G]>,
675        G: Into<E::Data> + Copy,
676    >(
677        vars: &ShapeVars<V>,
678    ) -> impl Fn(&mut [E::Data], VarIndex) -> Result<(), ShapeBulkEvalError>
679    {
680        |data: &mut [E::Data], i: VarIndex| {
681            let vars = vars.get(i).ok_or(MissingVar { var: i })?;
682            if vars.len() != data.len() {
683                return Err(ShapeBulkEvalError::MismatchedVarSlices {
684                    var_a: Var::X,
685                    length_a: data.len(),
686                    var_b: Var::V(i),
687                    length_b: vars.len(),
688                });
689            }
690            for (a, b) in data.iter_mut().zip(vars.deref().iter()) {
691                *a = (*b).into();
692            }
693            Ok(())
694        }
695    }
696
697    /// Helper to bind to a single-variable map
698    ///
699    /// This is a building block for working with [`eval_raw`](Self::eval_raw),
700    /// and you probably don't need it unless you're deep in the weeds.
701    #[inline]
702    pub fn var_value<G: Into<E::Data> + Copy>(
703        vars: &ShapeVars<G>,
704    ) -> impl Fn(&mut [E::Data], VarIndex) -> Result<(), ShapeBulkEvalError>
705    {
706        |data: &mut [E::Data], i: VarIndex| {
707            let value = vars.get(i).ok_or(MissingVar { var: i })?;
708            data.fill((*value).into());
709            Ok(())
710        }
711    }
712
713    /// Core implementation of evaluation
714    ///
715    /// `copy_vars` is a way to populate variable slices (other than the X, Y, Z
716    /// inputs), and may be generated from a helper like
717    /// [`var_value`](Self::var_value) or [`var_array`](Self::var_array).
718    #[inline]
719    pub fn eval_raw<F>(
720        &mut self,
721        tape: &ShapeTape<E::Tape>,
722        x: &[E::Data],
723        y: &[E::Data],
724        z: &[E::Data],
725        transform: Option<&Matrix4<f32>>,
726        copy_vars: F,
727    ) -> Result<&[E::Data], ShapeBulkEvalError>
728    where
729        F: Fn(&mut [E::Data], VarIndex) -> Result<(), ShapeBulkEvalError>,
730    {
731        assert_eq!(
732            tape.tape.output_count(),
733            1,
734            "ShapeTape has multiple outputs" // enforced by ShapeTape
735        );
736
737        // Make sure our scratch arrays are big enough for this evaluation
738        if x.len() != y.len() {
739            return Err(ShapeBulkEvalError::MismatchedVarSlices {
740                var_a: Var::X,
741                length_a: x.len(),
742                var_b: Var::Y,
743                length_b: y.len(),
744            });
745        }
746        if x.len() != z.len() {
747            return Err(ShapeBulkEvalError::MismatchedVarSlices {
748                var_a: Var::X,
749                length_a: x.len(),
750                var_b: Var::Z,
751                length_b: z.len(),
752            });
753        }
754        let n = x.len();
755
756        // We need at least one item in the scratch array to set evaluation
757        // size; otherwise, evaluating a single constant will return []
758        let vs = tape.vars();
759        self.scratch.resize_with(vs.len().max(1), Vec::new);
760        for s in &mut self.scratch {
761            s.resize(n, 0.0.into());
762        }
763
764        let mut axes = [None; 3];
765        for (var, index) in vs.iter() {
766            match var {
767                Var::X => axes[0] = Some(index),
768                Var::Y => axes[1] = Some(index),
769                Var::Z => axes[2] = Some(index),
770                Var::V(i) => {
771                    copy_vars(&mut self.scratch[index], i)?;
772                }
773            }
774        }
775
776        for i in 0..n {
777            let (x, y, z) = if let Some(t) = transform {
778                Transformable::transform(x[i], y[i], z[i], t)
779            } else {
780                (x[i], y[i], z[i])
781            };
782            if let Some(a) = axes[0] {
783                self.scratch[a][i] = x;
784            }
785            if let Some(b) = axes[1] {
786                self.scratch[b][i] = y;
787            }
788            if let Some(c) = axes[2] {
789                self.scratch[c][i] = z;
790            }
791        }
792
793        let out = match self.eval.eval(&tape.tape, &self.scratch) {
794            Ok(out) => out,
795            Err(BulkEvalError(e)) => match e {
796                // All of these conditions should be handled by `setup`
797                BulkArgError::BadVarSlice(..)
798                | BulkArgError::MismatchedSlices(..) => unreachable!(),
799            },
800        };
801        Ok(out.borrow(0))
802    }
803}
804
805/// Shape with bound variables
806///
807/// By construction, the [`vars`](Self::vars) member is guaranteed to be
808/// populated for every (non-XYZ) variable in the shape.
809#[derive(Clone)]
810pub struct BoundShape<'a, F, D> {
811    shape: Shape<F>,
812    vars: BorrowedOrDefault<'a, ShapeVars<D>>,
813}
814
815/// Helper object which either borrows or builds a default value
816enum BorrowedOrDefault<'a, T: Default> {
817    Borrowed(&'a T),
818    Default(T),
819}
820
821impl<'a, T: Default> Clone for BorrowedOrDefault<'a, T> {
822    fn clone(&self) -> Self {
823        match self {
824            Self::Borrowed(b) => Self::Borrowed(b),
825            Self::Default(..) => Self::Default(T::default()),
826        }
827    }
828}
829
830impl<'a, T: Default> std::ops::Deref for BorrowedOrDefault<'a, T> {
831    type Target = T;
832    fn deref(&self) -> &T {
833        match self {
834            Self::Borrowed(t) => t,
835            Self::Default(t) => t,
836        }
837    }
838}
839
840impl<T: Default> Default for BorrowedOrDefault<'_, T> {
841    fn default() -> Self {
842        Self::Default(T::default())
843    }
844}
845
846impl<'a, F: Function, D> BoundShape<'a, F, D> {
847    /// Builds a new object, checking that all variables are present
848    pub fn new(
849        shape: Shape<F>,
850        vars: &'a ShapeVars<D>,
851    ) -> Result<Self, MissingVar> {
852        vars.check(&shape)?;
853        Ok(Self {
854            shape,
855            vars: BorrowedOrDefault::Borrowed(vars),
856        })
857    }
858
859    /// Returns the shape handle
860    pub fn shape(&self) -> &Shape<F> {
861        &self.shape
862    }
863
864    /// Returns the variables map
865    pub fn vars(&self) -> &ShapeVars<D> {
866        &self.vars
867    }
868}
869
870impl<'a, F: Function> TryFrom<Shape<F>> for BoundShape<'a, F, f32> {
871    type Error = MissingVar;
872    fn try_from(shape: Shape<F>) -> Result<Self, Self::Error> {
873        if let Some(v) = shape
874            .inner()
875            .vars()
876            .iter()
877            .flat_map(|(v, _i)| match v {
878                Var::X | Var::Y | Var::Z => None,
879                Var::V(i) => Some(i),
880            })
881            .next()
882        {
883            Err(MissingVar { var: v })
884        } else {
885            Ok(BoundShape {
886                shape,
887                vars: BorrowedOrDefault::default(),
888            })
889        }
890    }
891}
892
893/// Trait for types that can be transformed by a 4x4 homogeneous transform matrix
894pub trait Transformable {
895    /// Apply the given transform to an `(x, y, z)` position
896    fn transform(
897        x: Self,
898        y: Self,
899        z: Self,
900        mat: &Matrix4<f32>,
901    ) -> (Self, Self, Self)
902    where
903        Self: Sized;
904}
905
906impl Transformable for f32 {
907    fn transform(
908        x: f32,
909        y: f32,
910        z: f32,
911        mat: &Matrix4<f32>,
912    ) -> (f32, f32, f32) {
913        let out = mat.transform_point(&Point3::new(x, y, z));
914        (out.x, out.y, out.z)
915    }
916}
917
918impl Transformable for Interval {
919    fn transform(
920        x: Interval,
921        y: Interval,
922        z: Interval,
923        mat: &Matrix4<f32>,
924    ) -> (Interval, Interval, Interval) {
925        let out = [0, 1, 2, 3].map(|i| {
926            let row = mat.row(i);
927            x * row[0] + y * row[1] + z * row[2] + Interval::from(row[3])
928        });
929
930        (out[0] / out[3], out[1] / out[3], out[2] / out[3])
931    }
932}
933
934impl Transformable for Grad {
935    fn transform(
936        x: Grad,
937        y: Grad,
938        z: Grad,
939        mat: &Matrix4<f32>,
940    ) -> (Grad, Grad, Grad) {
941        let out = [0, 1, 2, 3].map(|i| {
942            let row = mat.row(i);
943            x * row[0] + y * row[1] + z * row[2] + Grad::from(row[3])
944        });
945
946        (out[0] / out[3], out[1] / out[3], out[2] / out[3])
947    }
948}
949
950#[cfg(test)]
951mod test {
952    use super::*;
953    use crate::vm::VmShape;
954
955    #[test]
956    fn shape_vars() {
957        let v = Var::new();
958        let s = Tree::x() + Tree::y() + v;
959
960        let mut ctx = Context::new();
961        let s = ctx.import(&s);
962
963        let s = VmShape::new(&ctx, s).unwrap();
964        let vs = s.inner().vars();
965        assert_eq!(vs.len(), 3);
966
967        assert!(vs.get(&Var::X).is_some());
968        assert!(vs.get(&Var::Y).is_some());
969        assert!(vs.get(&Var::Z).is_none());
970        assert!(vs.get(&v).is_some());
971
972        let mut seen = [false; 3];
973        for v in [Var::X, Var::Y, v] {
974            seen[vs[&v]] = true;
975        }
976        assert!(seen.iter().all(|i| *i));
977    }
978
979    #[test]
980    fn shape_bind() {
981        let v = Var::new();
982        let s = Tree::x() + Tree::y() + v;
983
984        let mut ctx = Context::new();
985        let s = ctx.import(&s);
986
987        // Try the TryFrom conversion, which will fail
988        let s = VmShape::new(&ctx, s).unwrap();
989        let Err(e) = BoundShape::try_from(s.clone()) else {
990            panic!("expected error");
991        };
992        let vi = v.index().unwrap();
993        assert_eq!(e, MissingVar { var: vi });
994
995        // Bind to a map with no vars
996        let mut m = ShapeVars::new();
997        let Err(e) = s.bind(&m) else {
998            panic!("expected error");
999        };
1000        assert_eq!(e, MissingVar { var: vi });
1001
1002        // Add an unrelated var
1003        m.insert(Var::new().index().unwrap(), 1.0f32);
1004        let Err(e) = s.bind(&m) else {
1005            panic!("expected error");
1006        };
1007        assert_eq!(e, MissingVar { var: vi });
1008
1009        // Finally, bind to a map with the target var
1010        m.insert(vi, 2.0f32);
1011        assert!(s.bind(&m).is_ok());
1012    }
1013
1014    #[test]
1015    fn shape_eval_bulk_size() {
1016        let s = Tree::constant(1.0);
1017        let mut ctx = Context::new();
1018        let s = ctx.import(&s);
1019
1020        let s = VmShape::new(&ctx, s).unwrap();
1021        let tape = s.ez_float_slice_tape();
1022        let mut eval = VmShape::new_float_slice_eval();
1023        let out = eval
1024            .eval(&tape, &[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0], &[7.0, 8.0, 9.0])
1025            .unwrap();
1026        assert_eq!(out, [1.0, 1.0, 1.0]);
1027    }
1028}
fidget_core/shape/mod.rs

fidget_core/shape/
mod.rs
