Add affine transforms (#292)

When building math trees with transforms, merging affine transforms
helps to keep intervals tight.

Author: mkeeter

CHANGELOG.md

@@ -13,6 +13,10 @@
   a point with `i32` coordinates (instead of `f32`).  This helps us distinguish
   between screen (pixel) and world (floating-point) coordinates at the type
   level.
+- Add `Tree::remap_affine` (and `TreeOp::RemapAffine`) to perform affine
+  transformations on math expressions.  These transformations are composable;
+  two affine transforms will be combined into a single transform if stacked
+  together.
 
 # 0.3.5
 - Added `#[derive(Serialize, Deserialize)]` to `View2` and `View3`

fidget/src/core/context/mod.rs

@@ -30,6 +30,7 @@ use std::fmt::Write;
 use std::io::{BufRead, BufReader, Read};
 use std::sync::Arc;
 
+use nalgebra::Matrix4;
 use ordered_float::OrderedFloat;
 
 define_index!(Node, "An index in the `Context::ops` map");
@@ -1089,10 +1090,13 @@ impl Context {
             Up(&'a Arc<TreeOp>),
             /// Pops the latest axis frame
             Pop,
+            /// Pops the latest affine frame
+            PopAffine,
         }
         let mut axes = vec![(self.x(), self.y(), self.z())];
         let mut todo = vec![Action::Down(tree.arc())];
         let mut stack = vec![];
+        let mut affine: Vec<Matrix4<f64>> = vec![];
 
         // Cache of TreeOp -> Node mapping under a particular frame (axes)
         //
@@ -1145,11 +1149,43 @@ impl Context {
                             todo.push(Action::Down(y));
                             todo.push(Action::Down(z));
                         }
+                        TreeOp::RemapAffine { target, mat } => {
+                            let prev = affine
+                                .last()
+                                .cloned()
+                                .unwrap_or(Matrix4::identity());
+                            let mat = prev * mat.to_homogeneous();
+
+                            // Push either an affine frame or an axis frame,
+                            // depending on whether the target is also affine
+                            if matches!(&**target, TreeOp::RemapAffine { .. }) {
+                                affine.push(mat);
+                                todo.push(Action::PopAffine);
+                            } else {
+                                let (x, y, z) = axes.last().unwrap();
+                                let mut out = [None; 3];
+                                for i in 0..3 {
+                                    let a = self.mul(mat[(i, 0)], *x).unwrap();
+                                    let b = self.mul(mat[(i, 1)], *y).unwrap();
+                                    let c = self.mul(mat[(i, 2)], *z).unwrap();
+                                    let d = self.constant(mat[(i, 3)]);
+                                    let ab = self.add(a, b).unwrap();
+                                    let cd = self.add(c, d).unwrap();
+                                    out[i] = Some(self.add(ab, cd).unwrap());
+                                }
+                                let [x, y, z] = out.map(Option::unwrap);
+                                axes.push((x, y, z));
+                                todo.push(Action::Pop);
+                            }
+                            todo.push(Action::Down(target));
+                        }
                     }
                 }
                 Action::Up(t) => {
                     match t.as_ref() {
-                        TreeOp::Const(..) | TreeOp::Input(..) => unreachable!(),
+                        TreeOp::Const(..)
+                        | TreeOp::Input(..)
+                        | TreeOp::RemapAffine { .. } => unreachable!(),
                         TreeOp::Unary(op, ..) => {
                             let arg = stack.pop().unwrap();
                             let out = self.op_unary(arg, *op).unwrap();
@@ -1195,6 +1231,9 @@ impl Context {
                 Action::Pop => {
                     axes.pop().unwrap();
                 }
+                Action::PopAffine => {
+                    affine.pop().unwrap();
+                }
             }
         }
         assert_eq!(stack.len(), 1);

fidget/src/core/context/tree.rs

@@ -6,7 +6,8 @@ use std::sync::Arc;
 /// Opcode type for trees
 ///
 /// This is equivalent to [`Op`](crate::context::Op), but also includes the
-/// [`RemapAxes`](TreeOp::RemapAxes) operation for lazy remapping.
+/// [`RemapAxes`](TreeOp::RemapAxes) and [`TreeOp::RemapAffine`] operations for
+/// lazy remapping.
 #[derive(Debug)]
 #[allow(missing_docs)]
 pub enum TreeOp {
@@ -19,12 +20,23 @@ pub enum TreeOp {
     ///
     /// When imported into a `Context`, all `x/y/z` clauses within `target` will
     /// be replaced with the provided `x/y/z` trees.
+    ///
+    /// If the transform is affine, then `RemapAffine` should be preferred,
+    /// because it flattens sequences of affine transformations.
     RemapAxes {
         target: Arc<TreeOp>,
         x: Arc<TreeOp>,
         y: Arc<TreeOp>,
         z: Arc<TreeOp>,
     },
+    /// Lazy affine transforms
+    ///
+    /// When imported into a `Context`, the `x/y/z` clauses within `target` will
+    /// be transformed with the provided affine matrix.
+    RemapAffine {
+        target: Arc<TreeOp>,
+        mat: nalgebra::Affine3<f64>,
+    },
 }
 
 impl Drop for TreeOp {
@@ -66,6 +78,9 @@ impl TreeOp {
                     && matches!(**y, TreeOp::Const(..))
                     && matches!(**z, TreeOp::Const(..))
             }
+            TreeOp::RemapAffine { target, .. } => {
+                matches!(**target, TreeOp::Const(..))
+            }
         }
     }
 
@@ -77,6 +92,9 @@ impl TreeOp {
             TreeOp::RemapAxes { target, x, y, z } => {
                 [Some(target), Some(x), Some(y), Some(z)]
             }
+            TreeOp::RemapAffine { target, .. } => {
+                [Some(target), None, None, None]
+            }
         }
         .into_iter()
         .flatten()
@@ -155,6 +173,9 @@ impl Tree {
 
     /// Remaps the axes of the given tree
     ///
+    /// If the mapping is affine, then [`remap_affine`](Self::remap_affine)
+    /// should be preferred.
+    ///
     /// The remapping is lazy; it is not evaluated until the tree is imported
     /// into a `Context`.
     pub fn remap_xyz(&self, x: Tree, y: Tree, z: Tree) -> Tree {
@@ -166,6 +187,25 @@ impl Tree {
         }))
     }
 
+    /// Performs an affine remapping of the given tree
+    ///
+    /// The remapping is lazy; it is not evaluated until the tree is imported
+    /// into a `Context`.
+    pub fn remap_affine(&self, mat: nalgebra::Affine3<f64>) -> Tree {
+        // Flatten affine trees
+        let out = match &*self.0 {
+            TreeOp::RemapAffine { target, mat: next } => TreeOp::RemapAffine {
+                target: target.clone(),
+                mat: mat * next,
+            },
+            _ => TreeOp::RemapAffine {
+                target: self.0.clone(),
+                mat,
+            },
+        };
+        Self(out.into())
+    }
+
     /// Returns the inner [`Var`] if this is an input tree, or `None`
     pub fn var(&self) -> Option<Var> {
         if let TreeOp::Input(v) = &*self.0 {
@@ -399,6 +439,80 @@ mod test {
         assert_eq!(ctx.eval_xyz(v_, 0.0, 1.0, 0.0).unwrap(), 4.0);
     }
 
+    #[test]
+    fn test_remap_affine() {
+        let s = Tree::x();
+        // Two rotations by 45° -> 90°
+        let t = nalgebra::convert(nalgebra::Rotation3::<f64>::from_axis_angle(
+            &nalgebra::Vector3::<f64>::z_axis(),
+            -std::f64::consts::FRAC_PI_4,
+        ));
+        let s = s.remap_affine(t);
+        let s = s.remap_affine(t);
+
+        let TreeOp::RemapAffine { target, .. } = &*s else {
+            panic!("invalid shape");
+        };
+        assert!(matches!(&**target, TreeOp::Input(Var::X)));
+
+        let mut ctx = Context::new();
+        let v_ = ctx.import(&s);
+
+        assert!((ctx.eval_xyz(v_, 0.0, 1.0, 0.0).unwrap() - 1.0).abs() < 1e-6);
+        assert!(
+            (ctx.eval_xyz(v_, 0.0, -2.0, 0.0).unwrap() - -2.0).abs() < 1e-6
+        );
+    }
+
+    #[test]
+    fn test_remap_order() {
+        let translate = nalgebra::convert(nalgebra::Translation3::<f64>::new(
+            3.0, 0.0, 0.0,
+        ));
+        let scale =
+            nalgebra::convert(nalgebra::Scale3::<f64>::new(0.5, 0.5, 0.5));
+
+        let s = Tree::x();
+        let s = s.remap_affine(scale);
+        let s = s.remap_affine(translate);
+
+        // Confirm that we didn't stack up RemapAffine nodes
+        let TreeOp::RemapAffine { target, .. } = &*s else {
+            panic!("invalid shape");
+        };
+        assert!(matches!(&**target, TreeOp::Input(Var::X)));
+
+        // Basic evaluation testing
+        let mut ctx = Context::new();
+        let v_ = ctx.import(&s);
+        assert_eq!(ctx.eval_xyz(v_, 1.0, 0.0, 0.0).unwrap(), 3.5);
+        assert_eq!(ctx.eval_xyz(v_, 2.0, 0.0, 0.0).unwrap(), 4.0);
+
+        // Do the same thing but testing collapsing in `Context::import`
+        let manual = TreeOp::RemapAffine {
+            target: Arc::new(TreeOp::RemapAffine {
+                target: TreeOp::Input(Var::X).into(),
+                mat: scale,
+            }),
+            mat: translate,
+        }
+        .into();
+        let mut ctx = Context::new();
+        let v_ = ctx.import(&manual);
+        assert_eq!(ctx.eval_xyz(v_, 1.0, 0.0, 0.0).unwrap(), 3.5);
+        assert_eq!(ctx.eval_xyz(v_, 2.0, 0.0, 0.0).unwrap(), 4.0);
+
+        // Swap the order and make sure it still works
+        let s = Tree::x();
+        let s = s.remap_affine(translate);
+        let s = s.remap_affine(scale);
+
+        let mut ctx = Context::new();
+        let v_ = ctx.import(&s);
+        assert_eq!(ctx.eval_xyz(v_, 1.0, 0.0, 0.0).unwrap(), 2.0);
+        assert_eq!(ctx.eval_xyz(v_, 2.0, 0.0, 0.0).unwrap(), 2.5);
+    }
+
     #[test]
     fn deep_recursion_drop() {
         let mut x = Tree::x();