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        _id: AssetId<Self::SourceAsset>,
390        asset: Self::SourceAsset,
391        (binding_params, render_instance, gpu_buffers, gpu_textures): &mut SystemParamItem<
392            Self::Params
393        >
394    ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
395        let mut asset = asset.clone();
396
397        // Describe the asset
398        let mut asset_builder = RenderBindingBuilder::default();
399        T::describe(&mut asset, binding_params, &mut asset_builder);
400        if asset_builder.is_none {
401            // If any of the buffers or textures were not ready, retry in the next update
402            trace!(
403                "Not all resources for render binding {} are ready, retrying...",
404                asset.label()
405            );
406            return Err(PrepareAssetError::RetryNextUpdate(asset));
407        }
408
409        // Check every texture and buffers are ready
410        for (binding_type, idx) in &asset_builder.elements {
411            match binding_type {
412                RenderBindingType::Buffer => {
413                    if gpu_buffers
414                        .get(asset_builder.buffers[*idx as usize].unwrap())
415                        .is_none()
416                    {
417                        trace!(
418                            "Buffer for render binding {} is not ready, retrying...",
419                            asset.label()
420                        );
421                        return Err(PrepareAssetError::RetryNextUpdate(asset));
422                    }
423                }
424                RenderBindingType::TextureView
425                | RenderBindingType::TextureArrayView
426                | RenderBindingType::TextureSampler
427                | RenderBindingType::StorageTexture
428                | RenderBindingType::StorageTextureArray
429                | RenderBindingType::DepthTextureView => {
430                    if gpu_textures
431                        .get(asset_builder.textures[*idx as usize].unwrap())
432                        .is_none()
433                    {
434                        trace!(
435                            "Texture (View, ArrayView, Sampler, StorageTexture, or DepthView) {} for render binding {} is not ready, retrying...",
436                            asset_builder.textures[*idx as usize].unwrap(),
437                            asset.label()
438                        );
439                        return Err(PrepareAssetError::RetryNextUpdate(asset));
440                    };
441                }
442            }
443        }
444
445        // Create the bind group layout
446        let mut is_err = false;
447        let layout = BindGroupLayout::new(asset.label(), |builder| {
448            let vis = ShaderStages::all();
449            for (binding, (binding_type, idx)) in asset_builder.elements.iter().enumerate() {
450                match binding_type {
451                    RenderBindingType::Buffer => {
452                        let Some(buffer) =
453                            gpu_buffers.get(asset_builder.buffers[*idx as usize].unwrap())
454                        else {
455                            trace!(
456                                "Buffer for render binding {} is not ready, retrying...",
457                                asset.label()
458                            );
459                            is_err = true;
460                            return;
461                        };
462                        let binding_type = if buffer
463                            .buffer
464                            .buffer
465                            .usage()
466                            .contains(BufferUsage::UNIFORM)
467                        {
468                            BufferBindingType::Uniform
469                        } else if buffer.buffer.buffer.usage().contains(BufferUsage::STORAGE) {
470                            BufferBindingType::Storage { read_only: true }
471                        } else {
472                            error!(
473                                "Buffer has no usage flag, defaulting to UNIFORM for render binding {}.",
474                                asset.label()
475                            );
476                            BufferBindingType::Uniform
477                        };
478                        builder.add_buffer(binding as u32, vis, binding_type)
479                    }
480                    RenderBindingType::TextureView => {
481                        let Some(texture) =
482                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
483                        else {
484                            trace!(
485                                "Texture (View) {} for render binding {} is not ready, retrying...",
486                                asset_builder.textures[*idx as usize].unwrap(),
487                                asset.label()
488                            );
489                            is_err = true;
490                            return;
491                        };
492                        let multisampled = texture.texture.sample_count > 1;
493                        builder.add_texture_view_filterable(
494                            binding as u32,
495                            vis,
496                            multisampled,
497                            texture.texture.filterable
498                        )
499                    }
500                    RenderBindingType::TextureArrayView => {
501                        let Some(texture) =
502                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
503                        else {
504                            trace!(
505                                "Texture (ArrayView) {} for render binding {} is not ready, retrying...",
506                                asset_builder.textures[*idx as usize].unwrap(),
507                                asset.label()
508                            );
509                            is_err = true;
510                            return;
511                        };
512                        builder.add_texture_array_view(binding as u32, vis, texture.texture.filterable)
513                    }
514                    RenderBindingType::TextureSampler => {
515                        builder.add_texture_sampler(binding as u32, vis)
516                    }
517                    RenderBindingType::StorageTexture => {
518                        let Some(texture) =
519                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
520                        else {
521                            trace!(
522                                "Texture (Storage) {} for render binding {} is not ready, retrying...",
523                                asset_builder.textures[*idx as usize].unwrap(),
524                                asset.label()
525                            );
526                            is_err = true;
527                            return;
528                        };
529                        builder.add_storage_texture_view(binding as u32, texture.texture.format)
530                    }
531                    RenderBindingType::StorageTextureArray => {
532                        let Some(texture) =
533                            gpu_textures.get(asset_builder.textures[*idx as usize].unwrap())
534                        else {
535                            trace!(
536                                "Texture (StorageArray) {} for render binding {} is not ready, retrying...",
537                                asset_builder.textures[*idx as usize].unwrap(),
538                                asset.label()
539                            );
540                            is_err = true;
541                            return;
542                        };
543                        builder
544                            .add_storage_texture_array_view(binding as u32, texture.texture.format)
545                    }
546                    RenderBindingType::DepthTextureView => {
547                        builder.add_depth_texture_view(binding as u32, vis, false)
548                    }
549                };
550            }
551        });
552        if is_err {
553            return Err(PrepareAssetError::RetryNextUpdate(asset));
554        }
555
556        // Build the bind group layout
557        let render_instance = render_instance.0.read().unwrap();
558        let layout_built = match layout.build(&render_instance) {
559            Ok(layout) => layout,
560            Err(_) => return Err(PrepareAssetError::RetryNextUpdate(asset))
561        };
562
563        // Create the bind group
564        let mut bg_entries = vec![];
565        for (binding, (binding_type, idx)) in asset_builder.elements.iter().enumerate() {
566            match binding_type {
567                RenderBindingType::Buffer => {
568                    let buffer = gpu_buffers
569                        .get(asset_builder.buffers[*idx as usize].unwrap())
570                        .unwrap();
571                    bg_entries.push(BindGroupBuilder::buffer(binding as u32, &buffer.buffer));
572                }
573                RenderBindingType::TextureView
574                | RenderBindingType::TextureArrayView
575                | RenderBindingType::StorageTexture
576                | RenderBindingType::StorageTextureArray
577                | RenderBindingType::DepthTextureView => {
578                    let texture = gpu_textures
579                        .get(asset_builder.textures[*idx as usize].unwrap())
580                        .unwrap();
581                    bg_entries.push(BindGroupBuilder::texture_view(
582                        binding as u32,
583                        &texture.texture
584                    ));
585                }
586                RenderBindingType::TextureSampler => {
587                    let texture = gpu_textures
588                        .get(asset_builder.textures[*idx as usize].unwrap())
589                        .unwrap();
590                    bg_entries.push(BindGroupBuilder::texture_sampler(
591                        binding as u32,
592                        &texture.texture
593                    ));
594                }
595            }
596        }
597        let Ok(bind_group) =
598            BindGroupBuilder::build(asset.label(), &render_instance, &layout_built, &bg_entries)
599        else {
600            trace!(
601                "Not all resources for render binding {} are ready to create the bind group, retrying...",
602                asset.label()
603            );
604            return Err(PrepareAssetError::RetryNextUpdate(asset));
605        };
606
607        // Return the gpu asset
608        trace!("Prepared render binding {} GPU resources.", asset.label());
609        Ok(GpuRenderBinding {
610            _phantom: PhantomData,
611            bind_group,
612            layout,
613            asset
614        })
615    }
616}