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