Skip to main content

game_toolkit_gfx/
sprite.rs

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    /// Sample the `size` pixel rectangle at `min` (top-left, in pixels) of a texture that is
74    /// `tex_size` pixels large, e.g. one sprite of an atlas. Get `tex_size` from
75    /// `Graphics::texture_size` / `Painter::texture_size`, which never reports a zero side;
76    /// a zero `tex_size` side would produce non-finite UVs and is a debug-build panic.
77    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    /// Mirror horizontally by swapping the u bounds. Works with any UV range or region;
89    /// apply it after [`with_uv`](Self::with_uv) / [`with_region`](Self::with_region).
90    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    /// Mirror vertically by swapping the v bounds. Works with any UV range or region;
95    /// apply it after [`with_uv`](Self::with_uv) / [`with_region`](Self::with_region).
96    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    /// The wgpu blend state every batcher uses for this mode.
113    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    /// Record every layer that has pending sprites, for cross-batcher interleaving.
270    pub fn collect_layers(&self, out: &mut std::collections::BTreeSet<i16>) {
271        out.extend(self.pending.iter().map(|(k, _)| k.layer));
272    }
273
274    /// Sort all pending sprites by layer (then blend, then texture) and upload them to the
275    /// instance buffer in one write. Must run before any [`SpriteBatcher::draw_layer`]: the
276    /// buffer is written once per frame because `queue.write_buffer` does not interleave with
277    /// encoder passes, so a per-pass write would clobber the earlier layers' data.
278    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    /// Draw the sprites on `layer` (already uploaded by [`SpriteBatcher::upload`]) into the
302    /// already-cleared target, grouped by blend then texture to minimize state changes.
303    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        // `pending` is sorted by layer, so this layer's sprites are a contiguous range.
312        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}