enkou/enkou-shaders/src/transform.rs

//! # Transform
//!
//! Transforms are the "functions" in an iterated function system. They take in a point,
//! and generate a new point. For fractal flames, transforms are always affine,
//! but produce more interesting images once we add variations.
use crate::variation::Variation;
use bytemuck::{Pod, Zeroable};
use glam::{Affine2, FloatExt, UVec2, Vec2};
use rand::Rng;

/// Affine transform for use in the [`chaos_game`](crate::chaos_game).
#[derive(Copy, Clone, Pod, Zeroable)]
#[repr(C)]
pub struct Transform {
    coefficients: Affine2,
    variation_range: UVec2,
    color: Vec2,
}

impl Transform {
    /// Create a new transform from an affine transformation matrix
    ///
    /// 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,
        }
    }

    /// Apply this transform to a point in IFS coordinates, producing a new point
    pub fn transform_point<R: Rng>(
        &self,
        rng: &mut R,
        variations: &[Variation],
        point: Vec2,
    ) -> Vec2 {
        let point = self.coefficients.transform_point2(point);

        let mut point_output = Vec2::ZERO;

        let variation_start = self.variation_range.x;
        let variation_end = self.variation_range.y;
        for variation_index in variation_start..variation_end {
            let ref variation = variations[variation_index as usize];
            point_output += variation.transform_point(point, rng, &self.coefficients)
        }

        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);
    }
}