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