Move import_path definitions into shader source (#3976)

cart · cart · commit e9f52b9dd2b8 · 2022-02-18T21:54:03.000Z
This enables shaders to (optionally) define their import path inside their source. This has a number of benefits: 1. enables users to define their own custom paths directly in their assets 2. moves the import path "close" to the asset instead of centralized in the plugin definition, which seems "better" to me. 3. makes "internal hot shader reloading" way more reasonable (see #3966) 4. logically opens the door to importing "parts" of a shader by defining "import_path blocks". ```rust #define_import_path bevy_pbr::mesh_struct struct Mesh { model: mat4x4<f32>; inverse_transpose_model: mat4x4<f32>; // 'flags' is a bit field indicating various options. u32 is 32 bits so we have up to 32 options. flags: u32; }; let MESH_FLAGS_SHADOW_RECEIVER_BIT: u32 = 1u; ```
diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs
@@ -42,13 +42,11 @@ impl Plugin for MeshRenderPlugin {
         );
         shaders.set_untracked(
             MESH_STRUCT_HANDLE,
-            Shader::from_wgsl(include_str!("mesh_struct.wgsl"))
-                .with_import_path("bevy_pbr::mesh_struct"),
+            Shader::from_wgsl(include_str!("mesh_struct.wgsl")),
         );
         shaders.set_untracked(
             MESH_VIEW_BIND_GROUP_HANDLE,
-            Shader::from_wgsl(include_str!("mesh_view_bind_group.wgsl"))
-                .with_import_path("bevy_pbr::mesh_view_bind_group"),
+            Shader::from_wgsl(include_str!("mesh_view_bind_group.wgsl")),
         );
 
         app.add_plugin(UniformComponentPlugin::<MeshUniform>::default());
diff --git a/crates/bevy_pbr/src/render/mesh_struct.wgsl b/crates/bevy_pbr/src/render/mesh_struct.wgsl
@@ -1,3 +1,5 @@
+#define_import_path bevy_pbr::mesh_struct
+
 struct Mesh {
     model: mat4x4<f32>;
     inverse_transpose_model: mat4x4<f32>;
diff --git a/crates/bevy_pbr/src/render/mesh_view_bind_group.wgsl b/crates/bevy_pbr/src/render/mesh_view_bind_group.wgsl
@@ -1,3 +1,5 @@
+#define_import_path bevy_pbr::mesh_view_bind_group
+
 struct View {
     view_proj: mat4x4<f32>;
     view: mat4x4<f32>;
diff --git a/crates/bevy_render/src/render_resource/shader.rs b/crates/bevy_render/src/render_resource/shader.rs
@@ -45,27 +45,29 @@ pub struct Shader {
 impl Shader {
     pub fn from_wgsl(source: impl Into<Cow<'static, str>>) -> Shader {
         let source = source.into();
+        let shader_imports = SHADER_IMPORT_PROCESSOR.get_imports_from_str(&source);
         Shader {
-            imports: SHADER_IMPORT_PROCESSOR.get_imports_from_str(&source),
+            imports: shader_imports.imports,
+            import_path: shader_imports.import_path,
             source: Source::Wgsl(source),
-            import_path: None,
         }
     }
 
     pub fn from_glsl(source: impl Into<Cow<'static, str>>, stage: naga::ShaderStage) -> Shader {
         let source = source.into();
+        let shader_imports = SHADER_IMPORT_PROCESSOR.get_imports_from_str(&source);
         Shader {
-            imports: SHADER_IMPORT_PROCESSOR.get_imports_from_str(&source),
+            imports: shader_imports.imports,
+            import_path: shader_imports.import_path,
             source: Source::Glsl(source, stage),
-            import_path: None,
         }
     }
 
     pub fn from_spirv(source: impl Into<Cow<'static, [u8]>>) -> Shader {
         Shader {
             imports: Vec::new(),
-            source: Source::SpirV(source.into()),
             import_path: None,
+            source: Source::SpirV(source.into()),
         }
     }
 
@@ -238,12 +240,16 @@ impl AssetLoader for ShaderLoader {
                 _ => panic!("unhandled extension: {}", ext),
             };
 
-            shader.import_path = Some(ShaderImport::AssetPath(
-                load_context.path().to_string_lossy().to_string(),
-            ));
-            let imports = SHADER_IMPORT_PROCESSOR.get_imports(&shader);
+            let shader_imports = SHADER_IMPORT_PROCESSOR.get_imports(&shader);
+            if shader_imports.import_path.is_some() {
+                shader.import_path = shader_imports.import_path;
+            } else {
+                shader.import_path = Some(ShaderImport::AssetPath(
+                    load_context.path().to_string_lossy().to_string(),
+                ));
+            }
             let mut asset = LoadedAsset::new(shader);
-            for import in imports {
+            for import in shader_imports.imports {
                 if let ShaderImport::AssetPath(asset_path) = import {
                     let path = PathBuf::from_str(&asset_path)?;
                     asset.add_dependency(path.into());
@@ -281,6 +287,7 @@ pub enum ProcessShaderError {
 pub struct ShaderImportProcessor {
     import_asset_path_regex: Regex,
     import_custom_path_regex: Regex,
+    define_import_path_regex: Regex,
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Hash)]
@@ -292,34 +299,48 @@ pub enum ShaderImport {
 impl Default for ShaderImportProcessor {
     fn default() -> Self {
         Self {
-            import_asset_path_regex: Regex::new(r#"^\s*#\s*import\s*"(.+)""#).unwrap(),
-            import_custom_path_regex: Regex::new(r"^\s*#\s*import\s*(.+)").unwrap(),
+            import_asset_path_regex: Regex::new(r#"^\s*#\s*import\s+"(.+)""#).unwrap(),
+            import_custom_path_regex: Regex::new(r"^\s*#\s*import\s+(.+)").unwrap(),
+            define_import_path_regex: Regex::new(r"^\s*#\s*define_import_path\s+(.+)").unwrap(),
         }
     }
 }
 
+#[derive(Default)]
+pub struct ShaderImports {
+    imports: Vec<ShaderImport>,
+    import_path: Option<ShaderImport>,
+}
+
 impl ShaderImportProcessor {
-    pub fn get_imports(&self, shader: &Shader) -> Vec<ShaderImport> {
+    pub fn get_imports(&self, shader: &Shader) -> ShaderImports {
         match &shader.source {
             Source::Wgsl(source) => self.get_imports_from_str(source),
             Source::Glsl(source, _stage) => self.get_imports_from_str(source),
-            Source::SpirV(_source) => Vec::new(),
+            Source::SpirV(_source) => ShaderImports::default(),
         }
     }
 
-    pub fn get_imports_from_str(&self, shader: &str) -> Vec<ShaderImport> {
-        let mut imports = Vec::new();
+    pub fn get_imports_from_str(&self, shader: &str) -> ShaderImports {
+        let mut shader_imports = ShaderImports::default();
         for line in shader.lines() {
             if let Some(cap) = self.import_asset_path_regex.captures(line) {
                 let import = cap.get(1).unwrap();
-                imports.push(ShaderImport::AssetPath(import.as_str().to_string()));
+                shader_imports
+                    .imports
+                    .push(ShaderImport::AssetPath(import.as_str().to_string()));
             } else if let Some(cap) = self.import_custom_path_regex.captures(line) {
                 let import = cap.get(1).unwrap();
-                imports.push(ShaderImport::Custom(import.as_str().to_string()));
+                shader_imports
+                    .imports
+                    .push(ShaderImport::Custom(import.as_str().to_string()));
+            } else if let Some(cap) = self.define_import_path_regex.captures(line) {
+                let path = cap.get(1).unwrap();
+                shader_imports.import_path = Some(ShaderImport::Custom(path.as_str().to_string()));
             }
         }
 
-        imports
+        shader_imports
     }
 }
 
@@ -413,6 +434,11 @@ impl ShaderProcessor {
                     shader_defs,
                     &mut final_string,
                 )?;
+            } else if SHADER_IMPORT_PROCESSOR
+                .define_import_path_regex
+                .is_match(line)
+            {
+                // ignore import path lines
             } else if *scopes.last().unwrap() {
                 final_string.push_str(line);
                 final_string.push('\n');
diff --git a/crates/bevy_sprite/src/mesh2d/mesh.rs b/crates/bevy_sprite/src/mesh2d/mesh.rs
@@ -50,13 +50,11 @@ impl Plugin for Mesh2dRenderPlugin {
         );
         shaders.set_untracked(
             MESH2D_STRUCT_HANDLE,
-            Shader::from_wgsl(include_str!("mesh2d_struct.wgsl"))
-                .with_import_path("bevy_sprite::mesh2d_struct"),
+            Shader::from_wgsl(include_str!("mesh2d_struct.wgsl")),
         );
         shaders.set_untracked(
             MESH2D_VIEW_BIND_GROUP_HANDLE,
-            Shader::from_wgsl(include_str!("mesh2d_view_bind_group.wgsl"))
-                .with_import_path("bevy_sprite::mesh2d_view_bind_group"),
+            Shader::from_wgsl(include_str!("mesh2d_view_bind_group.wgsl")),
         );
 
         app.add_plugin(UniformComponentPlugin::<Mesh2dUniform>::default());
diff --git a/crates/bevy_sprite/src/mesh2d/mesh2d_struct.wgsl b/crates/bevy_sprite/src/mesh2d/mesh2d_struct.wgsl
@@ -1,3 +1,5 @@
+#define_import_path bevy_sprite::mesh2d_struct
+
 struct Mesh2d {
     model: mat4x4<f32>;
     inverse_transpose_model: mat4x4<f32>;
diff --git a/crates/bevy_sprite/src/mesh2d/mesh2d_view_bind_group.wgsl b/crates/bevy_sprite/src/mesh2d/mesh2d_view_bind_group.wgsl
@@ -1,3 +1,5 @@
+#define_import_path bevy_sprite::mesh2d_view_bind_group
+
 struct View {
     view_proj: mat4x4<f32>;
     view: mat4x4<f32>;

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#define_import_path bevy_pbr::mesh_struct`
	`2`	`+`
`1`	`3`	`struct Mesh {`
`2`	`4`	`model: mat4x4<f32>;`
`3`	`5`	`inverse_transpose_model: mat4x4<f32>;`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#define_import_path bevy_pbr::mesh_view_bind_group`
	`2`	`+`
`1`	`3`	`struct View {`
`2`	`4`	`view_proj: mat4x4<f32>;`
`3`	`5`	`view: mat4x4<f32>;`
Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,5 @@`
	`1`	`+#define_import_path bevy_sprite::mesh2d_struct`
	`2`	`+`
`1`	`3`	`struct Mesh2d {`
`2`	`4`	`model: mat4x4<f32>;`
`3`	`5`	`inverse_transpose_model: mat4x4<f32>;`