bevy_shader_plugin/lib.rs
1mod auto_array;
2mod bindings;
3mod compute;
4mod cpu_buffer;
5mod extra_bind_group;
6mod gpu_buffer;
7mod node;
8mod pipeline;
9mod readback;
10
11use bevy::{
12 core_pipeline::{Core3d, Core3dSystems},
13 prelude::*,
14 render::{
15 Render, RenderApp, RenderStartup, RenderSystems,
16 render_resource::{
17 BindGroup, BindGroupLayoutDescriptor, BindGroupLayoutEntries, ShaderStages, ShaderType,
18 binding_types::{storage_buffer_read_only_sized, storage_buffer_sized},
19 },
20 },
21};
22use encase::{
23 ShaderSize,
24 internal::{CreateFrom, WriteInto},
25};
26
27pub use auto_array::{ArrayBufferChanges, ArrayBufferState};
28pub use bindings::{
29 AutoBufferBindGroups, AutoBufferCompiledLayouts, AutoBufferKind, AutoBufferLayoutDescriptors,
30 AutoBufferLayouts, BindingTable, BoundResource,
31};
32pub use compute::{ComputePasses, ComputeShaderPlugin, Workgroups};
33pub use extra_bind_group::FragmentBindGroupBuilder;
34pub use gpu_buffer::GpuBuffer;
35pub use node::fullscreen_pass;
36pub use pipeline::{FullscreenPipeline, FullscreenPipelineConfig};
37pub use readback::{GpuReadback, ReadbackMode, StorageReadback};
38
39pub mod prelude {
40 pub use crate::{
41 ArrayBufferChanges, ComputePasses, ComputeShaderPlugin, FullscreenFragmentPlugin,
42 GpuBuffer, GpuReadback, ReadbackMode, ShaderAppExt, StorageReadback, Workgroups,
43 };
44 pub use bevy::render::render_resource::ShaderType;
45 pub use bevy::window::PrimaryWindow;
46}
47
48// ---------------------------------------------------------------------------
49// Re-exports for the fragment_layout! macro
50// ---------------------------------------------------------------------------
51
52#[doc(hidden)]
53pub mod __private {
54 pub use bevy::render::render_resource::{
55 BindGroupLayoutDescriptor, BindGroupLayoutEntries, ShaderStages,
56 };
57}
58
59// ---------------------------------------------------------------------------
60// Public types
61// ---------------------------------------------------------------------------
62
63/// System sets for the binding machinery shared by the fragment and compute plugins.
64#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
65pub enum ShaderSystems {
66 /// Runs in `RenderStartup`. Compiles the bind group layouts of all registered buffers.
67 /// Pipelines are created after this set.
68 CompileLayouts,
69 /// Runs in `Render`, inside [`RenderSystems::PrepareBindGroups`]. Uploads changed
70 /// buffers and records their handles in the [`BindingTable`].
71 PrepareBindings,
72 /// Runs in `Render`, inside [`RenderSystems::PrepareBindGroups`], after
73 /// [`PrepareBindings`](Self::PrepareBindings). Rebuilds bind groups whose bindings
74 /// changed.
75 FinalizeBindGroups,
76}
77
78/// System set label for the pipeline initialisation system so users can order
79/// their static-resource init systems before it.
80#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
81pub enum FragmentSystems {
82 /// Runs in `RenderStartup`. All user systems that populate
83 /// `FragmentExtraLayouts` must be ordered before this set.
84 InitPipeline,
85}
86
87/// Bind group layout descriptors for the manual extra bind groups.
88///
89/// Populate this resource in a layout system registered in `RenderStartup`
90/// before [`FragmentSystems::InitPipeline`].
91#[derive(Resource, Default)]
92pub struct FragmentExtraLayouts(pub Vec<BindGroupLayoutDescriptor>);
93
94impl FragmentExtraLayouts {
95 /// Push an arbitrary [`BindGroupLayoutDescriptor`].
96 pub fn push(&mut self, desc: BindGroupLayoutDescriptor) -> &mut Self {
97 self.0.push(desc);
98 self
99 }
100
101 /// Add a group with a single read-only storage buffer at binding 0.
102 pub fn storage_buffer_read_only(&mut self, label: &'static str) -> &mut Self {
103 self.0.push(BindGroupLayoutDescriptor::new(
104 label,
105 &BindGroupLayoutEntries::single(
106 ShaderStages::FRAGMENT,
107 storage_buffer_read_only_sized(false, None),
108 ),
109 ));
110 self
111 }
112
113 /// Add a group with a single read-write storage buffer at binding 0.
114 pub fn storage_buffer_read_write(&mut self, label: &'static str) -> &mut Self {
115 self.0.push(BindGroupLayoutDescriptor::new(
116 label,
117 &BindGroupLayoutEntries::single(
118 ShaderStages::FRAGMENT,
119 storage_buffer_sized(false, None),
120 ),
121 ));
122 self
123 }
124}
125
126/// Per-frame bind groups for manual extra groups.
127///
128/// Call [`clear`](Self::clear) at the start of your `PrepareBindGroups` system,
129/// then use [`push`](Self::push) or [`FragmentBindGroupBuilder`] to populate each
130/// group in order.
131#[derive(Resource, Default)]
132pub struct FragmentExtraBindGroups(pub Vec<BindGroup>);
133
134impl FragmentExtraBindGroups {
135 /// Push a pre-built bind group.
136 pub fn push(&mut self, bind_group: BindGroup) -> &mut Self {
137 self.0.push(bind_group);
138 self
139 }
140
141 /// Remove all bind groups. Call at the start of your `PrepareBindGroups` system.
142 pub fn clear(&mut self) -> &mut Self {
143 self.0.clear();
144 self
145 }
146}
147
148/// System sets for the GPU passes this crate records, so your own render systems can be
149/// ordered around them.
150#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
151pub enum ShaderPassSystems {
152 /// The [`ComputeShaderPlugin`] passes. Runs in the root
153 /// [`RenderGraph`] schedule, before any camera
154 /// renders.
155 Compute,
156 /// The [`FullscreenFragmentPlugin`] draw. Runs in each 3D camera's
157 /// [`Core3d`] schedule, after the prepass and before the main pass.
158 Fullscreen,
159 /// Buffer copies for readback. Runs in the root
160 /// [`RenderGraph`] schedule, after every camera
161 /// has rendered.
162 ReadbackCopy,
163}
164
165// ---------------------------------------------------------------------------
166// App extension trait
167// ---------------------------------------------------------------------------
168
169/// Extension methods on [`App`] for registering buffers shared by the fullscreen
170/// fragment shader and the compute shader.
171///
172/// Every registered buffer is visible to both stages at the same `@group`/`@binding`.
173/// Declare each binding with the same access mode (`read` or `read_write`) in every
174/// shader that uses it — wgpu rejects a mismatch.
175pub trait ShaderAppExt {
176 /// Register an auto-managed uniform buffer at `@group(group_index) @binding(binding_index)`.
177 ///
178 /// `U` must be inserted as a [`Resource`] in the main world. The library extracts
179 /// it to the render world and uploads it to a persistent uniform buffer whenever
180 /// the resource changes.
181 ///
182 /// Multiple calls with the same `group_index` but different `binding_index` values
183 /// pack several uniform bindings into one bind group.
184 ///
185 /// # Example
186 ///
187 /// ```rust,ignore
188 /// #[derive(Resource, Default, ShaderType, Clone)]
189 /// struct TimeUniforms { elapsed: f32 }
190 ///
191 /// app.insert_resource(TimeUniforms::default())
192 /// .register_uniform_buffer::<TimeUniforms>(0, 0);
193 /// // WGSL: @group(0) @binding(0) var<uniform> time_uniforms: TimeUniforms;
194 /// ```
195 fn register_uniform_buffer<U>(&mut self, group_index: u32, binding_index: u32) -> &mut Self
196 where
197 U: ShaderType + WriteInto + Default + Resource + Clone + Send + Sync + 'static;
198
199 /// Register an auto-managed storage buffer at `@group(group_index) @binding(binding_index)`.
200 ///
201 /// `S` must be inserted as a [`Resource`] in the main world. The library extracts
202 /// it to the render world and uploads it to a persistent storage buffer whenever
203 /// the resource changes.
204 ///
205 /// Multiple calls with the same `group_index` but different `binding_index` values
206 /// pack several storage bindings into one bind group.
207 /// `read_write`: `false` → `var<storage, read>`, `true` → `var<storage, read_write>`.
208 ///
209 /// The CPU owns this buffer: shader writes persist until the resource next changes,
210 /// then the CPU upload replaces the whole buffer. For data the GPU owns, use
211 /// [`register_gpu_buffer`](Self::register_gpu_buffer).
212 ///
213 /// # Example
214 ///
215 /// ```rust,ignore
216 /// #[derive(Resource, Default, ShaderType, Clone)]
217 /// struct Palette { colors: [Vec4; 16] }
218 ///
219 /// app.insert_resource(Palette::default())
220 /// .register_storage_buffer::<Palette>(1, 0, false);
221 /// // WGSL: @group(1) @binding(0) var<storage, read> palette: Palette;
222 /// ```
223 fn register_storage_buffer<S>(
224 &mut self,
225 group_index: u32,
226 binding_index: u32,
227 read_write: bool,
228 ) -> &mut Self
229 where
230 S: ShaderType + WriteInto + Default + Resource + Clone + Send + Sync + 'static;
231
232 /// Register a fixed-size array buffer at `@group(group_index) @binding(binding_index)`.
233 ///
234 /// Maps to WGSL `array<T, N>`. The buffer is initialized once with `T::default()` values
235 /// and persists across frames — only elements explicitly changed via [`ArrayBufferChanges`]
236 /// are uploaded each frame, batched into contiguous `write_buffer` runs.
237 ///
238 /// `Tag` is a user-defined zero-sized marker type. Define one per registration so that
239 /// multiple buffers of the same element type and length can coexist:
240 ///
241 /// ```rust,ignore
242 /// struct Colors;
243 /// struct Positions;
244 /// app.register_array_buffer::<Colors, Vec4, 64>(1, 0, false);
245 /// app.register_array_buffer::<Positions, Vec4, 64>(2, 0, false);
246 ///
247 /// // Each system names only its tag — no T or N required:
248 /// fn update_colors(mut changes: ResMut<ArrayBufferChanges<Colors>>) {
249 /// changes.set(0, Vec4::ONE);
250 /// }
251 /// ```
252 ///
253 /// `read_write`: `false` → `var<storage, read>`, `true` → `var<storage, read_write>`.
254 fn register_array_buffer<Tag, T, const N: usize>(
255 &mut self,
256 group_index: u32,
257 binding_index: u32,
258 read_write: bool,
259 ) -> &mut Self
260 where
261 Tag: Send + Sync + 'static,
262 T: ShaderSize + WriteInto + Default + Send + Sync + 'static;
263
264 /// Register a GPU-owned buffer of `len` elements at
265 /// `@group(group_index) @binding(binding_index)`.
266 ///
267 /// Maps to WGSL `var<storage, read_write> name: array<T>` (always `read_write`). The
268 /// buffer starts filled with `T::default()` and is never uploaded from the CPU again,
269 /// so shader writes persist across frames — use it for simulation state that compute
270 /// passes update. Its handle is available in the main world as [`GpuBuffer<Tag>`].
271 ///
272 /// `Tag` is a user-defined marker type naming the buffer, also used to size
273 /// dispatches with [`Workgroups::over`]:
274 ///
275 /// ```rust,ignore
276 /// #[derive(ShaderType, Clone, Copy, Default)]
277 /// struct Particle { pos: Vec2, vel: Vec2 }
278 /// struct Particles;
279 ///
280 /// app.register_gpu_buffer::<Particles, Particle>(1, 0, 1024);
281 /// // WGSL: @group(1) @binding(0) var<storage, read_write> particles: array<Particle>;
282 /// ```
283 fn register_gpu_buffer<Tag, T>(
284 &mut self,
285 group_index: u32,
286 binding_index: u32,
287 len: u32,
288 ) -> &mut Self
289 where
290 Tag: Send + Sync + 'static,
291 T: ShaderSize + WriteInto + Default + Send + Sync + 'static;
292
293 /// Enable reading the GPU buffer tagged `Tag` back to the CPU, decoded as `T` values
294 /// into the [`GpuReadback<Tag, T>`] resource. Call after
295 /// [`register_gpu_buffer`](Self::register_gpu_buffer), with the same `T`.
296 ///
297 /// With [`ReadbackMode::OnRequest`], call [`GpuReadback::request`] whenever you want a
298 /// copy. With [`ReadbackMode::EveryFrame`], the buffer is copied every frame. Results
299 /// arrive 1–3 frames after the copy is made.
300 ///
301 /// ```rust,ignore
302 /// app.register_gpu_buffer::<Particles, Particle>(1, 0, 1024)
303 /// .read_back_gpu_buffer::<Particles, Particle>(ReadbackMode::OnRequest);
304 ///
305 /// fn request(readback: Res<GpuReadback<Particles, Particle>>) {
306 /// readback.request();
307 /// }
308 ///
309 /// fn print(readback: Res<GpuReadback<Particles, Particle>>) {
310 /// if !readback.is_changed() { return; }
311 /// let Some(particles) = readback.latest() else { return };
312 /// info!("{} particles", particles.len());
313 /// }
314 /// ```
315 fn read_back_gpu_buffer<Tag, T>(&mut self, mode: ReadbackMode) -> &mut Self
316 where
317 Tag: Send + Sync + 'static,
318 T: ShaderSize + CreateFrom + Send + Sync + 'static;
319
320 /// Enable reading the array buffer tagged `Tag` back to the CPU, decoded as `T` values
321 /// into the [`GpuReadback<Tag, T>`] resource. Call after
322 /// [`register_array_buffer`](Self::register_array_buffer), with the same `T`.
323 ///
324 /// Works exactly like [`read_back_gpu_buffer`](Self::read_back_gpu_buffer). Useful
325 /// when a shader writes to a `read_write` array buffer.
326 fn read_back_array_buffer<Tag, T>(&mut self, mode: ReadbackMode) -> &mut Self
327 where
328 Tag: Send + Sync + 'static,
329 T: ShaderSize + CreateFrom + Send + Sync + 'static;
330
331 /// Enable reading the storage buffer of resource type `S` back to the CPU, decoded as
332 /// an `S` into the [`StorageReadback<S>`] resource. Call after
333 /// [`register_storage_buffer`](Self::register_storage_buffer). If `S` is registered at
334 /// several bindings, the first storage registration is read.
335 ///
336 /// Useful when a shader writes to a `read_write` storage buffer. With
337 /// [`ReadbackMode::OnRequest`], call [`StorageReadback::request`] whenever you want a
338 /// copy. With [`ReadbackMode::EveryFrame`], the buffer is copied every frame.
339 ///
340 /// ```rust,ignore
341 /// app.register_storage_buffer::<Stats>(1, 0, true)
342 /// .init_resource::<Stats>()
343 /// .read_back_storage_buffer::<Stats>(ReadbackMode::EveryFrame);
344 ///
345 /// fn print(readback: Res<StorageReadback<Stats>>) {
346 /// let Some(stats) = readback.latest() else { return };
347 /// info!("{:?}", stats.total);
348 /// }
349 /// ```
350 fn read_back_storage_buffer<S>(&mut self, mode: ReadbackMode) -> &mut Self
351 where
352 S: ShaderType + WriteInto + CreateFrom + Default + Resource + Clone + Send + Sync + 'static;
353}
354
355impl ShaderAppExt for App {
356 fn register_uniform_buffer<U>(&mut self, group_index: u32, binding_index: u32) -> &mut Self
357 where
358 U: ShaderType + WriteInto + Default + Resource + Clone + Send + Sync + 'static,
359 {
360 cpu_buffer::register::<U>(self, group_index, binding_index, AutoBufferKind::Uniform);
361 self
362 }
363
364 fn register_storage_buffer<S>(
365 &mut self,
366 group_index: u32,
367 binding_index: u32,
368 read_write: bool,
369 ) -> &mut Self
370 where
371 S: ShaderType + WriteInto + Default + Resource + Clone + Send + Sync + 'static,
372 {
373 cpu_buffer::register::<S>(
374 self,
375 group_index,
376 binding_index,
377 AutoBufferKind::Storage {
378 read_only: !read_write,
379 },
380 );
381 self
382 }
383
384 fn register_array_buffer<Tag, T, const N: usize>(
385 &mut self,
386 group_index: u32,
387 binding_index: u32,
388 read_write: bool,
389 ) -> &mut Self
390 where
391 Tag: Send + Sync + 'static,
392 T: ShaderSize + WriteInto + Default + Send + Sync + 'static,
393 {
394 auto_array::register::<Tag, T, N>(self, group_index, binding_index, read_write);
395 self
396 }
397
398 fn register_gpu_buffer<Tag, T>(
399 &mut self,
400 group_index: u32,
401 binding_index: u32,
402 len: u32,
403 ) -> &mut Self
404 where
405 Tag: Send + Sync + 'static,
406 T: ShaderSize + WriteInto + Default + Send + Sync + 'static,
407 {
408 gpu_buffer::register::<Tag, T>(self, group_index, binding_index, len);
409 self
410 }
411
412 fn read_back_gpu_buffer<Tag, T>(&mut self, mode: ReadbackMode) -> &mut Self
413 where
414 Tag: Send + Sync + 'static,
415 T: ShaderSize + CreateFrom + Send + Sync + 'static,
416 {
417 readback::register_gpu::<Tag, T>(self, mode);
418 self
419 }
420
421 fn read_back_array_buffer<Tag, T>(&mut self, mode: ReadbackMode) -> &mut Self
422 where
423 Tag: Send + Sync + 'static,
424 T: ShaderSize + CreateFrom + Send + Sync + 'static,
425 {
426 readback::register_array::<Tag, T>(self, mode);
427 self
428 }
429
430 fn read_back_storage_buffer<S>(&mut self, mode: ReadbackMode) -> &mut Self
431 where
432 S: ShaderType + WriteInto + CreateFrom + Default + Resource + Clone + Send + Sync + 'static,
433 {
434 readback::register_storage::<S>(self, mode);
435 self
436 }
437}
438
439// ---------------------------------------------------------------------------
440// fragment_layout! macro
441// ---------------------------------------------------------------------------
442
443/// Build a [`BindGroupLayoutDescriptor`] with sequential fragment-stage bindings.
444#[macro_export]
445macro_rules! fragment_layout {
446 ($label:expr, $($entry:expr),+ $(,)?) => {
447 $crate::__private::BindGroupLayoutDescriptor::new(
448 $label,
449 &$crate::__private::BindGroupLayoutEntries::sequential(
450 $crate::__private::ShaderStages::FRAGMENT,
451 ($($entry,)+),
452 ),
453 )
454 };
455}
456
457// ---------------------------------------------------------------------------
458// Plugins
459// ---------------------------------------------------------------------------
460
461/// Binding machinery shared by [`FullscreenFragmentPlugin`] and [`ComputeShaderPlugin`]:
462/// layout compilation and bind group assembly for every registered buffer.
463struct ShaderCorePlugin;
464
465impl Plugin for ShaderCorePlugin {
466 fn build(&self, app: &mut App) {
467 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
468 return;
469 };
470
471 bindings::init_resources(render_app);
472
473 render_app
474 .configure_sets(
475 Render,
476 (
477 ShaderSystems::PrepareBindings,
478 ShaderSystems::FinalizeBindGroups,
479 )
480 .chain()
481 .in_set(RenderSystems::PrepareBindGroups),
482 )
483 .add_systems(
484 Render,
485 bindings::finalize_bind_groups.in_set(ShaderSystems::FinalizeBindGroups),
486 )
487 .add_systems(
488 RenderStartup,
489 bindings::compile_layouts.in_set(ShaderSystems::CompileLayouts),
490 );
491 }
492}
493
494/// Adds [`ShaderCorePlugin`] unless another plugin already did.
495fn add_core_plugin(app: &mut App) {
496 if !app.is_plugin_added::<ShaderCorePlugin>() {
497 app.add_plugins(ShaderCorePlugin);
498 }
499}
500
501/// Bevy plugin that wires up a fullscreen fragment shader pipeline.
502///
503/// Call the [`ShaderAppExt`] registration methods on the [`App`] to bind data to your
504/// shader.
505///
506/// This plugin is not compatible with MSAA. Disable MSAA on all cameras.
507///
508/// # Example
509///
510/// ```rust,ignore
511/// #[derive(Resource, Default, ShaderType, Clone)]
512/// struct MyUniforms { time: f32 }
513///
514/// App::new()
515/// .add_plugins(FullscreenFragmentPlugin::new("shaders/my_effect.wgsl"))
516/// .insert_resource(MyUniforms::default())
517/// .register_uniform_buffer::<MyUniforms>(0, 0)
518/// .run();
519/// ```
520pub struct FullscreenFragmentPlugin {
521 /// Path to the fragment shader asset (e.g. `"shaders/effect.wgsl"` or
522 /// `"shaders/effect.spv"`).
523 pub shader_path: &'static str,
524 /// Fragment entry point name. `None` auto-detects the module's only one.
525 pub entry_point: Option<&'static str>,
526}
527
528impl FullscreenFragmentPlugin {
529 /// Creates a new plugin that renders the shader at the given asset path.
530 pub fn new(shader_path: &'static str) -> Self {
531 Self {
532 shader_path,
533 entry_point: None,
534 }
535 }
536
537 /// Names the fragment entry point. Required when the shader module has more
538 /// than one.
539 ///
540 /// ```rust,ignore
541 /// FullscreenFragmentPlugin::new("shaders/my_shader.spv")
542 /// .with_entry_point("main_fs")
543 /// ```
544 pub fn with_entry_point(mut self, entry_point: &'static str) -> Self {
545 self.entry_point = Some(entry_point);
546 self
547 }
548}
549
550impl Plugin for FullscreenFragmentPlugin {
551 fn build(&self, app: &mut App) {
552 add_core_plugin(app);
553
554 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
555 return;
556 };
557
558 render_app.insert_resource(FullscreenPipelineConfig {
559 shader_path: self.shader_path,
560 entry_point: self.entry_point,
561 });
562
563 render_app.init_resource::<FragmentExtraLayouts>();
564 render_app.init_resource::<FragmentExtraBindGroups>();
565
566 render_app
567 .configure_sets(
568 RenderStartup,
569 FragmentSystems::InitPipeline.after(ShaderSystems::CompileLayouts),
570 )
571 .add_systems(
572 RenderStartup,
573 pipeline::init_pipeline.in_set(FragmentSystems::InitPipeline),
574 )
575 .add_systems(
576 Render,
577 pipeline::queue_fullscreen_pipelines.in_set(RenderSystems::Prepare),
578 );
579
580 // After the prepass and before the main pass, which draws on top. Only one
581 // `FullscreenFragmentPlugin` is supported per app.
582 render_app.add_systems(
583 Core3d,
584 fullscreen_pass
585 .in_set(ShaderPassSystems::Fullscreen)
586 .after(Core3dSystems::Prepass)
587 .before(Core3dSystems::MainPass),
588 );
589 }
590}