renderling/atlas/
shader.rs

1use crabslab::{Id, Slab, SlabItem};
2use glam::{UVec2, UVec3, Vec2, Vec3, Vec3Swizzles, Vec4, Vec4Swizzles};
3use spirv_std::{image::Image2d, spirv, Sampler};
4
5/// Describes various qualities of the atlas, to be used on the GPU.
6#[derive(Clone, Copy, core::fmt::Debug, Default, PartialEq, SlabItem)]
7pub struct AtlasDescriptor {
8    pub size: UVec3,
9}
10
11/// A texture inside the atlas.
12#[cfg_attr(not(target_arch = "spirv"), derive(Debug))]
13#[derive(Clone, Copy, Default, PartialEq, SlabItem)]
14pub struct AtlasTextureDescriptor {
15    /// The top left offset of texture in the atlas.
16    pub offset_px: UVec2,
17    /// The size of the texture in the atlas.
18    pub size_px: UVec2,
19    /// The index of the layer within the atlas that this `AtlasTexture `belongs to.
20    pub layer_index: u32,
21    /// The index of this frame within the layer.
22    pub frame_index: u32,
23    /// Various toggles of texture modes.
24    pub modes: super::TextureModes,
25}
26
27impl AtlasTextureDescriptor {
28    /// Transform the given `uv` coordinates for this texture's address mode
29    /// and placement in the atlas of the given size.
30    pub fn uv(&self, mut uv: Vec2, atlas_size: UVec2) -> Vec3 {
31        uv.x = self.modes.s.wrap(uv.x);
32        uv.y = self.modes.t.wrap(uv.y);
33
34        // get the pixel index of the uv coordinate in terms of the original image
35        let mut px_index_s = (uv.x * self.size_px.x as f32) as u32;
36        let mut px_index_t = (uv.y * self.size_px.y as f32) as u32;
37
38        // convert the pixel index from image to atlas space
39        px_index_s += self.offset_px.x;
40        px_index_t += self.offset_px.y;
41
42        let sx = atlas_size.x as f32;
43        let sy = atlas_size.y as f32;
44        // normalize the pixels by dividing by the atlas size
45        let uv_s = px_index_s as f32 / sx;
46        let uv_t = px_index_t as f32 / sy;
47
48        Vec2::new(uv_s, uv_t).extend(self.layer_index as f32)
49    }
50
51    /// Constrain the input `clip_pos` to be within the bounds of this texture
52    /// within its atlas, in texture space.
53    pub fn constrain_clip_coords_to_texture_space(
54        &self,
55        clip_pos: Vec2,
56        atlas_size: UVec2,
57    ) -> Vec2 {
58        // Convert `clip_pos` into uv coords to figure out where in the texture
59        // this point lives
60        let input_uv = (clip_pos * Vec2::new(1.0, -1.0) + Vec2::splat(1.0)) * Vec2::splat(0.5);
61        self.uv(input_uv, atlas_size).xy()
62    }
63
64    /// Constrain the input `clip_pos` to be within the bounds of this texture
65    /// within its atlas.
66    pub fn constrain_clip_coords(&self, clip_pos: Vec2, atlas_size: UVec2) -> Vec2 {
67        let uv = self.constrain_clip_coords_to_texture_space(clip_pos, atlas_size);
68        // Convert `uv` back into clip space
69        (uv * Vec2::new(2.0, 2.0) - Vec2::splat(1.0)) * Vec2::new(1.0, -1.0)
70    }
71
72    #[cfg(cpu)]
73    /// Returns the frame of this texture as a [`wgpu::Origin3d`].
74    pub fn origin(&self) -> wgpu::Origin3d {
75        wgpu::Origin3d {
76            x: self.offset_px.x,
77            y: self.offset_px.y,
78            z: self.layer_index,
79        }
80    }
81
82    #[cfg(cpu)]
83    /// Returns the frame of this texture as a [`wgpu::Extent3d`].
84    pub fn size_as_extent(&self) -> wgpu::Extent3d {
85        wgpu::Extent3d {
86            width: self.size_px.x,
87            height: self.size_px.y,
88            depth_or_array_layers: 1,
89        }
90    }
91}
92
93#[derive(Clone, Copy, Default, SlabItem, core::fmt::Debug)]
94pub struct AtlasBlittingDescriptor {
95    pub atlas_texture_id: Id<AtlasTextureDescriptor>,
96    pub atlas_desc_id: Id<AtlasDescriptor>,
97}
98
99/// Vertex shader for blitting a texture into a the frame of an
100/// [`AtlasTextureDescriptor`].
101///
102/// This is useful for copying textures of unsupported formats, or
103/// textures of different sizes.
104#[spirv(vertex)]
105pub fn atlas_blit_vertex(
106    #[spirv(vertex_index)] vertex_id: u32,
107    #[spirv(instance_index)] atlas_blitting_desc_id: Id<AtlasBlittingDescriptor>,
108    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] slab: &[u32],
109    out_uv: &mut Vec2,
110    #[spirv(position)] out_pos: &mut Vec4,
111) {
112    let i = vertex_id as usize;
113    *out_uv = crate::math::UV_COORD_QUAD_CCW[i];
114
115    crate::println!("atlas_blitting_desc_id: {atlas_blitting_desc_id:?}");
116    let atlas_blitting_desc = slab.read_unchecked(atlas_blitting_desc_id);
117    crate::println!("atlas_blitting_desc: {atlas_blitting_desc:?}");
118    let atlas_texture = slab.read_unchecked(atlas_blitting_desc.atlas_texture_id);
119    crate::println!("atlas_texture: {atlas_texture:?}");
120    let atlas_desc = slab.read_unchecked(atlas_blitting_desc.atlas_desc_id);
121    crate::println!("atlas_desc: {atlas_desc:?}");
122    let clip_pos = crate::math::CLIP_SPACE_COORD_QUAD_CCW[i];
123    crate::println!("clip_pos: {clip_pos:?}");
124    *out_pos = atlas_texture
125        .constrain_clip_coords(clip_pos.xy(), atlas_desc.size.xy())
126        .extend(clip_pos.z)
127        .extend(clip_pos.w);
128    crate::println!("out_pos: {out_pos}");
129}
130
131/// Fragment shader for blitting a texture into a frame of an atlas.
132#[spirv(fragment)]
133pub fn atlas_blit_fragment(
134    #[spirv(descriptor_set = 0, binding = 1)] texture: &Image2d,
135    #[spirv(descriptor_set = 0, binding = 2)] sampler: &Sampler,
136    in_uv: Vec2,
137    frag_color: &mut Vec4,
138) {
139    *frag_color = texture.sample(*sampler, in_uv);
140}