Abstract export functionality behind plugin

ProfLander · ProfLander · commit 6530d7d03641 · 2023-03-04T01:34:52.000Z
diff --git a/src/entry_point_export.rs b/src/entry_point_export.rs
@@ -17,6 +17,16 @@ use bevy::{
 };
 use serde::{Deserialize, Serialize};
 
+#[cfg(feature = "hot-rebuild")]
+pub(crate) static EXPORT_HANDLES: once_cell::sync::Lazy<
+    std::sync::RwLock<HashMap<PathBuf, ExportHandle>>,
+> = once_cell::sync::Lazy::new(default);
+
+#[cfg(feature = "hot-rebuild")]
+pub(crate) static MATERIAL_EXPORTS: once_cell::sync::Lazy<
+    std::sync::RwLock<HashMap<std::any::TypeId, PathBuf>>,
+> = once_cell::sync::Lazy::new(default);
+
 /// Handles exporting known `RustGpuMaterial` permutations to a JSON file for static compilation.
 pub struct EntryPointExportPlugin;
 
@@ -25,6 +35,10 @@ impl Plugin for EntryPointExportPlugin {
         app.world.init_non_send_resource::<EntryPointExport>();
 
         app.add_system_to_stage(
+            CoreStage::Update,
+            EntryPointExport::create_export_containers_system,
+        )
+        .add_system_to_stage(
             CoreStage::Last,
             EntryPointExport::receive_entry_points_system,
         )
@@ -59,37 +73,36 @@ struct EntryPoints {
 /// Container for a set of entry points, with MPSC handles and change tracking
 #[derive(Debug)]
 struct EntryPointExportContainer {
-    tx: ExportHandle,
     rx: EntryPointReceiver,
     entry_points: EntryPoints,
     changed: bool,
 }
 
-impl Default for EntryPointExportContainer {
-    fn default() -> Self {
-        let (tx, rx) = std::sync::mpsc::sync_channel::<Export>(32);
-
-        EntryPointExportContainer {
-            tx,
-            rx,
-            entry_points: default(),
-            changed: default(),
-        }
-    }
-}
-
 /// Non-send resource used to register export files and aggregate their entry points.
-#[derive(Debug, Default)]
-pub struct EntryPointExport {
+#[derive(Debug, Default, Deref, DerefMut)]
+struct EntryPointExport {
     exports: HashMap<PathBuf, EntryPointExportContainer>,
 }
 
 impl EntryPointExport {
-    /// Registers a path to which entrypoints will be exported,
-    /// returning a corresponding [`ExportHandle`] that can be passed to a
-    /// [`RustGpu`](crate::rust_gpu::RustGpu) material.
-    pub fn export<T: Into<PathBuf>>(&mut self, path: T) -> ExportHandle {
-        self.exports.entry(path.into()).or_default().tx.clone()
+    /// System used to populate export containers for registered materials
+    pub fn create_export_containers_system(mut exports: NonSendMut<Self>) {
+        let material_exports = MATERIAL_EXPORTS.read().unwrap();
+        for (_, path) in material_exports.iter() {
+            if !exports.contains_key(path) {
+                let (tx, rx) = std::sync::mpsc::sync_channel::<Export>(32);
+
+                EXPORT_HANDLES.write().unwrap().insert(path.clone(), tx);
+
+                let container = EntryPointExportContainer {
+                    rx,
+                    entry_points: default(),
+                    changed: default(),
+                };
+
+                exports.insert(path.clone(), container);
+            }
+        }
     }
 
     /// System used to receive and store entry points sent from materials.
diff --git a/src/load_rust_gpu_shader.rs b/src/load_rust_gpu_shader.rs
@@ -2,8 +2,6 @@ use std::path::PathBuf;
 
 use bevy::prelude::{AssetServer, Handle, Shader};
 
-use crate::prelude::SHADER_META_MAP;
-
 #[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct RustGpuShader(pub(crate) Handle<Shader>);
 
@@ -30,7 +28,7 @@ impl LoadRustGpuShader for AssetServer {
 
         #[cfg(feature = "hot-reload")]
         {
-            let mut shader_meta_map = SHADER_META_MAP.write().unwrap();
+            let mut shader_meta_map = crate::prelude::SHADER_META_MAP.write().unwrap();
             shader_meta_map.add(
                 shader.clone_weak(),
                 self.load::<crate::prelude::ModuleMeta, _>(meta_path),
diff --git a/src/rust_gpu.rs b/src/rust_gpu.rs
@@ -1,22 +1,22 @@
 //! Wrapper for extending a `Material` with `rust-gpu` shader functionality.
 
-use std::{any::TypeId, marker::PhantomData, sync::RwLock};
+use std::{any::TypeId, marker::PhantomData, path::PathBuf, sync::RwLock};
 
 use bevy::{
     pbr::MaterialPipelineKey,
     prelude::{
         default, info, warn, CoreStage, Image, IntoSystemDescriptor, Material, MaterialPlugin,
-        Plugin,
+        Plugin, Handle, Shader, Deref, DerefMut, Resource,
     },
     reflect::TypeUuid,
     render::render_resource::{AsBindGroup, PreparedBindGroup, ShaderRef},
-    utils::HashMap,
+    utils::{HashMap, HashSet},
 };
 use once_cell::sync::Lazy;
 
 use crate::{
     load_rust_gpu_shader::RustGpuShader,
-    prelude::{EntryPoint, Export, ExportHandle, RustGpuMaterial, SHADER_META},
+    prelude::{EntryPoint, RustGpuMaterial},
     systems::{reload_materials, shader_events},
 };
 
@@ -68,6 +68,10 @@ where
     }
 }
 
+/// A resource to track `rust-gpu` shaders that have been reloaded on a given frame
+#[derive(Debug, Default, Clone, Deref, DerefMut, Resource)]
+pub struct ChangedShaders(pub HashSet<Handle<Shader>>);
+
 /// Type-level RustGpu material settings
 #[derive(Debug, Default, Copy, Clone)]
 pub struct RustGpuSettings {
@@ -86,8 +90,6 @@ where
     pub vertex_shader: Option<RustGpuShader>,
     pub fragment_shader: Option<RustGpuShader>,
     pub iteration: usize,
-    #[cfg(feature = "hot-rebuild")]
-    pub export_handle: Option<ExportHandle>,
 }
 
 impl<M> Clone for RustGpuKey<M>
@@ -101,7 +103,6 @@ where
             vertex_shader: self.vertex_shader.clone(),
             fragment_shader: self.fragment_shader.clone(),
             iteration: self.iteration.clone(),
-            export_handle: self.export_handle.clone(),
         }
     }
 }
@@ -154,10 +155,6 @@ pub struct RustGpu<M> {
 
     /// Current reload iteration, used to drive hot-reloading.
     pub iteration: usize,
-
-    /// If `Some`, active entry points will be reported to this handle.
-    #[cfg(feature = "hot-rebuild")]
-    pub export_handle: Option<ExportHandle>,
 }
 
 impl<M> PartialEq for RustGpu<M>
@@ -237,8 +234,6 @@ where
                     vertex_shader: self.vertex_shader.clone(),
                     fragment_shader: self.fragment_shader.clone(),
                     iteration: self.iteration,
-                    #[cfg(feature = "hot-rebuild")]
-                    export_handle: self.export_handle.clone(),
                 },
             })
     }
@@ -329,7 +324,7 @@ where
 
             #[cfg(feature = "hot-reload")]
             {
-                let metas = SHADER_META.read().unwrap();
+                let metas = crate::prelude::SHADER_META.read().unwrap();
                 if let Some(vertex_meta) = metas.get(&vertex_shader.0) {
                     info!("Vertex meta is valid");
                     info!("Checking entry point {entry_point:}");
@@ -343,15 +338,25 @@ where
             }
 
             #[cfg(feature = "hot-rebuild")]
-            if let Some(sender) = &key.bind_group_data.export_handle {
+            'hot_rebuild: {
+                let exports = crate::prelude::MATERIAL_EXPORTS.read().unwrap();
+                let Some(export) = exports.get(&std::any::TypeId::of::<Self>()) else {
+                    break 'hot_rebuild;
+                };
+
+                let handles = crate::prelude::EXPORT_HANDLES.read().unwrap();
+                let Some(handle) = handles.get(export) else {
+                    break 'hot_rebuild;
+                };
+
                 info!("Entrypoint sender is valid");
-                sender
-                    .send(Export {
+                handle
+                    .send(crate::prelude::Export {
                         shader: M::Vertex::NAME,
                         permutation: M::Vertex::permutation(&shader_defs),
                     })
                     .unwrap();
-            }
+            };
 
             if apply {
                 info!("Applying vertex shader and entry point");
@@ -382,7 +387,7 @@ where
                 #[cfg(feature = "hot-reload")]
                 {
                     info!("Fragment meta is present");
-                    let metas = SHADER_META.read().unwrap();
+                    let metas = crate::prelude::SHADER_META.read().unwrap();
                     if let Some(fragment_meta) = metas.get(&fragment_shader.0) {
                         info!("Fragment meta is valid");
                         info!("Checking entry point {entry_point:}");
@@ -396,14 +401,25 @@ where
                 }
 
                 #[cfg(feature = "hot-rebuild")]
-                if let Some(sender) = &key.bind_group_data.export_handle {
-                    sender
-                        .send(Export {
+                'hot_rebuild: {
+                    let exports = crate::prelude::MATERIAL_EXPORTS.read().unwrap();
+                    let Some(export) = exports.get(&std::any::TypeId::of::<Self>()) else {
+                        break 'hot_rebuild;
+                    };
+
+                    let handles = crate::prelude::EXPORT_HANDLES.read().unwrap();
+                    let Some(handle) = handles.get(export) else {
+                        break 'hot_rebuild;
+                    };
+
+                    info!("Entrypoint sender is valid");
+                    handle
+                        .send(crate::prelude::Export {
                             shader: M::Fragment::NAME,
                             permutation: M::Fragment::permutation(&shader_defs),
                         })
                         .unwrap();
-                }
+                };
 
                 if apply {
                     info!("Applying fragment shader and entry point");
@@ -431,4 +447,10 @@ where
         let mut settings = MATERIAL_SETTINGS.write().unwrap();
         f(&mut settings.entry(TypeId::of::<Self>()).or_default());
     }
+
+    #[cfg(feature = "hot-rebuild")]
+    pub fn export_to<P: Into<PathBuf>>(path: P) {
+        let mut handles = crate::prelude::MATERIAL_EXPORTS.write().unwrap();
+        handles.insert(std::any::TypeId::of::<Self>(), path.into());
+    }
 }
diff --git a/src/shader_meta.rs b/src/shader_meta.rs
@@ -18,7 +18,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     prelude::{reload_materials, RustGpuMaterial},
-    systems::shader_events,
+    systems::shader_events, ChangedShaders,
 };
 
 pub(crate) static SHADER_META: Lazy<RwLock<ShaderMeta>> = Lazy::new(Default::default);
@@ -103,10 +103,6 @@ pub struct ModuleMeta {
     pub module: String,
 }
 
-/// A resource to track `rust-gpu` shaders that have been reloaded on a given frame
-#[derive(Debug, Default, Clone, Deref, DerefMut, Resource)]
-pub struct ChangedShaders(pub HashSet<Handle<Shader>>);
-
 /// Listens for asset events, updates backend data, and triggers material re-specialization
 pub fn module_meta_events<M>(
     mut module_meta_events: EventReader<AssetEvent<ModuleMeta>>,
diff --git a/src/systems.rs b/src/systems.rs
@@ -3,7 +3,7 @@
 use bevy::prelude::{info, AssetEvent, Assets, EventReader, Res, ResMut, Shader};
 
 use crate::{
-    prelude::{ChangedShaders, ModuleMeta, RustGpuMaterial, SHADER_META, SHADER_META_MAP},
+    prelude::{ChangedShaders, RustGpuMaterial},
     RustGpu,
 };
 
@@ -25,7 +25,7 @@ pub fn shader_events<M>(
             } => {
                 #[cfg(feature = "hot-reload")]
                 // Remove meta in case the shader and meta load on different frames
-                SHADER_META.write().unwrap().remove(shader_handle);
+                crate::prelude::SHADER_META.write().unwrap().remove(shader_handle);
 
                 // Mark this shader for material reloading
                 changed_shaders.insert(shader_handle.clone_weak());
@@ -71,13 +71,13 @@ pub fn reload_materials<M>(
 /// Listens for [`ModuleMeta`] asset events, updates backend data,
 /// and aggregates change events for application in [`reload_materials`].
 pub fn module_meta_events<M>(
-    mut module_meta_events: EventReader<AssetEvent<ModuleMeta>>,
-    assets: Res<Assets<ModuleMeta>>,
+    mut module_meta_events: EventReader<AssetEvent<crate::prelude::ModuleMeta>>,
+    assets: Res<Assets<crate::prelude::ModuleMeta>>,
     mut changed_shaders: ResMut<ChangedShaders>,
 ) where
     M: RustGpuMaterial,
 {
-    let shader_meta_map = SHADER_META_MAP.read().unwrap();
+    let shader_meta_map = crate::prelude::SHADER_META_MAP.read().unwrap();
 
     for event in module_meta_events.iter() {
         match event {
@@ -89,7 +89,7 @@ pub fn module_meta_events<M>(
                     // Update module meta
                     if let Some(asset) = assets.get(handle) {
                         info!("Updating shader meta for {handle:?}");
-                        SHADER_META
+                        crate::prelude::SHADER_META
                             .write()
                             .unwrap()
                             .insert(shader.clone_weak(), asset.clone());
