bevy_shader_plugin/
auto_array.rs

1use std::{collections::BTreeMap, marker::PhantomData};
2
3use bevy::{
4    prelude::*,
5    render::{
6        ExtractSchedule, MainWorld, Render, RenderApp, RenderStartup,
7        render_resource::{Buffer, BufferUsages},
8        renderer::{RenderDevice, RenderQueue},
9    },
10};
11use encase::ShaderSize;
12use encase::internal::WriteInto;
13use wgpu::util::BufferInitDescriptor;
14
15use crate::{
16    ShaderSystems,
17    bindings::{AutoBufferKind, BindingTable, register_binding},
18};
19
20/// Main-world resource for queuing index-value changes to a fixed-size array buffer.
21///
22/// `Tag` is a zero-sized marker type that uniquely identifies this buffer, allowing
23/// multiple buffers of the same element type and length to coexist. The element type
24/// and length are specified at registration and do not appear in this type.
25///
26/// Changes are serialized to bytes immediately in [`set`](Self::set), then batched
27/// into contiguous `write_buffer` runs each frame. Indices need not be sorted or
28/// unique — duplicates use last-write-wins.
29///
30/// # Example
31///
32/// ```rust,ignore
33/// struct Colors;
34/// // registered via: app.register_array_buffer::<Colors, Vec4, 64>(0, 0, false)
35///
36/// fn update(mut changes: ResMut<ArrayBufferChanges<Colors>>) {
37///     changes.set(0, Vec4::new(1.0, 0.0, 0.0, 1.0)); // red at index 0
38///     changes.set_many([(1, Vec4::Y), (2, Vec4::Z)]); // green, blue
39/// }
40/// ```
41#[derive(Resource)]
42pub struct ArrayBufferChanges<Tag> {
43    pub(crate) changes: Vec<(usize, Vec<u8>)>,
44    pub(crate) len: usize,
45    pub(crate) _marker: PhantomData<Tag>,
46}
47
48impl<Tag> Default for ArrayBufferChanges<Tag> {
49    fn default() -> Self {
50        Self {
51            changes: Vec::new(),
52            len: 0,
53            _marker: PhantomData,
54        }
55    }
56}
57
58fn serialize_element<T: ShaderSize + WriteInto>(value: &T) -> Vec<u8> {
59    let el_size = T::SHADER_SIZE.get() as usize;
60    let mut bytes = vec![0u8; el_size];
61    encase::StorageBuffer::new(&mut bytes[..])
62        .write(value)
63        .unwrap();
64    bytes
65}
66
67/// WGSL array element stride of `T` in bytes (element size plus alignment padding).
68pub(crate) fn array_stride<T: ShaderSize>() -> usize {
69    // [T; 1]::SHADER_SIZE equals the WGSL array element stride for any length.
70    <[T; 1] as ShaderSize>::SHADER_SIZE.get() as usize
71}
72
73/// Bytes of a WGSL `array<T, len>` with every element set to `T::default()`.
74pub(crate) fn default_filled_bytes<T: ShaderSize + WriteInto + Default>(len: usize) -> Vec<u8> {
75    let stride = array_stride::<T>();
76    let el_bytes = serialize_element(&T::default());
77    let mut bytes = vec![0u8; stride * len];
78    for i in 0..len {
79        bytes[i * stride..i * stride + el_bytes.len()].copy_from_slice(&el_bytes);
80    }
81    bytes
82}
83
84/// Register a fixed-size array buffer. See
85/// [`ShaderAppExt::register_array_buffer`](crate::ShaderAppExt::register_array_buffer).
86pub(crate) fn register<Tag, T, const N: usize>(
87    app: &mut App,
88    group_index: u32,
89    binding_index: u32,
90    read_write: bool,
91) where
92    Tag: Send + Sync + 'static,
93    T: ShaderSize + WriteInto + Default + Send + Sync + 'static,
94{
95    app.insert_resource(ArrayBufferChanges::<Tag> {
96        changes: Vec::new(),
97        len: N,
98        _marker: PhantomData,
99    });
100
101    let render_app = app.sub_app_mut(RenderApp);
102    register_binding(
103        render_app,
104        group_index,
105        binding_index,
106        AutoBufferKind::Storage {
107            read_only: !read_write,
108        },
109    );
110
111    // RenderStartup: create the persistent GPU buffer pre-filled with T::default().
112    render_app.add_systems(
113        RenderStartup,
114        |mut commands: Commands, render_device: Res<RenderDevice>| {
115            let buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
116                label: Some("array_buffer"),
117                contents: &default_filled_bytes::<T>(N),
118                // COPY_SRC lets a future GPU→CPU readback copy out of this buffer.
119                usage: BufferUsages::STORAGE | BufferUsages::COPY_DST | BufferUsages::COPY_SRC,
120            });
121
122            commands.insert_resource(ArrayBufferState::<Tag> {
123                buffer,
124                stride: array_stride::<T>(),
125                _marker: PhantomData,
126            });
127        },
128    );
129
130    render_app.add_systems(ExtractSchedule, extract_array_changes::<Tag>);
131
132    render_app.add_systems(
133        Render,
134        (move |state: Option<Res<ArrayBufferState<Tag>>>,
135               changes: Option<Res<ArrayBufferChanges<Tag>>>,
136               render_queue: Res<RenderQueue>,
137               mut table: ResMut<BindingTable>| {
138            let (Some(state), Some(changes)) = (state, changes) else {
139                return;
140            };
141            apply_array_buffer_changes(
142                &state,
143                &changes,
144                &render_queue,
145                &mut table,
146                group_index,
147                binding_index,
148            );
149        })
150        .in_set(ShaderSystems::PrepareBindings),
151    );
152}
153
154impl<Tag> ArrayBufferChanges<Tag> {
155    /// Queue a change: element at `index` will be updated to `value` this frame.
156    /// The value is serialized to bytes immediately.
157    pub fn set<T: ShaderSize + WriteInto>(&mut self, index: usize, value: T) {
158        if index >= self.len {
159            panic!(
160                "ArrayBufferChanges::set index out of bounds: {} >= {}",
161                index, self.len
162            );
163        }
164        let bytes = serialize_element(&value);
165        self.changes.push((index, bytes));
166    }
167
168    /// Queue multiple changes at once from any iterable of `(index, value)` pairs.
169    ///
170    /// Equivalent to calling [`set`](Self::set) for each pair. Accepts arrays, vecs,
171    /// iterators, or anything that implements `IntoIterator<Item = (usize, T)>`:
172    ///
173    /// ```rust,ignore
174    /// changes.set_many([(8, color_a), (9, color_a), (24, color_b), (25, color_b)]);
175    /// changes.set_many((0..4).map(|i| (i, Vec4::ONE)));
176    /// ```
177    pub fn set_many<T, I>(&mut self, pairs: I)
178    where
179        T: ShaderSize + WriteInto,
180        I: IntoIterator<Item = (usize, T)>,
181    {
182        for (index, value) in pairs {
183            self.set(index, value);
184        }
185    }
186
187    /// Queue a change to set every element in the buffer to `value`.
188    /// The value is serialized once and the bytes are reused for each element.
189    pub fn set_all<T: ShaderSize + WriteInto>(&mut self, value: T) {
190        let bytes = serialize_element(&value);
191        self.changes.reserve(self.len);
192        for i in 0..self.len {
193            self.changes.push((i, bytes.clone()));
194        }
195    }
196}
197
198/// Render-world resource holding the persistent GPU buffer and element stride.
199#[derive(Resource)]
200pub struct ArrayBufferState<Tag> {
201    pub(crate) buffer: Buffer,
202    /// WGSL array element stride in bytes (includes any alignment padding).
203    pub(crate) stride: usize,
204    pub(crate) _marker: PhantomData<Tag>,
205}
206
207/// `ExtractSchedule` system: moves queued changes from main world to render world.
208///
209/// Uses `mem::take` so the main-world vec is left empty after extraction — no
210/// separate clear system is needed, and changes set during `Startup` survive to
211/// the first extraction without being wiped by a `First`-schedule clear.
212pub(crate) fn extract_array_changes<Tag: Send + Sync + 'static>(
213    mut commands: Commands,
214    mut main_world: ResMut<MainWorld>,
215) {
216    let extracted = {
217        let Some(mut changes) = main_world.get_resource_mut::<ArrayBufferChanges<Tag>>() else {
218            return;
219        };
220        (std::mem::take(&mut changes.changes), changes.len)
221    };
222    commands.insert_resource(ArrayBufferChanges::<Tag> {
223        changes: extracted.0,
224        len: extracted.1,
225        _marker: PhantomData,
226    });
227}
228
229/// Applies pending changes to the GPU buffer via contiguous-run `write_buffer` batching,
230/// then records the persistent buffer handle in the [`BindingTable`].
231fn apply_array_buffer_changes<Tag: Send + Sync + 'static>(
232    state: &ArrayBufferState<Tag>,
233    changes: &ArrayBufferChanges<Tag>,
234    render_queue: &RenderQueue,
235    table: &mut BindingTable,
236    group_index: u32,
237    binding_index: u32,
238) {
239    let stride = state.stride;
240
241    // Deduplicate: last .set() wins per index; BTreeMap gives sorted order.
242    let mut by_index: BTreeMap<usize, &[u8]> = BTreeMap::new();
243    for (i, bytes) in &changes.changes {
244        by_index.insert(*i, bytes);
245    }
246
247    // Group consecutive indices into runs; each run becomes one write_buffer call.
248    let mut runs: Vec<(u64, Vec<u8>)> = Vec::new();
249    let mut run_start: Option<usize> = None;
250    let mut run_bytes: Vec<u8> = Vec::new();
251    let mut prev: Option<usize> = None;
252
253    for (&index, &el_bytes) in &by_index {
254        // Pad element bytes to stride (tail bytes remain zero).
255        let mut padded = vec![0u8; stride];
256        padded[..el_bytes.len()].copy_from_slice(el_bytes);
257
258        match prev {
259            Some(p) if index == p + 1 => {
260                run_bytes.extend_from_slice(&padded);
261            }
262            _ => {
263                if let Some(start) = run_start {
264                    runs.push((start as u64 * stride as u64, std::mem::take(&mut run_bytes)));
265                }
266                run_start = Some(index);
267                run_bytes = padded;
268            }
269        }
270        prev = Some(index);
271    }
272    if let Some(start) = run_start {
273        runs.push((start as u64 * stride as u64, run_bytes));
274    }
275
276    for (offset, bytes) in runs {
277        render_queue.write_buffer(&state.buffer, offset, &bytes);
278    }
279
280    table.set_buffer(group_index, binding_index, &state.buffer);
281}
bevy_shader_plugin/auto_array.rs

bevy_shader_plugin/
auto_array.rs
