bevy_shader_plugin/
bindings.rs

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

bevy_shader_plugin/
bindings.rs
