Update to match bevy-pbr-rust 0.2

Author: ProfLander

src/bevy_pbr_rust.rs

@@ -2,7 +2,10 @@
 
 use bevy::{prelude::StandardMaterial, render::render_resource::ShaderDefVal};
 
-use crate::prelude::{EntryPoint, EntryPointName, EntryPointParameters, RustGpuMaterial};
+use crate::{
+    prelude::{EntryPoint, EntryPointName, EntryPointParameters, RustGpuMaterial},
+    EntryPointConstants,
+};
 
 /// `bevy_rust_gpu::mesh::entry_points::vertex`
 pub enum MeshVertex {}
@@ -14,6 +17,7 @@ impl EntryPoint for MeshVertex {
         (&[("VERTEX_COLORS", "some")], "none"),
         (&[("SKINNED", "some")], "none"),
     ];
+    const CONSTANTS: EntryPointConstants = &[];
 }
 
 /// `bevy_rust_gpu::mesh::entry_points::fragment`
@@ -22,6 +26,7 @@ pub enum MeshFragment {}
 impl EntryPoint for MeshFragment {
     const NAME: EntryPointName = "mesh::entry_points::fragment";
     const PARAMETERS: EntryPointParameters = &[];
+    const CONSTANTS: EntryPointConstants = &[];
 }
 
 /// `bevy_rust_gpu::pbr::entry_points::fragment`
@@ -37,8 +42,16 @@ impl EntryPoint for PbrFragment {
         (&[("STANDARDMATERIAL_NORMAL_MAP", "some")], "none"),
         (&[("SKINNED", "some")], "none"),
         (&[("TONEMAP_IN_SHADER", "some")], "none"),
-        (&[("PREMULTIPLY_ALPHA", "some")], "none"),
         (&[("DEBAND_DITHER", "some")], "none"),
+        (
+            &[
+                ("BLEND_MULTIPLY", "multiply"),
+                ("BLEND_PREMULTIPLIED_ALPHA", "blend_premultiplied_alpha"),
+            ],
+            "none",
+        ),
+        (&[("ENVIRONMENT_MAP", "some")], "none"),
+        (&[("PREMULTIPLY_ALPHA", "some")], "none"),
         (
             &[
                 ("CLUSTERED_FORWARD_DEBUG_Z_SLICES", "debug_z_slices"),
@@ -53,7 +66,12 @@ impl EntryPoint for PbrFragment {
             ],
             "none",
         ),
+        (
+            &[("DIRECTIONAL_LIGHT_SHADOW_MAP_DEBUG_CASCADES", "some")],
+            "none",
+        ),
     ];
+    const CONSTANTS: EntryPointConstants = &["MAX_DIRECTIONAL_LIGHTS", "MAX_CASCADES_PER_LIGHT"];
 
     fn permutation(shader_defs: &Vec<ShaderDefVal>) -> Vec<String> {
         let mut permutation = vec![];

src/entry_point.rs

@@ -1,14 +1,19 @@
 //! Trait representation of a `rust-gpu` entry point.
 
+use std::collections::BTreeMap;
+
 use bevy::render::render_resource::ShaderDefVal;
 
 /// An entry point name for use with the [`EntryPoint`] trait.
 pub type EntryPointName = &'static str;
 
-/// A set of entry compile parameters for use with the [`EntryPoint`] trait.
+/// A set of entry point compile parameters for use with the [`EntryPoint`] trait.
 pub type EntryPointParameters =
     &'static [(&'static [(&'static str, &'static str)], &'static str)];
 
+/// A set of entry point constants for use with the [`EntryPoint`] trait.
+pub type EntryPointConstants = &'static [&'static str];
+
 /// A `rust-gpu` entry point for use with [`RustGpuMaterial`](crate::rust_gpu_material::RustGpuMaterial).
 pub trait EntryPoint: 'static + Send + Sync {
     /// The entry point's base function name, including module path
@@ -31,6 +36,9 @@ pub trait EntryPoint: 'static + Send + Sync {
     /// ```
     const PARAMETERS: EntryPointParameters;
 
+    /// Mapping from bevy shader def values to `permutate-macro` constants
+    const CONSTANTS: EntryPointConstants;
+
     /// Constructs a permutation set from the provided shader defs
     fn permutation(shader_defs: &Vec<ShaderDefVal>) -> Vec<String> {
         let mut permutation = vec![];
@@ -52,19 +60,56 @@ pub trait EntryPoint: 'static + Send + Sync {
         permutation
     }
 
+    fn constants(shader_defs: &Vec<ShaderDefVal>) -> Vec<ShaderDefVal> {
+        shader_defs
+            .iter()
+            .filter(|def| match def {
+                ShaderDefVal::Bool(key, _)
+                | ShaderDefVal::Int(key, _)
+                | ShaderDefVal::UInt(key, _) => Self::CONSTANTS.contains(&key.as_str()),
+            })
+            .cloned()
+            .collect()
+    }
+
     /// Build an entry point name from the provided shader defs
     fn build(shader_defs: &Vec<ShaderDefVal>) -> String {
+        let constants = Self::constants(shader_defs)
+            .into_iter()
+            .map(|def| {
+                (
+                    match &def {
+                        ShaderDefVal::Bool(key, _)
+                        | ShaderDefVal::Int(key, _)
+                        | ShaderDefVal::UInt(key, _) => key.clone(),
+                    },
+                    match &def {
+                        ShaderDefVal::Bool(value, _) => value.to_string(),
+                        ShaderDefVal::Int(_, value) => value.to_string(),
+                        ShaderDefVal::UInt(_, value) => value.to_string(),
+                    },
+                )
+            })
+            .collect::<BTreeMap<_, _>>();
+
         std::iter::once(Self::NAME.to_string())
             .chain(
                 Self::permutation(shader_defs)
                     .into_iter()
                     .map(|variant| "__".to_string() + &variant),
             )
+            .chain(
+                constants
+                    .into_iter()
+                    .map(|(key, value)| key + "_" + &value)
+                    .map(|variant| "__".to_string() + &variant),
+            )
             .collect::<String>()
     }
 }
 
 impl EntryPoint for () {
     const NAME: &'static str = "";
     const PARAMETERS: EntryPointParameters = &[];
+    const CONSTANTS: EntryPointConstants = &[];
 }

src/entry_point_export.rs

@@ -10,6 +10,7 @@ use std::{
 
 use bevy::{
     prelude::{default, info, CoreSet, Deref, DerefMut, IntoSystemConfig, NonSendMut, Plugin},
+    render::render_resource::ShaderDefVal,
     tasks::IoTaskPool,
     utils::HashMap,
 };
@@ -53,13 +54,49 @@ type EntryPointReceiver = Receiver<Export>;
 pub struct Export {
     pub shader: &'static str,
     pub permutation: Vec<String>,
+    pub constants: Vec<ShaderDefVal>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
+#[serde(untagged)]
+enum PermutationConstant {
+    Bool(bool),
+    Uint(u32),
+    Int(i32),
+}
+
+#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, Deref, DerefMut)]
+struct PermutationConstants {
+    #[serde(flatten)]
+    constants: HashMap<String, PermutationConstant>,
+}
+
+impl From<Vec<ShaderDefVal>> for PermutationConstants {
+    fn from(value: Vec<ShaderDefVal>) -> Self {
+        PermutationConstants {
+            constants: value
+                .into_iter()
+                .map(|def| match def {
+                    ShaderDefVal::Bool(key, value) => (key, PermutationConstant::Bool(value)),
+                    ShaderDefVal::Int(key, value) => (key, PermutationConstant::Int(value)),
+                    ShaderDefVal::UInt(key, value) => (key, PermutationConstant::Uint(value)),
+                })
+                .collect(),
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize)]
+struct Permutation {
+    parameters: Vec<String>,
+    constants: PermutationConstants,
 }
 
 /// Serializable container for a single entry point
 #[derive(Debug, Default, Clone, Deref, DerefMut, Serialize, Deserialize)]
 struct EntryPoints {
     #[serde(flatten)]
-    entry_points: HashMap<String, Vec<Vec<String>>>,
+    entry_points: HashMap<String, Vec<Permutation>>,
 }
 
 /// Container for a set of entry points, with MPSC handles and change tracking
@@ -110,13 +147,18 @@ impl EntryPointExport {
                 }
 
                 let entry = &export.entry_points[entry_point.shader];
-                if !entry.contains(&entry_point.permutation) {
-                    info!("New permutation: {:?}", entry_point.permutation);
+                let permutation = Permutation {
+                    parameters: entry_point.permutation,
+                    constants: entry_point.constants.into(),
+                };
+
+                if !entry.contains(&permutation) {
+                    info!("New permutation: {:?}", permutation);
                     export
                         .entry_points
                         .get_mut(entry_point.shader)
                         .unwrap()
-                        .push(entry_point.permutation);
+                        .push(permutation);
                     export.changed = true;
                 }
             }

src/rust_gpu.rs

@@ -47,8 +47,12 @@ where
     M::Data: Clone + Eq + std::hash::Hash,
 {
     fn build(&self, app: &mut bevy::prelude::App) {
-        app.add_system(reload_materials::<M>.in_base_set(CoreSet::Last));
-        app.add_system(shader_events::<M>.before(reload_materials::<M>));
+        app.add_system(reload_materials::<M>.in_base_set(CoreSet::PreUpdate));
+        app.add_system(
+            shader_events::<M>
+                .in_base_set(CoreSet::PreUpdate)
+                .before(reload_materials::<M>),
+        );
 
         app.add_plugin(MaterialPlugin::<RustGpu<M>>::default());
 
@@ -335,6 +339,7 @@ where
                     .send(crate::prelude::Export {
                         shader: M::Vertex::NAME,
                         permutation: M::Vertex::permutation(&descriptor.vertex.shader_defs),
+                        constants: M::Vertex::constants(&descriptor.vertex.shader_defs),
                     })
                     .unwrap();
             };
@@ -394,6 +399,7 @@ where
                         .send(crate::prelude::Export {
                             shader: M::Fragment::NAME,
                             permutation: M::Fragment::permutation(&fragment_descriptor.shader_defs),
+                            constants: M::Fragment::constants(&fragment_descriptor.shader_defs),
                         })
                         .unwrap();
                 };

src/shader_meta.rs

@@ -50,7 +50,7 @@ where
 
         app.add_system(
             module_meta_events::<M>
-                .in_base_set(CoreSet::Last)
+                .in_base_set(CoreSet::PreUpdate)
                 .after(shader_events::<M>)
                 .before(reload_materials::<M>),
         );