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, 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
96fn remove_tracked_entity(batches: &mut BatchList, entity: Entity) -> bool {
97    let Some(tracked) = batches.tracked_entities.remove(&entity) else {
98        return false;
99    };
100
101    let mut removed_any = false;
102    let mut remove_batch_key = false;
103    if let Some(batch) = batches.batches.get_mut(&tracked.key) {
104        let len_before = batch.transform_ssbo_ids.len();
105        batch
106            .transform_ssbo_ids
107            .retain(|id| *id != tracked.transform_id);
108        removed_any = batch.transform_ssbo_ids.len() != len_before;
109        remove_batch_key = batch.transform_ssbo_ids.is_empty();
110    }
111
112    if remove_batch_key {
113        batches.batches.remove(&tracked.key);
114        batches.sorted_batches.retain(|key| *key != tracked.key);
115    }
116
117    removed_any
118}
119
120pub(crate) fn build_batches(
121    mut commands: Commands,
122    mut batches: ResMut<BatchList>,
123    mut removed_extracted_instances: RemovedComponents<ExtractedPbrInstance>,
124    mut removed_transform_uuid: RemovedComponents<PbrSsboTransformUuid>,
125    extracted_instances: Query<
126        (
127            Entity,
128            MainEntity,
129            &PbrSsboTransformUuid,
130            &ExtractedPbrInstance
131        ),
132        Changed<ExtractedPbrInstance>
133    >,
134    extracted_instances_to_retry: Query<
135        (
136            Entity,
137            MainEntity,
138            &PbrSsboTransformUuid,
139            &ExtractedPbrInstance
140        ),
141        With<ExtractedPbrInstanceToRetryMarker>
142    >,
143    transform_registry: Res<PbrUuidRegistry>
144) {
145    let mut changed = false;
146
147    for entity in removed_extracted_instances.read() {
148        changed |= remove_tracked_entity(&mut batches, entity);
149    }
150    for entity in removed_transform_uuid.read() {
151        changed |= remove_tracked_entity(&mut batches, entity);
152    }
153
154    // Add new or changed instances to the batches, and remove the retry marker if it exists
155    for (entity, main_entity, transform_uuid, instance) in extracted_instances
156        .into_iter()
157        .chain(extracted_instances_to_retry)
158    {
159        changed |= remove_tracked_entity(&mut batches, entity);
160
161        let transform_id = if let Some(id) = transform_registry.get(&transform_uuid.0) {
162            id
163        } else {
164            // Retry next frame
165            commands
166                .entity(entity)
167                .insert(ExtractedPbrInstanceToRetryMarker);
168            continue;
169        };
170
171        let key = BatchKey {
172            mesh: instance.mesh_id,
173            material: instance.material_id
174        };
175        batches
176            .batches
177            .entry(key)
178            .or_insert_with(|| Batch {
179                _key: key,
180                transform_ssbo_ids: Vec::new(),
181                instances_offset: 0
182            })
183            .transform_ssbo_ids
184            .push(transform_id);
185        batches.tracked_entities.insert(
186            entity,
187            TrackedBatchEntity {
188                key,
189                transform_id,
190                main_entity
191            }
192        );
193        if !batches.sorted_batches.contains(&key) {
194            batches.sorted_batches.push(key);
195        }
196        commands
197            .entity(entity)
198            .remove::<ExtractedPbrInstanceToRetryMarker>();
199        changed = true;
200    }
201    batches.dirty |= changed;
202
203    // Sort batches first by material and then by mesh
204    if changed {
205        batches.sorted_batches.sort_by(|a, b| {
206            a.material
207                .cmp(&b.material)
208                .then_with(|| a.mesh.cmp(&b.mesh))
209        });
210    }
211}
212
213/// Rebuilds `BatchList::shadow_casting_batches` every frame from render-world entities.
214pub(crate) fn update_shadow_casting_batches(
215    mut batches: ResMut<BatchList>,
216    cast_shadow_query: Query<MainEntity, With<CastShadow>>
217) {
218    let shadow_main_entities: HashSet<Entity> = cast_shadow_query.iter().collect();
219    let new_shadow_batches: HashSet<BatchKey> = batches
220        .tracked_entities
221        .iter()
222        .filter_map(|(_, tracked)| {
223            if shadow_main_entities.contains(&tracked.main_entity) {
224                Some(tracked.key)
225            } else {
226                None
227            }
228        })
229        .collect();
230    batches.shadow_casting_batches = new_shadow_batches;
231}
232
233pub(crate) struct BatchesPlugin;
234impl Plugin for BatchesPlugin {
235    fn build(&self, app: &mut App) {
236        app.add_plugins(ExtractComponentPlugin::<PbrBatchesMarker>::default());
237        app.get_sub_app_mut(RenderApp)
238            .unwrap()
239            .init_resource::<BatchList>();
240    }
241}
wde_pbr/deferred/batches/build_batches.rs

wde_pbr/deferred/batches/
build_batches.rs
