WESL: Suport importing shader from asset source

Author: beicause

crates/bevy_shader/src/shader_cache.rs

@@ -65,7 +65,7 @@ pub struct ShaderCache<ShaderModule, RenderDevice> {
         &ValidateShader,
     ) -> Result<ShaderModule, PipelineCacheError>,
     #[cfg(feature = "shader_format_wesl")]
-    asset_paths: HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
+    asset_paths: HashMap<wesl::ModulePath, AssetId<Shader>>,
     shaders: HashMap<AssetId<Shader>, Shader>,
     import_path_shaders: HashMap<ShaderImport, AssetId<Shader>>,
     waiting_on_import: HashMap<ShaderImport, Vec<AssetId<Shader>>>,
@@ -210,7 +210,7 @@ impl<ShaderModule, RenderDevice> ShaderCache<ShaderModule, RenderDevice> {
                         if let ShaderImport::AssetPath(path) = shader.import_path() {
                             let shader_resolver =
                                 ShaderResolver::new(&self.asset_paths, &self.shaders);
-                            let module_path = wesl::syntax::ModulePath::from_path(path);
+                            let module_path = wesl_module_path_from_asset_path(path);
                             let mut compiler_options = wesl::CompileOptions {
                                 imports: true,
                                 condcomp: true,
@@ -361,7 +361,7 @@ impl<ShaderModule, RenderDevice> ShaderCache<ShaderModule, RenderDevice> {
             && let ShaderImport::AssetPath(path) = shader.import_path()
         {
             self.asset_paths
-                .insert(wesl::syntax::ModulePath::from_path(path), id);
+                .insert(wesl_module_path_from_asset_path(path), id);
         }
         self.shaders.insert(id, shader);
         pipelines_to_queue
@@ -379,14 +379,14 @@ impl<ShaderModule, RenderDevice> ShaderCache<ShaderModule, RenderDevice> {
 
 #[cfg(feature = "shader_format_wesl")]
 pub struct ShaderResolver<'a> {
-    asset_paths: &'a HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
+    asset_paths: &'a HashMap<wesl::ModulePath, AssetId<Shader>>,
     shaders: &'a HashMap<AssetId<Shader>, Shader>,
 }
 
 #[cfg(feature = "shader_format_wesl")]
 impl<'a> ShaderResolver<'a> {
     pub fn new(
-        asset_paths: &'a HashMap<wesl::syntax::ModulePath, AssetId<Shader>>,
+        asset_paths: &'a HashMap<wesl::ModulePath, AssetId<Shader>>,
         shaders: &'a HashMap<AssetId<Shader>, Shader>,
     ) -> Self {
         Self {
@@ -400,7 +400,7 @@ impl<'a> ShaderResolver<'a> {
 impl<'a> wesl::Resolver for ShaderResolver<'a> {
     fn resolve_source(
         &self,
-        module_path: &wesl::syntax::ModulePath,
+        module_path: &wesl::ModulePath,
     ) -> Result<alloc::borrow::Cow<'_, str>, wesl::ResolveError> {
         let asset_id = self.asset_paths.get(module_path).ok_or_else(|| {
             wesl::ResolveError::ModuleNotFound(module_path.clone(), "Invalid asset id".to_string())
@@ -411,6 +411,27 @@ impl<'a> wesl::Resolver for ShaderResolver<'a> {
     }
 }
 
+#[cfg(feature = "shader_format_wesl")]
+fn wesl_module_path_from_asset_path(path: &String) -> wesl::ModulePath {
+    use bevy_asset::{io::AssetSourceId, AssetPath};
+    use std::path::PathBuf;
+
+    let asset_path = AssetPath::from(path);
+    if let AssetSourceId::Name(source) = asset_path.source() {
+        let mut comp = vec!["bevy_asset".to_string()];
+        let mut path = PathBuf::new();
+        path.push(source.as_ref());
+        path.push(asset_path.path());
+        comp.extend(wesl::ModulePath::from_path(path).components);
+        wesl::ModulePath {
+            origin: wesl::syntax::PathOrigin::Package,
+            components: comp,
+        }
+    } else {
+        wesl::ModulePath::from_path(asset_path.path())
+    }
+}
+
 /// Type of error returned by a `PipelineCache` when the creation of a GPU pipeline object failed.
 #[cfg_attr(
     not(target_arch = "wasm32"),

assets/shaders/custom_material.wesl renamed to examples/shader/files/custom_material.wesl

@@ -1,9 +1,5 @@
 import super::util::make_polka_dots;
-
-struct VertexOutput {
-    @builtin(position) position: vec4<f32>,
-    @location(2) uv: vec2<f32>,
-}
+import bevy_asset::embedded::shader_material_wesl::files::vertex_attr;
 
 struct CustomMaterial {
     // Needed for 16-bit alignment on WebGL2
@@ -14,7 +10,7 @@ struct CustomMaterial {
 
 @fragment
 fn fragment(
-    mesh: VertexOutput,
+    mesh: vertex_attr::VertexOutput,
 ) -> @location(0) vec4<f32> {
     return make_polka_dots(mesh.uv, material.time.x);
 }

assets/shaders/util.wesl renamed to examples/shader/files/util.wesl

@@ -0,0 +1,4 @@
+struct VertexOutput {
+    @builtin(position) position: vec4<f32>,
+    @location(2) uv: vec2<f32>,
+}

examples/shader/files/vertex_attr.wesl

@@ -1,5 +1,6 @@
 //! A shader that uses the WESL shading language.
 
+use bevy::asset::embedded_asset;
 use bevy::{
     mesh::MeshVertexBufferLayoutRef,
     pbr::{MaterialPipeline, MaterialPipelineKey},
@@ -12,7 +13,8 @@ use bevy::{
 };
 
 /// This example uses shader source files from the assets subdirectory
-const FRAGMENT_SHADER_ASSET_PATH: &str = "shaders/custom_material.wesl";
+const FRAGMENT_SHADER_ASSET_PATH: &str =
+    "embedded://shader_material_wesl/files/custom_material.wesl";
 
 fn main() {
     App::new()
@@ -37,13 +39,26 @@ pub struct CustomMaterialPlugin;
 #[derive(Resource)]
 struct UtilityShader(Handle<Shader>);
 
+#[expect(
+    dead_code,
+    reason = "used to kept a strong handle, shader is referenced by the material"
+)]
+#[derive(Resource)]
+struct VertexAttrShader(Handle<Shader>);
+
 impl Plugin for CustomMaterialPlugin {
     fn build(&self, app: &mut App) {
-        let handle = app
-            .world_mut()
-            .resource_mut::<AssetServer>()
-            .load::<Shader>("shaders/util.wesl");
-        app.insert_resource(UtilityShader(handle));
+        embedded_asset!(app, "examples/shader", "files/custom_material.wesl");
+        embedded_asset!(app, "examples/shader", "files/vertex_attr.wesl");
+        embedded_asset!(app, "examples/shader", "files/util.wesl");
+
+        let asset_server = app.world_mut().resource_mut::<AssetServer>();
+        let utils_handle =
+            asset_server.load::<Shader>("embedded://shader_material_wesl/files/util.wesl");
+        let vertex_attr_handle =
+            asset_server.load::<Shader>("embedded://shader_material_wesl/files/vertex_attr.wesl");
+        app.insert_resource(UtilityShader(utils_handle))
+            .insert_resource(VertexAttrShader(vertex_attr_handle));
     }
 }