Fix RNG transmutation

enkou-shaders/examples/gasket.rs

@@ -43,7 +43,7 @@ pub fn main() -> Result<()> {
 
     main_chaos_game(
         ITERATIONS_DISCARD,
-        &[4u8],
+        &[4u8; 32],
         &transforms,
         &weights,
         &variations,
@@ -77,10 +77,11 @@ pub fn main() -> Result<()> {
     let temp = NamedTempFile::with_suffix(".png").context("Unable to create file for image")?;
     image.save(temp.path()).context("Unable to save image")?;
 
-    let open_program = cfg_select! {
+    let open_program: &str = cfg_select! {
-        unix => "xdg-open",
+        unix => Some("xdg-open"),
-        _ => panic!("Unknown system"),
+        _ => None,
-    };
+    }
+    .expect("No available program to open images");
 
     Command::new(open_program)
         .arg(temp.path())

enkou-shaders/src/camera.rs

@@ -90,11 +90,13 @@ 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 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,

enkou-shaders/src/chaos_game.rs

@@ -107,21 +107,19 @@ impl<'a, R: Rng> Iterator for ChaosGame<'a, R> {
     }
 }
 
+/// Shader entry point for running the chaos game to produce new IFS coordinates
 pub mod entry {
     use crate::chaos_game::ChaosGame;
+    use crate::rng::xoshiro256starstar_from_seed;
     use crate::transform::Transform;
     use crate::variation::Variation;
     use glam::Vec2;
-    use rand_xoshiro::Xoshiro256StarStar;
     use spirv_std::spirv;
 
-    fn xoshiro_from_state(rng_state: [u8; 32]) -> Xoshiro256StarStar {
+    /// Given a set of fractal flame parameters, generate new IFS coordinates
-        let mut rng_state_actual = [1u64, 2u64, 3u64, 4u64];
+    /// and store them in the output array.
-        unsafe { core::mem::transmute(rng_state_actual) }
-    }
-
     #[spirv(compute(entry_point_name = "main_chaos_game", threads(1)))]
-    pub extern "C" fn main_chaos_game(
+    pub fn main_chaos_game(
         #[spirv(spec_constant(id = 1, default = 20))] iteration_discard: u32,
         #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] rng_seed: &[u8],
         #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] transforms: &[Transform],
@@ -129,9 +127,10 @@ pub mod entry {
         #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] variations: &[Variation],
         #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] output: &mut [Vec2],
     ) {
-        let rng_seed_actual = [0u8; 32];
+        let mut rng_seed_actual = [0u8; 32];
+        (0..32).for_each(|i| rng_seed_actual[i] = rng_seed[i]);
 
-        let mut rng = xoshiro_from_state(rng_seed_actual);
+        let mut rng = xoshiro256starstar_from_seed(rng_seed_actual);
         let mut chaos_game = ChaosGame::new(&mut rng, transforms, weights, variations);
 
         for _ in 0..iteration_discard {

enkou-shaders/src/lib.rs

@@ -4,6 +4,7 @@
 
 pub mod camera;
 pub mod chaos_game;
+mod rng;
 pub mod transform;
 pub mod variation;
 

enkou-shaders/src/rng.rs

@@ -0,0 +1,55 @@
+use rand::SeedableRng;
+use rand_xoshiro::Xoshiro256StarStar;
+
+/// Convert an RNG state buffer to an instance of [`Xoshiro256StarStar`].
+///
+/// While [`SeedableRng::from_seed`] is an infallible function,
+/// it relies on some methods that can't be compiled by the SPIR-V
+/// backend (specifically, formatting functions in the core crate).
+///
+/// In practice, the xoshiro RNG state is entirely defined by its seed,
+/// so this function does the work of [`SeedableRng::from_seed`] by
+/// transmuting the seed value to an RNG instance.
+///
+/// This function assumes a properly-initialized state array;
+/// output may silently degenerate if the initial state is all zeros,
+/// so this module is private to the crate.
+pub(crate) fn xoshiro256starstar_from_seed(
+    rng_state: <Xoshiro256StarStar as SeedableRng>::Seed,
+) -> Xoshiro256StarStar {
+    let mut rng_state_actual = [0u64; 4];
+
+    // NOTE: Bit shifting is bad, 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() {
+        for j in 0..size_of::<u64>() {
+            rng_state_actual[i] |= (rng_state[i * size_of::<u64>() + j] as u64) << j * 8;
+        }
+    }
+
+    unsafe { core::mem::transmute(rng_state_actual) }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::rng::xoshiro256starstar_from_seed;
+    use core::iter::zip;
+    use rand::{RngExt, SeedableRng};
+    use rand_xoshiro::Xoshiro256StarStar;
+
+    #[test]
+    fn match_seeded() {
+        let mut seed: <Xoshiro256StarStar as SeedableRng>::Seed = [0u8; 32];
+        for i in 0..seed.len() {
+            seed[i] = i as u8;
+        }
+
+        let rng1 = Xoshiro256StarStar::from_seed(seed).random_iter::<u64>();
+        let rng2 = xoshiro256starstar_from_seed(seed).random_iter::<u64>();
+
+        zip(rng1, rng2)
+            .take(100)
+            .for_each(|(rng1_value, rng2_value)| assert_eq!(rng1_value, rng2_value));
+    }
+}

enkou-shaders/src/variation.rs

@@ -1,4 +1,8 @@
 //! # Variation
+//!
+//! Variations extend the fractal flame iterated function system
+//! with non-linear transforms (as opposed to [`Transform`]s,
+//! which are strictly affine transformations).
 use crate::Coefficients2;
 use bytemuck::{Pod, Zeroable};
 use core::f32::consts::PI;
@@ -7,14 +11,25 @@ use libm::{atan2f, cosf, powf, sinf, sqrtf, tanf};
 use rand::distr::StandardUniform;
 use rand::{Rng, RngExt};
 
+/// Generic variation parameters
+///
+/// Not all variations will use these parameters, but passing them
+/// as an array per variation allows shaders to use a consistent struct size
+/// no matter what the variation actually needs.
 #[derive(Copy, Clone, Pod, Zeroable)]
 #[repr(C)]
 pub struct VariationParams([f32; 4]);
 
+/// Enum for all supported variation types
+///
+/// ID numbers are chosen to match the variation identifier also used by `flam3`
 #[derive(Copy, Clone)]
 #[repr(u32)]
+#[allow(missing_docs)]
 pub enum VariationKind {
+    /// Identity variation, returns the point as-is
     Linear = 0,
+
     Julia = 13,
     Popcorn = 17,
     Pdj = 24,
@@ -25,6 +40,10 @@ unsafe impl bytemuck::Zeroable for VariationKind {}
 // UNSAFE: Sound because enum has guaranteed layout (u32) and defined zero-value
 unsafe impl bytemuck::Pod for VariationKind {}
 
+/// Parameters required for shaders to run the variation function.
+///
+/// Not all variations use the [`VariationParams`], but using the struct
+/// makes it easy to provide parameters to the shader.
 #[derive(Copy, Clone, Pod, Zeroable)]
 #[repr(C)]
 pub struct Variation {
@@ -34,12 +53,15 @@ pub struct Variation {
 }
 
 impl Variation {
+    /// Identity variation; calling [`transform_point`] will yield
+    /// the same point as the input.
     pub const IDENTITY: Variation = Variation {
         kind: VariationKind::Linear,
         weight: 1.0,
         params: VariationParams([0f32; 4]),
     };
 
+    /// Create a new variation by providing the variation kind, weight, and parameters.
     pub fn new(kind: VariationKind, weight: f32, params: VariationParams) -> Variation {
         Variation {
             kind,
@@ -48,6 +70,9 @@ impl Variation {
         }
     }
 
+    /// Transform a point by applying this variation.
+    ///
+    /// Output points are scaled by this variation's weight.
     pub fn transform_point<R: Rng>(
         &self,
         point: Vec2,