Skip to main content

wde_pbr/deferred/batches/
ssbo_batches.rs

1use bevy::{ecs::system::SystemParamItem, prelude::*};
2use wde_renderer::prelude::*;
3
4use crate::prelude::BatchList;
5
6/// The maximum number of batches in the ssbo.
7pub const SSBO_MAX_BATCHES: usize = 100_000;
8
9pub(crate) struct SsboInstancesToTransformPlugin;
10impl Plugin for SsboInstancesToTransformPlugin {
11    fn build(&self, app: &mut App) {
12        app.add_plugins(RenderDataRegisterPlugin::<PbrSsboInstanceToTransform>::default());
13    }
14}
15
16/// Contains pointers to the transforms of the entities in the SSBO, as well as the offset in the "instance to transform" buffer of the first entity of this batch. This is used to bind the correct transforms to the instances of each batch when rendering.
17#[derive(Asset, Clone, TypePath, Default)]
18pub struct PbrSsboInstanceToTransform;
19impl PbrSsboInstanceToTransform {
20    pub const INSTANCE_TO_TRANSFORM_IDX: u32 = 0;
21    pub const INSTANCE_TO_TRANSFORM_STAGING_IDX: u32 = 1;
22}
23impl RenderData for PbrSsboInstanceToTransform {
24    type Params = ();
25
26    fn describe(_params: &mut SystemParamItem<Self::Params>, builder: &mut RenderDataBuilder) {
27        builder
28            .add_buffer(
29                Self::INSTANCE_TO_TRANSFORM_IDX,
30                Buffer {
31                    label: "pbr-instance-to-transform-buffer-gpu".to_string(),
32                    size: std::mem::size_of::<u32>() * SSBO_MAX_BATCHES,
33                    usage: BufferUsage::STORAGE | BufferUsage::COPY_DST,
34                    content: None
35                }
36            )
37            .add_buffer(
38                Self::INSTANCE_TO_TRANSFORM_STAGING_IDX,
39                Buffer {
40                    label: "pbr-instance-to-transform-staging".to_string(),
41                    size: std::mem::size_of::<u32>() * SSBO_MAX_BATCHES,
42                    usage: BufferUsage::COPY_SRC | BufferUsage::COPY_DST,
43                    content: None
44                }
45            );
46    }
47}
48
49pub(crate) fn set_instances_to_transform(
50    instance_to_transform: ResRenderData<PbrSsboInstanceToTransform>,
51    mut batches: ResMut<BatchList>,
52    buffers: Res<RenderAssets<GpuBuffer>>,
53    render_instance: Res<RenderInstance>
54) {
55    if !batches.dirty {
56        return;
57    }
58
59    // Get buffers
60    let instance_to_transform = match instance_to_transform.iter().next() {
61        Some((_, data)) => data,
62        None => return
63    };
64    let (staging_buffer, gpu_buffer) = match (
65        buffers.get(
66            &instance_to_transform
67                .get_buffer(PbrSsboInstanceToTransform::INSTANCE_TO_TRANSFORM_STAGING_IDX)
68                .unwrap()
69        ),
70        buffers.get(
71            &instance_to_transform
72                .get_buffer(PbrSsboInstanceToTransform::INSTANCE_TO_TRANSFORM_IDX)
73                .unwrap()
74        )
75    ) {
76        (Some(staging), Some(gpu)) => (staging, gpu),
77        _ => return
78    };
79
80    // Update the staging buffer with the new instance to transform mappings for each batch
81    let mut staging_content = vec![0u8; staging_buffer.buffer.buffer.size() as usize];
82    let mut offset = 0;
83    for batch_key in batches.sorted_batches.clone().iter() {
84        let batch = batches.batches.get_mut(batch_key).unwrap();
85        batch.instances_offset = offset;
86        for transform_id in &batch.transform_ssbo_ids {
87            let bytes = transform_id.to_ne_bytes();
88            staging_content[(offset as usize * 4)..((offset as usize + 1) * 4)]
89                .copy_from_slice(&bytes);
90            offset += 1;
91        }
92    }
93    let render_instance = render_instance.0.read().unwrap();
94    staging_buffer
95        .buffer
96        .write(&render_instance, &staging_content, 0);
97
98    // Copy the staging buffer to the GPU buffer
99    gpu_buffer
100        .buffer
101        .copy_from_buffer(&render_instance, &staging_buffer.buffer);
102
103    // Clear dirty batches
104    batches.dirty = false;
105}