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

fidget_core/shape/
mod.rs
