Track shader defs by source, handle asset changes, add test

yrns · yrns · commit d7b37f51c10d · 2020-12-18T13:17:31.000-07:00
diff --git a/crates/bevy_render/src/pipeline/pipeline_compiler.rs b/crates/bevy_render/src/pipeline/pipeline_compiler.rs
@@ -2,13 +2,12 @@ use super::{state_descriptors::PrimitiveTopology, IndexFormat, PipelineDescripto
 use crate::{
     pipeline::{BindType, InputStepMode, VertexBufferDescriptor},
     renderer::RenderResourceContext,
-    shader::{Shader, ShaderError, ShaderSource},
+    shader::{Shader, ShaderDefSource, ShaderError, ShaderSource},
 };
 use bevy_asset::{Assets, Handle};
 use bevy_reflect::Reflect;
 use bevy_utils::{HashMap, HashSet};
 use once_cell::sync::Lazy;
-use serde::{Deserialize, Serialize};
 
 #[derive(Clone, Eq, PartialEq, Debug, Reflect)]
 pub struct PipelineSpecialization {
@@ -40,9 +39,17 @@ impl PipelineSpecialization {
     }
 }
 
-#[derive(Clone, Eq, PartialEq, Debug, Default, Reflect, Serialize, Deserialize)]
+#[derive(Clone, Eq, PartialEq, Debug, Default, Reflect)]
 pub struct ShaderSpecialization {
-    pub shader_defs: HashSet<String>,
+    /// ShaderDefs tracked per-asset/component.
+    pub shader_defs: HashMap<ShaderDefSource, Vec<String>>,
+}
+
+impl ShaderSpecialization {
+    // get_specialized_shader should take a &[&str]?
+    pub fn get_shader_defs(&self) -> Vec<String> {
+        self.shader_defs.values().flatten().cloned().collect()
+    }
 }
 
 #[derive(Debug)]
@@ -95,13 +102,8 @@ impl PipelineCompiler {
             Ok(specialized_shader.shader.clone_weak())
         } else {
             // if no shader exists with the current configuration, create new shader and compile
-            let shader_def_vec = shader_specialization
-                .shader_defs
-                .iter()
-                .cloned()
-                .collect::<Vec<String>>();
-            let compiled_shader =
-                render_resource_context.get_specialized_shader(shader, Some(&shader_def_vec))?;
+            let compiled_shader = render_resource_context
+                .get_specialized_shader(shader, Some(&shader_specialization.get_shader_defs()))?;
             let specialized_handle = shaders.add(compiled_shader);
             let weak_specialized_handle = specialized_handle.clone_weak();
             specialized_shaders.push(SpecializedShader {
@@ -316,17 +318,10 @@ impl PipelineCompiler {
         if let Some(specialized_shaders) = self.specialized_shaders.get_mut(shader) {
             for specialized_shader in specialized_shaders {
                 // Recompile specialized shader. If it fails, we bail immediately.
-                let shader_def_vec = specialized_shader
-                    .specialization
-                    .shader_defs
-                    .iter()
-                    .cloned()
-                    .collect::<Vec<String>>();
-                let new_handle =
-                    shaders.add(render_resource_context.get_specialized_shader(
-                        shaders.get(shader).unwrap(),
-                        Some(&shader_def_vec),
-                    )?);
+                let new_handle = shaders.add(render_resource_context.get_specialized_shader(
+                    shaders.get(shader).unwrap(),
+                    Some(&specialized_shader.specialization.get_shader_defs()),
+                )?);
 
                 // Replace handle and remove old from assets.
                 let old_handle = std::mem::replace(&mut specialized_shader.shader, new_handle);
diff --git a/crates/bevy_render/src/shader/shader_defs.rs b/crates/bevy_render/src/shader/shader_defs.rs
@@ -1,8 +1,10 @@
-use bevy_asset::{Asset, Assets, Handle};
-
 use crate::{pipeline::RenderPipelines, Texture};
+use bevy_app::{EventReader, Events};
+use bevy_asset::{Asset, AssetEvent, Assets, Handle, HandleUntyped};
 pub use bevy_derive::ShaderDefs;
-use bevy_ecs::{Changed, Mut, Query, Res};
+use bevy_ecs::{Changed, Local, Mut, Query, QuerySet, Res};
+use bevy_reflect::{Reflect, TypeUuid};
+use bevy_utils::{HashSet, Uuid};
 
 /// Something that can either be "defined" or "not defined". This is used to determine if a "shader def" should be considered "defined"
 pub trait ShaderDef {
@@ -51,46 +53,181 @@ impl ShaderDef for Option<Handle<Texture>> {
     }
 }
 
+#[derive(Debug, Clone, Hash, Eq, PartialEq, Reflect)]
+pub enum ShaderDefSource {
+    Component(Uuid),
+    Asset(HandleUntyped),
+}
+
+impl<T: Asset> From<&Handle<T>> for ShaderDefSource {
+    fn from(h: &Handle<T>) -> Self {
+        Self::Asset(h.clone_weak_untyped())
+    }
+}
+
+impl From<Uuid> for ShaderDefSource {
+    fn from(uuid: Uuid) -> Self {
+        Self::Component(uuid)
+    }
+}
+
 /// Updates [RenderPipelines] with the latest [ShaderDefs]
 pub fn shader_defs_system<T>(mut query: Query<(&T, &mut RenderPipelines), Changed<T>>)
 where
-    T: ShaderDefs + Send + Sync + 'static,
+    T: ShaderDefs + TypeUuid + Send + Sync + 'static,
 {
-    query.iter_mut().for_each(update_render_pipelines)
+    query
+        .iter_mut()
+        .map(|(s, p)| (s, (T::TYPE_UUID).into(), p))
+        .for_each(update_render_pipelines)
 }
 
-/// Insert defined shader defs and remove undefined ones from render pipelines.
-fn update_render_pipelines<T>(q: (&T, Mut<RenderPipelines>))
+fn update_render_pipelines<T>(q: (&T, ShaderDefSource, Mut<RenderPipelines>))
 where
     T: ShaderDefs + Send + Sync + 'static,
 {
-    let (shader_defs, mut render_pipelines) = q;
-    for (shader_def, is_defined) in shader_defs.iter_shader_defs() {
-        for render_pipeline in render_pipelines.pipelines.iter_mut() {
-            let shader_defs = &mut render_pipeline
-                .specialization
-                .shader_specialization
-                .shader_defs;
-            let s = shader_def.to_string();
-            if is_defined {
-                shader_defs.insert(s);
-            } else {
-                shader_defs.remove(&s);
-            }
+    let (shader_defs, src, mut render_pipelines) = q;
+
+    let new_defs = shader_defs
+        .iter_shader_defs()
+        // FIX: revert macro
+        .filter_map(|(def, defined)| if defined { Some(def.to_string()) } else { None })
+        .collect::<Vec<_>>();
+    render_pipelines.pipelines.iter_mut().for_each(|p| {
+        *(p.specialization
+            .shader_specialization
+            .shader_defs
+            .entry(src.clone())
+            .or_default()) = new_defs.clone();
+    });
+}
+
+// FIX: track entities or clean this up
+//#[derive(Default)]
+pub struct AssetShaderDefsState<T: Asset> {
+    event_reader: EventReader<AssetEvent<T>>,
+    //entities: HashMap<Handle<T>, HashSet<Entity>>,
+}
+
+impl<T: Asset> Default for AssetShaderDefsState<T> {
+    fn default() -> Self {
+        Self {
+            event_reader: Default::default(),
+            //entities: Default::default(),
         }
     }
 }
 
 /// Updates [RenderPipelines] with the latest [ShaderDefs] from a given asset type
-pub fn asset_shader_defs_system<T: Asset>(
+pub fn asset_shader_defs_system<T>(
+    mut state: Local<AssetShaderDefsState<T>>,
     assets: Res<Assets<T>>,
-    mut query: Query<(&Handle<T>, &mut RenderPipelines), Changed<Handle<T>>>,
+    events: Res<Events<AssetEvent<T>>>,
+    mut queries: QuerySet<(
+        Query<(&Handle<T>, &mut RenderPipelines)>,
+        Query<(&Handle<T>, &mut RenderPipelines), Changed<Handle<T>>>,
+    )>,
 ) where
-    T: ShaderDefs + Send + Sync + 'static,
+    T: Default + Asset + ShaderDefs + Send + Sync + 'static,
 {
-    query
+    let changed = state
+        .event_reader
+        .iter(&events)
+        .fold(HashSet::default(), |mut set, event| {
+            match event {
+                AssetEvent::Created { handle } | AssetEvent::Modified { handle } => {
+                    set.insert(handle.clone_weak());
+                }
+                AssetEvent::Removed { handle } => {
+                    set.remove(&handle);
+                }
+            }
+            set
+        });
+
+    // Update for changed assets.
+    if changed.len() > 0 {
+        queries
+            .q0_mut()
+            .iter_mut()
+            .filter(|(h, _)| changed.contains(h))
+            .filter_map(|(h, p)| assets.get(h).map(|a| (a, h.into(), p)))
+            .for_each(update_render_pipelines);
+    }
+
+    // Update for changed asset handles.
+    queries
+        .q1_mut()
         .iter_mut()
-        // (Handle<T>, _) -> (&T, _)
-        .filter_map(|(h, p)| assets.get(h).map(|a| (a, p)))
+        // Not worth?
+        //.filter(|(h, _)| !changed.contains(h))
+        .filter_map(|(h, p)| assets.get(h).map(|a| (a, h.into(), p)))
         .for_each(update_render_pipelines);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::asset_shader_defs_system;
+    use super::ShaderDefs;
+    use crate::{self as bevy_render, pipeline::RenderPipeline, prelude::RenderPipelines};
+    use bevy_app::App;
+    use bevy_asset::{AddAsset, AssetPlugin, AssetServer, Assets, HandleId};
+    use bevy_core::CorePlugin;
+    use bevy_ecs::{Commands, ResMut};
+    use bevy_reflect::{ReflectPlugin, TypeUuid};
+
+    #[derive(Debug, Default, ShaderDefs, TypeUuid)]
+    #[uuid = "3130b0bf-46a6-42f2-8556-c1a04da20b7e"]
+    struct A {
+        #[shader_def]
+        d: bool,
+    }
+
+    fn shader_def_len(app: &App) -> usize {
+        app.world
+            .query::<&RenderPipelines>()
+            .next()
+            .unwrap()
+            .pipelines[0]
+            .specialization
+            .shader_specialization
+            .shader_defs
+            .len()
+    }
+
+    #[test]
+    fn empty_handle() {
+        // Insert an empty asset handle, and empty render pipelines.
+        let handle_id = HandleId::random::<A>();
+        let setup = move |commands: &mut Commands, asset_server: ResMut<AssetServer>| {
+            let h = asset_server.get_handle::<A, HandleId>(handle_id);
+            let render_pipelines = RenderPipelines::from_pipelines(vec![RenderPipeline::default()]);
+            commands.spawn((h, render_pipelines));
+        };
+
+        App::build()
+            .add_plugin(ReflectPlugin::default())
+            .add_plugin(CorePlugin::default())
+            .add_plugin(AssetPlugin::default())
+            .add_asset::<A>()
+            .add_system(asset_shader_defs_system::<A>)
+            .add_startup_system(setup)
+            .set_runner(move |mut app: App| {
+                app.initialize();
+                app.update();
+                assert_eq!(shader_def_len(&app), 0);
+                {
+                    let mut asset_server = app.resources.get_mut::<Assets<A>>().unwrap();
+                    asset_server.set(handle_id, A { d: true });
+                }
+
+                // Asset changed events are sent post-update, so we
+                // have to update twice to see the change.
+                app.update();
+                app.update();
+
+                assert_eq!(shader_def_len(&app), 1);
+            })
+            .run();
+    }
+}
