Implement RustGpuShader to ensure shader meta load

Author: ProfLander

README.md

@@ -25,9 +25,9 @@ Currently this must occur before `RenderPlugin` is added (most often via `Defaul
 
 For each `RustGpuMaterial` implementor, add a `RustGpuMaterialPlugin::<M>` to your app to setup rendering machinery and hot-reload / hot-rebuild support if the respective features are enabled (see below.)
 
-When instantiating `RustGpu` materials, handles to the `Shader` representing your `rust-gpu` SPIR-V will be required.
-These handles should be acquired via the `AssetServer::load_rust_gpu_shader` extension method, which is provided by the `LoadRustGpuShader` trait.
-This will allow the plugin to automatically load the `.spv.json` metadata asset required for hot-reloading support.
+When instantiating `RustGpu` materials, `RustGpuShader` handles will be required.
+These are equivalent to `Handle<Shader>` with some extra hot-reloading machinery,
+and can be acquired via the `AssetServer::load_rust_gpu_shader` extension method provided by the `LoadRustGpuShader` trait.
 
 ## Feature Flags
 

src/load_rust_gpu_shader.rs

@@ -1,17 +1,20 @@
 use std::path::PathBuf;
 
-use bevy::prelude::{Shader, Handle, AssetServer};
+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>);
+
 /// Loads a `Shader` asset,
 /// and optionally its metadata if the corresponding feature flag is enabled.
 pub trait LoadRustGpuShader {
-    fn load_rust_gpu_shader<'a, P: Into<PathBuf>>(&self, path: P) -> Handle<Shader>;
+    fn load_rust_gpu_shader<'a, P: Into<PathBuf>>(&self, path: P) -> RustGpuShader;
 }
 
 impl LoadRustGpuShader for AssetServer {
-    fn load_rust_gpu_shader<'a, P: Into<PathBuf>>(&self, path: P) -> Handle<Shader> {
+    fn load_rust_gpu_shader<'a, P: Into<PathBuf>>(&self, path: P) -> RustGpuShader {
         let path = path.into();
 
         let mut meta_path: PathBuf;
@@ -23,7 +26,6 @@ impl LoadRustGpuShader for AssetServer {
             meta_path.push(last + ".json");
         }
 
-        #[allow(unused_variables)]
         let shader = self.load(path);
 
         #[cfg(feature = "hot-reload")]
@@ -35,7 +37,6 @@ impl LoadRustGpuShader for AssetServer {
             );
         }
 
-        shader
+        RustGpuShader(shader)
     }
 }
-

src/rust_gpu.rs

@@ -5,8 +5,8 @@ use std::{any::TypeId, marker::PhantomData, sync::RwLock};
 use bevy::{
     pbr::MaterialPipelineKey,
     prelude::{
-        default, info, warn, CoreStage, Handle, Image, IntoSystemDescriptor, Material,
-        MaterialPlugin, Plugin, Shader,
+        default, info, warn, CoreStage, Image, IntoSystemDescriptor, Material, MaterialPlugin,
+        Plugin,
     },
     reflect::TypeUuid,
     render::render_resource::{AsBindGroup, PreparedBindGroup, ShaderRef},
@@ -15,6 +15,7 @@ use bevy::{
 use once_cell::sync::Lazy;
 
 use crate::{
+    load_rust_gpu_shader::RustGpuShader,
     prelude::{EntryPoint, Export, ExportHandle, RustGpuMaterial, SHADER_META},
     systems::{reload_materials, shader_events},
 };
@@ -82,8 +83,8 @@ where
     M: AsBindGroup,
 {
     pub base: M::Data,
-    pub vertex_shader: Option<Handle<Shader>>,
-    pub fragment_shader: Option<Handle<Shader>>,
+    pub vertex_shader: Option<RustGpuShader>,
+    pub fragment_shader: Option<RustGpuShader>,
     pub iteration: usize,
     #[cfg(feature = "hot-rebuild")]
     pub export_handle: Option<ExportHandle>,
@@ -146,10 +147,10 @@ pub struct RustGpu<M> {
     pub base: M,
 
     /// If `Some`, overrides [`Material::vertex_shader`] during specialization.
-    pub vertex_shader: Option<Handle<Shader>>,
+    pub vertex_shader: Option<RustGpuShader>,
 
     /// If `Some`, overrides [`Material::fragment_shader`] during specialization.
-    pub fragment_shader: Option<Handle<Shader>>,
+    pub fragment_shader: Option<RustGpuShader>,
 
     /// Current reload iteration, used to drive hot-reloading.
     pub iteration: usize,
@@ -329,7 +330,7 @@ where
             #[cfg(feature = "hot-reload")]
             {
                 let metas = SHADER_META.read().unwrap();
-                if let Some(vertex_meta) = metas.get(&vertex_shader) {
+                if let Some(vertex_meta) = metas.get(&vertex_shader.0) {
                     info!("Vertex meta is valid");
                     info!("Checking entry point {entry_point:}");
                     if !vertex_meta.entry_points.contains(&entry_point) {
@@ -354,7 +355,7 @@ where
 
             if apply {
                 info!("Applying vertex shader and entry point");
-                descriptor.vertex.shader = vertex_shader;
+                descriptor.vertex.shader = vertex_shader.0;
                 descriptor.vertex.entry_point = entry_point.into();
             } else {
                 warn!("Falling back to default vertex shader.");
@@ -382,7 +383,7 @@ where
                 {
                     info!("Fragment meta is present");
                     let metas = SHADER_META.read().unwrap();
-                    if let Some(fragment_meta) = metas.get(&fragment_shader) {
+                    if let Some(fragment_meta) = metas.get(&fragment_shader.0) {
                         info!("Fragment meta is valid");
                         info!("Checking entry point {entry_point:}");
                         if !fragment_meta.entry_points.contains(&entry_point) {
@@ -406,7 +407,7 @@ where
 
                 if apply {
                     info!("Applying fragment shader and entry point");
-                    fragment_descriptor.shader = fragment_shader;
+                    fragment_descriptor.shader = fragment_shader.0;
                     fragment_descriptor.entry_point = entry_point.into();
                 } else {
                     warn!("Falling back to default fragment shader.");

src/rust_gpu_material.rs

@@ -1,18 +1,8 @@
 //! Trait
 
-use std::marker::PhantomData;
+use bevy::prelude::Material;
 
-use bevy::{
-    prelude::{
-        default, CoreStage, Handle, IntoSystemDescriptor, Material, MaterialPlugin, Plugin, Shader,
-    },
-    render::render_resource::AsBindGroup,
-};
-
-use crate::prelude::{reload_materials, shader_events, EntryPoint, RustGpu};
-
-#[cfg(feature = "hot-rebuild")]
-use crate::prelude::ExportHandle;
+use crate::prelude::EntryPoint;
 
 /// A [`Material`] type with statically-known `rust-gpu` vertex and fragment entry points.
 pub trait RustGpuMaterial: Material {

src/systems.rs

@@ -48,13 +48,13 @@ pub fn reload_materials<M>(
         let mut reload = false;
 
         if let Some(vertex_shader) = &material.vertex_shader {
-            if changed_shaders.contains(vertex_shader) {
+            if changed_shaders.contains(&vertex_shader.0) {
                 reload = true;
             }
         }
 
         if let Some(fragment_shader) = &material.fragment_shader {
-            if changed_shaders.contains(fragment_shader) {
+            if changed_shaders.contains(&fragment_shader.0) {
                 reload = true;
             }
         }