Skip to main content

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}