wde_renderer/assets/bindings/
render_binding.rs

1use wde_logger::prelude::*;
2
3use crate::{assets::bindings::render_data::RenderDataRecreated, prelude::*, sync::{ExtractResource, ExtractResourcePlugin}};
4use bevy::{
5    ecs::system::{
6        ReadOnlySystemParam, SystemParamItem, SystemState,
7        lifetimeless::{SRes, SResMut}
8    },
9    prelude::*
10};
11use std::{any::TypeId, collections::HashSet, marker::PhantomData};
12use wde_wgpu::pipelines::{BindGroup, BindGroupBuilder, BindGroupLayout};
13
14// Reexport wgpu types
15pub use wde_wgpu::buffer::{BufferBindingType, BufferUsage};
16
17// ============= RENDER DATA PLUGIN REGISTER =============
18#[derive(Resource)]
19pub struct RenderBindingHolder<R: RenderBinding + TypePath + Sync + Send + Asset>(pub Handle<R>);
20impl<R: RenderBinding + TypePath + Sync + Send + Asset> Clone for RenderBindingHolder<R> {
21    fn clone(&self) -> Self {
22        Self(self.0.clone())
23    }
24}
25impl<R: RenderBinding + TypePath + Sync + Send + Asset> ExtractResource for RenderBindingHolder<R> {
26    type Source = Self;
27
28    fn extract(source: &Self::Source) -> Self {
29        source.clone()
30    }
31}
32
33/// A plugin to register a render binding.
34///
35/// This should be added to the app to initialize the [`RenderBinding`] asset and create the
36/// corresponding [`GpuRenderBinding`].
37///
38/// It also tracks [`RenderData`](crate::assets::bindings::RenderData) dependencies declared in
39/// [`RenderBinding::describe`], and recreates the binding asset when one of those dependencies
40/// is recreated.
41///
42/// The render binding type must implement [`RenderBinding`] and usually derive [`Asset`],
43/// [`TypePath`], [`Clone`] and [`Default`].
44pub struct RenderBindingRegisterPlugin<R: RenderBinding>(std::marker::PhantomData<R>);
45impl<R: RenderBinding> Default for RenderBindingRegisterPlugin<R> {
46    fn default() -> Self {
47        RenderBindingRegisterPlugin(std::marker::PhantomData)
48    }
49}
50impl<R: RenderBinding + TypePath + Sync + Send + Clone + Asset + Default> Plugin
51    for RenderBindingRegisterPlugin<R>
52{
53    fn build(&self, app: &mut App) {
54        app.init_asset::<R>()
55            .add_plugins((
56                RenderAssetsPlugin::<GpuRenderBinding<R>>::default(),
57                ExtractResourcePlugin::<RenderBindingHolder<R>>::default()
58            ))
59            .add_systems(Update, on_dependency_recreate::<R>);
60    }
61
62    fn finish(&self, app: &mut App) {
63        let mut default_res = R::default();
64
65        // Register dependencies
66        let builder = {
67            let world = app.get_sub_app_mut(RenderApp).unwrap().world_mut();
68            let mut builder = RenderBindingBuilder::default();
69            let mut state: SystemState<R::Params> = SystemState::new(world);
70            let params = state.get_mut(world);
71            R::describe(&mut default_res, &params, &mut builder);
72            builder
73        };
74        app.world_mut()
75            .insert_resource(RenderBindingDependencies::<R> {
76                _phantom: PhantomData,
77                dependencies: builder.dependencies
78            });
79
80        // Create the asset and insert the handle in the resource
81        let binding: Handle<R> = app.world_mut().add_asset(default_res);
82        app.world_mut()
83            .insert_resource(RenderBindingHolder(binding));
84    }
85}
86
87#[derive(Resource)]
88struct RenderBindingDependencies<R: RenderBinding> {
89    _phantom: PhantomData<R>,
90    dependencies: HashSet<TypeId>
91}
92fn on_dependency_recreate<R: RenderBinding + Asset + Default>(
93    asset_server: Res<AssetServer>,
94    mut dependency_update: MessageReader<RenderDataRecreated>,
95    mut render_binding_holder: ResMut<RenderBindingHolder<R>>,
96    dependencies: Res<RenderBindingDependencies<R>>
97) {
98    for message in dependency_update.read() {
99        let type_id = message.0;
100
101        // Check if the render binding depends on the render data that have been recreated
102        if !dependencies.dependencies.contains(&type_id) {
103            return;
104        }
105
106        // Recreate asset
107        debug!(
108            "Recreating render binding {} because of dependency render data change.",
109            std::any::type_name::<R>()
110        );
111        render_binding_holder.0 = asset_server.add(R::default());
112    }
113}
114
115// ============= RENDER BINDING AND BUILDER METHODS =============
116enum RenderBindingType {
117    Buffer,
118    TextureView,
119    TextureArrayView,
120    TextureSampler,
121    StorageTexture
122}
123
124/// A builder for the render binding. This is used in the description of the render binding to specify the buffers and textures that should be created on the gpu, and to build the bind group layout and bind group entries.
125#[derive(Default)]
126pub struct RenderBindingBuilder {
127    is_none: bool, // True if there is at least one buffer or texture that is not ready
128    elements: Vec<(RenderBindingType, u32)>, // (binding type, buffer/texture index in the corresponding vector)
129    buffers: Vec<Option<AssetId<Buffer>>>,
130    textures: Vec<Option<AssetId<Texture>>>,
131    dependencies: HashSet<TypeId>
132}
133impl RenderBindingBuilder {
134    /// Adds a buffer to the render binding. The `render_data` parameter is the [`RenderData`](RenderData) that contains the buffer, and the `render_data_idx` parameter is the index of the buffer in this render data.
135    pub fn add_buffer<R>(
136        &mut self,
137        render_data: &RenderAssets<GpuRenderData<R>>,
138        render_data_idx: u32
139    ) -> &mut Self
140    where
141        R: RenderData + Clone + Asset + 'static
142    {
143        self.dependencies.insert(TypeId::of::<R>());
144        let buffer = render_data
145            .iter()
146            .next()
147            .and_then(|(_, d)| d.get_buffer(render_data_idx))
148            .map(|b| b.id());
149        if buffer.is_none() {
150            self.is_none = true;
151        }
152        self.buffers.push(buffer);
153        self.elements
154            .push((RenderBindingType::Buffer, self.buffers.len() as u32 - 1));
155        self
156    }
157    /// Adds a buffer directly from an asset id.
158    pub fn add_buffer_from_id(&mut self, buffer: Option<AssetId<Buffer>>) -> &mut Self {
159        if buffer.is_none() {
160            self.is_none = true;
161        }
162        self.buffers.push(buffer);
163        self.elements
164            .push((RenderBindingType::Buffer, self.buffers.len() as u32 - 1));
165        self
166    }
167    /// Adds a texture view to the render binding. The `render_data` parameter is the [`RenderData`](RenderData) that contains the texture, and the `render_data_idx` parameter is the index of the texture in this render data.
168    pub fn add_texture_view<R>(
169        &mut self,
170        render_data: &RenderAssets<GpuRenderData<R>>,
171        render_data_idx: u32
172    ) -> &mut Self
173    where
174        R: RenderData + Clone + Asset + 'static
175    {
176        self.add_texture(render_data, render_data_idx, RenderBindingType::TextureView)
177    }
178    /// Adds a texture view directly from an asset id.
179    pub fn add_texture_view_from_id(&mut self, texture: Option<AssetId<Texture>>) -> &mut Self {
180        self.add_texture_from_id(texture, RenderBindingType::TextureView)
181    }
182    /// Adds a texture array view to the render binding. The `render_data` parameter is the [`RenderData`](RenderData) that contains the texture, and the `render_data_idx` parameter is the index of the texture in this render data.
183    pub fn add_texture_array_view<R>(
184        &mut self,
185        render_data: &RenderAssets<GpuRenderData<R>>,
186        render_data_idx: u32
187    ) -> &mut Self
188    where
189        R: RenderData + Clone + Asset + 'static
190    {
191        self.add_texture(
192            render_data,
193            render_data_idx,
194            RenderBindingType::TextureArrayView
195        )
196    }
197    /// Adds a texture array view directly from an asset id.
198    pub fn add_texture_array_view_from_id(
199        &mut self,
200        texture: Option<AssetId<Texture>>
201    ) -> &mut Self {
202        self.add_texture_from_id(texture, RenderBindingType::TextureArrayView)
203    }
204    /// Adds a texture sampler to the render binding. The `render_data` parameter is the [`RenderData`](RenderData) that contains the texture, and the `render_data_idx` parameter is the index of the texture in this render data.
205    pub fn add_texture_sampler<R>(
206        &mut self,
207        render_data: &RenderAssets<GpuRenderData<R>>,
208        render_data_idx: u32
209    ) -> &mut Self
210    where
211        R: RenderData + Clone + Asset + 'static
212    {
213        self.add_texture(
214            render_data,
215            render_data_idx,
216            RenderBindingType::TextureSampler
217        )
218    }
219    /// Adds a texture sampler directly from an asset id.
220    pub fn add_texture_sampler_from_id(&mut self, texture: Option<AssetId<Texture>>) -> &mut Self {
221        self.add_texture_from_id(texture, RenderBindingType::TextureSampler)
222    }
223    /// Adds a storage texture view to the render binding. The `render_data` parameter is the [`RenderData`](RenderData) that contains the texture, and the `render_data_idx` parameter is the index of the texture in this render data.
224    pub fn add_storage_texture<R>(
225        &mut self,
226        render_data: &RenderAssets<GpuRenderData<R>>,
227        render_data_idx: u32
228    ) -> &mut Self
229    where
230        R: RenderData + Clone + Asset + 'static
231    {
232        self.add_texture(
233            render_data,
234            render_data_idx,
235            RenderBindingType::StorageTexture
236        )
237    }
238    /// Adds a storage texture view directly from an asset id.
239    pub fn add_storage_texture_from_id(&mut self, texture: Option<AssetId<Texture>>) -> &mut Self {
240        self.add_texture_from_id(texture, RenderBindingType::StorageTexture)
241    }
242
243    fn add_texture<R>(
244        &mut self,
245        render_data: &RenderAssets<GpuRenderData<R>>,
246        render_data_idx: u32,
247        binding_type: RenderBindingType
248    ) -> &mut Self
249    where
250        R: RenderData + Clone + Asset + 'static
251    {
252        self.dependencies.insert(TypeId::of::<R>());
253        let texture = render_data
254            .iter()
255            .next()
256            .and_then(|(_, d)| d.get_texture(render_data_idx))
257            .map(|t| t.id());
258        if texture.is_none() {
259            self.is_none = true;
260        }
261        self.textures.push(texture);
262        self.elements
263            .push((binding_type, self.textures.len() as u32 - 1));
264        self
265    }
266    fn add_texture_from_id(
267        &mut self,
268        texture: Option<AssetId<Texture>>,
269        binding_type: RenderBindingType
270    ) -> &mut Self {
271        if texture.is_none() {
272            self.is_none = true;
273        }
274        self.textures.push(texture);
275        self.elements
276            .push((binding_type, self.textures.len() as u32 - 1));
277        self
278    }
279}
280
281/// A trait for describing a render binding.
282///
283/// This trait is used to declare bind-group entries by referencing resources from
284/// [`RenderData`](crate::assets::bindings::RenderData). Register the type using
285/// [`RenderBindingRegisterPlugin`] and retrieve the prepared GPU side through
286/// [`GpuRenderBinding`] in render systems.
287pub trait RenderBinding {
288    type Params: ReadOnlySystemParam;
289
290    /// Describes the render binding. This is used to specify the [`RenderData`](RenderData) that should be used to create the bind group layout and bind group entries, and to specify the binding index for each buffer and texture.
291    fn describe(
292        &mut self,
293        params: &SystemParamItem<Self::Params>,
294        builder: &mut RenderBindingBuilder
295    );
296    /// Whether the gpu asset should be recreated. This is called every frame. The default is to never recreate (None).
297    fn label(&self) -> &str;
298}
299
300// ============= RENDER BINDING RENDER ASSET GPU CREATION =============
301/// The gpu asset that is created from the render binding holder. This is the asset that is actually used in the render pass.
302pub type SBinding<T> = SRes<RenderAssets<GpuRenderBinding<T>>>;
303/// The gpu asset that is created from the render binding holder. This is the asset that is actually used in the render pass.
304pub type SBindingMut<T> = SResMut<RenderAssets<GpuRenderBinding<T>>>;
305/// The gpu asset that is created from the render binding holder. This is the asset that is actually used in the render pass.
306pub type Binding<'w, T> = Res<'w, RenderAssets<GpuRenderBinding<T>>>;
307/// The gpu asset that is created from the render binding holder. This is the asset that is actually used in the render pass.
308pub type BindingMut<'w, T> = ResMut<'w, RenderAssets<GpuRenderBinding<T>>>;
309
310/// The gpu asset that is created from the render binding holder. This is the asset that is actually used in the render pass.
311/// It contains the bind group and bind group layout that are created according to the description in the [`RenderBinding`](RenderBinding) implementation, and can be retrieved using the binding index specified in the description.
312pub struct GpuRenderBinding<T> {
313    _phantom: PhantomData<T>,
314    pub layout: BindGroupLayout,
315    pub bind_group: BindGroup
316}
317impl<T> RenderAsset for GpuRenderBinding<T>
318where
319    T: RenderBinding + TypePath + Sync + Send + Clone + Asset
320{
321    type SourceAsset = T;
322    type Params = (
323        T::Params,
324        SRes<RenderInstance>,
325        SRes<RenderAssets<GpuBuffer>>,
326        SRes<RenderAssets<GpuTexture>>
327    );
328
329    fn prepare(
330        asset: Self::SourceAsset,
331        (binding_params, render_instance, gpu_buffers, gpu_textures): &mut SystemParamItem<
332            Self::Params
333        >
334    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
335        let mut asset = asset.clone();
336
337        // Describe the asset
338        let mut asset_builder = RenderBindingBuilder::default();
339        T::describe(&mut asset, binding_params, &mut asset_builder);
340        if asset_builder.is_none {
341            // If any of the buffers or textures were not ready, retry in the next update
342            trace!(
343                "Not all resources for render binding {} are ready, retrying...",
344                asset.label()
345            );
346            return Err(PrepareAssetError::RetryNextUpdate(asset));
347        }
348
349        // Check every texture and buffers are ready
350        for (binding_type, idx) in &asset_builder.elements {
351            match binding_type {
352                RenderBindingType::Buffer => {
353                    if gpu_buffers
354                        .get(asset_builder.buffers[*idx as usize].unwrap())
355                        .is_none()
356                    {
357                        trace!(
358                            "Buffer for render binding {} is not ready, retrying...",
359                            asset.label()
360                        );
361                        return Err(PrepareAssetError::RetryNextUpdate(asset));
362                    }
363                }
364                RenderBindingType::TextureView
365                | RenderBindingType::TextureArrayView
366                | RenderBindingType::TextureSampler
367                | RenderBindingType::StorageTexture => {
368                    if gpu_textures
369                        .get(asset_builder.textures[*idx as usize].unwrap())
370                        .is_none()
371                    {
372                        trace!(
373                            "Texture for render binding {} is not ready, retrying...",
374                            asset.label()
375                        );
376                        return Err(PrepareAssetError::RetryNextUpdate(asset));
377                    };
378                }
379            }
380        }
381
382        // Create the bind group layout
383        let mut is_err = false;
384        let layout = BindGroupLayout::new(asset.label(), |builder| {
385            let vis = ShaderStages::all();
386            for (binding, (binding_type, idx)) in asset_builder.elements.iter().enumerate() {
387                match binding_type {
388                    RenderBindingType::Buffer => {
389                        let Some(buffer) =
390                            gpu_buffers.get(asset_builder.buffers[*idx as usize].unwrap())
391                        else {
392                            trace!(
393                                "Buffer for render binding {} is not ready, retrying...",
394                                asset.label()
395                            );
396                            is_err = true;
397                            return;
398                        };
399                        let binding_type = if buffer
400                            .buffer
401                            .buffer
402                            .usage()
403                            .contains(BufferUsage::UNIFORM)
404                        {
405                            BufferBindingType::Uniform
406                        } else if buffer.buffer.buffer.usage().contains(BufferUsage::STORAGE) {
407                            BufferBindingType::Storage { read_only: true }
408                        } else {
409                            error!(
410                                "Buffer has no usage flag, defaulting to UNIFORM for render binding {}.",
411                                asset.label()
412                            );
413                            BufferBindingType::Uniform
414                        };
415                        builder.add_buffer(binding as u32, vis, binding_type)
416                    }
417                    RenderBindingType::TextureView => {
418                        let Some(texture) =
419                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
420                        else {
421                            trace!(
422                                "Texture for render binding {} is not ready, retrying...",
423                                asset.label()
424                            );
425                            is_err = true;
426                            return;
427                        };
428                        let multisampled = texture.texture.sample_count > 1;
429                        builder.add_texture_view(binding as u32, vis, multisampled)
430                    }
431                    RenderBindingType::TextureArrayView => {
432                        builder.add_texture_array_view(binding as u32, vis)
433                    }
434                    RenderBindingType::TextureSampler => {
435                        builder.add_texture_sampler(binding as u32, vis)
436                    }
437                    RenderBindingType::StorageTexture => {
438                        let Some(texture) =
439                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
440                        else {
441                            trace!(
442                                "Texture for render binding {} is not ready, retrying...",
443                                asset.label()
444                            );
445                            is_err = true;
446                            return;
447                        };
448                        builder.add_storage_texture_view(binding as u32, texture.texture.format)
449                    }
450                };
451            }
452        });
453        if is_err {
454            return Err(PrepareAssetError::RetryNextUpdate(asset));
455        }
456
457        // Build the bind group layout
458        let render_instance = render_instance.0.read().unwrap();
459        let layout_built = match layout.build(&render_instance) {
460            Ok(layout) => layout,
461            Err(_) => return Err(PrepareAssetError::RetryNextUpdate(asset))
462        };
463
464        // Create the bind group
465        let mut bg_entries = vec![];
466        for (binding, (binding_type, idx)) in asset_builder.elements.iter().enumerate() {
467            match binding_type {
468                RenderBindingType::Buffer => {
469                    let buffer = gpu_buffers
470                        .get(asset_builder.buffers[*idx as usize].unwrap())
471                        .unwrap();
472                    bg_entries.push(BindGroupBuilder::buffer(binding as u32, &buffer.buffer));
473                }
474                RenderBindingType::TextureView
475                | RenderBindingType::TextureArrayView
476                | RenderBindingType::StorageTexture => {
477                    let texture = gpu_textures
478                        .get(asset_builder.textures[*idx as usize].unwrap())
479                        .unwrap();
480                    bg_entries.push(BindGroupBuilder::texture_view(
481                        binding as u32,
482                        &texture.texture
483                    ));
484                }
485                RenderBindingType::TextureSampler => {
486                    let texture = gpu_textures
487                        .get(asset_builder.textures[*idx as usize].unwrap())
488                        .unwrap();
489                    bg_entries.push(BindGroupBuilder::texture_sampler(
490                        binding as u32,
491                        &texture.texture
492                    ));
493                }
494            }
495        }
496        let Ok(bind_group) =
497            BindGroupBuilder::build(asset.label(), &render_instance, &layout_built, &bg_entries)
498        else {
499            trace!(
500                "Not all resources for render binding {} are ready to create the bind group, retrying...",
501                asset.label()
502            );
503            return Err(PrepareAssetError::RetryNextUpdate(asset));
504        };
505
506        // Return the gpu asset
507        debug!("Prepared render binding {} GPU resources.", asset.label());
508        Ok(GpuRenderBinding {
509            _phantom: PhantomData,
510            bind_group,
511            layout
512        })
513    }
514}