Implement transform colors

enkou-shaders/src/transform.rs

@@ -5,7 +5,7 @@
 //! but produce more interesting images once we add variations.
 use crate::variation::Variation;
 use bytemuck::{Pod, Zeroable};
-use glam::{Affine2, UVec2, Vec2};
+use glam::{Affine2, FloatExt, UVec2, Vec2};
 use rand::Rng;
 
 /// Affine transform for use in the [`chaos_game`](crate::chaos_game).
@@ -14,14 +14,21 @@ use rand::Rng;
 pub struct Transform {
     coefficients: Affine2,
     variation_range: UVec2,
+    color: Vec2,
 }
 
 impl Transform {
     /// Create a new transform from an affine transformation matrix
-    pub fn new(coefficients: Affine2, variation_range: UVec2) -> Self {
+    ///
+    /// Arguments:
+    /// * `coefficients` - Affine transform coefficients for this transformation. Applied prior to variations
+    /// * `variation_range` - (half-open) range of variations to apply during [`Self::transform_point`]
+    /// * `color` - Color value and speed to apply during [`Self::transform_color`]
+    pub fn new(coefficients: Affine2, variation_range: UVec2, color: Vec2) -> Self {
         Transform {
             coefficients,
             variation_range,
+            color,
         }
     }
 
@@ -45,4 +52,105 @@ impl Transform {
 
         point_output
     }
+
+    /// Mix an existing color with this transform's color
+    pub fn transform_color(&self, color: f32) -> f32 {
+        color.lerp(self.color.x, self.color.y)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::Coefficients2;
+    use crate::transform::Transform;
+    use crate::variation::{Variation, VariationKind};
+    use core::convert::Infallible;
+    use glam::{Affine2, Vec2, uvec2, vec2};
+    use rand::TryRng;
+
+    struct NullRng;
+
+    impl NullRng {
+        pub fn new() -> Self {
+            NullRng
+        }
+    }
+
+    impl TryRng for NullRng {
+        type Error = Infallible;
+
+        fn try_next_u32(&mut self) -> Result<u32, Self::Error> {
+            Ok(0)
+        }
+
+        fn try_next_u64(&mut self) -> Result<u64, Self::Error> {
+            Ok(0)
+        }
+
+        fn try_fill_bytes(&mut self, dst: &mut [u8]) -> Result<(), Self::Error> {
+            dst.iter_mut().for_each(|b| *b = 0);
+            Ok(())
+        }
+    }
+
+    #[test]
+    fn test_transform_point_identity() {
+        let transform = Transform::new(Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 0.0));
+        let variations = [Variation::IDENTITY];
+
+        let transform_point =
+            |point: Vec2| transform.transform_point(&mut NullRng::new(), &variations, point);
+        for (input, expected) in [
+            (vec2(0.0, 1.0), vec2(0.0, 1.0)),
+            (vec2(1.0, 0.0), vec2(1.0, 0.0)),
+            (vec2(1.0, 1.0), vec2(1.0, 1.0)),
+        ] {
+            assert_eq!(transform_point(input), expected);
+        }
+    }
+
+    #[test]
+    fn test_transform_point_scaling() {
+        let transform = Transform::new(
+            Affine2::from_coefficients(0.5, 0.0, 0.0, 0.0, 0.5, 0.0),
+            uvec2(0, 1),
+            vec2(0.0, 0.0),
+        );
+        let variations = [Variation::IDENTITY];
+
+        let transform_point =
+            |point: Vec2| transform.transform_point(&mut NullRng::new(), &variations, point);
+        for (input, expected) in [
+            (vec2(0.0, 1.0), vec2(0.0, 0.5)),
+            (vec2(1.0, 0.0), vec2(0.5, 0.0)),
+            (vec2(1.0, 1.0), vec2(0.5, 0.5)),
+        ] {
+            assert_eq!(transform_point(input), expected);
+        }
+    }
+
+    #[test]
+    fn test_transform_point_scaling_variation() {
+        let transform = Transform::new(Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 0.0));
+        let variations = [Variation::new(VariationKind::Linear, 2.0, [0.0; 4].into())];
+
+        let transform_point =
+            |point: Vec2| transform.transform_point(&mut NullRng::new(), &variations, point);
+        for (input, expected) in [
+            (vec2(0.0, 1.0), vec2(0.0, 2.0)),
+            (vec2(1.0, 0.0), vec2(2.0, 0.0)),
+            (vec2(1.0, 1.0), vec2(2.0, 2.0)),
+        ] {
+            assert_eq!(transform_point(input), expected);
+        }
+    }
+
+    #[test]
+    fn test_color_mixing() {
+        let transform = Transform::new(Affine2::IDENTITY, uvec2(0, 1), vec2(0.0, 0.5));
+
+        assert_eq!(transform.transform_color(0.0), 0.0);
+        assert_eq!(transform.transform_color(1.0), 0.5);
+        assert_eq!(transform.transform_color(0.5), 0.25);
+    }
 }