1import bevy_pbr::render::pbr_functions::calculate_tbn_mikktspace;
2import bevy_render::maths::affine3_to_square;
3
4enable wgpu_ray_query;
5enable wgpu_binding_array;
6
7struct InstanceGeometryIds {
8    vertex_buffer_id: u32,
9    vertex_buffer_offset: u32,
10    index_buffer_id: u32,
11    index_buffer_offset: u32,
12    triangle_count: u32,
13}
14
15struct VertexBuffer { vertices: array<PackedVertex> }
16
17struct IndexBuffer { indices: array<u32> }
18
19struct PackedVertex {
20    a: vec4<f32>,
21    b: vec4<f32>,
22    tangent: vec4<f32>,
23}
24
25struct Vertex {
26    position: vec3<f32>,
27    normal: vec3<f32>,
28    uv: vec2<f32>,
29    tangent: vec4<f32>,
30}
31
32fn unpack_vertex(packed: PackedVertex) -> Vertex {
33    var vertex: Vertex;
34    vertex.position = packed.a.xyz;
35    vertex.normal = vec3(packed.a.w, packed.b.xy);
36    vertex.uv = packed.b.zw;
37    vertex.tangent = packed.tangent;
38    return vertex;
39}
40
41struct Material {
42    normal_map_texture_id: u32,
43    base_color_texture_id: u32,
44    emissive_texture_id: u32,
45    metallic_roughness_texture_id: u32,
46
47    base_color: vec3<f32>,
48    perceptual_roughness: f32,
49    emissive: vec3<f32>,
50    metallic: f32,
51    _padding: vec3<f32>,
52    reflectance: f32,
53}
54
55const TEXTURE_MAP_NONE = 0xFFFFFFFFu;
56
57const MIRROR_ROUGHNESS_THRESHOLD = 0.001f;
58
59struct LightSource {
60    kind: u32, // 1 bit for kind, 31 bits for extra data
61    id: u32,
62}
63
64const LIGHT_SOURCE_KIND_EMISSIVE_MESH = 0u;
65const LIGHT_SOURCE_KIND_DIRECTIONAL = 1u;
66
67struct DirectionalLight {
68    direction_to_light: vec3<f32>,
69    cos_theta_max: f32,
70    luminance: vec3<f32>,
71    inverse_pdf: f32,
72}
73
74struct EnvironmentMapLight {
75    light_from_world: mat3x3<f32>,
76    intensity: f32,
77}
78
79const LIGHT_NOT_PRESENT_THIS_FRAME = 0xFFFFFFFFu;
80
81@group(0) @binding(0) var<storage> vertex_buffers: binding_array<VertexBuffer>;
82@group(0) @binding(1) var<storage> index_buffers: binding_array<IndexBuffer>;
83@group(0) @binding(2) var textures: binding_array<texture_2d<f32>>;
84@group(0) @binding(3) var samplers: binding_array<sampler>;
85@group(0) @binding(4) var<storage> materials: array<Material>;
86@group(0) @binding(5) var tlas: acceleration_structure;
87@group(0) @binding(6) var previous_frame_tlas: acceleration_structure;
88@group(0) @binding(7) var<storage> transforms: array<mat3x4<f32>>;
89@group(0) @binding(8) var<storage> previous_frame_transforms: array<mat3x4<f32>>;
90@group(0) @binding(9) var<storage> geometry_ids: array<InstanceGeometryIds>;
91@group(0) @binding(10) var<storage> material_ids: array<u32>;
92@group(0) @binding(11) var<storage> light_sources: array<LightSource>;
93@group(0) @binding(12) var<storage> directional_lights: array<DirectionalLight>;
94@group(0) @binding(13) var<storage> previous_frame_light_id_translations: array<u32>;
95@group(0) @binding(14) var brdf_dfg_lut: texture_2d<f32>;
96@group(0) @binding(15) var brdf_dfg_lut_sampler: sampler;
97@group(0) @binding(16) var environment_map: texture_cube<f32>;
98@group(0) @binding(17) var environment_map_sampler: sampler;
99@group(0) @binding(18) var<storage> environment_map_light: EnvironmentMapLight;
100
101const RAY_T_MIN = 0.001f;
102const RAY_T_MAX = 100000.0f;
103
104const RAY_NO_CULL = 0xFFu;
105
106fn trace_ray(ray_origin: vec3<f32>, ray_direction: vec3<f32>, ray_t_min: f32, ray_t_max: f32, ray_flag: u32) -> RayIntersection {
107    let ray = RayDesc(ray_flag, RAY_NO_CULL, ray_t_min, ray_t_max, ray_origin, ray_direction);
108    var rq: ray_query;
109    rayQueryInitialize(&rq, tlas, ray);
110    rayQueryProceed(&rq);
111    return rayQueryGetCommittedIntersection(&rq);
112}
113
114fn trace_ray_previous_frame(ray_origin: vec3<f32>, ray_direction: vec3<f32>, ray_t_min: f32, ray_t_max: f32, ray_flag: u32) -> RayIntersection {
115    let ray = RayDesc(ray_flag, RAY_NO_CULL, ray_t_min, ray_t_max, ray_origin, ray_direction);
116    var rq: ray_query;
117    rayQueryInitialize(&rq, previous_frame_tlas, ray);
118    rayQueryProceed(&rq);
119    return rayQueryGetCommittedIntersection(&rq);
120}
121
122fn sample_texture(id: u32, uv: vec2<f32>) -> vec3<f32> {
123    return textureSampleLevel(textures[id], samplers[id], uv, 0.0).rgb; // TODO: Mipmap
124}
125
126struct ResolvedMaterial {
127    base_color: vec3<f32>,
128    emissive: vec3<f32>,
129    reflectance: f32,
130    perceptual_roughness: f32,
131    roughness: f32,
132    metallic: f32,
133}
134
135struct ResolvedRayHitFull {
136    world_position: vec3<f32>,
137    previous_frame_world_position: vec3<f32>,
138    world_normal: vec3<f32>,
139    geometric_world_normal: vec3<f32>,
140    world_tangent: vec4<f32>,
141    uv: vec2<f32>,
142    triangle_area: f32,
143    triangle_count: u32,
144    material: ResolvedMaterial,
145}
146
147fn resolve_material(material: Material, uv: vec2<f32>) -> ResolvedMaterial {
148    var m: ResolvedMaterial;
149
150    m.base_color = material.base_color.rgb;
151    if material.base_color_texture_id != TEXTURE_MAP_NONE {
152        m.base_color *= sample_texture(material.base_color_texture_id, uv);
153    }
154
155    m.emissive = material.emissive.rgb;
156    if material.emissive_texture_id != TEXTURE_MAP_NONE {
157        m.emissive *= sample_texture(material.emissive_texture_id, uv);
158    }
159
160    m.reflectance = material.reflectance;
161
162    m.perceptual_roughness = material.perceptual_roughness;
163    m.metallic = material.metallic;
164    if material.metallic_roughness_texture_id != TEXTURE_MAP_NONE {
165        let metallic_roughness = sample_texture(material.metallic_roughness_texture_id, uv);
166        m.perceptual_roughness *= metallic_roughness.g;
167        m.metallic *= metallic_roughness.b;
168    }
169
170    m.roughness = m.perceptual_roughness * m.perceptual_roughness;
171
172    return m;
173}
174
175fn resolve_ray_hit_full(ray_hit: RayIntersection) -> ResolvedRayHitFull {
176    let barycentrics = vec3(1.0 - ray_hit.barycentrics.x - ray_hit.barycentrics.y, ray_hit.barycentrics);
177    return resolve_triangle_data_full(ray_hit.instance_custom_data, ray_hit.primitive_index, barycentrics);
178}
179
180fn load_vertices(instance_geometry_ids: InstanceGeometryIds, triangle_id: u32) -> array<Vertex, 3> {
181    let index_buffer = &index_buffers[instance_geometry_ids.index_buffer_id].indices;
182    let vertex_buffer = &vertex_buffers[instance_geometry_ids.vertex_buffer_id].vertices;
183
184    let indices_i = (triangle_id * 3u) + vec3(0u, 1u, 2u) + instance_geometry_ids.index_buffer_offset;
185    let indices = vec3((*index_buffer)[indices_i.x], (*index_buffer)[indices_i.y], (*index_buffer)[indices_i.z]) + instance_geometry_ids.vertex_buffer_offset;
186
187    return array<Vertex, 3>(
188        unpack_vertex((*vertex_buffer)[indices.x]),
189        unpack_vertex((*vertex_buffer)[indices.y]),
190        unpack_vertex((*vertex_buffer)[indices.z])
191    );
192}
193
194fn transform_positions(transform: mat4x4<f32>, vertices: array<Vertex, 3>) -> array<vec3<f32>, 3> {
195    return array<vec3<f32>, 3>(
196        (transform * vec4(vertices[0].position, 1.0)).xyz,
197        (transform * vec4(vertices[1].position, 1.0)).xyz,
198        (transform * vec4(vertices[2].position, 1.0)).xyz
199    );
200}
201
202fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics: vec3<f32>) -> ResolvedRayHitFull {
203    let material_id = material_ids[instance_id];
204    let material = materials[material_id];
205
206    let transform = affine3_to_square(transforms[instance_id]);
207    let previous_frame_transform = affine3_to_square(previous_frame_transforms[instance_id]);
208
209    let instance_geometry_ids = geometry_ids[instance_id];
210    let vertices = load_vertices(instance_geometry_ids, triangle_id);
211
212    let world_vertices = transform_positions(transform, vertices);
213    let world_position = mat3x3(world_vertices[0], world_vertices[1], world_vertices[2]) * barycentrics;
214
215    let previous_frame_world_vertices = transform_positions(previous_frame_transform, vertices);
216    let previous_frame_world_position = mat3x3(previous_frame_world_vertices[0], previous_frame_world_vertices[1], previous_frame_world_vertices[2]) * barycentrics;
217
218    let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics;
219
220    let transform_3x3 = mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz);
221
222    let cofactor = mat3x3(cross(transform_3x3[1], transform_3x3[2]), cross(transform_3x3[2], transform_3x3[0]), cross(transform_3x3[0], transform_3x3[1]));
223    let determinant = dot(transform_3x3[0], cofactor[0]);
224    let determinant_sign = select(sign(determinant), 1.0, determinant == 0.0);
225
226    let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics;
227    let world_tangent = vec4(normalize(transform_3x3 * local_tangent), vertices[0].tangent.w * determinant_sign);
228
229    let local_normal = mat3x3(vertices[0].normal, vertices[1].normal, vertices[2].normal) * barycentrics; // TODO: Use barycentric lerp, ray_hit.object_to_world, cross product geo normal
230    var world_normal = normalize(determinant_sign * (cofactor * local_normal));
231
232    let geometric_world_normal = world_normal;
233    if material.normal_map_texture_id != TEXTURE_MAP_NONE {
234        let TBN = calculate_tbn_mikktspace(world_normal, world_tangent);
235        let T = TBN[0];
236        let B = TBN[1];
237        let N = TBN[2];
238        var Nt = sample_texture(material.normal_map_texture_id, uv) * 2.0 - 1.0;
239        Nt.z = sqrt(max(1.0 - dot(Nt.xy, Nt.xy), 0.0)); // Reconstruct Z to support two-channel normal maps
240        world_normal = normalize(Nt.x * T + Nt.y * B + Nt.z * N);
241    }
242
243    let triangle_edge0 = world_vertices[0] - world_vertices[1];
244    let triangle_edge1 = world_vertices[0] - world_vertices[2];
245    let triangle_area = length(cross(triangle_edge0, triangle_edge1)) / 2.0;
246
247    let resolved_material = resolve_material(material, uv);
248
249    return ResolvedRayHitFull(
250        world_position,
251        previous_frame_world_position,
252        world_normal,
253        geometric_world_normal,
254        world_tangent,
255        uv,
256        triangle_area,
257        instance_geometry_ids.triangle_count,
258        resolved_material,
259    );
260}