renderling/cubemap/
shader.rs

1use crabslab::{Array, Id, Slab};
2use glam::{Mat4, Vec2, Vec3, Vec3Swizzles, Vec4};
3use spirv_std::{num_traits::Zero, spirv};
4
5use crate::{
6    atlas::shader::{AtlasDescriptor, AtlasTextureDescriptor},
7    math::{IsSampler, Sample2dArray},
8};
9
10/// Vertex shader for testing cubemap sampling.
11#[spirv(vertex)]
12pub fn cubemap_sampling_test_vertex(
13    #[spirv(vertex_index)] vertex_index: u32,
14    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] uv: &Vec3,
15    out_uv: &mut Vec3,
16    #[spirv(position)] out_clip_coords: &mut Vec4,
17) {
18    let vertex_index = vertex_index as usize % 6;
19    *out_clip_coords = crate::math::CLIP_SPACE_COORD_QUAD_CCW[vertex_index];
20    *out_uv = *uv;
21}
22
23/// Vertex shader for testing cubemap sampling.
24#[spirv(fragment)]
25pub fn cubemap_sampling_test_fragment(
26    #[spirv(descriptor_set = 0, binding = 1)] cubemap: &spirv_std::image::Cubemap,
27    #[spirv(descriptor_set = 0, binding = 2)] sampler: &spirv_std::Sampler,
28    in_uv: Vec3,
29    frag_color: &mut Vec4,
30) {
31    *frag_color = cubemap.sample(*sampler, in_uv);
32}
33
34/// Represents one side of a cubemap.
35///
36/// Assumes the camera is at the origin, inside the cube, with
37/// a left-handed coordinate system (+Z going into the screen).
38#[derive(Clone, Copy)]
39pub struct CubemapFaceDirection {
40    /// Where is the camera
41    pub eye: Vec3,
42    /// Where is the camera looking
43    pub dir: Vec3,
44    /// Which direction is up
45    pub up: Vec3,
46}
47
48impl CubemapFaceDirection {
49    pub const X: Self = Self {
50        eye: Vec3::ZERO,
51        dir: Vec3::X,
52        up: Vec3::Y,
53    };
54    pub const NEG_X: Self = Self {
55        eye: Vec3::ZERO,
56        dir: Vec3::NEG_X,
57        up: Vec3::Y,
58    };
59
60    pub const Y: Self = Self {
61        eye: Vec3::ZERO,
62        dir: Vec3::Y,
63        up: Vec3::NEG_Z,
64    };
65    pub const NEG_Y: Self = Self {
66        eye: Vec3::ZERO,
67        dir: Vec3::NEG_Y,
68        up: Vec3::Z,
69    };
70
71    pub const Z: Self = Self {
72        eye: Vec3::ZERO,
73        dir: Vec3::Z,
74        up: Vec3::Y,
75    };
76    pub const NEG_Z: Self = Self {
77        eye: Vec3::ZERO,
78        dir: Vec3::NEG_Z,
79        up: Vec3::Y,
80    };
81
82    pub const FACES: [Self; 6] = [
83        CubemapFaceDirection::X,
84        CubemapFaceDirection::NEG_X,
85        CubemapFaceDirection::Y,
86        CubemapFaceDirection::NEG_Y,
87        CubemapFaceDirection::Z,
88        CubemapFaceDirection::NEG_Z,
89    ];
90
91    pub fn right(&self) -> Vec3 {
92        -self.dir.cross(self.up)
93    }
94
95    /// The view from _inside_ the cube.
96    pub fn view(&self) -> Mat4 {
97        Mat4::look_at_lh(self.eye, self.eye + self.dir, self.up)
98    }
99}
100
101pub struct CubemapDescriptor {
102    atlas_descriptor_id: Id<AtlasDescriptor>,
103    faces: Array<AtlasTextureDescriptor>,
104}
105
106impl CubemapDescriptor {
107    /// Return the face index and UV coordinates that can be used to sample
108    /// a cubemap from the given directional coordinate.
109    pub fn get_face_index_and_uv(coord: Vec3) -> (usize, Vec2) {
110        let abs_x = coord.x.abs();
111        let abs_y = coord.y.abs();
112        let abs_z = coord.z.abs();
113
114        let (face_index, uv) = if abs_x >= abs_y && abs_x >= abs_z {
115            if coord.x > 0.0 {
116                (0, Vec2::new(-coord.z, -coord.y) / abs_x)
117            } else {
118                (1, Vec2::new(coord.z, -coord.y) / abs_x)
119            }
120        } else if abs_y >= abs_x && abs_y >= abs_z {
121            if coord.y > 0.0 {
122                (2, Vec2::new(coord.x, coord.z) / abs_y)
123            } else {
124                (3, Vec2::new(coord.x, -coord.z) / abs_y)
125            }
126        } else if coord.z > 0.0 {
127            (4, Vec2::new(coord.x, -coord.y) / abs_z)
128        } else {
129            (5, Vec2::new(-coord.x, -coord.y) / abs_z)
130        };
131
132        (face_index, (uv + Vec2::ONE) / 2.0)
133    }
134
135    /// Sample the cubemap with a directional coordinate.
136    pub fn sample<A, S>(&self, coord: Vec3, slab: &[u32], atlas: &A, sampler: &S) -> Vec4
137    where
138        A: Sample2dArray<Sampler = S>,
139        S: IsSampler,
140    {
141        let coord = if coord.length().is_zero() {
142            Vec3::X
143        } else {
144            coord.normalize()
145        };
146        let (face_index, uv) = Self::get_face_index_and_uv(coord);
147        let atlas_image = slab.read_unchecked(self.faces.at(face_index));
148        let atlas_desc = slab.read_unchecked(self.atlas_descriptor_id);
149        let uv = atlas_image.uv(uv, atlas_desc.size.xy());
150        atlas.sample_by_lod(*sampler, uv, 0.0)
151    }
152}
153
154#[cfg(test)]
155mod test {
156    use super::*;
157
158    #[test]
159    fn cubemap_right() {
160        assert_eq!(Vec3::NEG_Z, CubemapFaceDirection::X.right());
161        assert_eq!(Vec3::Z, CubemapFaceDirection::NEG_X.right());
162        assert_eq!(Vec3::X, CubemapFaceDirection::Y.right());
163        assert_eq!(Vec3::X, CubemapFaceDirection::NEG_Y.right());
164        assert_eq!(Vec3::X, CubemapFaceDirection::Z.right());
165        assert_eq!(Vec3::NEG_X, CubemapFaceDirection::NEG_Z.right());
166
167        assert_eq!(
168            (1, Vec2::new(0.0, 1.0)),
169            CubemapDescriptor::get_face_index_and_uv(Vec3::NEG_ONE)
170        );
171    }
172
173    #[test]
174    fn cubemap_face_index() {
175        let center = Vec2::splat(0.5);
176        let data = [
177            (Vec3::X, 0, center),
178            (Vec3::NEG_X, 1, center),
179            (Vec3::Y, 2, center),
180            (Vec3::NEG_Y, 3, center),
181            (Vec3::Z, 4, center),
182            (Vec3::NEG_Z, 5, center),
183        ];
184        for (coord, expected_face_index, expected_uv) in data {
185            let (seen_face_index, seen_uv) = CubemapDescriptor::get_face_index_and_uv(coord);
186            dbg!((coord, seen_face_index, seen_uv));
187            assert_eq!(expected_face_index, seen_face_index);
188            assert_eq!(expected_uv, seen_uv);
189        }
190    }
191}