1import package::meshlet::bindings::{
2    MeshletAabb,
3    DispatchIndirectArgs,
4    InstancedOffset,
5    depth_pyramid,
6    view,
7    previous_view,
8    meshlet_instance_uniforms,
9};
10import bevy_render::maths::affine3_to_square;
11
12/// https://github.com/zeux/meshoptimizer/blob/1e48e96c7e8059321de492865165e9ef071bffba/demo/nanite.cpp#L115
13fn lod_error_is_imperceptible(lod_sphere: vec4<f32>, simplification_error: f32, instance_id: u32) -> bool {
14    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
15    let world_scale = max(length(world_from_local[0]), max(length(world_from_local[1]), length(world_from_local[2])));
16    let camera_pos = view.world_position;
17
18    let projection = view.clip_from_view;
19    if projection[3][3] == 1.0 {
20        // Orthographic
21        let world_error = simplification_error * world_scale;
22        let proj = projection[1][1];
23        let height = 2.0 / proj;
24        let norm_error = world_error / height;
25        return norm_error * view.viewport.w < 1.0;
26    } else {
27        // Perspective
28        var near = projection[3][2];
29        let world_sphere_center = (world_from_local * vec4<f32>(lod_sphere.xyz, 1.0)).xyz;
30        let world_sphere_radius = lod_sphere.w * world_scale;
31        let d_pos = world_sphere_center - camera_pos;
32        let d = sqrt(dot(d_pos, d_pos)) - world_sphere_radius;
33        let norm_error = simplification_error / max(d, near) * projection[1][1] * 0.5;
34        return norm_error * view.viewport.w < 1.0;
35    }
36}
37
38fn normalize_plane(p: vec4<f32>) -> vec4<f32> {
39    return p / length(p.xyz);
40}
41
42/// https://fgiesen.wordpress.com/2012/08/31/frustum-planes-from-the-projection-matrix/
43/// https://fgiesen.wordpress.com/2010/10/17/view-frustum-culling/
44fn aabb_in_frustum(aabb: MeshletAabb, instance_id: u32) -> bool {
45    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
46    let clip_from_local = view.clip_from_world * world_from_local;
47    let row_major = transpose(clip_from_local);
48    let planes = array(
49        row_major[3] + row_major[0],
50        row_major[3] - row_major[0],
51        row_major[3] + row_major[1],
52        row_major[3] - row_major[1],
53        row_major[2],
54    );
55
56    for (var i = 0; i < 5; i++) {
57        let plane = normalize_plane(planes[i]);
58        let flipped = aabb.half_extent * sign(plane.xyz);
59        if dot(aabb.center + flipped, plane.xyz) <= -plane.w {
60            return false;
61        }
62    }
63    return true;
64}
65
66struct ScreenAabb {
67    min: vec3<f32>,
68    max: vec3<f32>,
69}
70
71fn min8(a: vec3<f32>, b: vec3<f32>, c: vec3<f32>, d: vec3<f32>, e: vec3<f32>, f: vec3<f32>, g: vec3<f32>, h: vec3<f32>) -> vec3<f32> {
72    return min(min(min(a, b), min(c, d)), min(min(e, f), min(g, h)));
73}
74
75fn max8(a: vec3<f32>, b: vec3<f32>, c: vec3<f32>, d: vec3<f32>, e: vec3<f32>, f: vec3<f32>, g: vec3<f32>, h: vec3<f32>) -> vec3<f32> {
76    return max(max(max(a, b), max(c, d)), max(max(e, f), max(g, h)));
77}
78
79fn min8_4(a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>, e: vec4<f32>, f: vec4<f32>, g: vec4<f32>, h: vec4<f32>) -> vec4<f32> {
80    return min(min(min(a, b), min(c, d)), min(min(e, f), min(g, h)));
81}
82
83/// https://zeux.io/2023/01/12/approximate-projected-bounds/
84fn project_aabb(clip_from_local: mat4x4<f32>, near: f32, aabb: MeshletAabb, out: ptr<function, ScreenAabb>) -> bool {
85    let extent = aabb.half_extent * 2.0;
86    let sx = clip_from_local * vec4<f32>(extent.x, 0.0, 0.0, 0.0);
87    let sy = clip_from_local * vec4<f32>(0.0, extent.y, 0.0, 0.0);
88    let sz = clip_from_local * vec4<f32>(0.0, 0.0, extent.z, 0.0);
89
90    let p0 = clip_from_local * vec4<f32>(aabb.center - aabb.half_extent, 1.0);
91    let p1 = p0 + sz;
92    let p2 = p0 + sy;
93    let p3 = p2 + sz;
94    let p4 = p0 + sx;
95    let p5 = p4 + sz;
96    let p6 = p4 + sy;
97    let p7 = p6 + sz;
98
99    let depth = min8_4(p0, p1, p2, p3, p4, p5, p6, p7).w;
100    // do not occlusion cull if we are inside the aabb
101    if depth < near {
102        return false;
103    }
104
105    let dp0 = p0.xyz / p0.w;
106    let dp1 = p1.xyz / p1.w;
107    let dp2 = p2.xyz / p2.w;
108    let dp3 = p3.xyz / p3.w;
109    let dp4 = p4.xyz / p4.w;
110    let dp5 = p5.xyz / p5.w;
111    let dp6 = p6.xyz / p6.w;
112    let dp7 = p7.xyz / p7.w;
113    let min = min8(dp0, dp1, dp2, dp3, dp4, dp5, dp6, dp7);
114    let max = max8(dp0, dp1, dp2, dp3, dp4, dp5, dp6, dp7);
115    var vaabb = vec4<f32>(min.xy, max.xy);
116    // convert ndc to texture coordinates by rescaling and flipping Y
117    vaabb = vaabb.xwzy * vec4<f32>(0.5, -0.5, 0.5, -0.5) + 0.5;
118    (*out).min = vec3<f32>(vaabb.xy, min.z);
119    (*out).max = vec3<f32>(vaabb.zw, max.z);
120    return true;
121}
122
123fn sample_hzb(smin: vec2<u32>, smax: vec2<u32>, mip: i32) -> f32 {
124    let texel = vec4<u32>(0, 1, 2, 3);
125    let sx = min(smin.x + texel, smax.xxxx);
126    let sy = min(smin.y + texel, smax.yyyy);
127    // TODO: switch to min samplers when wgpu has them
128    // sampling 16 times a finer mip is worth the extra cost for better culling
129    let a = sample_hzb_row(sx, sy.x, mip);
130    let b = sample_hzb_row(sx, sy.y, mip);
131    let c = sample_hzb_row(sx, sy.z, mip);
132    let d = sample_hzb_row(sx, sy.w, mip);
133    return min(min(a, b), min(c, d));
134}
135
136fn sample_hzb_row(sx: vec4<u32>, sy: u32, mip: i32) -> f32 {
137    let a = textureLoad(depth_pyramid, vec2(sx.x, sy), mip).x;
138    let b = textureLoad(depth_pyramid, vec2(sx.y, sy), mip).x;
139    let c = textureLoad(depth_pyramid, vec2(sx.z, sy), mip).x;
140    let d = textureLoad(depth_pyramid, vec2(sx.w, sy), mip).x;
141    return min(min(a, b), min(c, d));
142}
143
144fn occlusion_cull_screen_aabb(aabb: ScreenAabb, screen: vec2<f32>) -> bool {
145    let hzb_size = vec2<f32>(textureDimensions(depth_pyramid).xy);
146    let aabb_min = aabb.min.xy * hzb_size;
147    let aabb_max = aabb.max.xy * hzb_size;
148
149    let min_texel = vec2<u32>(max(aabb_min, vec2<f32>(0.0)));
150    let max_texel = vec2<u32>(min(aabb_max, hzb_size - 1.0));
151    let size = max_texel - min_texel;
152    let max_size = max(size.x, size.y);
153
154    // note: add 1 before max because the unsigned overflow behavior is intentional
155    // it wraps around firstLeadingBit(0) = ~0 to 0
156    // TODO: we actually sample a 4x4 block, so ideally this would be `max(..., 3u) - 3u`.
157    var mip = max(firstLeadingBit(max_size) + 1u, 2u) - 2u;
158
159    if any((max_texel >> vec2(mip)) > (min_texel >> vec2(mip)) + 3) {
160        mip += 1u;
161    }
162
163    let smin = min_texel >> vec2<u32>(mip);
164    let smax = max_texel >> vec2<u32>(mip);
165
166    let curr_depth = sample_hzb(smin, smax, i32(mip));
167    return aabb.max.z <= curr_depth;
168}
169
170fn occlusion_cull_projection() -> mat4x4<f32> {
171    @if(MESHLET_FIRST_CULLING_PASS)
172    return previous_view.clip_from_world;
173    @else
174    return view.clip_from_world;
175}
176
177fn occlusion_cull_clip_from_local(instance_id: u32) -> mat4x4<f32> {
178    @if(MESHLET_FIRST_CULLING_PASS)
179    let prev_world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].previous_world_from_local);
180    @if(MESHLET_FIRST_CULLING_PASS)
181    return previous_view.clip_from_world * prev_world_from_local;
182    @if(!MESHLET_FIRST_CULLING_PASS)
183    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
184    @if(!MESHLET_FIRST_CULLING_PASS)
185    return view.clip_from_world * world_from_local;
186}
187
188fn should_occlusion_cull_aabb(aabb: MeshletAabb, instance_id: u32) -> bool {
189    let projection = occlusion_cull_projection();
190    var near: f32;
191    if projection[3][3] == 1.0 {
192        near = projection[3][2] / projection[2][2];
193    } else {
194        near = projection[3][2];
195    }
196
197    let clip_from_local = occlusion_cull_clip_from_local(instance_id);
198    var screen_aabb = ScreenAabb(vec3<f32>(0.0), vec3<f32>(0.0));
199    if project_aabb(clip_from_local, near, aabb, &screen_aabb) {
200        return occlusion_cull_screen_aabb(screen_aabb, view.viewport.zw);
201    }
202    return false;
203}