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