Add computeparticles rust shader

Author: LegNeato

shaders/rust/Cargo.toml

@@ -74,6 +74,7 @@ members = [
     "pushdescriptors/cube",
     "descriptorbuffer/cube",
     "dynamicrendering/texture",
+    "computeparticles/particle",
 ]
 
 [workspace.dependencies]

shaders/rust/computeparticles/particle.comp.spv

@@ -0,0 +1,10 @@
+[package]
+name = "computeparticles-particle"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+spirv-std = { workspace = true }
+
+[lib]
+crate-type = ["dylib"]

shaders/rust/computeparticles/particle.frag.spv

@@ -0,0 +1,102 @@
+#![cfg_attr(target_arch = "spirv", no_std)]
+#![allow(clippy::missing_safety_doc)]
+
+use spirv_std::{spirv, glam::{vec2, vec4, Vec2, Vec4, UVec3}, Image, num_traits::Float};
+use spirv_std::image::SampledImage;
+
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct Particle {
+    pub pos: Vec2,
+    pub vel: Vec2,
+    pub gradient_pos: Vec4,
+}
+
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub struct Ubo {
+    pub delta_t: f32,
+    pub dest_x: f32,
+    pub dest_y: f32,
+    pub particle_count: i32,
+}
+
+fn attraction(pos: Vec2, attract_pos: Vec2) -> Vec2 {
+    let delta = attract_pos - pos;
+    const DAMP: f32 = 0.5;
+    let d_damped_dot = delta.dot(delta) + DAMP;
+    let inv_dist = 1.0 / d_damped_dot.sqrt();
+    let inv_dist_cubed = inv_dist * inv_dist * inv_dist;
+    delta * inv_dist_cubed * 0.0035
+}
+
+fn repulsion(pos: Vec2, attract_pos: Vec2) -> Vec2 {
+    let delta = attract_pos - pos;
+    let target_distance = delta.dot(delta).sqrt();
+    delta * (1.0 / (target_distance * target_distance * target_distance)) * -0.000035
+}
+
+#[spirv(compute(threads(256, 1, 1)))]
+pub fn main_cs(
+    #[spirv(global_invocation_id)] global_id: UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] particles: &mut [Particle],
+    #[spirv(uniform, descriptor_set = 0, binding = 1)] ubo: &Ubo,
+) {
+    let index = global_id.x;
+    if index >= ubo.particle_count as u32 {
+        return;
+    }
+
+    let idx = index as usize;
+    let mut vel = particles[idx].vel;
+    let mut pos = particles[idx].pos;
+
+    let dest_pos = vec2(ubo.dest_x, ubo.dest_y);
+
+    let delta = dest_pos - pos;
+    let target_distance = delta.dot(delta).sqrt();
+    vel += repulsion(pos, dest_pos) * 0.05;
+
+    pos += vel * ubo.delta_t;
+
+    if pos.x < -1.0 || pos.x > 1.0 || pos.y < -1.0 || pos.y > 1.0 {
+        vel = (-vel * 0.1) + attraction(pos, dest_pos) * 12.0;
+    } else {
+        particles[idx].pos = pos;
+    }
+
+    particles[idx].vel = vel;
+    particles[idx].gradient_pos.x += 0.02 * ubo.delta_t;
+    if particles[idx].gradient_pos.x > 1.0 {
+        particles[idx].gradient_pos.x -= 1.0;
+    }
+}
+
+#[spirv(vertex)]
+pub fn main_vs(
+    in_pos: Vec2,
+    in_gradient_pos: Vec4,
+    #[spirv(position)] out_position: &mut Vec4,
+    #[spirv(point_size)] out_point_size: &mut f32,
+    out_color: &mut Vec4,
+    out_gradient_pos: &mut f32,
+) {
+    *out_point_size = 8.0;
+    *out_color = vec4(0.035, 0.035, 0.035, 1.0);
+    *out_gradient_pos = in_gradient_pos.x;
+    *out_position = vec4(in_pos.x, in_pos.y, 1.0, 1.0);
+}
+
+#[spirv(fragment)]
+pub fn main_fs(
+    _in_color: Vec4,
+    in_gradient_pos: f32,
+    #[spirv(point_coord)] point_coord: Vec2,
+    #[spirv(descriptor_set = 0, binding = 0)] sampler_color_map: &SampledImage<Image!(2D, type=f32, sampled)>,
+    #[spirv(descriptor_set = 0, binding = 1)] sampler_gradient_ramp: &SampledImage<Image!(2D, type=f32, sampled)>,
+    out_frag_color: &mut Vec4,
+) {
+    let color = sampler_gradient_ramp.sample(vec2(in_gradient_pos, 0.0)).truncate();
+    let tex_color = sampler_color_map.sample(point_coord).truncate();
+    *out_frag_color = vec4(tex_color.x * color.x, tex_color.y * color.y, tex_color.z * color.z, 1.0);
+}