[WIP] Migrate to using a compute shader

abiogenesis/Cargo.toml

@@ -20,6 +20,7 @@ bevy = { version = "0.16", default-features = false, features = [
   "png",
   "std",
   "webgl2",
+  "shader_format_wesl"
 ] }
 rand = "0.8"
 log = { version = "0.4", features = [
@@ -38,6 +39,7 @@ bevy_tweening = { version = "0.13.0", default-features = false, features = [
 itertools = "0.14.0"
 rand_distr = "0.4.3"
 serde = { version = "1.0.219", features = ["derive"] }
+wgpu-types = "24.0.0"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 wasm-bindgen = { version = "0.2.100" }

abiogenesis/assets/shaders/compute_shader.wgsl

@@ -0,0 +1,131 @@
+#import utils::colours::NUM_COLOURS
+#import utils::math::{
+    lcg, random_float, remap,
+    toroidal_displacement, toroidal_wrap,
+    Rect, clamp_length, safe_normalize
+}
+#import bevy_render::view::View
+#import bevy_render::globals::Globals
+
+@group(0) @binding(0)
+var<uniform> view: View;
+
+@group(0) @binding(1)
+var<uniform> globals: Globals;
+
+@group(0) @binding(2)
+var<storage, read_write> particle_colours: array<u32>;
+
+@group(0) @binding(3)
+var<storage, read> in_position: array<vec2<f32>>;
+
+@group(0) @binding(4)
+var<storage, read> in_velocity: array<vec2<f32>>;
+
+@group(0) @binding(5)
+var<storage, read_write> out_position: array<vec2<f32>>;
+
+@group(0) @binding(6)
+var<storage, read_write> out_velocity: array<vec2<f32>>;
+
+override NUM_PARTICLES = 400u * 64u;
+
+override x_dim = 4.0 * 1920.0;
+override y_dim = 4.0 * 1080.0;
+
+
+@compute @workgroup_size(64)
+fn init(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+
+    particle_colours[index] = lcg(index) % 3;
+
+    out_position[index] = vec2<f32>(
+        mix(-x_dim / 2.0, x_dim / 2.0, random_float(index + 1283718)),
+        mix(-y_dim / 2.0, y_dim / 2.0, random_float(index + 3879349)),
+    );
+
+    out_velocity[index] = vec2<f32>(0.0, 0.0);
+}
+
+override friction = 2.0;
+const model = array<f32, 9>(
+    0.3, 0.4, 0.5,
+    0.7, -0.4, 0.3,
+    -0.5, 0.5, 0.0
+);
+
+fn get_model_value(a: u32, b: u32) -> f32 {
+    return model[a * 3 + b];
+}
+
+const BOUNDS: Rect = Rect(
+    vec2<f32>(-4.0 * 1920.0 / 2.0, -4.0 * 1080.0 / 2.0),
+    vec2<f32>(4.0 * 1920.0 / 2.0, 4.0 * 1080.0 / 2.0)
+);
+
+@compute @workgroup_size(64)
+fn update(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+
+    var position = in_position[index];
+
+    let lower_index = u32(((globals.time - globals.delta_time) % 60.0 / 60.0) * f32(NUM_PARTICLES));
+    let upper_index = u32((globals.time % 60.0 / 60.0) * f32(NUM_PARTICLES));
+
+    if lower_index <= index && index < upper_index {
+        particle_colours[index] = lcg(index + u32(globals.time)) % 3;
+        position = vec2<f32>(
+            mix(-x_dim / 2.0, x_dim / 2.0, random_float(index + 1283718)),
+            mix(-y_dim / 2.0, y_dim / 2.0, random_float(index + 3879349)),
+        );
+    }
+
+    let friction = exp(-friction * globals.delta_time);
+
+    var force = vec2<f32>(0.0, 0.0);
+
+    for (var other = 0u; other < NUM_PARTICLES; other++) {
+        if other == index { continue; }
+
+        let displacement = toroidal_displacement(BOUNDS, position, in_position[other]);
+        // let displacement = in_position[other] - position;
+
+        let magnitude = influence(get_model_value(particle_colours[index], particle_colours[other]), length(displacement));
+
+        force += 100.0 * magnitude * safe_normalize(displacement);
+    }
+
+
+    let velocity = in_velocity[index] + clamp_length(force, 0.0, 1000.0) * globals.delta_time;
+    out_velocity[index] = clamp_length(velocity * friction, 0.0, 400.0);
+    out_position[index] = toroidal_wrap(BOUNDS, position + velocity * globals.delta_time);
+}
+
+override replusion_radius = 25.0;
+override peak_attraction_radius = 50.0;
+override attraction_radius = 75.0;
+
+fn influence(factor: f32, distance: f32) -> f32 {
+    if distance <= replusion_radius {
+        return remap(distance, 0.0, replusion_radius, -1.0, 0.0);
+    } else if distance <= peak_attraction_radius {
+        return remap(
+            distance,
+            replusion_radius,
+            peak_attraction_radius,
+            0.0,
+            factor,
+        );
+    } else if distance <= attraction_radius {
+        return remap(
+            distance,
+            66.6,
+            attraction_radius,
+            factor,
+            0.0,
+        );
+    } else {
+        return 0.0;
+    }
+}

abiogenesis/assets/shaders/particle_render.wgsl

@@ -0,0 +1,84 @@
+#import bevy_render::view::View
+#import utils::colours::COLOURS
+#import utils::math::{Rect, toroidal_displacement}
+
+struct VertexInput {
+    @builtin(instance_index) instance_index: u32,
+    @location(0) position: vec3<f32>,
+}
+
+struct VertexOutput {
+    @builtin(position) clip_position: vec4<f32>,
+    @location(0) fragUV: vec2<f32>,
+    @location(1) color: vec3<f32>,
+}
+
+@group(0) @binding(0) var<uniform> view: View;
+@group(0) @binding(1) var<storage, read> particle_positions: array<vec2<f32>>;
+@group(0) @binding(2) var<storage, read> particle_colours: array<u32>;
+
+override size = 1.0;
+override sharpness = 1.0;
+
+const BOUNDS: Rect = Rect(
+    vec2<f32>(-4.0 * 1920.0 / 2.0, -4.0 * 1080.0 / 2.0),
+    vec2<f32>(4.0 * 1920.0 / 2.0, 4.0 * 1080.0 / 2.0)
+);
+
+fn wrap_particle_position(particle_pos: vec2<f32>, camera_pos: vec2<f32>) -> vec2<f32> {
+    let width = BOUNDS.max.x - BOUNDS.min.x;
+    let height = BOUNDS.max.y - BOUNDS.min.y;
+
+    var wrapped_pos = particle_pos;
+
+    // Calculate the displacement from camera to particle
+    let displacement = toroidal_displacement(BOUNDS, camera_pos, particle_pos);
+    let direct_displacement = particle_pos - camera_pos;
+
+    // If the shortest path crosses a boundary, render at the wrapped position
+    if abs(displacement.x) < abs(direct_displacement.x) {
+        if direct_displacement.x > 0.0 && displacement.x < 0.0 {
+            wrapped_pos.x -= width;
+        } else if direct_displacement.x < 0.0 && displacement.x > 0.0 {
+            wrapped_pos.x += width;
+        }
+    }
+
+    if abs(displacement.y) < abs(direct_displacement.y) {
+        if direct_displacement.y > 0.0 && displacement.y < 0.0 {
+            wrapped_pos.y -= height;
+        } else if direct_displacement.y < 0.0 && displacement.y > 0.0 {
+            wrapped_pos.y += height;
+        }
+    }
+
+    return wrapped_pos;
+}
+
+@vertex
+fn vertex(input: VertexInput) -> VertexOutput {
+    var out: VertexOutput;
+
+    let raw_particle_pos = particle_positions[input.instance_index];
+
+    // Extract camera position from the existing View uniform
+    let camera_pos = view.world_position.xy;
+
+    // Wrap the particle position for seamless toroidal rendering
+    let wrapped_particle_pos = wrap_particle_position(raw_particle_pos, camera_pos);
+
+    let world_position = vec4<f32>(input.position.xy * size + wrapped_particle_pos, 0.0, 1.0);
+    out.clip_position = view.clip_from_world * world_position;
+
+    out.color = COLOURS[particle_colours[input.instance_index]];
+    out.fragUV = input.position.xy;
+
+    return out;
+}
+
+@fragment
+fn fragment(input: VertexOutput) -> @location(0) vec4<f32> {
+    let alpha = clamp(smoothstep(1.0, 0.0, sharpness * (dot(input.fragUV, input.fragUV) * 2.0 - 1.0)), 0.0, 1.0);
+
+    return vec4<f32>(input.color, alpha);
+}

abiogenesis/assets/shaders/utils/colours.wgsl

@@ -0,0 +1,21 @@
+#define_import_path utils::colours
+
+const RED: vec3<f32> = vec3<f32>(172.0 / 255.0, 40.0 / 255.0, 71.0 / 255.0);
+const GREEN: vec3<f32> = vec3<f32>(90.0 / 255.0, 181.0 / 255.0, 82.0 / 255.0);
+const BLUE: vec3<f32> = vec3<f32>(51.0 / 255.0, 136.0 / 255.0, 222.0 / 255.0);
+const ORANGE: vec3<f32> = vec3<f32>(255.0 / 255.0, 155.0 / 255.0, 37.0 / 255.0);
+const PINK: vec3<f32> = vec3<f32>(233.0 / 255.0, 75.0 / 255.0, 234.0 / 255.0);
+const AQUA: vec3<f32> = vec3<f32>(57.0 / 255.0, 247.0 / 255.0, 241.0 / 255.0);
+
+const COLOURS: array<vec3<f32>, NUM_COLOURS> = array<vec3<f32>, NUM_COLOURS>(RED, GREEN, BLUE, ORANGE, PINK, AQUA);
+
+const NUM_COLOURS: u32 = 6;
+
+const RED_ID: u32 = 0;
+const GREEN_ID: u32 = 1;
+const BLUE_ID: u32 = 2;
+const ORANGE_ID: u32 = 3;
+const PINK_ID: u32 = 4;
+const AQUA_ID: u32 = 5;
+
+const EMPTY_ID: u32 = 0 << 1;

abiogenesis/assets/shaders/utils/math.wgsl

@@ -0,0 +1,94 @@
+#define_import_path utils::math
+
+fn lcg(seed: u32) -> u32 {
+    var result = seed;
+
+    result ^= 2747636419u;
+    result *= 2654435769u;
+    result ^= result >> 16;
+    result *= 2654435769u;
+    result ^= result >> 16;
+    result *= 2654435769u;
+
+    return result;
+}
+
+fn random_float(seed: u32) -> f32 {
+    let rand = lcg(seed);
+    return f32(rand) / f32(0xffffffffu);
+}
+
+fn clamp_length(value: vec2<f32>, min: f32, max: f32) -> vec2<f32> {
+    let len = length(value);
+
+    if len == 0.0 {
+        return vec2f(0.0);
+    } else {
+        return value * clamp(len, min, max) / len;
+    }
+}
+
+fn lerp(a: f32, b: f32, t: f32) -> f32 {
+    return a + (b - a) * t;
+}
+
+fn inverse_lerp(a: f32, b: f32, value: f32) -> f32 {
+    return (value - a) / (b - a);
+}
+
+fn remap(value: f32, a: f32, b: f32, c: f32, d: f32) -> f32 {
+    return lerp(c, d, inverse_lerp(a, b, value));
+}
+
+struct Rect {
+    min: vec2f,
+    max: vec2f,
+}
+
+fn toroidal_displacement(bounds: Rect, a: vec2f, b: vec2f) -> vec2f {
+    let width = bounds.max.x - bounds.min.x;
+    let height = bounds.max.y - bounds.min.y;
+
+    var dx = b.x - a.x;
+    var dy = b.y - a.y;
+
+    // Adjust x displacement for wrapping
+    if dx > width / 2.0 {
+        dx -= width;
+    } else if dx < -width / 2.0 {
+        dx += width;
+    }
+
+    // Adjust y displacement for wrapping
+    if dy > height / 2.0 {
+        dy -= height;
+    } else if dy < -height / 2.0 {
+        dy += height;
+    }
+
+    return vec2f(dx, dy);
+}
+
+fn toroidal_wrap(bounds: Rect, pos: vec2f) -> vec2f {
+    let width = bounds.max.x - bounds.min.x;
+    let height = bounds.max.y - bounds.min.y;
+
+    // Normalize position relative to bounds origin
+    let relative_x = pos.x - bounds.min.x;
+    let relative_y = pos.y - bounds.min.y;
+
+    // Use modulo operation to wrap coordinates
+    let wrapped_x = relative_x - floor(relative_x / width) * width;
+    let wrapped_y = relative_y - floor(relative_y / height) * height;
+
+    // Convert back to absolute coordinates
+    return vec2f(
+        wrapped_x + bounds.min.x,
+        wrapped_y + bounds.min.y
+    );
+}
+
+fn safe_normalize(v: vec2<f32>) -> vec2<f32> {
+    let len = length(v);
+    return select(v / len, vec2f(0.0), len == 0.0);
+}

abiogenesis/src/bundle_fn.rs

@@ -5,18 +5,18 @@ use bevy::prelude::*;
 pub trait Thunk: FnOnce(&mut EntityWorldMut) + Send + Sync + 'static {}
 impl<F: FnOnce(&mut EntityWorldMut) + Send + Sync + 'static> Thunk for F {}
 
-pub struct BundleFn<F: Thunk>(pub F);
+pub struct BundleFn<T: Thunk>(pub T);
 
-unsafe impl<F: Thunk> Bundle for BundleFn<F> {
+unsafe impl<T: Thunk> Bundle for BundleFn<T> {
     fn component_ids(_: &mut ComponentsRegistrator, _: &mut impl FnMut(ComponentId)) {}
 
     fn get_component_ids(_: &Components, _: &mut impl FnMut(Option<ComponentId>)) {}
 
     fn register_required_components(_: &mut ComponentsRegistrator, _: &mut RequiredComponents) {}
 }
 
-impl<F: Thunk> DynamicBundle for BundleFn<F> {
-    type Effect = BundleFn<F>;
+impl<T: Thunk> DynamicBundle for BundleFn<T> {
+    type Effect = BundleFn<T>;
 
     fn get_components(
         self,
@@ -26,7 +26,7 @@ impl<F: Thunk> DynamicBundle for BundleFn<F> {
     }
 }
 
-impl<F: Thunk> BundleEffect for BundleFn<F> {
+impl<T: Thunk> BundleEffect for BundleFn<T> {
     fn apply(self, entity: &mut EntityWorldMut) {
         (self.0)(entity);
     }