1use std::collections::HashMap;
2
3use bytemuck::{Pod, Zeroable};
4use wgpu::util::DeviceExt;
5
6use crate::target::Targets;
7use crate::texture::{TextureId, TextureRegistry};
8
9#[repr(C)]
10#[derive(Copy, Clone, Pod, Zeroable)]
11struct QuadVertex {
12 pos: [f32; 2],
13 uv: [f32; 2],
14}
15
16const QUAD_VERTS: &[QuadVertex] = &[
17 QuadVertex {
18 pos: [0.0, 0.0],
19 uv: [0.0, 0.0],
20 },
21 QuadVertex {
22 pos: [1.0, 0.0],
23 uv: [1.0, 0.0],
24 },
25 QuadVertex {
26 pos: [1.0, 1.0],
27 uv: [1.0, 1.0],
28 },
29 QuadVertex {
30 pos: [0.0, 1.0],
31 uv: [0.0, 1.0],
32 },
33];
34const QUAD_INDICES: &[u16] = &[0, 1, 2, 0, 2, 3];
35
36#[repr(C)]
37#[derive(Copy, Clone, Pod, Zeroable, Debug)]
38pub struct SpriteInstance {
39 pub pos: [f32; 2],
40 pub size: [f32; 2],
41 pub uv_min: [f32; 2],
42 pub uv_max: [f32; 2],
43 pub color: [f32; 4],
44 pub rotation: f32,
45 pub _pad: [f32; 3],
46}
47
48impl SpriteInstance {
49 pub fn at(pos: [f32; 2], size: [f32; 2]) -> Self {
50 Self {
51 pos,
52 size,
53 uv_min: [0.0, 0.0],
54 uv_max: [1.0, 1.0],
55 color: [1.0; 4],
56 rotation: 0.0,
57 _pad: [0.0; 3],
58 }
59 }
60 pub fn with_color(mut self, c: [f32; 4]) -> Self {
61 self.color = c;
62 self
63 }
64 pub fn with_rotation(mut self, r: f32) -> Self {
65 self.rotation = r;
66 self
67 }
68 pub fn with_uv(mut self, min: [f32; 2], max: [f32; 2]) -> Self {
69 self.uv_min = min;
70 self.uv_max = max;
71 self
72 }
73 pub fn with_region(self, min: [f32; 2], size: [f32; 2], tex_size: [u32; 2]) -> Self {
78 debug_assert!(
79 tex_size[0] > 0 && tex_size[1] > 0,
80 "with_region: zero-sized texture {tex_size:?}"
81 );
82 let [tw, th] = [tex_size[0] as f32, tex_size[1] as f32];
83 self.with_uv(
84 [min[0] / tw, min[1] / th],
85 [(min[0] + size[0]) / tw, (min[1] + size[1]) / th],
86 )
87 }
88 pub fn flip_x(mut self) -> Self {
91 std::mem::swap(&mut self.uv_min[0], &mut self.uv_max[0]);
92 self
93 }
94 pub fn flip_y(mut self) -> Self {
97 std::mem::swap(&mut self.uv_min[1], &mut self.uv_max[1]);
98 self
99 }
100}
101
102#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
103pub enum BlendMode {
104 Alpha,
105 Additive,
106 Premultiplied,
107}
108
109impl BlendMode {
110 pub(crate) const ALL: [Self; 3] = [Self::Alpha, Self::Additive, Self::Premultiplied];
111
112 pub(crate) fn state(self) -> wgpu::BlendState {
114 match self {
115 Self::Alpha => wgpu::BlendState::ALPHA_BLENDING,
116 Self::Premultiplied => wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING,
117 Self::Additive => wgpu::BlendState {
118 color: wgpu::BlendComponent {
119 src_factor: wgpu::BlendFactor::SrcAlpha,
120 dst_factor: wgpu::BlendFactor::One,
121 operation: wgpu::BlendOperation::Add,
122 },
123 alpha: wgpu::BlendComponent::OVER,
124 },
125 }
126 }
127
128 pub(crate) fn label(self) -> &'static str {
129 match self {
130 Self::Alpha => "alpha",
131 Self::Additive => "add",
132 Self::Premultiplied => "premul",
133 }
134 }
135}
136
137#[derive(Copy, Clone, PartialEq, Eq, Hash)]
138struct BatchKey {
139 texture: TextureId,
140 layer: i16,
141 blend: BlendMode,
142}
143
144pub(crate) struct SpriteBatcher {
145 quad_vb: wgpu::Buffer,
146 quad_ib: wgpu::Buffer,
147 instance_vb: wgpu::Buffer,
148 instance_capacity: usize,
149 pipelines: HashMap<BlendMode, wgpu::RenderPipeline>,
150 pending: Vec<(BatchKey, SpriteInstance)>,
151}
152
153impl SpriteBatcher {
154 pub fn new(
155 device: &wgpu::Device,
156 surface_format: wgpu::TextureFormat,
157 camera_bgl: &wgpu::BindGroupLayout,
158 texture_bgl: &wgpu::BindGroupLayout,
159 sample_count: u32,
160 depth_format: Option<wgpu::TextureFormat>,
161 ) -> Self {
162 let multisample = crate::target::multisample(sample_count);
163 let depth_stencil = depth_format.map(crate::target::no_write_depth);
164 let quad_vb = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
165 label: Some("sprite.quad_vb"),
166 contents: bytemuck::cast_slice(QUAD_VERTS),
167 usage: wgpu::BufferUsages::VERTEX,
168 });
169 let quad_ib = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
170 label: Some("sprite.quad_ib"),
171 contents: bytemuck::cast_slice(QUAD_INDICES),
172 usage: wgpu::BufferUsages::INDEX,
173 });
174
175 let instance_capacity = 4096usize;
176 let instance_vb = device.create_buffer(&wgpu::BufferDescriptor {
177 label: Some("sprite.instances"),
178 size: (instance_capacity * size_of::<SpriteInstance>()) as u64,
179 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
180 mapped_at_creation: false,
181 });
182
183 let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
184 label: Some("sprite.shader"),
185 source: wgpu::ShaderSource::Wgsl(include_str!("sprite.wgsl").into()),
186 });
187
188 let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
189 label: Some("sprite.layout"),
190 bind_group_layouts: &[Some(camera_bgl), Some(texture_bgl)],
191 immediate_size: 0,
192 });
193
194 let make_pipeline = |blend: wgpu::BlendState, label: &str| {
195 device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
196 label: Some(label),
197 layout: Some(&layout),
198 vertex: wgpu::VertexState {
199 module: &shader,
200 entry_point: Some("vs_main"),
201 compilation_options: Default::default(),
202 buffers: &[
203 Some(wgpu::VertexBufferLayout {
204 array_stride: size_of::<QuadVertex>() as u64,
205 step_mode: wgpu::VertexStepMode::Vertex,
206 attributes: &wgpu::vertex_attr_array![0 => Float32x2, 1 => Float32x2],
207 }),
208 Some(wgpu::VertexBufferLayout {
209 array_stride: size_of::<SpriteInstance>() as u64,
210 step_mode: wgpu::VertexStepMode::Instance,
211 attributes: &wgpu::vertex_attr_array![
212 2 => Float32x2,
213 3 => Float32x2,
214 4 => Float32x2,
215 5 => Float32x2,
216 6 => Float32x4,
217 7 => Float32,
218 ],
219 }),
220 ],
221 },
222 fragment: Some(wgpu::FragmentState {
223 module: &shader,
224 entry_point: Some("fs_main"),
225 compilation_options: Default::default(),
226 targets: &[Some(wgpu::ColorTargetState {
227 format: surface_format,
228 blend: Some(blend),
229 write_mask: wgpu::ColorWrites::ALL,
230 })],
231 }),
232 primitive: wgpu::PrimitiveState::default(),
233 depth_stencil: depth_stencil.clone(),
234 multisample,
235 multiview_mask: None,
236 cache: None,
237 })
238 };
239
240 let pipelines = BlendMode::ALL
241 .into_iter()
242 .map(|blend| {
243 let label = format!("sprite.{}", blend.label());
244 (blend, make_pipeline(blend.state(), &label))
245 })
246 .collect();
247
248 Self {
249 quad_vb,
250 quad_ib,
251 instance_vb,
252 instance_capacity,
253 pipelines,
254 pending: Vec::new(),
255 }
256 }
257
258 pub fn draw(&mut self, tex: TextureId, layer: i16, blend: BlendMode, inst: SpriteInstance) {
259 self.pending.push((
260 BatchKey {
261 texture: tex,
262 layer,
263 blend,
264 },
265 inst,
266 ));
267 }
268
269 pub fn collect_layers(&self, out: &mut std::collections::BTreeSet<i16>) {
271 out.extend(self.pending.iter().map(|(k, _)| k.layer));
272 }
273
274 pub fn upload(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
279 if self.pending.is_empty() {
280 return;
281 }
282 self.pending.sort_by(|a, b| {
283 a.0.layer
284 .cmp(&b.0.layer)
285 .then((a.0.blend as u8).cmp(&(b.0.blend as u8)))
286 .then(a.0.texture.0.cmp(&b.0.texture.0))
287 });
288 if self.pending.len() > self.instance_capacity {
289 self.instance_capacity = self.pending.len().next_power_of_two();
290 self.instance_vb = device.create_buffer(&wgpu::BufferDescriptor {
291 label: Some("sprite.instances"),
292 size: (self.instance_capacity * size_of::<SpriteInstance>()) as u64,
293 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
294 mapped_at_creation: false,
295 });
296 }
297 let flat: Vec<SpriteInstance> = self.pending.iter().map(|(_, i)| *i).collect();
298 queue.write_buffer(&self.instance_vb, 0, bytemuck::cast_slice(&flat));
299 }
300
301 pub fn draw_layer(
304 &self,
305 layer: i16,
306 encoder: &mut wgpu::CommandEncoder,
307 targets: &Targets<'_>,
308 camera_bg: &wgpu::BindGroup,
309 textures: &TextureRegistry,
310 ) {
311 let lo = self.pending.partition_point(|(k, _)| k.layer < layer);
313 let hi = self.pending.partition_point(|(k, _)| k.layer <= layer);
314 if lo == hi {
315 return;
316 }
317
318 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
319 label: Some("sprite.pass"),
320 color_attachments: &[Some(targets.color_attachment(wgpu::LoadOp::Load))],
321 depth_stencil_attachment: targets.depth_attachment(wgpu::LoadOp::Load),
322 occlusion_query_set: None,
323 timestamp_writes: None,
324 multiview_mask: None,
325 });
326
327 pass.set_bind_group(0, camera_bg, &[]);
328 pass.set_vertex_buffer(0, self.quad_vb.slice(..));
329 pass.set_index_buffer(self.quad_ib.slice(..), wgpu::IndexFormat::Uint16);
330 pass.set_vertex_buffer(1, self.instance_vb.slice(..));
331
332 let mut i = lo;
333 while i < hi {
334 let key = self.pending[i].0;
335 let start = i;
336 while i < hi
337 && self.pending[i].0.blend == key.blend
338 && self.pending[i].0.texture == key.texture
339 {
340 i += 1;
341 }
342 let count = (i - start) as u32;
343 pass.set_pipeline(&self.pipelines[&key.blend]);
344 pass.set_bind_group(1, textures.bind_group(key.texture), &[]);
345 pass.draw_indexed(0..6, 0, (start as u32)..(start as u32 + count));
346 }
347 }
348
349 pub fn clear(&mut self) {
350 self.pending.clear();
351 }
352}
353
354#[cfg(test)]
355mod tests {
356 use super::*;
357
358 #[test]
359 fn with_region_maps_pixels_to_uvs() {
360 let inst = SpriteInstance::at([0.0, 0.0], [32.0, 16.0]).with_region(
361 [64.0, 32.0],
362 [32.0, 16.0],
363 [256, 128],
364 );
365 assert_eq!(inst.uv_min, [0.25, 0.25]);
366 assert_eq!(inst.uv_max, [0.375, 0.375]);
367 }
368
369 #[test]
370 #[cfg(debug_assertions)]
371 #[should_panic(expected = "zero-sized texture")]
372 fn with_region_rejects_a_zero_sized_texture_in_debug_builds() {
373 let _ =
374 SpriteInstance::at([0.0, 0.0], [1.0, 1.0]).with_region([0.0, 0.0], [1.0, 1.0], [0, 64]);
375 }
376
377 #[test]
378 fn flip_x_swaps_u_bounds_only() {
379 let inst = SpriteInstance::at([0.0, 0.0], [1.0, 1.0])
380 .with_uv([0.1, 0.2], [0.3, 0.4])
381 .flip_x();
382 assert_eq!(inst.uv_min, [0.3, 0.2]);
383 assert_eq!(inst.uv_max, [0.1, 0.4]);
384 }
385
386 #[test]
387 fn flip_y_swaps_v_bounds_only() {
388 let inst = SpriteInstance::at([0.0, 0.0], [1.0, 1.0])
389 .with_uv([0.1, 0.2], [0.3, 0.4])
390 .flip_y();
391 assert_eq!(inst.uv_min, [0.1, 0.4]);
392 assert_eq!(inst.uv_max, [0.3, 0.2]);
393 }
394
395 #[test]
396 fn flipping_twice_restores_the_region() {
397 let inst = SpriteInstance::at([0.0, 0.0], [1.0, 1.0]).with_region(
398 [8.0, 8.0],
399 [16.0, 16.0],
400 [64, 64],
401 );
402 let back = inst.flip_x().flip_y().flip_x().flip_y();
403 assert_eq!(back.uv_min, inst.uv_min);
404 assert_eq!(back.uv_max, inst.uv_max);
405 }
406}