1#define_import_path bevy_pbr::pbr_functions
2
3#import bevy_pbr::pbr_types as PbrTypes
4#import bevy_pbr::mesh_types as MeshTypes
5#import bevy_pbr::mesh_bindings as MeshBindings
6#import bevy_pbr::mesh_view_types as ViewTypes
7#import bevy_pbr::mesh_view_bindings as ViewBindings
8#import bevy_pbr::lighting as Lighting
9#import bevy_pbr::clustered_forward as Clustering
10#import bevy_pbr::shadows as Shadows
11
12// NOTE: This ensures that the world_normal is normalized and if
13// vertex tangents and normal maps then normal mapping may be applied.
14fn prepare_normal(
15    standard_material_flags: u32,
16    world_normal: vec3<f32>,
17#ifdef VERTEX_TANGENTS
18#ifdef STANDARDMATERIAL_NORMAL_MAP
19    world_tangent: vec4<f32>,
20#endif
21#endif
22#ifdef VERTEX_UVS
23    uv: vec2<f32>,
24#endif
25    is_front: bool,
26) -> vec3<f32> {
27    var N: vec3<f32> = normalize(world_normal);
28
29#ifdef VERTEX_TANGENTS
30#ifdef STANDARDMATERIAL_NORMAL_MAP
31    // NOTE: The mikktspace method of normal mapping explicitly requires that these NOT be
32    // normalized nor any Gram-Schmidt applied to ensure the vertex normal is orthogonal to the
33    // vertex tangent! Do not change this code unless you really know what you are doing.
34    // http://www.mikktspace.com/
35    var T: vec3<f32> = world_tangent.xyz;
36    var B: vec3<f32> = world_tangent.w * cross(N, T);
37#endif
38#endif
39
40    if ((standard_material_flags & PbrTypes::STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u) {
41        if (!is_front) {
42            N = -N;
43#ifdef VERTEX_TANGENTS
44#ifdef STANDARDMATERIAL_NORMAL_MAP
45            T = -T;
46            B = -B;
47#endif
48#endif
49        }
50    }
51
52#ifdef VERTEX_TANGENTS
53#ifdef VERTEX_UVS
54#ifdef STANDARDMATERIAL_NORMAL_MAP
55    // Nt is the tangent-space normal.
56    var Nt = textureSample(normal_map_texture, normal_map_sampler, uv).rgb;
57    if ((standard_material_flags & PbrTypes::STANDARD_MATERIAL_FLAGS_TWO_COMPONENT_NORMAL_MAP) != 0u) {
58        // Only use the xy components and derive z for 2-component normal maps.
59        Nt = vec3<f32>(Nt.rg * 2.0 - 1.0, 0.0);
60        Nt.z = sqrt(1.0 - Nt.x * Nt.x - Nt.y * Nt.y);
61    } else {
62        Nt = Nt * 2.0 - 1.0;
63    }
64    // Normal maps authored for DirectX require flipping the y component
65    if ((standard_material_flags & PbrTypes::STANDARD_MATERIAL_FLAGS_FLIP_NORMAL_MAP_Y) != 0u) {
66        Nt.y = -Nt.y;
67    }
68    // NOTE: The mikktspace method of normal mapping applies maps the tangent-space normal from
69    // the normal map texture in this way to be an EXACT inverse of how the normal map baker
70    // calculates the normal maps so there is no error introduced. Do not change this code
71    // unless you really know what you are doing.
72    // http://www.mikktspace.com/
73    N = normalize(Nt.x * T + Nt.y * B + Nt.z * N);
74#endif
75#endif
76#endif
77
78    return N;
79}
80
81// NOTE: Correctly calculates the view vector depending on whether
82// the projection is orthographic or perspective.
83fn calculate_view(
84    world_position: vec4<f32>,
85    is_orthographic: bool,
86) -> vec3<f32> {
87    var V: vec3<f32>;
88    if (is_orthographic) {
89        // Orthographic view vector
90        V = normalize(vec3<f32>(ViewBindings::view.view_proj[0].z, ViewBindings::view.view_proj[1].z, ViewBindings::view.view_proj[2].z));
91    } else {
92        // Only valid for a perpective projection
93        V = normalize(ViewBindings::view.world_position.xyz - world_position.xyz);
94    }
95    return V;
96}
97
98struct PbrInput {
99    material: PbrTypes::StandardMaterial,
100    occlusion: f32,
101    frag_coord: vec4<f32>,
102    world_position: vec4<f32>,
103    // Normalized world normal used for shadow mapping as normal-mapping is not used for shadow
104    // mapping
105    world_normal: vec3<f32>,
106    // Normalized normal-mapped world normal used for lighting
107    N: vec3<f32>,
108    // Normalized view vector in world space, pointing from the fragment world position toward the
109    // view world position
110    V: vec3<f32>,
111    is_orthographic: bool,
112};
113
114// Creates a PbrInput with default values
115fn pbr_input_new() -> PbrInput {
116    var pbr_input: PbrInput;
117
118    pbr_input.material = PbrTypes::standard_material_new();
119    pbr_input.occlusion = 1.0;
120
121    pbr_input.frag_coord = vec4<f32>(0.0, 0.0, 0.0, 1.0);
122    pbr_input.world_position = vec4<f32>(0.0, 0.0, 0.0, 1.0);
123    pbr_input.world_normal = vec3<f32>(0.0, 0.0, 1.0);
124
125    pbr_input.is_orthographic = false;
126
127    pbr_input.N = vec3<f32>(0.0, 0.0, 1.0);
128    pbr_input.V = vec3<f32>(1.0, 0.0, 0.0);
129
130    return pbr_input;
131}
132
133fn pbr(
134    in: PbrInput,
135) -> vec4<f32> {
136    var output_color: vec4<f32> = in.material.base_color;
137
138    // TODO use .a for exposure compensation in HDR
139    let emissive = in.material.emissive;
140
141    // calculate non-linear roughness from linear perceptualRoughness
142    let metallic = in.material.metallic;
143    let perceptual_roughness = in.material.perceptual_roughness;
144    let roughness = Lighting::perceptualRoughnessToRoughness(perceptual_roughness);
145
146    let occlusion = in.occlusion;
147
148    if ((in.material.flags & PbrTypes::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE) != 0u) {
149        // NOTE: If rendering as opaque, alpha should be ignored so set to 1.0
150        output_color.a = 1.0;
151    } else if ((in.material.flags & PbrTypes::STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK) != 0u) {
152        if (output_color.a >= in.material.alpha_cutoff) {
153            // NOTE: If rendering as masked alpha and >= the cutoff, render as fully opaque
154            output_color.a = 1.0;
155        } else {
156            // NOTE: output_color.a < in.material.alpha_cutoff should not is not rendered
157            // NOTE: This and any other discards mean that early-z testing cannot be done!
158            discard;
159        }
160    }
161
162    // Neubelt and Pettineo 2013, "Crafting a Next-gen Material Pipeline for The Order: 1886"
163    let NdotV = max(dot(in.N, in.V), 0.0001);
164
165    // Remapping [0,1] reflectance to F0
166    // See https://google.github.io/filament/Filament.html#materialsystem/parameterization/remapping
167    let reflectance = in.material.reflectance;
168    let F0 = 0.16 * reflectance * reflectance * (1.0 - metallic) + output_color.rgb * metallic;
169
170    // Diffuse strength inversely related to metallicity
171    let diffuse_color = output_color.rgb * (1.0 - metallic);
172
173    let R = reflect(-in.V, in.N);
174
175    // accumulate color
176    var light_accum: vec3<f32> = vec3<f32>(0.0);
177
178    let view_z = dot(vec4<f32>(
179        ViewBindings::view.inverse_view[0].z,
180        ViewBindings::view.inverse_view[1].z,
181        ViewBindings::view.inverse_view[2].z,
182        ViewBindings::view.inverse_view[3].z
183    ), in.world_position);
184    let cluster_index = Clustering::fragment_cluster_index(in.frag_coord.xy, view_z, in.is_orthographic);
185    let offset_and_counts = Clustering::unpack_offset_and_counts(cluster_index);
186
187    // point lights
188    for (var i: u32 = offset_and_counts[0]; i < offset_and_counts[0] + offset_and_counts[1]; i = i + 1u) {
189        let light_id = Clustering::get_light_id(i);
190        let light = ViewBindings::point_lights.data[light_id];
191        var shadow: f32 = 1.0;
192        if ((MeshBindings::mesh.flags & MeshTypes::MESH_FLAGS_SHADOW_RECEIVER_BIT) != 0u
193                && (light.flags & ViewTypes::POINT_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u) {
194            shadow = Shadows::fetch_point_shadow(light_id, in.world_position, in.world_normal);
195        }
196        let light_contrib = Lighting::point_light(in.world_position.xyz, light, roughness, NdotV, in.N, in.V, R, F0, diffuse_color);
197        light_accum = light_accum + light_contrib * shadow;
198    }
199
200    // spot lights
201    for (var i: u32 = offset_and_counts[0] + offset_and_counts[1]; i < offset_and_counts[0] + offset_and_counts[1] + offset_and_counts[2]; i = i + 1u) {
202        let light_id = Clustering::get_light_id(i);
203        let light = ViewBindings::point_lights.data[light_id];
204        var shadow: f32 = 1.0;
205        if ((MeshBindings::mesh.flags & MeshTypes::MESH_FLAGS_SHADOW_RECEIVER_BIT) != 0u
206                && (light.flags & ViewTypes::POINT_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u) {
207            shadow = Shadows::fetch_spot_shadow(light_id, in.world_position, in.world_normal);
208        }
209        let light_contrib = Lighting::spot_light(in.world_position.xyz, light, roughness, NdotV, in.N, in.V, R, F0, diffuse_color);
210        light_accum = light_accum + light_contrib * shadow;
211    }
212
213    let n_directional_lights = ViewBindings::lights.n_directional_lights;
214    for (var i: u32 = 0u; i < n_directional_lights; i = i + 1u) {
215        let light = ViewBindings::lights.directional_lights[i];
216        var shadow: f32 = 1.0;
217        if ((MeshBindings::mesh.flags & MeshTypes::MESH_FLAGS_SHADOW_RECEIVER_BIT) != 0u
218                && (light.flags & ViewTypes::DIRECTIONAL_LIGHT_FLAGS_SHADOWS_ENABLED_BIT) != 0u) {
219            shadow = Shadows::fetch_directional_shadow(i, in.world_position, in.world_normal);
220        }
221        let light_contrib = Lighting::directional_light(light, roughness, NdotV, in.N, in.V, R, F0, diffuse_color);
222        light_accum = light_accum + light_contrib * shadow;
223    }
224
225    let diffuse_ambient = Lighting::EnvBRDFApprox(diffuse_color, 1.0, NdotV);
226    let specular_ambient = Lighting::EnvBRDFApprox(F0, perceptual_roughness, NdotV);
227
228    output_color = vec4<f32>(
229        light_accum +
230            (diffuse_ambient + specular_ambient) * ViewBindings::lights.ambient_color.rgb * occlusion +
231            emissive.rgb * output_color.a,
232        output_color.a);
233
234    output_color = Clustering::cluster_debug_visualization(
235        output_color,
236        view_z,
237        in.is_orthographic,
238        offset_and_counts,
239        cluster_index,
240    );
241
242    return output_color;
243}
244
245fn tone_mapping(in: vec4<f32>) -> vec4<f32> {
246    // tone_mapping
247    return vec4<f32>(Lighting::reinhard_luminance(in.rgb), in.a);
248    
249    // Gamma correction.
250    // Not needed with sRGB buffer
251    // output_color.rgb = pow(output_color.rgb, vec3(1.0 / 2.2));
252}