diff --git a/assets/shaders/recolor.wgsl b/assets/shaders/recolor.wgsl new file mode 100644 index 0000000..d3158fe --- /dev/null +++ b/assets/shaders/recolor.wgsl @@ -0,0 +1,19 @@ +struct Context { + color: vec4, +} + +@group(0) @binding(0) +var texture: texture_storage_2d; + +@group(0) @binding(1) +var context: Context; + +@compute @workgroup_size(8, 8, 1) +fn recolor(@builtin(global_invocation_id) id: vec3) { + let location = vec2(id.xy); + let value = textureLoad(texture, location); + let recolored_value = vec4(context.color[0], context.color[1], context.color[2], value[3]); + + storageBarrier(); + textureStore(texture, location, recolored_value); +} diff --git a/assets/shaders/simulation.wgsl b/assets/shaders/simulation.wgsl index 7b1c7ad..669c36b 100644 --- a/assets/shaders/simulation.wgsl +++ b/assets/shaders/simulation.wgsl @@ -1,5 +1,5 @@ struct Context { - pause: u32, + pause: u32, // align(4) width: u32, height: u32, speed: f32, @@ -9,6 +9,9 @@ struct Context { senseDistance: f32, turnSpeed: f32, turnRandomness: f32, + r: f32, + g: f32, + b: f32, } struct Agent { @@ -96,7 +99,7 @@ fn update(@builtin(global_invocation_id) id: vec3) { agents[id.x].position = newPosition; let location = vec2(agents[id.x].position); - let color = vec4(1.0, 1.0, 1.0, 1.0); + let color = vec4(context.r, context.g, context.b, 1.0); storageBarrier(); textureStore(textureOut, location, color); diff --git a/src/pipeline/fade.rs b/src/pipeline/fade.rs index ab643d1..1eedce1 100644 --- a/src/pipeline/fade.rs +++ b/src/pipeline/fade.rs @@ -48,7 +48,7 @@ impl SubShaderPipeline for FadeShaderPipeline { usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST, mapped_at_creation: false, }, - ), + ) ); } diff --git a/src/pipeline.rs b/src/pipeline/mod.rs similarity index 97% rename from src/pipeline.rs rename to src/pipeline/mod.rs index feae215..bcdca80 100644 --- a/src/pipeline.rs +++ b/src/pipeline/mod.rs @@ -9,11 +9,13 @@ use bevy::render::renderer::{RenderContext, RenderDevice, RenderQueue}; use crate::pipeline::blur::BlurShaderPipeline; use crate::pipeline::fade::FadeShaderPipeline; +use crate::pipeline::recolor::RecolorShaderPipeline; use crate::pipeline::simulation::SimulationShaderPipeline; use crate::plugin::{PluginSettings, PluginTime}; pub mod blur; pub mod fade; +pub mod recolor; pub mod simulation; pub struct MainShaderPipeline { @@ -27,6 +29,7 @@ impl FromWorld for MainShaderPipeline { Box::new(SimulationShaderPipeline::new(world)), Box::new(FadeShaderPipeline::new(world)), Box::new(BlurShaderPipeline::new(world)), + Box::new(RecolorShaderPipeline::new(world)), ], }; diff --git a/src/pipeline/recolor.rs b/src/pipeline/recolor.rs new file mode 100644 index 0000000..d99ddc0 --- /dev/null +++ b/src/pipeline/recolor.rs @@ -0,0 +1,148 @@ +use bevy::asset::Handle; +use bevy::core::{Pod, Zeroable}; +use bevy::prelude::{AssetServer, Image, World}; +use bevy::render::render_asset::RenderAssets; +use bevy::render::render_resource::*; +use bevy::render::renderer::{RenderDevice, RenderQueue}; + +use crate::pipeline::{get_compute_pipeline_id, PipelineData, SubShaderPipeline, WorkgroupSize}; +use crate::plugin::{PluginSettings, PluginTime}; +use crate::SETTINGS; + +pub struct RecolorShaderPipeline { + bind_group_layout: BindGroupLayout, + bind_group: Option, + pipeline: CachedComputePipelineId, + context: PipelineData, +} + +impl RecolorShaderPipeline { + pub fn new(world: &mut World) -> Self { + let bind_group_layout = get_bind_group_layout( + world.resource::(), + ); + + let shader = world.resource::().load("shaders/recolor.wgsl"); + + Self { + pipeline: get_compute_pipeline_id( + shader, + world.resource_mut::().as_mut(), + bind_group_layout.clone(), + "recolor shader update".to_string(), + "recolor".to_string(), + ), + bind_group_layout, + bind_group: None, + context: PipelineData::default(), + } + } +} + +impl SubShaderPipeline for RecolorShaderPipeline { + fn init_data(&mut self, render_device: &RenderDevice, _settings: &PluginSettings) { + self.context.buffer = Some(render_device + .create_buffer( + &BufferDescriptor { + label: Some("recolor context uniform buffer"), + size: std::mem::size_of::() as u64, + usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST, + mapped_at_creation: false, + }, + ) + ); + } + + fn prepare_data(&mut self, render_queue: &RenderQueue, settings: &PluginSettings, _time: &PluginTime) { + self.context.data = Some(RecolorPipelineContext { + color: settings.color.as_rgba_f32(), + }); + + render_queue.write_buffer( + self.context.buffer.as_ref().expect("context buffer to exist"), + 0, + bevy::core::cast_slice(&[ + self.context.data.expect("context data to exist"), + ]), + ); + } + + fn queue_bind_groups(&mut self, render_device: &RenderDevice, gpu_images: &RenderAssets, images: &Vec>) { + self.bind_group = Some( + render_device.create_bind_group( + &BindGroupDescriptor { + label: Some("recolor bind group"), + layout: &self.bind_group_layout, + entries: &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::TextureView( + &gpu_images[images.get(1).unwrap()].texture_view, + ), + }, + BindGroupEntry { + binding: 1, + resource: self.context.buffer + .as_ref() + .expect("context buffer to exist") + .as_entire_binding(), + }, + ], + }, + ) + ) + } + + fn get_pipeline(&self) -> CachedComputePipelineId { + self.pipeline + } + + fn get_bind_group(&self) -> Option<&BindGroup> { + self.bind_group.as_ref() + } + + fn get_workgroup_size(&self, _settings: &PluginSettings) -> WorkgroupSize { + WorkgroupSize { + x: SETTINGS.texture_size.0 / 8, + y: SETTINGS.texture_size.1 / 8, + z: 1, + } + } +} + +fn get_bind_group_layout(render_device: &RenderDevice) -> BindGroupLayout { + render_device + .create_bind_group_layout( + &BindGroupLayoutDescriptor { + label: Some("recolor bind group layout"), + entries: &[ + BindGroupLayoutEntry { + binding: 0, + visibility: ShaderStages::COMPUTE, + ty: BindingType::StorageTexture { + access: StorageTextureAccess::ReadWrite, + format: TextureFormat::Rgba8Unorm, + view_dimension: TextureViewDimension::D2, + }, + count: None, + }, + BindGroupLayoutEntry { + binding: 1, + visibility: ShaderStages::COMPUTE, + ty: BindingType::Buffer { + ty: BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: BufferSize::new(std::mem::size_of::() as u64), + }, + count: None, + }, + ], + }, + ) +} + +#[repr(C)] +#[derive(Copy, Clone, Default, Pod, Zeroable)] +struct RecolorPipelineContext { + color: [f32; 4], +} diff --git a/src/pipeline/simulation.rs b/src/pipeline/simulation.rs index c07079f..352faae 100644 --- a/src/pipeline/simulation.rs +++ b/src/pipeline/simulation.rs @@ -101,6 +101,9 @@ impl SubShaderPipeline for SimulationShaderPipeline { sense_distance: settings.agent_sense_distance, turn_speed: settings.agent_turn_speed, turn_randomness: settings.agent_turn_randomness, + r: settings.color.r(), + g: settings.color.g(), + b: settings.color.b(), }); render_queue.write_buffer( @@ -239,6 +242,9 @@ struct SimulationPipelineContext { sense_distance: f32, turn_speed: f32, turn_randomness: f32, + r: f32, + g: f32, + b: f32, } diff --git a/src/plugin.rs b/src/plugin.rs index 2386231..6497a31 100644 --- a/src/plugin.rs +++ b/src/plugin.rs @@ -59,7 +59,7 @@ fn create_images(mut commands: Commands, mut images: ResMut>) { depth_or_array_layers: 1, }, TextureDimension::D2, - &[0, 0, 0, 255], + &[0, 0, 0, 0], TextureFormat::Rgba8Unorm, ); @@ -104,6 +104,7 @@ pub struct PluginSettings { pub agent_turn_speed: f32, #[inspectable(min = 0.0, max = 2.0, speed = 0.05)] pub agent_turn_randomness: f32, + pub color: Color, pub has_trails: bool, #[inspectable(min = 0.0, max = 5.0, speed = 0.005)] pub fade_rate: f32, @@ -121,6 +122,7 @@ impl FromWorld for PluginSettings { agent_sense_distance: 20.0, agent_turn_speed: 1.0, agent_turn_randomness: 0.1, + color: Color::WHITE, has_trails: true, fade_rate: 0.15, blur_radius: 1,