1//! Vertex/fragment shader for rasterizing large clusters into a visibility buffer.
2
3import package::{
4    meshlet::bindings::{
5        meshlet_cluster_meshlet_ids,
6        meshlets,
7        meshlet_cluster_instance_ids,
8        meshlet_instance_uniforms,
9        meshlet_raster_clusters,
10        meshlet_previous_raster_counts,
11        meshlet_visibility_buffer,
12        view,
13        get_meshlet_triangle_count,
14        get_meshlet_vertex_id,
15        get_meshlet_vertex_position,
16    },
17    render::mesh_functions::mesh_position_local_to_world,
18};
19import bevy_render::maths::affine3_to_square;
20var<immediate> meshlet_raster_cluster_rightmost_slot: u32;
21
22struct VertexOutput {
23    @builtin(position) position: vec4<f32>,
24    @if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
25    @location(0) @interpolate(flat) packed_ids: u32,
26}
27
28@vertex
29fn vertex(@builtin(instance_index) instance_index: u32, @builtin(vertex_index) vertex_index: u32) -> VertexOutput {
30    let cluster_in_draw = meshlet_previous_raster_counts[1] + instance_index;
31    let cluster_id = meshlet_raster_cluster_rightmost_slot - cluster_in_draw;
32    let instanced_offset = meshlet_raster_clusters[cluster_id];
33    var meshlet = meshlets[instanced_offset.offset];
34
35    let triangle_id = vertex_index / 3u;
36    if triangle_id >= get_meshlet_triangle_count(&meshlet) { return dummy_vertex(); }
37    let index_id = vertex_index;
38    let vertex_id = get_meshlet_vertex_id(meshlet.start_index_id + index_id);
39
40    let instance_uniform = meshlet_instance_uniforms[instanced_offset.instance_id];
41
42    let vertex_position = get_meshlet_vertex_position(&meshlet, vertex_id);
43    let world_from_local = affine3_to_square(instance_uniform.world_from_local);
44    let world_position = mesh_position_local_to_world(world_from_local, vec4(vertex_position, 1.0));
45    let clip_position = view.clip_from_world * vec4(world_position.xyz, 1.0);
46
47    @if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
48    return VertexOutput(
49        clip_position,
50        (cluster_id << 7u) | triangle_id,
51    );
52    @else
53    return VertexOutput(
54        clip_position,
55    );
56}
57
58@fragment
59fn fragment(vertex_output: VertexOutput) {
60    let depth = bitcast<u32>(vertex_output.position.z);
61    @if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
62    let visibility = (u64(depth) << 32u) | u64(vertex_output.packed_ids);
63    @else
64    let visibility = depth;
65    textureAtomicMax(meshlet_visibility_buffer, vec2<u32>(vertex_output.position.xy), visibility);
66}
67
68fn dummy_vertex() -> VertexOutput {
69    @if(MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT)
70    return VertexOutput(
71        vec4(divide(0.0, 0.0)), // NaN vertex position
72        0u,
73    );
74    @else
75    return VertexOutput(
76        vec4(divide(0.0, 0.0)),
77    );
78}
79
80/// Naga doesn't allow divide by zero literals, but this lets us work around it
81fn divide(a: f32, b: f32) -> f32 {
82    return a / b;
83}