renderling/primitive/
shader.rs

1//! Shader support for rendering primitives.
2use crabslab::{Array, Id, Slab, SlabItem};
3use glam::{Mat4, Vec2, Vec3, Vec4, Vec4Swizzles};
4use spirv_std::{
5    image::{Cubemap, Image2d, Image2dArray},
6    spirv, Image, Sampler,
7};
8
9// use glam::Mat4;
10// #[cfg(not(target_arch = "spirv"))]
11// use glam::UVec2;
12
13// #[allow(unused_imports)]
14// use spirv_std::num_traits::Float;
15
16use crate::{
17    bvol::BoundingSphere,
18    geometry::{
19        shader::{GeometryDescriptor, SkinDescriptor},
20        MorphTarget, Vertex,
21    },
22    material::shader::MaterialDescriptor,
23    math::IsVector,
24    transform::shader::TransformDescriptor,
25};
26
27#[allow(unused_imports)]
28use spirv_std::num_traits::Float;
29
30/// Returned by [`PrimitiveDescriptor::get_vertex_info`].
31pub struct VertexInfo {
32    pub vertex: Vertex,
33    pub transform: TransformDescriptor,
34    pub model_matrix: Mat4,
35    pub world_pos: Vec3,
36}
37
38/// A draw call used to render some geometry.
39#[derive(Clone, Copy, PartialEq, SlabItem, Debug)]
40#[offsets]
41pub struct PrimitiveDescriptor {
42    pub visible: bool,
43    pub vertices_array: Array<Vertex>,
44    /// Bounding sphere of the entire primitive, in local space.
45    pub bounds: BoundingSphere,
46    pub indices_array: Array<u32>,
47    pub transform_id: Id<TransformDescriptor>,
48    pub material_id: Id<MaterialDescriptor>,
49    pub skin_id: Id<SkinDescriptor>,
50    pub morph_targets: Array<Array<MorphTarget>>,
51    pub morph_weights: Array<f32>,
52    pub geometry_descriptor_id: Id<GeometryDescriptor>,
53}
54
55impl Default for PrimitiveDescriptor {
56    fn default() -> Self {
57        PrimitiveDescriptor {
58            visible: true,
59            vertices_array: Array::default(),
60            bounds: BoundingSphere::default(),
61            indices_array: Array::default(),
62            transform_id: Id::NONE,
63            material_id: Id::NONE,
64            skin_id: Id::NONE,
65            morph_targets: Array::default(),
66            morph_weights: Array::default(),
67            geometry_descriptor_id: Id::new(0),
68        }
69    }
70}
71
72impl PrimitiveDescriptor {
73    /// Returns the vertex at the given index and its related values.
74    ///
75    /// These values are often used in shaders, so they are grouped together.
76    pub fn get_vertex_info(&self, vertex_index: u32, geometry_slab: &[u32]) -> VertexInfo {
77        let vertex = self.get_vertex(vertex_index, geometry_slab);
78        let transform = self.get_transform(vertex, geometry_slab);
79        let model_matrix = Mat4::from(transform);
80        let world_pos = model_matrix.transform_point3(vertex.position);
81        VertexInfo {
82            vertex,
83            transform,
84            model_matrix,
85            world_pos,
86        }
87    }
88    /// Retrieve the transform of this `primitive`.
89    ///
90    /// This takes into consideration all skinning matrices.
91    pub fn get_transform(&self, vertex: Vertex, slab: &[u32]) -> TransformDescriptor {
92        let config = slab.read_unchecked(self.geometry_descriptor_id);
93        if config.has_skinning && self.skin_id.is_some() {
94            let skin = slab.read(self.skin_id);
95            TransformDescriptor::from(skin.get_skinning_matrix(vertex, slab))
96        } else {
97            slab.read(self.transform_id)
98        }
99    }
100
101    /// Retrieve the vertex from the slab, calculating any displacement due to
102    /// morph targets.
103    pub fn get_vertex(&self, vertex_index: u32, slab: &[u32]) -> Vertex {
104        let index = if self.indices_array.is_null() {
105            vertex_index as usize
106        } else {
107            slab.read(self.indices_array.at(vertex_index as usize)) as usize
108        };
109        let vertex_id = self.vertices_array.at(index);
110        let mut vertex = slab.read_unchecked(vertex_id);
111        for i in 0..self.morph_targets.len() {
112            let morph_target_array = slab.read(self.morph_targets.at(i));
113            let morph_target = slab.read(morph_target_array.at(index));
114            let weight = slab.read(self.morph_weights.at(i));
115            vertex.position += weight * morph_target.position;
116            vertex.normal += weight * morph_target.normal;
117            vertex.tangent += weight * morph_target.tangent.extend(0.0);
118        }
119        vertex
120    }
121
122    pub fn get_vertex_count(&self) -> u32 {
123        if self.indices_array.is_null() {
124            self.vertices_array.len() as u32
125        } else {
126            self.indices_array.len() as u32
127        }
128    }
129}
130
131#[cfg(test)]
132/// A helper struct that contains all outputs of the primitive's PBR vertex
133/// shader.
134#[derive(Default, Debug, Clone, Copy, PartialEq)]
135pub struct PrimitivePbrVertexInfo {
136    pub primitive: PrimitiveDescriptor,
137    pub primitive_id: Id<PrimitiveDescriptor>,
138    pub vertex_index: u32,
139    pub vertex: Vertex,
140    pub transform: TransformDescriptor,
141    pub model_matrix: Mat4,
142    pub view_projection: Mat4,
143    pub out_color: Vec4,
144    pub out_uv0: Vec2,
145    pub out_uv1: Vec2,
146    pub out_norm: Vec3,
147    pub out_tangent: Vec3,
148    pub out_bitangent: Vec3,
149    pub out_pos: Vec3,
150    pub out_clip_pos: Vec4,
151}
152
153/// primitive vertex shader.
154#[spirv(vertex)]
155#[allow(clippy::too_many_arguments)]
156pub fn primitive_vertex(
157    // Points at a `primitive`
158    #[spirv(instance_index)] primitive_id: Id<PrimitiveDescriptor>,
159    // Which vertex within the primitive are we rendering
160    #[spirv(vertex_index)] vertex_index: u32,
161    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] geometry_slab: &[u32],
162
163    #[spirv(flat)] out_primitive: &mut Id<PrimitiveDescriptor>,
164    // TODO: Think about placing all these out values in a G-Buffer
165    // But do we have enough buffers + enough space on web?
166    // ...and can we write to buffers from vertex shaders on web?
167    out_color: &mut Vec4,
168    out_uv0: &mut Vec2,
169    out_uv1: &mut Vec2,
170    out_norm: &mut Vec3,
171    out_tangent: &mut Vec3,
172    out_bitangent: &mut Vec3,
173    out_world_pos: &mut Vec3,
174    #[spirv(position)] out_clip_pos: &mut Vec4,
175    // test-only info struct
176    #[cfg(test)] out_info: &mut PrimitivePbrVertexInfo,
177) {
178    let primitive = geometry_slab.read_unchecked(primitive_id);
179    if !primitive.visible {
180        // put it outside the clipping frustum
181        *out_clip_pos = Vec4::new(10.0, 10.0, 10.0, 1.0);
182        return;
183    }
184
185    *out_primitive = primitive_id;
186
187    let VertexInfo {
188        vertex,
189        transform,
190        model_matrix,
191        world_pos,
192    } = primitive.get_vertex_info(vertex_index, geometry_slab);
193    *out_color = vertex.color;
194    *out_uv0 = vertex.uv0;
195    *out_uv1 = vertex.uv1;
196    *out_world_pos = world_pos;
197
198    let scale2 = transform.scale * transform.scale;
199    let normal = vertex.normal.alt_norm_or_zero();
200    let tangent = vertex.tangent.xyz().alt_norm_or_zero();
201    let normal_w: Vec3 = (model_matrix * (normal / scale2).extend(0.0))
202        .xyz()
203        .alt_norm_or_zero();
204    *out_norm = normal_w;
205
206    let tangent_w: Vec3 = (model_matrix * tangent.extend(0.0))
207        .xyz()
208        .alt_norm_or_zero();
209    *out_tangent = tangent_w;
210
211    let bitangent_w = normal_w.cross(tangent_w) * if vertex.tangent.w >= 0.0 { 1.0 } else { -1.0 };
212    *out_bitangent = bitangent_w;
213
214    let camera_id = geometry_slab
215        .read_unchecked(primitive.geometry_descriptor_id + GeometryDescriptor::OFFSET_OF_CAMERA_ID);
216    let camera = geometry_slab.read(camera_id);
217    let clip_pos = camera.view_projection() * world_pos.extend(1.0);
218    *out_clip_pos = clip_pos;
219    #[cfg(test)]
220    {
221        *out_info = PrimitivePbrVertexInfo {
222            primitive_id,
223            vertex_index,
224            vertex,
225            transform,
226            model_matrix,
227            view_projection: camera.view_projection(),
228            out_clip_pos: clip_pos,
229            primitive,
230            out_color: *out_color,
231            out_uv0: *out_uv0,
232            out_uv1: *out_uv1,
233            out_norm: *out_norm,
234            out_tangent: *out_tangent,
235            out_bitangent: *out_bitangent,
236            out_pos: *out_world_pos,
237        };
238    }
239}
240
241/// primitive fragment shader
242#[allow(clippy::too_many_arguments, dead_code)]
243#[spirv(fragment)]
244pub fn primitive_fragment(
245    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] geometry_slab: &[u32],
246    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] material_slab: &[u32],
247    #[spirv(descriptor_set = 0, binding = 2)] atlas: &Image2dArray,
248    #[spirv(descriptor_set = 0, binding = 3)] atlas_sampler: &Sampler,
249    #[spirv(descriptor_set = 0, binding = 4)] irradiance: &Cubemap,
250    #[spirv(descriptor_set = 0, binding = 5)] irradiance_sampler: &Sampler,
251    #[spirv(descriptor_set = 0, binding = 6)] prefiltered: &Cubemap,
252    #[spirv(descriptor_set = 0, binding = 7)] prefiltered_sampler: &Sampler,
253    #[spirv(descriptor_set = 0, binding = 8)] brdf: &Image2d,
254    #[spirv(descriptor_set = 0, binding = 9)] brdf_sampler: &Sampler,
255    #[spirv(storage_buffer, descriptor_set = 0, binding = 10)] light_slab: &[u32],
256    #[spirv(descriptor_set = 0, binding = 11)] shadow_map: &Image!(2D, type=f32, sampled, arrayed),
257    #[spirv(descriptor_set = 0, binding = 12)] shadow_map_sampler: &Sampler,
258    #[cfg(feature = "debug-slab")]
259    #[spirv(storage_buffer, descriptor_set = 0, binding = 13)]
260    _debug_slab: &mut [u32],
261
262    #[spirv(flat)] primitive_id: Id<PrimitiveDescriptor>,
263    #[spirv(frag_coord)] frag_coord: Vec4,
264    in_color: Vec4,
265    in_uv0: Vec2,
266    in_uv1: Vec2,
267    in_norm: Vec3,
268    in_tangent: Vec3,
269    in_bitangent: Vec3,
270    world_pos: Vec3,
271    output: &mut Vec4,
272) {
273    // proxy to a separate impl that allows us to test on CPU
274    crate::pbr::shader::fragment_impl(
275        atlas,
276        atlas_sampler,
277        irradiance,
278        irradiance_sampler,
279        prefiltered,
280        prefiltered_sampler,
281        brdf,
282        brdf_sampler,
283        shadow_map,
284        shadow_map_sampler,
285        geometry_slab,
286        material_slab,
287        light_slab,
288        primitive_id,
289        frag_coord,
290        in_color,
291        in_uv0,
292        in_uv1,
293        in_norm,
294        in_tangent,
295        in_bitangent,
296        world_pos,
297        output,
298    );
299}
300
301#[cfg(feature = "test_i8_16_extraction")]
302#[spirv(compute(threads(32)))]
303/// A shader to ensure that we can extract i8 and i16 values from a storage
304/// buffer.
305pub fn test_i8_i16_extraction(
306    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] slab: &mut [u32],
307    #[spirv(global_invocation_id)] global_id: UVec3,
308) {
309    let index = global_id.x as usize;
310    let (value, _, _) = crate::bits::extract_i8(index, 2, slab);
311    if value > 0 {
312        slab[index] = value as u32;
313    }
314    let (value, _, _) = crate::bits::extract_i16(index, 2, slab);
315    if value > 0 {
316        slab[index] = value as u32;
317    }
318}
renderling/primitive/shader.rs

renderling/primitive/
shader.rs
