wde_pbr/deferred/transform/
registry.rs

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

wde_pbr/deferred/transform/
registry.rs
