wde_pbr/deferred/transform/
registry.rs

1use std::collections::HashMap;
2use wde_logger::prelude::*;
3
4use bevy::prelude::*;
5use wde_camera::prelude::*;
6use wde_renderer::prelude::*;
7
8use crate::prelude::{PbrSsboTransform, SSBO_TRANSFORM_MAX_ENTITY};
9
10/// Marker component to indicate that an entity should be included in the SSBO transform updates.
11#[derive(Component, Default)]
12pub struct PbrSsboTransformMarker;
13
14pub(crate) struct SsboTransformRegistryPlugin;
15impl Plugin for SsboTransformRegistryPlugin {
16    fn build(&self, app: &mut App) {
17        app
18            .init_resource::<PbrSsboTransformRegistry>()
19            .add_systems(Update, set_dirty_transforms);
20        app.get_sub_app_mut(RenderApp)
21            .unwrap()
22            .init_resource::<PbrSsboTransformRegistry>()
23            .add_systems(Extract, extract_registry)
24            .add_systems(Render, update_ssbo_transforms.in_set(RenderSet::Prepare));
25    }
26}
27
28/// Resource to track the mapping between entities and their transform IDs in the SSBO, as well as the list of dirty transforms that need to be updated in the SSBO. This is used to update the [PbrSsboTransform] with the new transforms of the entities with a PBR material.
29#[derive(Resource, Default, Clone)]
30pub struct PbrSsboTransformRegistry {
31    pub entity_to_transform: HashMap<Entity, u32>, // (entity, transform_id in the SSBO)
32    pub available_transform_ids: Vec<u32>,         // List of available transform IDs in the SSBO
33    pub next_transform_id: u32,                    // Next available transform ID in the SSBO
34    pub dirty_transforms: Vec<(u32, TransformUniform)>, // List of dirty transforms that need to be updated in the SSBO
35}
36impl PbrSsboTransformRegistry {
37    /// Register an entity with its transform in the registry, returning the assigned transform ID in the SSBO. If the entity is already registered, returns the existing transform ID.
38    pub fn add_entity(&mut self, entity: Entity, transform: GlobalTransform) -> u32 {
39        if let Some(&transform_id) = self.entity_to_transform.get(&entity) {
40            transform_id
41        } else {
42            let transform_id = if let Some(free_id) = self.available_transform_ids.pop() {
43                free_id
44            } else if self.next_transform_id < SSBO_TRANSFORM_MAX_ENTITY as u32 {
45                let id = self.next_transform_id;
46                self.next_transform_id += 1;
47                id
48            } else {
49                error!(
50                    "PbrSsboTransformRegistry: Maximum number of transform IDs reached ({})",
51                    SSBO_TRANSFORM_MAX_ENTITY
52                );
53                0
54            };
55            self.entity_to_transform.insert(entity, transform_id);
56            self.dirty_transforms
57                .push((transform_id, TransformUniform::new(&transform)));
58            transform_id
59        }
60    }
61
62    /// Get the transform ID of an entity in the registry, if it exists.
63    pub fn get_transform_id(&self, entity: Entity) -> Option<u32> {
64        self.entity_to_transform.get(&entity).copied()
65    }
66
67    /// Update the transform of an entity in the registry, marking it as dirty for the next SSBO update. If the entity is not registered, it will be added to the registry.
68    pub fn update_entity(&mut self, entity: Entity, transform: GlobalTransform) {
69        let transform_id = self.add_entity(entity, transform); // Make sure the entity is registered and get its transform ID
70        self.dirty_transforms
71            .push((transform_id, TransformUniform::new(&transform)));
72    }
73
74    /// Unregister an entity from the registry, freeing its transform ID for future use.
75    pub fn remove_entity(&mut self, entity: Entity) {
76        if let Some(transform_id) = self.entity_to_transform.remove(&entity) {
77            self.available_transform_ids.push(transform_id);
78        }
79    }
80}
81
82// Update the list of dirty transforms for this frame
83#[allow(clippy::type_complexity)]
84fn set_dirty_transforms(
85    mut registry: ResMut<PbrSsboTransformRegistry>,
86    new_transforms: Query<
87        (Entity, &GlobalTransform),
88        Or<(
89            (With<PbrSsboTransformMarker>, Added<GlobalTransform>),
90            (With<GlobalTransform>, Added<PbrSsboTransformMarker>),
91        )>,
92    >,
93    changed_transforms: Query<
94        (Entity, &GlobalTransform),
95        (With<PbrSsboTransformMarker>, Changed<GlobalTransform>),
96    >,
97    mut removed_transforms: RemovedComponents<PbrSsboTransformMarker>,
98) {
99    // Clear the list of dirty transforms from the previous frame
100    registry.dirty_transforms.clear();
101
102    // Add new or changed transforms to the dirty list
103    for (entity, transform) in new_transforms.iter().chain(changed_transforms.iter()) {
104        registry.update_entity(entity, *transform);
105    }
106
107    // Remove transforms that are no longer needed
108    for entity in removed_transforms.read() {
109        registry.remove_entity(entity);
110    }
111}
112
113fn extract_registry(
114    main_registry: ExtractWorld<Res<PbrSsboTransformRegistry>>,
115    mut registry: ResMut<PbrSsboTransformRegistry>
116) {
117    registry.entity_to_transform = main_registry.entity_to_transform.clone();
118    registry.dirty_transforms.extend(main_registry.dirty_transforms.clone());
119}
120
121// Update the ssbo with the dirty transforms from the registry
122fn update_ssbo_transforms(
123    ssbo: ResRenderData<PbrSsboTransform>,
124    buffers: Res<RenderAssets<GpuBuffer>>,
125    mut registry: ResMut<PbrSsboTransformRegistry>,
126    render_instance: Res<RenderInstance>,
127) {
128    // Return early if there are no dirty transforms to update
129    if registry.dirty_transforms.is_empty() {
130        return;
131    }
132
133    // Get the ssbo cpu buffer and gpu buffer
134    let (ssbo_staging, ssbo_gpu) = match (
135        ssbo.iter().next().and_then(|(_, d)| buffers.get(&d.get_buffer(PbrSsboTransform::TRANSFORM_STAGING_IDX)?)),
136        ssbo.iter().next().and_then(|(_, d)| buffers.get(&d.get_buffer(PbrSsboTransform::TRANSFORM_IDX)?))
137    ) {
138        (Some(staging), Some(gpu)) => (&staging.buffer, &gpu.buffer),
139        _ => return,
140    };
141
142    // Write the dirty transforms to the staging buffer
143    let render_instance = render_instance.0.read().unwrap();
144    for (transform_id, transform) in &registry.dirty_transforms {
145        let offset = (*transform_id as usize) * std::mem::size_of::<TransformUniform>();
146        ssbo_staging.write(&render_instance, bytemuck::cast_slice(&[*transform]), offset);
147    }
148
149    // Copy the staging buffer to the GPU buffer
150    ssbo_gpu.copy_from_buffer(&render_instance, ssbo_staging);
151
152    // Only then, clear the list of dirty transforms
153    registry.dirty_transforms.clear();
154}
wde_pbr/deferred/transform/registry.rs

wde_pbr/deferred/transform/
registry.rs
