wde_pbr/deferred/batches/
build_batches.rs

1use wde_logger::prelude::*;
2use std::collections::HashMap;
3
4use bevy::prelude::*;
5use wde_renderer::prelude::*;
6
7use crate::prelude::{PbrMaterial, PbrMaterial3d, PbrSsboTransformMarker, PbrSsboTransformRegistry};
8
9#[derive(Component)]
10pub struct ExtractedPbrInstance {
11    mesh_id: AssetId<Mesh>,
12    material_id: AssetId<PbrMaterial>
13}
14
15/// Marker component to indicate that an entity should be included in the PBR render batches.
16/// This will automatically add the [PbrSsboTransformMarker] to the entity, so that its transform will be included in the SSBO updates for rendering.
17#[derive(Component, Default)]
18#[require(PbrSsboTransformMarker)]
19pub struct PbrBatchesMarker;
20impl SyncComponent for PbrBatchesMarker {
21    type Target = Self;
22}
23impl ExtractComponent for PbrBatchesMarker {
24    type QueryData = (&'static Mesh3d, &'static PbrMaterial3d);
25    type QueryFilter = (With<PbrBatchesMarker>, With<Transform>, Or<(Changed<Mesh3d>, Changed<PbrMaterial3d>)>);
26    type Out = ExtractedPbrInstance;
27
28    fn extract_component((mesh, material): QueryItem<'_, '_, Self::QueryData>) -> Option<Self::Out> {
29        Some(ExtractedPbrInstance {
30            mesh_id: mesh.0.id(),
31            material_id: material.0.id()
32        })
33    }
34}
35
36// Key for batching, based on mesh and material. Entities with the same mesh and material will be batched together.
37#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
38pub(crate) struct BatchKey {
39    pub(crate) mesh: AssetId<Mesh>,
40    pub(crate) material: AssetId<PbrMaterial>
41}
42// Batch of entities that share the same mesh and material, represented by their transform IDs in the SSBO.
43pub(crate) struct Batch {
44    pub _key: BatchKey,
45    pub transform_ssbo_ids: Vec<u32>, // Offset in the SSBO "transform" buffer of the first transform of this batch
46    pub instances_offset: u32,        // Offset in the SSBO "instance to transform" buffer of the first entity of this batch
47}
48// 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.
49#[derive(Resource, Default)]
50pub(crate) struct BatchList {
51    pub batches: HashMap<BatchKey, Batch>,
52    pub sorted_batches: Vec<BatchKey>, // List of batch keys sorted first by material and then by mesh
53    pub dirty_batches: Vec<BatchKey>   // List of batch keys that have been modified and need to be updated in the GPU instance to transform buffer
54}
55impl std::fmt::Debug for BatchList {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        let mut debug_struct = f.debug_struct("BatchList");
58        for batch_key in &self.sorted_batches {
59            if let Some(batch) = self.batches.get(batch_key) {
60                debug_struct.field(
61                    &format!("Batch(mesh: {:?}, material: {:?})", batch_key.mesh, batch_key.material),
62                    &format!("Batch {{ transform_ssbo_ids: {:?}, instances_offset: {} }}", batch.transform_ssbo_ids, batch.instances_offset)
63                );
64            }
65        }
66        debug_struct.finish()
67    }
68}
69
70pub(crate) fn build_batches(
71    mut batches: ResMut<BatchList>,
72    extracted_instances: Query<(MainEntity, &ExtractedPbrInstance), Changed<ExtractedPbrInstance>>,
73    transform_registry: Res<PbrSsboTransformRegistry>
74) {
75    for (entity, instance) in &extracted_instances {
76        let transform_id = if let Some(&id) = transform_registry.entity_to_transform.get(&entity) {
77            id
78        } else {
79            warn!("Entity {:?} with mesh {:?} and material {:?} is not registered in the PbrSsboTransformRegistry, skipping it for batching", entity, instance.mesh_id, instance.material_id);
80            continue; // If the entity doesn't have a transform ID, skip it
81        };
82
83        let key = BatchKey {
84            mesh: instance.mesh_id,
85            material: instance.material_id
86        };
87        batches.batches.entry(key).or_insert_with(|| Batch {
88            _key: key,
89            transform_ssbo_ids: Vec::new(),
90            instances_offset: 0
91        }).transform_ssbo_ids.push(transform_id);
92        batches.sorted_batches.push(key);
93        batches.dirty_batches.push(key);
94    }
95
96    // Sort batches first by material and then by mesh
97    if extracted_instances.iter().next().is_some() {
98        batches.sorted_batches.sort_by(|a, b| {
99            a.material.cmp(&b.material).then_with(|| a.mesh.cmp(&b.mesh))
100        });
101    }
102}
103
104pub(crate) struct BatchesPlugin;
105impl Plugin for BatchesPlugin {
106    fn build(&self, app: &mut App) {
107        app.add_plugins(ExtractComponentPlugin::<PbrBatchesMarker>::default());
108        app.get_sub_app_mut(RenderApp).unwrap()
109            .init_resource::<BatchList>();
110    }
111}
wde_pbr/deferred/batches/build_batches.rs

wde_pbr/deferred/batches/
build_batches.rs
