Skip to main content

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