Skip to main content

wde_renderer/assets/bindings/
render_binding.rs

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