bevy_shader_plugin/
bindings.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    num::NonZeroU64,
4};
5
6use bevy::{
7    app::SubApp,
8    prelude::*,
9    render::{
10        render_resource::{
11            BindGroup, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor,
12            BindGroupLayoutEntry, BindingResource, BindingType, Buffer, BufferBindingType,
13            PipelineCache, ShaderStages,
14        },
15        renderer::RenderDevice,
16    },
17};
18
19/// Shader stages that can see auto-managed bindings. Every registered binding is shared
20/// by the fullscreen fragment pipeline and all compute pipelines.
21pub(crate) const AUTO_BINDING_VISIBILITY: ShaderStages =
22    ShaderStages::FRAGMENT.union(ShaderStages::COMPUTE);
23
24/// Whether a registered auto-managed buffer is a uniform or a storage buffer.
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum AutoBufferKind {
27    Uniform,
28    Storage { read_only: bool },
29}
30
31/// (group → binding → kind) for all auto-managed buffers.
32///
33/// Populated eagerly at app-build time by the [`ShaderAppExt`](crate::ShaderAppExt)
34/// registration methods. Read at startup to compile per-group layouts.
35#[derive(Resource, Default)]
36pub struct AutoBufferLayouts(pub BTreeMap<u32, BTreeMap<u32, AutoBufferKind>>);
37
38/// Minimum byte size of each registered `(group, binding)`, where known. Put in the
39/// layouts as `min_binding_size`, which lets wgpu check buffer sizes once when a bind
40/// group is created instead of on every draw and dispatch.
41#[derive(Resource, Default)]
42pub(crate) struct AutoBufferMinSizes(BTreeMap<(u32, u32), NonZeroU64>);
43
44/// Compiled [`BindGroupLayout`] per group, populated at startup.
45///
46/// Used in [`ShaderSystems::FinalizeBindGroups`](crate::ShaderSystems::FinalizeBindGroups) to create bind groups.
47#[derive(Resource, Default)]
48pub struct AutoBufferCompiledLayouts(pub BTreeMap<u32, BindGroupLayout>);
49
50/// Layout descriptors for every auto-managed group, in group order, populated at startup.
51///
52/// Both the fragment pipeline and the compute pipelines build their pipeline layouts
53/// from this list, so they agree on every group.
54#[derive(Resource, Default)]
55pub struct AutoBufferLayoutDescriptors(pub Vec<BindGroupLayoutDescriptor>);
56
57/// Assembled bind groups for all auto-managed buffers, keyed by WGSL group index.
58///
59/// Rebuilt in [`ShaderSystems::FinalizeBindGroups`](crate::ShaderSystems::FinalizeBindGroups) only when a binding in the group changes.
60/// Read by [`fullscreen_pass`](crate::fullscreen_pass) and the compute pass.
61#[derive(Resource, Default)]
62pub struct AutoBufferBindGroups(pub BTreeMap<u32, BindGroup>);
63
64/// A GPU resource bound at one `(group, binding)` slot.
65#[derive(Clone)]
66pub enum BoundResource {
67    Buffer(Buffer),
68    // Storage textures will add a `TextureView` variant here.
69}
70
71impl BoundResource {
72    fn binding(&self) -> BindingResource<'_> {
73        match self {
74            BoundResource::Buffer(buffer) => buffer.as_entire_binding(),
75        }
76    }
77}
78
79/// Persistent `(group, binding)` → GPU resource table for all auto-managed bindings.
80///
81/// Each registered buffer records its GPU handle here every frame. A group is only marked
82/// dirty (and its bind group rebuilt) when one of its handles actually changes, e.g. the
83/// first frame or after a buffer is reallocated to a larger size.
84#[derive(Resource, Default)]
85pub struct BindingTable {
86    entries: BTreeMap<(u32, u32), BoundResource>,
87    dirty_groups: BTreeSet<u32>,
88}
89
90impl BindingTable {
91    /// Record `buffer` at `(group, binding)`. Does nothing if that exact buffer is
92    /// already recorded there.
93    pub fn set_buffer(&mut self, group: u32, binding: u32, buffer: &Buffer) {
94        let unchanged = matches!(
95            self.entries.get(&(group, binding)),
96            Some(BoundResource::Buffer(existing)) if existing.id() == buffer.id()
97        );
98        if !unchanged {
99            self.entries
100                .insert((group, binding), BoundResource::Buffer(buffer.clone()));
101            self.dirty_groups.insert(group);
102        }
103    }
104}
105
106/// Initialise all binding resources in the render world. Idempotent.
107pub(crate) fn init_resources(render_app: &mut SubApp) {
108    render_app.init_resource::<AutoBufferLayouts>();
109    render_app.init_resource::<AutoBufferCompiledLayouts>();
110    render_app.init_resource::<AutoBufferLayoutDescriptors>();
111    render_app.init_resource::<AutoBufferBindGroups>();
112    render_app.init_resource::<BindingTable>();
113    render_app.init_resource::<AutoBufferMinSizes>();
114}
115
116/// Record that `(group_index, binding_index)` holds a buffer of `kind`, at least
117/// `min_size` bytes long.
118pub(crate) fn register_binding(
119    render_app: &mut SubApp,
120    group_index: u32,
121    binding_index: u32,
122    kind: AutoBufferKind,
123    min_size: Option<NonZeroU64>,
124) {
125    init_resources(render_app);
126    let world = render_app.world_mut();
127    world
128        .resource_mut::<AutoBufferLayouts>()
129        .0
130        .entry(group_index)
131        .or_default()
132        .insert(binding_index, kind);
133    let mut min_sizes = world.resource_mut::<AutoBufferMinSizes>();
134    match min_size {
135        Some(size) => min_sizes.0.insert((group_index, binding_index), size),
136        None => min_sizes.0.remove(&(group_index, binding_index)),
137    };
138}
139
140/// `RenderStartup` system. Builds one bind group layout per registered group.
141pub(crate) fn compile_layouts(
142    pipeline_cache: Res<PipelineCache>,
143    auto_buffer_layouts: Res<AutoBufferLayouts>,
144    min_sizes: Res<AutoBufferMinSizes>,
145    mut compiled_layouts: ResMut<AutoBufferCompiledLayouts>,
146    mut descriptors: ResMut<AutoBufferLayoutDescriptors>,
147) {
148    // Validate: registered group indices must be contiguous (no gaps).
149    let keys: Vec<u32> = auto_buffer_layouts.0.keys().cloned().collect();
150    debug_assert!(
151        keys.windows(2).all(|w| w[1] == w[0] + 1),
152        "register_uniform_buffer/register_storage_buffer/register_array_buffer/register_gpu_buffer group indices must be contiguous (no gaps)"
153    );
154
155    for (&group_index, binding_map) in auto_buffer_layouts.0.iter() {
156        let entries: Vec<BindGroupLayoutEntry> = binding_map
157            .iter()
158            .map(|(&binding, &kind)| BindGroupLayoutEntry {
159                binding,
160                visibility: AUTO_BINDING_VISIBILITY,
161                ty: BindingType::Buffer {
162                    ty: match kind {
163                        AutoBufferKind::Uniform => BufferBindingType::Uniform,
164                        AutoBufferKind::Storage { read_only } => {
165                            BufferBindingType::Storage { read_only }
166                        }
167                    },
168                    has_dynamic_offset: false,
169                    min_binding_size: min_sizes.0.get(&(group_index, binding)).copied(),
170                },
171                count: None,
172            })
173            .collect();
174        let desc = BindGroupLayoutDescriptor::new("auto_buffer_layout", &entries);
175        compiled_layouts
176            .0
177            .insert(group_index, pipeline_cache.get_bind_group_layout(&desc));
178        descriptors.0.push(desc);
179    }
180}
181
182/// Rebuilds the bind group of every dirty group whose registered bindings are all present.
183/// Groups that are still missing a binding stay dirty and are retried next frame.
184pub(crate) fn finalize_bind_groups(
185    mut auto_bind_groups: ResMut<AutoBufferBindGroups>,
186    mut table: ResMut<BindingTable>,
187    compiled_layouts: Res<AutoBufferCompiledLayouts>,
188    auto_layouts: Res<AutoBufferLayouts>,
189    render_device: Res<RenderDevice>,
190) {
191    if table.dirty_groups.is_empty() {
192        return;
193    }
194    let BindingTable {
195        entries,
196        dirty_groups,
197    } = &mut *table;
198
199    dirty_groups.retain(|&group_index| {
200        let (Some(expected), Some(layout)) = (
201            auto_layouts.0.get(&group_index),
202            compiled_layouts.0.get(&group_index),
203        ) else {
204            return true;
205        };
206        let Some(bind_entries) = expected
207            .keys()
208            .map(|&binding| {
209                entries
210                    .get(&(group_index, binding))
211                    .map(|resource| BindGroupEntry {
212                        binding,
213                        resource: resource.binding(),
214                    })
215            })
216            .collect::<Option<Vec<_>>>()
217        else {
218            return true;
219        };
220
221        let bind_group =
222            render_device.create_bind_group("auto_buffer_bind_group", layout, &bind_entries);
223        auto_bind_groups.0.insert(group_index, bind_group);
224        false
225    });
226}
227
228/// `true` once every registered group has a bind group.
229pub(crate) fn all_groups_bound(
230    layouts: &AutoBufferLayouts,
231    bind_groups: &AutoBufferBindGroups,
232) -> bool {
233    layouts
234        .0
235        .keys()
236        .all(|group| bind_groups.0.contains_key(group))
237}
bevy_shader_plugin/bindings.rs

bevy_shader_plugin/
bindings.rs
