6 changed files with 152 additions and 20 deletions
@@ -5,7 +5,7 @@ use enkou_shaders::camera::entry::main_camera;
 use enkou_shaders::chaos_game::entry::main_chaos_game;
 use enkou_shaders::transform::Transform;
 use enkou_shaders::variation::Variation;
-use glam::{Affine2, IVec2, UVec2, Vec2, uvec2};
+use glam::{Affine2, IVec2, UVec2, Vec2, Vec4, uvec2, vec2};
 use image::{GrayImage, Luma};
 use std::mem;
 use std::process::Command;
@@ -21,16 +21,19 @@ pub fn main() -> Result<()> {
        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),
        ),
        // F_1: ((x + 1) / 2, y / 2)
        Transform::new(
            Affine2::from_coefficients(0.5, 0.0, 0.5, 0.0, 0.5, 0.0),
            uvec2(0, 1),
+            vec2(0.0, 0.0),
        ),
        // F_2: (x / 2, (y + 1) / 2)
        Transform::new(
            Affine2::from_coefficients(0.5, 0.0, 0.0, 0.0, 0.5, 0.5),
            uvec2(0, 1),
+            vec2(0.0, 0.0),
        ),
    ];

@@ -39,7 +42,7 @@ pub fn main() -> Result<()> {
    let variations = [Variation::IDENTITY];

    let mut output_points_ifs = Vec::new();
-    output_points_ifs.resize(ITERATIONS as usize, Vec2::ZERO);
+    output_points_ifs.resize(ITERATIONS as usize, Vec4::ZERO);

    main_chaos_game(
        ITERATIONS_DISCARD,
@@ -93,18 +93,19 @@ impl Camera {
 /// Shader entry point for running the camera transformation over a list of IFS coordinates
 pub mod entry {
    use crate::camera::Camera;
-    use spirv_std::glam::{IVec2, Vec2};
+    use glam::{Vec4, Vec4Swizzles};
+    use spirv_std::glam::IVec2;
    use spirv_std::spirv;

    /// Transform IFS coordinates to pixel coordinates
    #[spirv(compute(entry_point_name = "main_camera", threads(1)))]
    pub fn main_camera(
        #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] camera: &Camera,
-        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] coordinates_ifs: &[Vec2],
+        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] coordinates_ifs: &[Vec4],
        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] coordinates_pixel: &mut [IVec2],
    ) {
        for i in 0..coordinates_ifs.len() {
-            coordinates_pixel[i] = camera.transform_point(coordinates_ifs[i])
+            coordinates_pixel[i] = camera.transform_point(coordinates_ifs[i].xy())
        }
    }
 }
@@ -12,7 +12,6 @@
 //!
 //! This algorithm is also known as the ["chaos game"](https://en.wikipedia.org/wiki/Chaos_game),
 //! and it forms the basic system for producing images.
-
 use crate::transform::Transform;
 use crate::variation::Variation;
 use rand::distr::{Distribution, StandardUniform};
@@ -36,13 +35,14 @@ impl Distribution<f32> for BiUnit {
 /// * `weights` - Weights are assumed to be normalized; adding all elements together should return the value 1
 pub fn step_chaos_game<R: Rng>(
    point: Vec2,
+    color: f32,
    rng: &mut R,
    transforms: &[Transform],
    weights: &[f32],
    variations: &[Variation],
-) -> (Vec2, u32) {
+) -> (Vec2, f32, usize) {
    let mut choice_weight = rng.sample::<f32, _>(StandardUniform);
-    let mut transform_index: u32 = 0;
+    let mut transform_index: usize = 0;

    for i in 0..weights.len() {
        choice_weight -= weights[i];
@@ -53,8 +53,10 @@ pub fn step_chaos_game<R: Rng>(
        transform_index += 1;
    }

+    let ref transform = transforms[transform_index];
    (
-        transforms[transform_index as usize].transform_point(rng, variations, point),
+        transform.transform_point(rng, variations, point),
+        transform.transform_color(color),
        transform_index,
    )
 }
@@ -65,6 +67,7 @@ pub fn step_chaos_game<R: Rng>(
 /// New points in the chaos game are produced by iterating on the chaos game.
 pub struct ChaosGame<'a, R: Rng> {
    current_point: Vec2,
+    current_color: f32,
    rng: &'a mut R,
    transforms: &'a [Transform],
    weights: &'a [f32],
@@ -79,9 +82,9 @@ impl<'a, R: Rng> ChaosGame<'a, R> {
        weights: &'a [f32],
        variations: &'a [Variation],
    ) -> Self {
-        let current_point = vec2(rng.sample(BiUnit), rng.sample(BiUnit));
        ChaosGame {
-            current_point,
+            current_point: vec2(rng.sample(BiUnit), rng.sample(BiUnit)),
+            current_color: rng.sample(StandardUniform),
            rng,
            transforms,
            weights,
@@ -91,19 +94,21 @@ impl<'a, R: Rng> ChaosGame<'a, R> {
 }

 impl<'a, R: Rng> Iterator for ChaosGame<'a, R> {
-    type Item = Vec2;
+    type Item = (Vec2, f32);

    fn next(&mut self) -> Option<Self::Item> {
-        let (next_point, _) = step_chaos_game(
+        let (next_point, next_color, _) = step_chaos_game(
            self.current_point,
+            self.current_color,
            self.rng,
            self.transforms,
            self.weights,
            self.variations,
        );
        self.current_point = next_point;
+        self.current_color = next_color;

-        Some(next_point)
+        Some((next_point, next_color))
    }
 }

@@ -113,11 +118,17 @@ pub mod entry {
    use crate::rng::xoshiro256starstar_from_seed;
    use crate::transform::Transform;
    use crate::variation::Variation;
-    use glam::Vec2;
+    use glam::Vec4;
    use spirv_std::spirv;

    /// Given a set of fractal flame parameters, generate new IFS coordinates
    /// and store them in the output array.
+    ///
+    /// Arguments:
+    /// * `iteration_discard` - Choas game steps to discard prior to recording into the output buffer
+    /// * `output` - Output buffer to record chaos game steps into. Because of alignment issues,
+    ///     the output is recorded as a [`Vec4`]; the IFS (x, y) coordinate is in `x` and `y`,
+    ///     and color is in `w`
    #[spirv(compute(entry_point_name = "main_chaos_game", threads(1)))]
    pub fn main_chaos_game(
        #[spirv(spec_constant(id = 1, default = 20))] iteration_discard: u32,
@@ -125,7 +136,7 @@ pub mod entry {
        #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] transforms: &[Transform],
        #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] weights: &[f32],
        #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] variations: &[Variation],
-        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] output: &mut [Vec2],
+        #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] output: &mut [Vec4],
    ) {
        let mut rng_seed_actual = [0u8; 32];
        (0..32).for_each(|i| rng_seed_actual[i] = rng_seed[i]);
@@ -138,7 +149,10 @@ pub mod entry {
        }

        for i in 0..output.len() {
-            output[i] = chaos_game.next().unwrap();
+            output[i] = chaos_game
+                .next()
+                .map(|output| (output.0, 0.0, output.1).into())
+                .unwrap();
        }
    }
 }
@@ -19,7 +19,7 @@ pub(crate) fn xoshiro256starstar_from_seed(
 ) -> Xoshiro256StarStar {
    let mut rng_state_actual = [0u64; 4];

-    // NOTE: Bit shifting is bad, but we don't have great alternatives:
+    // NOTE: Bit shifting is tedious, but we don't have great alternatives:
    //  - `chunks_exact` has issues with pointer casting
    //  - `u64::from_le_bytes` has issues with `OpBitcast` in SPIR-V validation
    for i in 0..rng_state_actual.len() {
@@ -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);
+    }
 }
@@ -20,6 +20,12 @@ use rand::{Rng, RngExt};
 #[repr(C)]
 pub struct VariationParams([f32; 4]);

+impl From<[f32; 4]> for VariationParams {
+    fn from(value: [f32; 4]) -> Self {
+        VariationParams(value)
+    }
+}
+
 /// Enum for all supported variation types
 ///
 /// ID numbers are chosen to match the variation identifier also used by `flam3`