Shader defs can now have a value (#5900)

# Objective

- shaders defs can now have a `bool` or `int` value
- `#if SHADER_DEF <operator> 3`
  - ok if `SHADER_DEF` is defined, has the correct type and pass the comparison
  - `==`, `!=`, `>=`, `>`, `<`, `<=` supported
- `#SHADER_DEF` or `#{SHADER_DEF}`
  - will be replaced by the value in the shader code
---

## Migration Guide

- replace `shader_defs.push(String::from("NAME"));` by `shader_defs.push("NAME".into());`
- if you used shader def `NO_STORAGE_BUFFERS_SUPPORT`, check how `AVAILABLE_STORAGE_BUFFER_BINDINGS` is now used in Bevy default shaders

Author: mockersf

crates/bevy_core_pipeline/src/fxaa/mod.rs

@@ -203,8 +203,8 @@ impl SpecializedRenderPipeline for FxaaPipeline {
             fragment: Some(FragmentState {
                 shader: FXAA_SHADER_HANDLE.typed(),
                 shader_defs: vec![
-                    format!("EDGE_THRESH_{}", key.edge_threshold.get_str()),
-                    format!("EDGE_THRESH_MIN_{}", key.edge_threshold_min.get_str()),
+                    format!("EDGE_THRESH_{}", key.edge_threshold.get_str()).into(),
+                    format!("EDGE_THRESH_MIN_{}", key.edge_threshold_min.get_str()).into(),
                 ],
                 entry_point: "fragment".into(),
                 targets: vec![Some(ColorTargetState {

crates/bevy_core_pipeline/src/tonemapping/mod.rs

@@ -66,7 +66,7 @@ impl SpecializedRenderPipeline for TonemappingPipeline {
     fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
         let mut shader_defs = Vec::new();
         if key.deband_dither {
-            shader_defs.push("DEBAND_DITHER".to_string());
+            shader_defs.push("DEBAND_DITHER".into());
         }
         RenderPipelineDescriptor {
             label: Some("tonemapping pipeline".into()),

crates/bevy_pbr/src/pbr_material.rs

@@ -419,7 +419,7 @@ impl Material for StandardMaterial {
                 .as_mut()
                 .unwrap()
                 .shader_defs
-                .push(String::from("STANDARDMATERIAL_NORMAL_MAP"));
+                .push("STANDARDMATERIAL_NORMAL_MAP".into());
         }
         descriptor.primitive.cull_mode = key.bind_group_data.cull_mode;
         if let Some(label) = &mut descriptor.label {

crates/bevy_pbr/src/render/clustered_forward.wgsl

@@ -29,7 +29,9 @@ fn fragment_cluster_index(frag_coord: vec2<f32>, view_z: f32, is_orthographic: b
 // this must match CLUSTER_COUNT_SIZE in light.rs
 let CLUSTER_COUNT_SIZE = 9u;
 fn unpack_offset_and_counts(cluster_index: u32) -> vec3<u32> {
-#ifdef NO_STORAGE_BUFFERS_SUPPORT
+#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
+    return cluster_offsets_and_counts.data[cluster_index].xyz;
+#else
     let offset_and_counts = cluster_offsets_and_counts.data[cluster_index >> 2u][cluster_index & ((1u << 2u) - 1u)];
     //  [ 31     ..     18 | 17      ..      9 | 8       ..     0 ]
     //  [      offset      | point light count | spot light count ]
@@ -38,20 +40,18 @@ fn unpack_offset_and_counts(cluster_index: u32) -> vec3<u32> {
         (offset_and_counts >> CLUSTER_COUNT_SIZE)        & ((1u << CLUSTER_COUNT_SIZE) - 1u),
         offset_and_counts                                & ((1u << CLUSTER_COUNT_SIZE) - 1u),
     );
-#else
-    return cluster_offsets_and_counts.data[cluster_index].xyz;
 #endif
 }
 
 fn get_light_id(index: u32) -> u32 {
-#ifdef NO_STORAGE_BUFFERS_SUPPORT
+#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
+    return cluster_light_index_lists.data[index];
+#else
     // The index is correct but in cluster_light_index_lists we pack 4 u8s into a u32
     // This means the index into cluster_light_index_lists is index / 4
     let indices = cluster_light_index_lists.data[index >> 4u][(index >> 2u) & ((1u << 2u) - 1u)];
     // And index % 4 gives the sub-index of the u8 within the u32 so we shift by 8 * sub-index
     return (indices >> (8u * (index & ((1u << 2u) - 1u)))) & ((1u << 8u) - 1u);
-#else
-    return cluster_light_index_lists.data[index];
 #endif
 }
 

crates/bevy_pbr/src/render/light.rs

@@ -321,11 +321,15 @@ impl SpecializedMeshPipeline for ShadowPipeline {
 
         let mut bind_group_layout = vec![self.view_layout.clone()];
         let mut shader_defs = Vec::new();
+        shader_defs.push(ShaderDefVal::Int(
+            "MAX_DIRECTIONAL_LIGHTS".to_string(),
+            MAX_DIRECTIONAL_LIGHTS as i32,
+        ));
 
         if layout.contains(Mesh::ATTRIBUTE_JOINT_INDEX)
             && layout.contains(Mesh::ATTRIBUTE_JOINT_WEIGHT)
         {
-            shader_defs.push(String::from("SKINNED"));
+            shader_defs.push("SKINNED".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_INDEX.at_shader_location(4));
             vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_WEIGHT.at_shader_location(5));
             bind_group_layout.push(self.skinned_mesh_layout.clone());

crates/bevy_pbr/src/render/mesh.rs

@@ -1,7 +1,7 @@
 use crate::{
     GlobalLightMeta, GpuLights, GpuPointLights, LightMeta, NotShadowCaster, NotShadowReceiver,
     ShadowPipeline, ViewClusterBindings, ViewLightsUniformOffset, ViewShadowBindings,
-    CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT,
+    CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT, MAX_DIRECTIONAL_LIGHTS,
 };
 use bevy_app::Plugin;
 use bevy_asset::{load_internal_asset, Assets, Handle, HandleUntyped};
@@ -581,35 +581,40 @@ impl SpecializedMeshPipeline for MeshPipeline {
         let mut vertex_attributes = Vec::new();
 
         if layout.contains(Mesh::ATTRIBUTE_POSITION) {
-            shader_defs.push(String::from("VERTEX_POSITIONS"));
+            shader_defs.push("VERTEX_POSITIONS".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_POSITION.at_shader_location(0));
         }
 
         if layout.contains(Mesh::ATTRIBUTE_NORMAL) {
-            shader_defs.push(String::from("VERTEX_NORMALS"));
+            shader_defs.push("VERTEX_NORMALS".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_NORMAL.at_shader_location(1));
         }
 
+        shader_defs.push(ShaderDefVal::Int(
+            "MAX_DIRECTIONAL_LIGHTS".to_string(),
+            MAX_DIRECTIONAL_LIGHTS as i32,
+        ));
+
         if layout.contains(Mesh::ATTRIBUTE_UV_0) {
-            shader_defs.push(String::from("VERTEX_UVS"));
+            shader_defs.push("VERTEX_UVS".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_UV_0.at_shader_location(2));
         }
 
         if layout.contains(Mesh::ATTRIBUTE_TANGENT) {
-            shader_defs.push(String::from("VERTEX_TANGENTS"));
+            shader_defs.push("VERTEX_TANGENTS".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_TANGENT.at_shader_location(3));
         }
 
         if layout.contains(Mesh::ATTRIBUTE_COLOR) {
-            shader_defs.push(String::from("VERTEX_COLORS"));
+            shader_defs.push("VERTEX_COLORS".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_COLOR.at_shader_location(4));
         }
 
         let mut bind_group_layout = vec![self.view_layout.clone()];
         if layout.contains(Mesh::ATTRIBUTE_JOINT_INDEX)
             && layout.contains(Mesh::ATTRIBUTE_JOINT_WEIGHT)
         {
-            shader_defs.push(String::from("SKINNED"));
+            shader_defs.push("SKINNED".into());
             vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_INDEX.at_shader_location(5));
             vertex_attributes.push(Mesh::ATTRIBUTE_JOINT_WEIGHT.at_shader_location(6));
             bind_group_layout.push(self.skinned_mesh_layout.clone());
@@ -636,11 +641,11 @@ impl SpecializedMeshPipeline for MeshPipeline {
         }
 
         if key.contains(MeshPipelineKey::TONEMAP_IN_SHADER) {
-            shader_defs.push("TONEMAP_IN_SHADER".to_string());
+            shader_defs.push("TONEMAP_IN_SHADER".into());
 
             // Debanding is tied to tonemapping in the shader, cannot run without it.
             if key.contains(MeshPipelineKey::DEBAND_DITHER) {
-                shader_defs.push("DEBAND_DITHER".to_string());
+                shader_defs.push("DEBAND_DITHER".into());
             }
         }
 

crates/bevy_pbr/src/render/mesh_view_bindings.wgsl

@@ -25,20 +25,20 @@ var directional_shadow_textures: texture_depth_2d_array;
 @group(0) @binding(5)
 var directional_shadow_textures_sampler: sampler_comparison;
 
-#ifdef NO_STORAGE_BUFFERS_SUPPORT
+#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
 @group(0) @binding(6)
-var<uniform> point_lights: PointLights;
+var<storage> point_lights: PointLights;
 @group(0) @binding(7)
-var<uniform> cluster_light_index_lists: ClusterLightIndexLists;
+var<storage> cluster_light_index_lists: ClusterLightIndexLists;
 @group(0) @binding(8)
-var<uniform> cluster_offsets_and_counts: ClusterOffsetsAndCounts;
+var<storage> cluster_offsets_and_counts: ClusterOffsetsAndCounts;
 #else
 @group(0) @binding(6)
-var<storage> point_lights: PointLights;
+var<uniform> point_lights: PointLights;
 @group(0) @binding(7)
-var<storage> cluster_light_index_lists: ClusterLightIndexLists;
+var<uniform> cluster_light_index_lists: ClusterLightIndexLists;
 @group(0) @binding(8)
-var<storage> cluster_offsets_and_counts: ClusterOffsetsAndCounts;
+var<uniform> cluster_offsets_and_counts: ClusterOffsetsAndCounts;
 #endif
 
 @group(0) @binding(9)

crates/bevy_pbr/src/render/mesh_view_types.wgsl

@@ -42,7 +42,7 @@ let DIRECTIONAL_LIGHT_FLAGS_SHADOWS_ENABLED_BIT: u32 = 1u;
 
 struct Lights {
     // NOTE: this array size must be kept in sync with the constants defined in bevy_pbr/src/render/light.rs
-    directional_lights: array<DirectionalLight, 10u>,
+    directional_lights: array<DirectionalLight, #{MAX_DIRECTIONAL_LIGHTS}u>,
     ambient_color: vec4<f32>,
     // x/y/z dimensions and n_clusters in w
     cluster_dimensions: vec4<u32>,
@@ -61,28 +61,28 @@ struct Lights {
     spot_light_shadowmap_offset: i32,
 };
 
-#ifdef NO_STORAGE_BUFFERS_SUPPORT
+#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
 struct PointLights {
-    data: array<PointLight, 256u>,
+    data: array<PointLight>,
 };
 struct ClusterLightIndexLists {
-    // each u32 contains 4 u8 indices into the PointLights array
-    data: array<vec4<u32>, 1024u>,
+    data: array<u32>,
 };
 struct ClusterOffsetsAndCounts {
-    // each u32 contains a 24-bit index into ClusterLightIndexLists in the high 24 bits
-    // and an 8-bit count of the number of lights in the low 8 bits
-    data: array<vec4<u32>, 1024u>,
+    data: array<vec4<u32>>,
 };
 #else
 struct PointLights {
-    data: array<PointLight>,
+    data: array<PointLight, 256u>,
 };
 struct ClusterLightIndexLists {
-    data: array<u32>,
+    // each u32 contains 4 u8 indices into the PointLights array
+    data: array<vec4<u32>, 1024u>,
 };
 struct ClusterOffsetsAndCounts {
-    data: array<vec4<u32>>,
+    // each u32 contains a 24-bit index into ClusterLightIndexLists in the high 24 bits
+    // and an 8-bit count of the number of lights in the low 8 bits
+    data: array<vec4<u32>, 1024u>,
 };
 #endif
 

crates/bevy_render/src/render_resource/pipeline.rs

@@ -7,6 +7,7 @@ use wgpu::{
     VertexAttribute, VertexFormat, VertexStepMode,
 };
 
+use super::ShaderDefVal;
 use crate::render_resource::resource_macros::*;
 
 /// A [`RenderPipeline`] identifier.
@@ -115,7 +116,7 @@ pub struct RenderPipelineDescriptor {
 pub struct VertexState {
     /// The compiled shader module for this stage.
     pub shader: Handle<Shader>,
-    pub shader_defs: Vec<String>,
+    pub shader_defs: Vec<ShaderDefVal>,
     /// The name of the entry point in the compiled shader. There must be a
     /// function with this name in the shader.
     pub entry_point: Cow<'static, str>,
@@ -167,7 +168,7 @@ impl VertexBufferLayout {
 pub struct FragmentState {
     /// The compiled shader module for this stage.
     pub shader: Handle<Shader>,
-    pub shader_defs: Vec<String>,
+    pub shader_defs: Vec<ShaderDefVal>,
     /// The name of the entry point in the compiled shader. There must be a
     /// function with this name in the shader.
     pub entry_point: Cow<'static, str>,
@@ -182,7 +183,7 @@ pub struct ComputePipelineDescriptor {
     pub layout: Option<Vec<BindGroupLayout>>,
     /// The compiled shader module for this stage.
     pub shader: Handle<Shader>,
-    pub shader_defs: Vec<String>,
+    pub shader_defs: Vec<ShaderDefVal>,
     /// The name of the entry point in the compiled shader. There must be a
     /// function with this name in the shader.
     pub entry_point: Cow<'static, str>,

crates/bevy_render/src/render_resource/pipeline_cache.rs

@@ -107,7 +107,7 @@ impl CachedPipelineState {
 #[derive(Default)]
 struct ShaderData {
     pipelines: HashSet<CachedPipelineId>,
-    processed_shaders: HashMap<Vec<String>, ErasedShaderModule>,
+    processed_shaders: HashMap<Vec<ShaderDefVal>, ErasedShaderModule>,
     resolved_imports: HashMap<ShaderImport, Handle<Shader>>,
     dependents: HashSet<Handle<Shader>>,
 }
@@ -121,13 +121,31 @@ struct ShaderCache {
     processor: ShaderProcessor,
 }
 
+#[derive(Clone, PartialEq, Eq, Debug, Hash)]
+pub enum ShaderDefVal {
+    Bool(String, bool),
+    Int(String, i32),
+}
+
+impl From<&str> for ShaderDefVal {
+    fn from(key: &str) -> Self {
+        ShaderDefVal::Bool(key.to_string(), true)
+    }
+}
+
+impl From<String> for ShaderDefVal {
+    fn from(key: String) -> Self {
+        ShaderDefVal::Bool(key, true)
+    }
+}
+
 impl ShaderCache {
     fn get(
         &mut self,
         render_device: &RenderDevice,
         pipeline: CachedPipelineId,
         handle: &Handle<Shader>,
-        shader_defs: &[String],
+        shader_defs: &[ShaderDefVal],
     ) -> Result<ErasedShaderModule, PipelineCacheError> {
         let shader = self
             .shaders
@@ -156,21 +174,20 @@ impl ShaderCache {
                 let mut shader_defs = shader_defs.to_vec();
                 #[cfg(feature = "webgl")]
                 {
-                    shader_defs.push(String::from("NO_ARRAY_TEXTURES_SUPPORT"));
-                    shader_defs.push(String::from("SIXTEEN_BYTE_ALIGNMENT"));
+                    shader_defs.push("NO_ARRAY_TEXTURES_SUPPORT".into());
+                    shader_defs.push("SIXTEEN_BYTE_ALIGNMENT".into());
                 }
 
-                // TODO: 3 is the value from CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT declared in bevy_pbr
-                // consider exposing this in shaders in a more generally useful way, such as:
-                // # if AVAILABLE_STORAGE_BUFFER_BINDINGS == 3
-                // /* use storage buffers here */
-                // # elif
-                // /* use uniforms here */
-                if !matches!(
+                // 3 is the value from CLUSTERED_FORWARD_STORAGE_BUFFER_COUNT declared in bevy_pbr
+                // Using the value directly here to avoid the cyclic dependency
+                if matches!(
                     render_device.get_supported_read_only_binding_type(3),
                     BufferBindingType::Storage { .. }
                 ) {
-                    shader_defs.push(String::from("NO_STORAGE_BUFFERS_SUPPORT"));
+                    shader_defs.push(ShaderDefVal::Int(
+                        String::from("AVAILABLE_STORAGE_BUFFER_BINDINGS"),
+                        3,
+                    ));
                 }
 
                 debug!(