Dynamic UI Materials

crates/bevy_dev_tools/src/frame_time_graph/mod.rs

@@ -84,7 +84,7 @@ pub struct FrametimeGraphMaterial {
 }
 
 impl UiMaterial for FrametimeGraphMaterial {
-    fn fragment_shader() -> ShaderRef {
+    fn fragment_shader(&self) -> ShaderRef {
         FRAME_TIME_GRAPH_SHADER_HANDLE.into()
     }
 }

crates/bevy_feathers/src/alpha_pattern.rs

@@ -18,7 +18,7 @@ use bevy_ui_render::ui_material::{MaterialNode, UiMaterial};
 pub(crate) struct AlphaPatternMaterial {}
 
 impl UiMaterial for AlphaPatternMaterial {
-    fn fragment_shader() -> ShaderRef {
+    fn fragment_shader(&self) -> ShaderRef {
         "embedded://bevy_feathers/assets/shaders/alpha_pattern.wgsl".into()
     }
 }

crates/bevy_ui_render/src/ui_material.rs

@@ -7,7 +7,7 @@ use bevy_render::{
     extract_component::ExtractComponent,
     render_resource::{AsBindGroup, RenderPipelineDescriptor},
 };
-use bevy_shader::ShaderRef;
+use bevy_shader::{Shader, ShaderRef};
 use derive_more::derive::From;
 
 /// Materials are used alongside [`UiMaterialPlugin`](crate::UiMaterialPlugin) and [`MaterialNode`]
@@ -54,7 +54,7 @@ use derive_more::derive::From;
 /// // All functions on `UiMaterial` have default impls. You only need to implement the
 /// // functions that are relevant for your material.
 /// impl UiMaterial for CustomMaterial {
-///     fn fragment_shader() -> ShaderRef {
+///     fn fragment_shader(&self) -> ShaderRef {
 ///         "shaders/custom_material.wgsl".into()
 ///     }
 /// }
@@ -99,16 +99,16 @@ use derive_more::derive::From;
 ///
 /// }
 /// ```
-pub trait UiMaterial: AsBindGroup + Asset + Clone + Sized {
+pub trait UiMaterial: AsBindGroup + Asset + Clone + Sized + 'static {
     /// Returns this materials vertex shader. If [`ShaderRef::Default`] is returned, the default UI
     /// vertex shader will be used.
-    fn vertex_shader() -> ShaderRef {
+    fn vertex_shader(&self) -> ShaderRef {
         ShaderRef::Default
     }
 
     /// Returns this materials fragment shader. If [`ShaderRef::Default`] is returned, the default
     /// UI fragment shader will be used.
-    fn fragment_shader() -> ShaderRef {
+    fn fragment_shader(&self) -> ShaderRef {
         ShaderRef::Default
     }
 
@@ -122,6 +122,8 @@ pub trait UiMaterial: AsBindGroup + Asset + Clone + Sized {
 
 pub struct UiMaterialKey<M: UiMaterial> {
     pub hdr: bool,
+    pub vertex_shader: Handle<Shader>,
+    pub fragment_shader: Handle<Shader>,
     pub bind_group_data: M::Data,
 }
 
@@ -132,7 +134,10 @@ where
     M::Data: PartialEq,
 {
     fn eq(&self, other: &Self) -> bool {
-        self.hdr == other.hdr && self.bind_group_data == other.bind_group_data
+        self.hdr == other.hdr
+            && self.vertex_shader == other.vertex_shader
+            && self.fragment_shader == other.fragment_shader
+            && self.bind_group_data == other.bind_group_data
     }
 }
 
@@ -143,6 +148,8 @@ where
     fn clone(&self) -> Self {
         Self {
             hdr: self.hdr,
+            vertex_shader: self.vertex_shader.clone(),
+            fragment_shader: self.fragment_shader.clone(),
             bind_group_data: self.bind_group_data.clone(),
         }
     }
@@ -154,6 +161,8 @@ where
 {
     fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
         self.hdr.hash(state);
+        self.vertex_shader.hash(state);
+        self.fragment_shader.hash(state);
         self.bind_group_data.hash(state);
     }
 }

crates/bevy_ui_render/src/ui_material_pipeline.rs

@@ -115,8 +115,7 @@ pub struct UiMaterialBatch<M: UiMaterial> {
 pub struct UiMaterialPipeline<M: UiMaterial> {
     pub ui_layout: BindGroupLayout,
     pub view_layout: BindGroupLayout,
-    pub vertex_shader: Handle<Shader>,
-    pub fragment_shader: Handle<Shader>,
+    pub default_shader: Handle<Shader>,
     marker: PhantomData<M>,
 }
 
@@ -146,13 +145,13 @@ where
 
         let mut descriptor = RenderPipelineDescriptor {
             vertex: VertexState {
-                shader: self.vertex_shader.clone(),
+                shader: key.vertex_shader.clone(),
                 shader_defs: shader_defs.clone(),
                 buffers: vec![vertex_layout],
                 ..default()
             },
             fragment: Some(FragmentState {
-                shader: self.fragment_shader.clone(),
+                shader: key.fragment_shader.clone(),
                 shader_defs,
                 targets: vec![Some(ColorTargetState {
                     format: if key.hdr {
@@ -195,21 +194,10 @@ pub fn init_ui_material_pipeline<M: UiMaterial>(
         ),
     );
 
-    let load_default = || load_embedded_asset!(asset_server.as_ref(), "ui_material.wgsl");
-
     commands.insert_resource(UiMaterialPipeline::<M> {
         ui_layout,
         view_layout,
-        vertex_shader: match M::vertex_shader() {
-            ShaderRef::Default => load_default(),
-            ShaderRef::Handle(handle) => handle,
-            ShaderRef::Path(path) => asset_server.load(path),
-        },
-        fragment_shader: match M::fragment_shader() {
-            ShaderRef::Default => load_default(),
-            ShaderRef::Handle(handle) => handle,
-            ShaderRef::Path(path) => asset_server.load(path),
-        },
+        default_shader: load_embedded_asset!(asset_server.as_ref(), "ui_material.wgsl"),
         marker: PhantomData,
     });
 }
@@ -306,6 +294,8 @@ pub struct ExtractedUiMaterialNode<M: UiMaterial> {
     // Nodes with ambiguous camera will be ignored.
     pub extracted_camera_entity: Entity,
     pub main_entity: MainEntity,
+    pub vertex_shader: Handle<Shader>,
+    pub fragment_shader: Handle<Shader>,
     pub render_entity: Entity,
 }
 
@@ -337,6 +327,8 @@ pub fn extract_ui_material_nodes<M: UiMaterial>(
             &ComputedUiTargetCamera,
         )>,
     >,
+    ui_material_pipeline: Res<UiMaterialPipeline<M>>,
+    asset_server: Extract<Res<AssetServer>>,
     camera_map: Extract<UiCameraMap>,
 ) {
     let mut camera_mapper = camera_map.get_mapper();
@@ -349,10 +341,21 @@ pub fn extract_ui_material_nodes<M: UiMaterial>(
             continue;
         }
 
-        // Skip loading materials
-        if !materials.contains(handle) {
+        let Some(material) = materials.get(handle) else {
+            // Skip if the material isn't loaded yet
             continue;
-        }
+        };
+
+        let vertex_shader = match material.vertex_shader() {
+            ShaderRef::Default => ui_material_pipeline.default_shader.clone(),
+            ShaderRef::Handle(handle) => handle,
+            ShaderRef::Path(path) => asset_server.load(path),
+        };
+        let fragment_shader = match material.fragment_shader() {
+            ShaderRef::Default => ui_material_pipeline.default_shader.clone(),
+            ShaderRef::Handle(handle) => handle,
+            ShaderRef::Path(path) => asset_server.load(path),
+        };
 
         let Some(extracted_camera_entity) = camera_mapper.map(camera) else {
             continue;
@@ -372,6 +375,8 @@ pub fn extract_ui_material_nodes<M: UiMaterial>(
             clip: clip.map(|clip| clip.clip),
             extracted_camera_entity,
             main_entity: entity.into(),
+            vertex_shader,
+            fragment_shader,
         });
     }
 }
@@ -611,6 +616,8 @@ pub fn queue_ui_material_nodes<M: UiMaterial>(
             &ui_material_pipeline,
             UiMaterialKey {
                 hdr: view.hdr,
+                vertex_shader: extracted_uinode.vertex_shader.clone(),
+                fragment_shader: extracted_uinode.fragment_shader.clone(),
                 bind_group_data: material.key.clone(),
             },
         );

examples/ui/ui_material.rs

@@ -80,7 +80,7 @@ struct CustomUiMaterial {
 }
 
 impl UiMaterial for CustomUiMaterial {
-    fn fragment_shader() -> ShaderRef {
+    fn fragment_shader(&self) -> ShaderRef {
         SHADER_ASSET_PATH.into()
     }
 }

release-content/migration-guides/UiMaterial_functions_param.md

@@ -0,0 +1,6 @@
+---
+title: UiMaterial shader functions param change
+pull_requests: [20895]
+---
+
+`fn fragment_shader()` is now `fn fragment_shader(&self)` and `fn vertex_shader()` is now `fn vertex_shader(&self)` in `impl UiMaterial`, to allow accessing `&self`, which allows making the shaders dynamic