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