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.01f;
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    m.roughness = clamp(m.perceptual_roughness * m.perceptual_roughness, 0.001, 1.0);
146
147    return m;
148}
149
150fn resolve_ray_hit_full(ray_hit: RayIntersection) -> ResolvedRayHitFull {
151    let barycentrics = vec3(1.0 - ray_hit.barycentrics.x - ray_hit.barycentrics.y, ray_hit.barycentrics);
152    return resolve_triangle_data_full(ray_hit.instance_index, ray_hit.primitive_index, barycentrics);
153}
154
155fn load_vertices(instance_geometry_ids: InstanceGeometryIds, triangle_id: u32) -> array<Vertex, 3> {
156    let index_buffer = &index_buffers[instance_geometry_ids.index_buffer_id].indices;
157    let vertex_buffer = &vertex_buffers[instance_geometry_ids.vertex_buffer_id].vertices;
158
159    let indices_i = (triangle_id * 3u) + vec3(0u, 1u, 2u) + instance_geometry_ids.index_buffer_offset;
160    let indices = vec3((*index_buffer)[indices_i.x], (*index_buffer)[indices_i.y], (*index_buffer)[indices_i.z]) + instance_geometry_ids.vertex_buffer_offset;
161
162    return array<Vertex, 3>(
163        unpack_vertex((*vertex_buffer)[indices.x]),
164        unpack_vertex((*vertex_buffer)[indices.y]),
165        unpack_vertex((*vertex_buffer)[indices.z])
166    );
167}
168
169fn transform_positions(transform: mat4x4<f32>, vertices: array<Vertex, 3>) -> array<vec3<f32>, 3> {
170    return array<vec3<f32>, 3>(
171        (transform * vec4(vertices[0].position, 1.0)).xyz,
172        (transform * vec4(vertices[1].position, 1.0)).xyz,
173        (transform * vec4(vertices[2].position, 1.0)).xyz
174    );
175}
176
177fn resolve_triangle_data_full(instance_id: u32, triangle_id: u32, barycentrics: vec3<f32>) -> ResolvedRayHitFull {
178    let material_id = material_ids[instance_id];
179    let material = materials[material_id];
180
181    let instance_geometry_ids = geometry_ids[instance_id];
182    let vertices = load_vertices(instance_geometry_ids, triangle_id);
183    let transform = transforms[instance_id];
184    let world_vertices = transform_positions(transform, vertices);
185
186    let world_position = mat3x3(world_vertices[0], world_vertices[1], world_vertices[2]) * barycentrics;
187
188    let uv = mat3x2(vertices[0].uv, vertices[1].uv, vertices[2].uv) * barycentrics;
189
190    let local_tangent = mat3x3(vertices[0].tangent.xyz, vertices[1].tangent.xyz, vertices[2].tangent.xyz) * barycentrics;
191    let world_tangent = vec4(
192        normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_tangent),
193        vertices[0].tangent.w,
194    );
195
196    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
197    var world_normal = normalize(mat3x3(transform[0].xyz, transform[1].xyz, transform[2].xyz) * local_normal);
198    let geometric_world_normal = world_normal;
199    if material.normal_map_texture_id != TEXTURE_MAP_NONE {
200        let TBN = calculate_tbn_mikktspace(world_normal, world_tangent);
201        let T = TBN[0];
202        let B = TBN[1];
203        let N = TBN[2];
204        let Nt = sample_texture(material.normal_map_texture_id, uv);
205        world_normal = normalize(Nt.x * T + Nt.y * B + Nt.z * N);
206    }
207
208    let triangle_edge0 = world_vertices[0] - world_vertices[1];
209    let triangle_edge1 = world_vertices[0] - world_vertices[2];
210    let triangle_area = length(cross(triangle_edge0, triangle_edge1)) / 2.0;
211
212    let resolved_material = resolve_material(material, uv);
213
214    return ResolvedRayHitFull(world_position, world_normal, geometric_world_normal, world_tangent, uv, triangle_area, instance_geometry_ids.triangle_count, resolved_material);
215}