Live reloading of shaders

assets/shaders/custom.frag

@@ -0,0 +1,11 @@
+#version 450
+
+layout(location = 0) out vec4 o_Target;
+
+layout(set = 1, binding = 1) uniform MyMaterial_color {
+    vec4 color;
+};
+
+void main() {
+    o_Target = color * 0.5;
+}

crates/bevy_render/src/draw.rs

@@ -6,8 +6,9 @@ use crate::{
     },
     shader::Shader,
 };
-use bevy_asset::{Assets, Handle};
-use bevy_ecs::{Query, Res, ResMut, SystemParam};
+use bevy_app::{EventReader, Events};
+use bevy_asset::{AssetEvent, Assets, Handle};
+use bevy_ecs::{Local, Query, Res, ResMut, SystemParam};
 use bevy_reflect::Reflect;
 use std::{ops::Range, sync::Arc};
 use thiserror::Error;
@@ -123,6 +124,8 @@ pub enum DrawError {
 pub struct DrawContext<'a> {
     pub pipelines: ResMut<'a, Assets<PipelineDescriptor>>,
     pub shaders: ResMut<'a, Assets<Shader>>,
+    pub shader_events: ResMut<'a, Events<AssetEvent<Shader>>>,
+    pub shader_event_reader: Local<'a, EventReader<AssetEvent<Shader>>>,
     pub pipeline_compiler: ResMut<'a, PipelineCompiler>,
     pub render_resource_context: Res<'a, Box<dyn RenderResourceContext>>,
     pub shared_buffers: Res<'a, SharedBuffers>,

crates/bevy_render/src/lib.rs

@@ -44,6 +44,7 @@ use render_graph::{
     RenderGraph,
 };
 use renderer::{AssetRenderResourceBindings, RenderResourceBindings};
+use shader::ShaderLoader;
 #[cfg(feature = "hdr")]
 use texture::HdrTextureLoader;
 #[cfg(feature = "png")]
@@ -87,6 +88,8 @@ impl Plugin for RenderPlugin {
             app.init_asset_loader::<HdrTextureLoader>();
         }
 
+        app.init_asset_loader::<ShaderLoader>();
+
         if app.resources().get::<ClearColor>().is_none() {
             app.resources_mut().insert(ClearColor::default());
         }

crates/bevy_render/src/pipeline/pipeline_compiler.rs

@@ -282,4 +282,22 @@ impl PipelineCompiler {
             })
             .flatten()
     }
+
+    /// Remove any specialized shaders or pipelines that match the
+    /// changed shader set.
+    pub fn remove_shaders(
+        &mut self,
+        shaders: HashSet<&Handle<Shader>>,
+        pipelines: &Assets<PipelineDescriptor>,
+    ) {
+        self.specialized_shaders.retain(|k, _| !shaders.contains(k));
+        self.specialized_pipelines.retain(|k, _| {
+            pipelines
+                .get(k)
+                .unwrap()
+                .shader_stages
+                .iter()
+                .all(|s| !shaders.contains(&s))
+        });
+    }
 }

crates/bevy_render/src/pipeline/render_pipelines.rs

@@ -5,9 +5,10 @@ use crate::{
     prelude::Msaa,
     renderer::RenderResourceBindings,
 };
-use bevy_asset::{Assets, Handle};
+use bevy_asset::{AssetEvent, Assets, Handle};
 use bevy_ecs::{Query, Res, ResMut};
 use bevy_reflect::Reflect;
+use bevy_utils::HashSet;
 
 #[derive(Debug, Default, Clone, Reflect)]
 pub struct RenderPipeline {
@@ -82,6 +83,27 @@ pub fn draw_render_pipelines_system(
     meshes: Res<Assets<Mesh>>,
     mut query: Query<(&mut Draw, &mut RenderPipelines, &Handle<Mesh>)>,
 ) {
+    // Pipelines will be rebuilt for shaders that have been modified.
+    let mut changed_shaders = HashSet::default();
+    for event in draw_context
+        .shader_event_reader
+        .iter(&draw_context.shader_events)
+    {
+        match event {
+            AssetEvent::Modified { handle } => {
+                changed_shaders.insert(handle);
+            }
+            // Creating shaders on the fly is unhandled since they
+            // have to exist already when assigned to a pipeline. If a
+            // shader is removed the pipeline keeps using its
+            // specialized version. Maybe this should be a warning?
+            AssetEvent::Created { .. } | AssetEvent::Removed { .. } => (),
+        }
+    }
+    draw_context
+        .pipeline_compiler
+        .remove_shaders(changed_shaders, &draw_context.pipelines);
+
     for (mut draw, mut render_pipelines, mesh_handle) in query.iter_mut() {
         if !draw.is_visible {
             continue;

crates/bevy_render/src/shader/shader.rs

@@ -1,6 +1,7 @@
 use super::ShaderLayout;
-use bevy_asset::Handle;
+use bevy_asset::{AssetLoader, Handle, LoadContext, LoadedAsset};
 use bevy_reflect::TypeUuid;
+use bevy_utils::BoxedFuture;
 use std::marker::Copy;
 
 /// The stage of a shader
@@ -188,3 +189,31 @@ impl ShaderStages {
         }
     }
 }
+
+#[derive(Default)]
+pub struct ShaderLoader;
+
+impl AssetLoader for ShaderLoader {
+    fn load<'a>(
+        &'a self,
+        bytes: &'a [u8],
+        load_context: &'a mut LoadContext,
+    ) -> BoxedFuture<'a, Result<(), anyhow::Error>> {
+        Box::pin(async move {
+            let ext = load_context.path().extension().unwrap().to_str().unwrap();
+
+            let shader = match ext {
+                "vert" => Shader::from_glsl(ShaderStage::Vertex, std::str::from_utf8(bytes)?),
+                "frag" => Shader::from_glsl(ShaderStage::Fragment, std::str::from_utf8(bytes)?),
+                _ => panic!("unhandled extension: {}", ext),
+            };
+
+            load_context.set_default_asset(LoadedAsset::new(shader));
+            Ok(())
+        })
+    }
+
+    fn extensions(&self) -> &[&str] {
+        &["vert", "frag"]
+    }
+}

examples/shader/shader_custom_material.rs

@@ -39,29 +39,21 @@ void main() {
 }
 "#;
 
-const FRAGMENT_SHADER: &str = r#"
-#version 450
-layout(location = 0) out vec4 o_Target;
-layout(set = 1, binding = 1) uniform MyMaterial_color {
-    vec4 color;
-};
-void main() {
-    o_Target = color;
-}
-"#;
-
 fn setup(
     commands: &mut Commands,
+    asset_server: ResMut<AssetServer>,
     mut pipelines: ResMut<Assets<PipelineDescriptor>>,
     mut shaders: ResMut<Assets<Shader>>,
     mut meshes: ResMut<Assets<Mesh>>,
     mut materials: ResMut<Assets<MyMaterial>>,
     mut render_graph: ResMut<RenderGraph>,
 ) {
+    asset_server.watch_for_changes().unwrap();
+
     // Create a new shader pipeline
     let pipeline_handle = pipelines.add(PipelineDescriptor::default_config(ShaderStages {
         vertex: shaders.add(Shader::from_glsl(ShaderStage::Vertex, VERTEX_SHADER)),
-        fragment: Some(shaders.add(Shader::from_glsl(ShaderStage::Fragment, FRAGMENT_SHADER))),
+        fragment: Some(asset_server.load::<Shader, _>("shaders/custom.frag")),
     }));
 
     // Add an AssetRenderResourcesNode to our Render Graph. This will bind MyMaterial resources to our shader