wde_pbr/deferred/batches/
build_batches.rs

1use std::collections::{HashMap, HashSet};
2
3use bevy::prelude::*;
4use wde_renderer::prelude::*;
5
6use crate::prelude::{
7    CastShadow, PbrAlphaMode, PbrMaterial, PbrMaterial3d, PbrSsboTransformUuid, PbrUuidRegistry
8};
9
10#[derive(Component)]
11pub struct ExtractedPbrInstance {
12    mesh_id: AssetId<Mesh>,
13    material_id: AssetId<PbrMaterial>
14}
15
16#[derive(Clone, Copy)]
17pub(crate) struct TrackedBatchEntity {
18    pub(crate) key: BatchKey,
19    transform_id: u32,
20    pub(crate) main_entity: Entity
21}
22
23#[derive(Component)]
24pub(crate) struct ExtractedPbrInstanceToRetryMarker;
25
26/// Marker component to indicate that an entity should be included in the PBR render batches.
27/// This will automatically add the [PbrSsboTransformMarker] to the entity, so that its transform will be included in the SSBO updates for rendering.
28#[derive(Component, Default, Reflect)]
29#[reflect(Component)]
30#[require(PbrSsboTransformUuid)]
31pub struct PbrBatchesMarker;
32impl SyncComponent for PbrBatchesMarker {
33    type Target = Self;
34}
35impl ExtractComponent for PbrBatchesMarker {
36    type QueryData = (&'static Mesh3d, &'static PbrMaterial3d<PbrMaterial>);
37    type QueryFilter = (
38        With<PbrBatchesMarker>,
39        With<Transform>,
40        Or<(Changed<Mesh3d>, Changed<PbrMaterial3d<PbrMaterial>>)>
41    );
42    type Out = ExtractedPbrInstance;
43
44    fn extract_component(
45        (mesh, material): QueryItem<'_, '_, Self::QueryData>
46    ) -> Option<Self::Out> {
47        Some(ExtractedPbrInstance {
48            mesh_id: mesh.0.id(),
49            material_id: material.0.id()
50        })
51    }
52}
53
54// Key for batching, based on mesh and material. Entities with the same mesh and material will be batched together.
55#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
56pub(crate) struct BatchKey {
57    pub(crate) mesh: AssetId<Mesh>,
58    pub(crate) material: AssetId<PbrMaterial>
59}
60// Batch of entities that share the same mesh and material, represented by their transform IDs in the SSBO.
61pub(crate) struct Batch {
62    pub _key: BatchKey,
63    pub transform_ssbo_ids: Vec<u32>, // Offset in the SSBO "transform" buffer of the first transform of this batch
64    pub instances_offset: u32 // Offset in the SSBO "instance to transform" buffer of the first entity of this batch
65}
66// Resource to store the batches of entities for rendering. The key is the combination of mesh and material, and the value is the batch of transform IDs for entities that share that mesh and material.
67#[derive(Resource, Default)]
68pub(crate) struct BatchList {
69    pub batches: HashMap<BatchKey, Batch>,
70    pub sorted_batches: Vec<BatchKey>, // List of batch keys sorted first by material and then by mesh
71    pub(crate) tracked_entities: HashMap<Entity, TrackedBatchEntity>,
72    pub shadow_casting_batches: HashSet<BatchKey>, // Subset of sorted_batches whose entities have CastShadow
73    pub dirty: bool // Did anything change? If true, rebuild ssbo batches
74}
75impl std::fmt::Debug for BatchList {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        let mut debug_struct = f.debug_struct("BatchList");
78        for batch_key in &self.sorted_batches {
79            if let Some(batch) = self.batches.get(batch_key) {
80                debug_struct.field(
81                    &format!(
82                        "Batch(mesh: {:?}, material: {:?})",
83                        batch_key.mesh, batch_key.material
84                    ),
85                    &format!(
86                        "Batch {{ transform_ssbo_ids: {:?}, instances_offset: {} }}",
87                        batch.transform_ssbo_ids, batch.instances_offset
88                    )
89                );
90            }
91        }
92        debug_struct.finish()
93    }
94}
95
96/// Whether the batch material is alpha-blended, i.e. drawn by the transparent pass only.
97pub(crate) fn is_batch_blended(
98    materials: &RenderAssets<GpuRenderBinding<PbrMaterial>>,
99    key: &BatchKey
100) -> bool {
101    materials
102        .get(key.material)
103        .is_some_and(|m| m.asset.alpha_mode == PbrAlphaMode::Blend)
104}
105
106fn remove_tracked_entity(batches: &mut BatchList, entity: Entity) -> bool {
107    let Some(tracked) = batches.tracked_entities.remove(&entity) else {
108        return false;
109    };
110
111    let mut removed_any = false;
112    let mut remove_batch_key = false;
113    if let Some(batch) = batches.batches.get_mut(&tracked.key) {
114        let len_before = batch.transform_ssbo_ids.len();
115        batch
116            .transform_ssbo_ids
117            .retain(|id| *id != tracked.transform_id);
118        removed_any = batch.transform_ssbo_ids.len() != len_before;
119        remove_batch_key = batch.transform_ssbo_ids.is_empty();
120    }
121
122    if remove_batch_key {
123        batches.batches.remove(&tracked.key);
124        batches.sorted_batches.retain(|key| *key != tracked.key);
125    }
126
127    removed_any
128}
129
130pub(crate) fn build_batches(
131    mut commands: Commands,
132    mut batches: ResMut<BatchList>,
133    mut removed_extracted_instances: RemovedComponents<ExtractedPbrInstance>,
134    mut removed_transform_uuid: RemovedComponents<PbrSsboTransformUuid>,
135    extracted_instances: Query<
136        (
137            Entity,
138            MainEntity,
139            &PbrSsboTransformUuid,
140            &ExtractedPbrInstance
141        ),
142        Changed<ExtractedPbrInstance>
143    >,
144    extracted_instances_to_retry: Query<
145        (
146            Entity,
147            MainEntity,
148            &PbrSsboTransformUuid,
149            &ExtractedPbrInstance
150        ),
151        With<ExtractedPbrInstanceToRetryMarker>
152    >,
153    transform_registry: Res<PbrUuidRegistry>
154) {
155    let mut changed = false;
156
157    for entity in removed_extracted_instances.read() {
158        changed |= remove_tracked_entity(&mut batches, entity);
159    }
160    for entity in removed_transform_uuid.read() {
161        changed |= remove_tracked_entity(&mut batches, entity);
162    }
163
164    // Add new or changed instances to the batches, and remove the retry marker if it exists
165    for (entity, main_entity, transform_uuid, instance) in extracted_instances
166        .into_iter()
167        .chain(extracted_instances_to_retry)
168    {
169        changed |= remove_tracked_entity(&mut batches, entity);
170
171        let transform_id = if let Some(id) = transform_registry.get(&transform_uuid.0) {
172            id
173        } else {
174            // Retry next frame
175            commands
176                .entity(entity)
177                .insert(ExtractedPbrInstanceToRetryMarker);
178            continue;
179        };
180
181        let key = BatchKey {
182            mesh: instance.mesh_id,
183            material: instance.material_id
184        };
185        batches
186            .batches
187            .entry(key)
188            .or_insert_with(|| Batch {
189                _key: key,
190                transform_ssbo_ids: Vec::new(),
191                instances_offset: 0
192            })
193            .transform_ssbo_ids
194            .push(transform_id);
195        batches.tracked_entities.insert(
196            entity,
197            TrackedBatchEntity {
198                key,
199                transform_id,
200                main_entity
201            }
202        );
203        if !batches.sorted_batches.contains(&key) {
204            batches.sorted_batches.push(key);
205        }
206        commands
207            .entity(entity)
208            .remove::<ExtractedPbrInstanceToRetryMarker>();
209        changed = true;
210    }
211    batches.dirty |= changed;
212
213    // Sort batches first by material and then by mesh
214    if changed {
215        batches.sorted_batches.sort_by(|a, b| {
216            a.material
217                .cmp(&b.material)
218                .then_with(|| a.mesh.cmp(&b.mesh))
219        });
220    }
221}
222
223/// Rebuilds `BatchList::shadow_casting_batches` every frame from render-world entities.
224pub(crate) fn update_shadow_casting_batches(
225    mut batches: ResMut<BatchList>,
226    cast_shadow_query: Query<MainEntity, With<CastShadow>>
227) {
228    let shadow_main_entities: HashSet<Entity> = cast_shadow_query.iter().collect();
229    let new_shadow_batches: HashSet<BatchKey> = batches
230        .tracked_entities
231        .iter()
232        .filter_map(|(_, tracked)| {
233            if shadow_main_entities.contains(&tracked.main_entity) {
234                Some(tracked.key)
235            } else {
236                None
237            }
238        })
239        .collect();
240    batches.shadow_casting_batches = new_shadow_batches;
241}
242
243pub(crate) struct BatchesPlugin;
244impl Plugin for BatchesPlugin {
245    fn build(&self, app: &mut App) {
246        app.add_plugins(ExtractComponentPlugin::<PbrBatchesMarker>::default());
247        app.get_sub_app_mut(RenderApp)
248            .unwrap()
249            .init_resource::<BatchList>();
250    }
251}
wde_pbr/deferred/batches/build_batches.rs

wde_pbr/deferred/batches/
build_batches.rs
