From 30cb95d96e0b19ca44f92b30d638a9c1fc62a83f Mon Sep 17 00:00:00 2001
From: robtfm <50659922+robtfm@users.noreply.github.com>
Date: Fri, 6 Oct 2023 02:34:57 +0100
Subject: [PATCH] fix custom shader imports (#10030)

# Objective

assets v2 broke custom shader imports. fix them

## Solution

store handles of any file dependencies in the `Shader` to avoid them
being immediately dropped.
also added a use into the `shader_material` example so that it'll be
harder to break support in future.
---
 assets/shaders/custom_material.wgsl              |  4 +++-
 assets/shaders/custom_material_import.wgsl       |  2 ++
 crates/bevy_render/src/render_resource/shader.rs | 13 +++++++++----
 3 files changed, 14 insertions(+), 5 deletions(-)
 create mode 100644 assets/shaders/custom_material_import.wgsl

diff --git a/assets/shaders/custom_material.wgsl b/assets/shaders/custom_material.wgsl
index c8e5a50ad3..34c2fc3f28 100644
--- a/assets/shaders/custom_material.wgsl
+++ b/assets/shaders/custom_material.wgsl
@@ -1,4 +1,6 @@
 #import bevy_pbr::mesh_vertex_output MeshVertexOutput
+// we can import items from shader modules in the assets folder with a quoted path
+#import "shaders/custom_material_import.wgsl" COLOR_MULTIPLIER
 
 struct CustomMaterial {
     color: vec4<f32>,
@@ -12,5 +14,5 @@ struct CustomMaterial {
 fn fragment(
     mesh: MeshVertexOutput,
 ) -> @location(0) vec4<f32> {
-    return material.color * textureSample(base_color_texture, base_color_sampler, mesh.uv);
+    return material.color * textureSample(base_color_texture, base_color_sampler, mesh.uv) * COLOR_MULTIPLIER;
 }
diff --git a/assets/shaders/custom_material_import.wgsl b/assets/shaders/custom_material_import.wgsl
new file mode 100644
index 0000000000..ac9f2eb673
--- /dev/null
+++ b/assets/shaders/custom_material_import.wgsl
@@ -0,0 +1,2 @@
+// this is made available to the importing module
+const COLOR_MULTIPLIER: vec4<f32> = vec4<f32>(1.0, 1.0, 1.0, 0.5);
diff --git a/crates/bevy_render/src/render_resource/shader.rs b/crates/bevy_render/src/render_resource/shader.rs
index acf46932b4..782a206ac8 100644
--- a/crates/bevy_render/src/render_resource/shader.rs
+++ b/crates/bevy_render/src/render_resource/shader.rs
@@ -34,6 +34,9 @@ pub struct Shader {
     pub additional_imports: Vec<naga_oil::compose::ImportDefinition>,
     // any shader defs that will be included when this module is used
     pub shader_defs: Vec<ShaderDefVal>,
+    // we must store strong handles to our dependencies to stop them
+    // from being immediately dropped if we are the only user.
+    pub file_dependencies: Vec<Handle<Shader>>,
 }
 
 impl Shader {
@@ -75,6 +78,7 @@ impl Shader {
             source: Source::Wgsl(source),
             additional_imports: Default::default(),
             shader_defs: Default::default(),
+            file_dependencies: Default::default(),
         }
     }
 
@@ -104,6 +108,7 @@ impl Shader {
             source: Source::Glsl(source, stage),
             additional_imports: Default::default(),
             shader_defs: Default::default(),
+            file_dependencies: Default::default(),
         }
     }
 
@@ -116,6 +121,7 @@ impl Shader {
             source: Source::SpirV(source.into()),
             additional_imports: Default::default(),
             shader_defs: Default::default(),
+            file_dependencies: Default::default(),
         }
     }
 
@@ -246,7 +252,7 @@ impl AssetLoader for ShaderLoader {
 
             let mut bytes = Vec::new();
             reader.read_to_end(&mut bytes).await?;
-            let shader = match ext {
+            let mut shader = match ext {
                 "spv" => Shader::from_spirv(bytes, load_context.path().to_string_lossy()),
                 "wgsl" => Shader::from_wgsl(
                     String::from_utf8(bytes)?,
@@ -270,11 +276,10 @@ impl AssetLoader for ShaderLoader {
                 _ => panic!("unhandled extension: {ext}"),
             };
 
-            // collect file dependencies
+            // collect and store file dependencies
             for import in &shader.imports {
                 if let ShaderImport::AssetPath(asset_path) = import {
-                    // TODO: should we just allow this handle to be dropped?
-                    let _handle: Handle<Shader> = load_context.load(asset_path);
+                    shader.file_dependencies.push(load_context.load(asset_path));
                 }
             }
             Ok(shader)