wde_pbr/deferred/transform/
registry.rs

1use std::collections::HashMap;
2use uuid::Uuid;
3
4use bevy::prelude::*;
5use wde_camera::prelude::*;
6use wde_renderer::prelude::*;
7
8use crate::prelude::PbrSsboTransform;
9
10pub(crate) struct SsboTransformRegistryPlugin;
11impl Plugin for SsboTransformRegistryPlugin {
12    fn build(&self, app: &mut App) {
13        app.add_plugins((
14            ExtractComponentPlugin::<PbrSsboTransformUuid>::default(),
15            ExtractResourcePlugin::<PbrSsboTransformPendingRemovals>::default()
16        ))
17        .init_resource::<PbrSsboTransformPendingRemovals>()
18        .init_resource::<MainWorldEntityUuidMap>()
19        .add_systems(
20            Update,
21            (MainWorldEntityUuidMap::add, MainWorldEntityUuidMap::remove).chain()
22        );
23        app.get_sub_app_mut(RenderApp)
24            .unwrap()
25            .init_resource::<PbrUuidRegistry>()
26            .add_systems(Render, update_registry.in_set(RenderSet::Prepare));
27    }
28}
29
30/// Marker component to track the UUID of an entity's transform in the SSBO registry
31#[derive(Component, Clone, Reflect)]
32#[reflect(Component)]
33pub struct PbrSsboTransformUuid(pub Uuid);
34impl Default for PbrSsboTransformUuid {
35    fn default() -> Self {
36        PbrSsboTransformUuid(Uuid::new_v4())
37    }
38}
39impl SyncComponent for PbrSsboTransformUuid {
40    type Target = Self;
41}
42impl ExtractComponent for PbrSsboTransformUuid {
43    type QueryData = (&'static GlobalTransform, &'static Self);
44    type QueryFilter = Changed<GlobalTransform>;
45    type Out = (GlobalTransform, Self);
46
47    fn extract_component(
48        (transform, uuid): QueryItem<'_, '_, Self::QueryData>
49    ) -> Option<Self::Out> {
50        Some((*transform, uuid.clone()))
51    }
52}
53
54/// Similar to `PbrSsboTransformUuid`, but tracks the entity that were removed last frame
55#[derive(Resource, Default, Clone, ExtractResource)]
56struct PbrSsboTransformPendingRemovals(pub Vec<Uuid>);
57#[derive(Resource, Default)]
58struct MainWorldEntityUuidMap(HashMap<Entity, Uuid>);
59impl MainWorldEntityUuidMap {
60    pub fn add(
61        mut map: ResMut<Self>,
62        added: Query<(Entity, &PbrSsboTransformUuid), Added<PbrSsboTransformUuid>>
63    ) {
64        for (entity, uuid) in added.iter() {
65            map.0.insert(entity, uuid.0);
66        }
67    }
68
69    pub fn remove(
70        mut map: ResMut<Self>,
71        mut pending_removals: ResMut<PbrSsboTransformPendingRemovals>,
72        mut removed: RemovedComponents<PbrSsboTransformUuid>
73    ) {
74        pending_removals.0.clear();
75        for entity in removed.read() {
76            if let Some(uuid) = map.0.remove(&entity) {
77                pending_removals.0.push(uuid);
78            }
79        }
80    }
81}
82
83/// Resource to track the mapping between render entities and their transform IDs in the SSBO.
84#[derive(Resource, Default, Clone)]
85pub struct PbrUuidRegistry {
86    pub uuid_to_transform: HashMap<Uuid, u32>, // (uuid, transform_id in the SSBO)
87    pub available_transform_ids: Vec<u32>,     // List of available transform IDs in the SSBO
88    pub next_transform_id: u32                 // Next available transform ID in the SSBO
89}
90impl PbrUuidRegistry {
91    /// Get the transform ID for a given UUID, adding it to the registry if it doesn't exist.
92    pub fn get_or_add(&mut self, uuid: Uuid) -> u32 {
93        if let Some(&transform_id) = self.uuid_to_transform.get(&uuid) {
94            transform_id
95        } else {
96            let transform_id = self.available_transform_ids.pop().unwrap_or_else(|| {
97                let id = self.next_transform_id;
98                self.next_transform_id += 1;
99                id
100            });
101            self.uuid_to_transform.insert(uuid, transform_id);
102            transform_id
103        }
104    }
105
106    /// Remove a UUID from the registry, freeing its transform ID for future use.
107    pub fn remove(&mut self, uuid: Uuid) {
108        if let Some(transform_id) = self.uuid_to_transform.remove(&uuid) {
109            self.available_transform_ids.push(transform_id);
110        }
111    }
112
113    /// Get the transform ID for a given UUID, if it exists.
114    pub fn get(&self, uuid: &Uuid) -> Option<u32> {
115        self.uuid_to_transform.get(uuid).copied()
116    }
117}
118
119#[derive(Component, Default)]
120struct ChangedTransformToRetryMarker;
121#[allow(clippy::type_complexity, clippy::too_many_arguments)]
122fn update_registry(
123    mut commands: Commands,
124    mut registry: ResMut<PbrUuidRegistry>,
125    pending_removals: Res<PbrSsboTransformPendingRemovals>,
126    changed_transforms: Query<
127        (Entity, &PbrSsboTransformUuid, &GlobalTransform),
128        (With<PbrSsboTransformUuid>, Changed<GlobalTransform>)
129    >,
130    changeds_to_retry: Query<
131        (Entity, &PbrSsboTransformUuid, &GlobalTransform),
132        (
133            With<PbrSsboTransformUuid>,
134            With<ChangedTransformToRetryMarker>
135        )
136    >,
137    ssbo: ResRenderData<PbrSsboTransform>,
138    buffers: Res<RenderAssets<GpuBuffer>>,
139    render_instance: Res<RenderInstance>
140) {
141    // Free transform IDs for removed entities
142    for uuid in &pending_removals.0 {
143        registry.remove(*uuid);
144    }
145
146    // Return early if there are no changed transforms to update
147    if changed_transforms.is_empty() && changeds_to_retry.is_empty() {
148        return;
149    }
150
151    // Update transform IDs for changed entities
152    let (ssbo_staging, ssbo_gpu) = match (
153        ssbo.iter().next().and_then(|(_, d)| {
154            buffers.get(&d.get_buffer(PbrSsboTransform::TRANSFORM_STAGING_IDX)?)
155        }),
156        ssbo.iter()
157            .next()
158            .and_then(|(_, d)| buffers.get(&d.get_buffer(PbrSsboTransform::TRANSFORM_IDX)?))
159    ) {
160        (Some(staging), Some(gpu)) => (&staging.buffer, &gpu.buffer),
161        _ => {
162            for (entity, _, _) in changed_transforms.iter().chain(changeds_to_retry.iter()) {
163                commands
164                    .entity(entity)
165                    .insert(ChangedTransformToRetryMarker);
166            }
167            return;
168        }
169    };
170    let render_instance = render_instance.0.read().unwrap();
171    for (entity, uuid, transform) in changed_transforms.iter().chain(changeds_to_retry.iter()) {
172        let transform_id = registry.get_or_add(uuid.0);
173        let offset = (transform_id as usize) * std::mem::size_of::<TransformUniform>();
174        ssbo_staging.write(
175            &render_instance,
176            bytemuck::cast_slice(&[TransformUniform::new(transform)]),
177            offset
178        );
179        commands
180            .entity(entity)
181            .remove::<ChangedTransformToRetryMarker>();
182    }
183    ssbo_gpu.copy_from_buffer(&render_instance, ssbo_staging);
184}
wde_pbr/deferred/transform/registry.rs

wde_pbr/deferred/transform/
registry.rs
