1use crate::{GpuBuffer, GpuSlot};
2use wgpu::{BufferAddress, BufferUsages, CommandEncoder, Device, Queue};
3
4pub const STREAM: BufferUsages = BufferUsages::STORAGE
5 .union(BufferUsages::COPY_DST)
6 .union(BufferUsages::COPY_SRC);
7pub const UNIFORM: BufferUsages = BufferUsages::UNIFORM.union(BufferUsages::COPY_DST);
8pub const PACK: BufferUsages = BufferUsages::COPY_DST.union(BufferUsages::COPY_SRC);
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub struct StreamElement {
12 wgsl: &'static str,
13 bytes: u64,
14}
15
16impl StreamElement {
17 pub const fn new(wgsl: &'static str, bytes: u64) -> Self {
18 assert!(bytes > 0, "a stream element must occupy at least one byte");
19 Self { wgsl, bytes }
20 }
21
22 pub const fn wgsl(self) -> &'static str {
23 self.wgsl
24 }
25
26 pub const fn bytes(self) -> u64 {
27 self.bytes
28 }
29}
30
31#[derive(Clone, Copy, PartialEq, Eq)]
32pub enum Contents {
33 Scratch,
34 Durable,
35 Seeded(u32),
36}
37
38impl Contents {
39 pub const fn durable(self) -> bool {
40 matches!(self, Self::Durable | Self::Seeded(_))
41 }
42
43 pub const fn seed(self) -> Option<u32> {
44 match self {
45 Self::Seeded(word) => Some(word),
46 Self::Scratch | Self::Durable => None,
47 }
48 }
49}
50
51pub struct Stream {
52 label: &'static str,
53 buffer: GpuBuffer,
54 slots: u32,
55 stride: u64,
56 element: StreamElement,
57 usage: BufferUsages,
58 contents: Contents,
59}
60
61fn assert_fits(device: &Device, label: &str, bytes: BufferAddress) {
62 assert!(bytes > 0, "stream {label:?} requires at least one byte");
63 let limits = device.limits();
64 assert!(
65 bytes <= limits.max_storage_buffer_binding_size,
66 "stream {label:?} requires {bytes} bytes but the device storage binding limit is {}",
67 limits.max_storage_buffer_binding_size
68 );
69 assert!(
70 bytes <= limits.max_buffer_size,
71 "stream {label:?} requires {bytes} bytes but the device buffer limit is {}",
72 limits.max_buffer_size
73 );
74}
75
76pub struct StreamDesc {
77 pub label: &'static str,
78 pub slots: u32,
79 pub element: StreamElement,
80 pub elements_per_slot: u64,
81 pub usage: BufferUsages,
82 pub contents: Contents,
83}
84
85#[derive(Clone, Copy)]
86pub struct TypedSlot<'a> {
87 slot: GpuSlot<'a>,
88 element: StreamElement,
89}
90
91impl<'a> TypedSlot<'a> {
92 pub fn new(slot: GpuSlot<'a>, element: StreamElement) -> Self {
93 Self { slot, element }
94 }
95
96 pub fn slot(self) -> GpuSlot<'a> {
97 self.slot
98 }
99
100 pub fn element(self) -> StreamElement {
101 self.element
102 }
103
104 pub fn identity(self) -> (u64, u64, u64) {
105 self.slot.identity()
106 }
107}
108
109impl Stream {
110 pub fn new(device: &Device, queue: &Queue, desc: StreamDesc) -> Self {
111 let StreamDesc {
112 label,
113 slots,
114 element,
115 elements_per_slot,
116 usage,
117 contents,
118 } = desc;
119 assert!(slots > 0, "stream {label:?} requires at least one slot");
120 assert!(
121 elements_per_slot > 0,
122 "stream {label:?} requires at least one element per slot"
123 );
124 assert!(
125 element.bytes().is_multiple_of(4),
126 "stream {label:?} requires a word aligned record of {}",
127 element.wgsl()
128 );
129 let stride = element.bytes() * elements_per_slot;
130 let bytes = slots as BufferAddress * stride;
131 assert_fits(device, label, bytes);
132 let stream = Self {
133 label,
134 buffer: GpuBuffer::new(device, label, bytes, usage),
135 slots,
136 stride,
137 element,
138 usage,
139 contents,
140 };
141 if let Some(word) = contents.seed() {
142 stream.write_at(queue, 0, &word.to_le_bytes());
143 }
144 stream
145 }
146
147 pub fn slots(&self) -> u32 {
148 self.slots
149 }
150
151 pub fn stride(&self) -> u64 {
152 self.stride
153 }
154
155 pub fn element(&self) -> StreamElement {
156 self.element
157 }
158
159 pub fn size(&self) -> BufferAddress {
160 self.buffer.size()
161 }
162
163 pub fn buffer(&self) -> &wgpu::Buffer {
164 self.buffer.buffer()
165 }
166
167 pub fn gpu(&self) -> &GpuBuffer {
168 &self.buffer
169 }
170
171 pub fn slot(&self) -> GpuSlot<'_> {
172 GpuSlot::whole(&self.buffer)
173 }
174
175 pub fn write(&self, queue: &Queue, bytes: &[u8]) {
176 self.buffer.write(queue, bytes);
177 }
178
179 pub fn write_at(&self, queue: &Queue, offset: BufferAddress, bytes: &[u8]) {
180 self.buffer.write_at(queue, offset, bytes);
181 }
182
183 pub fn reserve(&mut self, device: &Device, encoder: &mut CommandEncoder, slots: u32) -> bool {
184 assert!(
185 slots > 0,
186 "stream {:?} requires at least one slot",
187 self.label
188 );
189 if slots == self.slots {
190 return false;
191 }
192 let bytes = slots as BufferAddress * self.stride;
193 assert_fits(device, self.label, bytes);
194 let next = GpuBuffer::new(device, self.label, bytes, self.usage);
195 if self.contents.durable() {
196 encoder.copy_buffer_to_buffer(
197 self.buffer.buffer(),
198 0,
199 next.buffer(),
200 0,
201 self.buffer.size().min(next.size()),
202 );
203 }
204 self.buffer = next;
205 self.slots = slots;
206 true
207 }
208}
209
210impl<'a> From<&'a Stream> for GpuSlot<'a> {
211 fn from(stream: &'a Stream) -> Self {
212 stream.slot()
213 }
214}
215
216impl<'a> From<&'a Stream> for TypedSlot<'a> {
217 fn from(stream: &'a Stream) -> Self {
218 Self::new(stream.slot(), stream.element())
219 }
220}