Var API cleanup

fidget/src/core/eval/mod.rs

@@ -18,15 +18,18 @@ pub use tracing::TracingEvaluator;
 
 /// A tape represents something that can be evaluated by an evaluator
 ///
-/// The only property enforced on the trait is that we must have some way to
-/// recycle its internal storage.  This matters most for JIT evaluators, whose
-/// tapes are regions of executable memory-mapped RAM (which is expensive to map
-/// and unmap).
+/// It includes some kind of storage and the ability to look up variable
+/// mapping.  The variable mapping should be identical to calling
+/// [`Function::vars`] on the `Function` which produced this tape, but it's
+/// convenient to be able to look up vars locally.
 pub trait Tape {
     /// Associated type for this tape's data storage
     type Storage: Default;
 
     /// Retrieves the internal storage from this tape
+    ///
+    /// This matters most for JIT evaluators, whose tapes are regions of
+    /// executable memory-mapped RAM (which is expensive to map and unmap).
     fn recycle(self) -> Self::Storage;
 
     /// Returns a mapping from [`Var`](crate::var::Var) to evaluation index
@@ -169,9 +172,6 @@ pub trait Function: Send + Sync + Clone {
     /// This is underspecified and only used for unit testing; for tape-based
     /// functions, it's typically the length of the tape,
     fn size(&self) -> usize;
-
-    /// Returns a map from variable to index
-    fn vars(&self) -> &VarMap;
 }
 
 /// A [`Function`] which can be built from a math expression

fidget/src/core/shape/mod.rs

@@ -51,12 +51,8 @@ pub struct Shape<F> {
     /// Wrapped function
     f: F,
 
-    /// Index of x, y, z axes within the function's variable list (if present)
-    ///
-    /// We could instead store an array of [`Var`]s and look them in with
-    /// [`self.f.vars()[&v]`](Function::vars), but it's more efficient to cache
-    /// them upon construction (because they never change).
-    axes: [Option<usize>; 3],
+    /// Variables representing x, y, z axes
+    axes: [Var; 3],
 
     /// Optional transform to apply to the shape
     transform: Option<Matrix4<f32>>,
@@ -100,9 +96,12 @@ impl<F: Function + Clone> Shape<F> {
         &self,
         storage: F::TapeStorage,
     ) -> ShapeTape<<F::PointEval as TracingEvaluator>::Tape> {
+        let tape = self.f.point_tape(storage);
+        let vars = tape.vars();
+        let axes = self.axes.map(|v| vars.get(&v));
         ShapeTape {
-            tape: self.f.point_tape(storage),
-            axes: self.axes,
+            tape,
+            axes,
             transform: self.transform,
         }
     }
@@ -112,9 +111,12 @@ impl<F: Function + Clone> Shape<F> {
         &self,
         storage: F::TapeStorage,
     ) -> ShapeTape<<F::IntervalEval as TracingEvaluator>::Tape> {
+        let tape = self.f.interval_tape(storage);
+        let vars = tape.vars();
+        let axes = self.axes.map(|v| vars.get(&v));
         ShapeTape {
-            tape: self.f.interval_tape(storage),
-            axes: self.axes,
+            tape,
+            axes,
             transform: self.transform,
         }
     }
@@ -124,9 +126,12 @@ impl<F: Function + Clone> Shape<F> {
         &self,
         storage: F::TapeStorage,
     ) -> ShapeTape<<F::FloatSliceEval as BulkEvaluator>::Tape> {
+        let tape = self.f.float_slice_tape(storage);
+        let vars = tape.vars();
+        let axes = self.axes.map(|v| vars.get(&v));
         ShapeTape {
-            tape: self.f.float_slice_tape(storage),
-            axes: self.axes,
+            tape,
+            axes,
             transform: self.transform,
         }
     }
@@ -136,9 +141,12 @@ impl<F: Function + Clone> Shape<F> {
         &self,
         storage: F::TapeStorage,
     ) -> ShapeTape<<F::GradSliceEval as BulkEvaluator>::Tape> {
+        let tape = self.f.grad_slice_tape(storage);
+        let vars = tape.vars();
+        let axes = self.axes.map(|v| vars.get(&v));
         ShapeTape {
-            tape: self.f.grad_slice_tape(storage),
-            axes: self.axes,
+            tape,
+            axes,
             transform: self.transform,
         }
     }
@@ -185,12 +193,12 @@ impl<F> Shape<F> {
     }
 
     /// Borrows the inner axis mapping
-    pub fn axes(&self) -> &[Option<usize>; 3] {
+    pub fn axes(&self) -> &[Var; 3] {
         &self.axes
     }
 
     /// Raw constructor
-    pub fn new_raw(f: F, axes: [Option<usize>; 3]) -> Self {
+    pub fn new_raw(f: F, axes: [Var; 3]) -> Self {
         Self {
             f,
             axes,
@@ -302,8 +310,6 @@ impl<F: MathFunction> Shape<F> {
         axes: [Var; 3],
     ) -> Result<Self, Error> {
         let f = F::new(ctx, node)?;
-        let vars = f.vars();
-        let axes = axes.map(|v| vars.get(&v));
         Ok(Self {
             f,
             axes,
@@ -642,7 +648,7 @@ mod test {
         let s = ctx.import(&s);
 
         let s = VmShape::new(&mut ctx, s).unwrap();
-        let vs = Function::vars(s.inner());
+        let vs = s.inner().vars();
         assert_eq!(vs.len(), 3);
 
         assert!(vs.get(&Var::X).is_some());

fidget/src/core/var/mod.rs

@@ -17,7 +17,18 @@ use std::collections::HashMap;
 /// same thing.  To generate a "local" variable, [`Var::new`] picks a random
 /// 64-bit value, which is very unlikely to collide with anything else.
 #[allow(missing_docs)]
-#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Hash,
+    Eq,
+    PartialEq,
+    Ord,
+    PartialOrd,
+    Serialize,
+    Deserialize,
+)]
 pub enum Var {
     X,
     Y,

fidget/src/core/vm/mod.rs

@@ -181,10 +181,6 @@ impl<const N: usize> Function for GenericVmFunction<N> {
     fn size(&self) -> usize {
         GenericVmFunction::size(self)
     }
-
-    fn vars(&self) -> &VarMap {
-        &self.0.vars
-    }
 }
 
 impl<const N: usize> RenderHints for GenericVmFunction<N> {

fidget/src/jit/mod.rs

@@ -900,10 +900,6 @@ impl Function for JitFunction {
     fn size(&self) -> usize {
         self.0.size()
     }
-
-    fn vars(&self) -> &VarMap {
-        Function::vars(&self.0)
-    }
 }
 
 impl RenderHints for JitFunction {

wasm-demo/src/lib.rs

@@ -2,6 +2,7 @@ use fidget::{
     context::{Context, Tree},
     render::{BitRenderMode, RenderConfig},
     shape::Bounds,
+    var::Var,
     vm::{VmData, VmShape},
     Error,
 };
@@ -37,7 +38,7 @@ pub fn serialize_into_tape(t: JsTree) -> Result<Vec<u8>, String> {
 /// Deserialize a `bincode`-packed `VmData` into a `VmShape`
 #[wasm_bindgen]
 pub fn deserialize_tape(data: Vec<u8>) -> Result<JsVmShape, String> {
-    let (d, axes): (VmData<255>, [Option<usize>; 3]) =
+    let (d, axes): (VmData<255>, [Var; 3]) =
         bincode::deserialize(&data).map_err(|e| format!("{e}"))?;
     Ok(JsVmShape(VmShape::new_raw(d.into(), axes)))
 }