From e4203c39258a55ec3f598cce845e4dcdfcdd69bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois?= <mockersf@gmail.com>
Date: Tue, 22 Feb 2022 20:21:04 +0000
Subject: [PATCH] shader preprocessor - do not import if scope is not valid
 (#4012)

# Objective

- fix #4011
- imports are not limited by the current `ifdef` they are in

## Solution

- process imports only if the current scope is enabled
---
 .../bevy_render/src/render_resource/shader.rs | 131 +++++++++++++-----
 1 file changed, 98 insertions(+), 33 deletions(-)

diff --git a/crates/bevy_render/src/render_resource/shader.rs b/crates/bevy_render/src/render_resource/shader.rs
index 77c6ac1bcf..b55ef0dab2 100644
--- a/crates/bevy_render/src/render_resource/shader.rs
+++ b/crates/bevy_render/src/render_resource/shader.rs
@@ -408,40 +408,42 @@ impl ShaderProcessor {
                 if scopes.is_empty() {
                     return Err(ProcessShaderError::TooManyEndIfs);
                 }
-            } else if let Some(cap) = SHADER_IMPORT_PROCESSOR
-                .import_asset_path_regex
-                .captures(line)
-            {
-                let import = ShaderImport::AssetPath(cap.get(1).unwrap().as_str().to_string());
-                self.apply_import(
-                    import_handles,
-                    shaders,
-                    &import,
-                    shader,
-                    shader_defs,
-                    &mut final_string,
-                )?;
-            } else if let Some(cap) = SHADER_IMPORT_PROCESSOR
-                .import_custom_path_regex
-                .captures(line)
-            {
-                let import = ShaderImport::Custom(cap.get(1).unwrap().as_str().to_string());
-                self.apply_import(
-                    import_handles,
-                    shaders,
-                    &import,
-                    shader,
-                    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');
+                if let Some(cap) = SHADER_IMPORT_PROCESSOR
+                    .import_asset_path_regex
+                    .captures(line)
+                {
+                    let import = ShaderImport::AssetPath(cap.get(1).unwrap().as_str().to_string());
+                    self.apply_import(
+                        import_handles,
+                        shaders,
+                        &import,
+                        shader,
+                        shader_defs,
+                        &mut final_string,
+                    )?;
+                } else if let Some(cap) = SHADER_IMPORT_PROCESSOR
+                    .import_custom_path_regex
+                    .captures(line)
+                {
+                    let import = ShaderImport::Custom(cap.get(1).unwrap().as_str().to_string());
+                    self.apply_import(
+                        import_handles,
+                        shaders,
+                        &import,
+                        shader,
+                        shader_defs,
+                        &mut final_string,
+                    )?;
+                } else if SHADER_IMPORT_PROCESSOR
+                    .define_import_path_regex
+                    .is_match(line)
+                {
+                    // ignore import path lines
+                } else {
+                    final_string.push_str(line);
+                    final_string.push('\n');
+                }
             }
         }
 
@@ -1231,4 +1233,67 @@ fn in_main() { }
             .unwrap();
         assert_eq!(result.get_wgsl_source().unwrap(), EXPECTED);
     }
+
+    #[test]
+    fn process_import_in_ifdef() {
+        #[rustfmt::skip]
+        const BAR: &str = r"
+fn bar() { }
+";
+        #[rustfmt::skip]
+        const BAZ: &str = r"
+fn baz() { }
+";
+        #[rustfmt::skip]
+        const INPUT: &str = r"
+#ifdef FOO
+    #import BAR
+#else
+    #import BAZ
+#endif
+";
+        #[rustfmt::skip]
+        const EXPECTED_FOO: &str = r"
+
+fn bar() { }
+";
+        #[rustfmt::skip]
+        const EXPECTED: &str = r"
+
+fn baz() { }
+";
+        let processor = ShaderProcessor::default();
+        let mut shaders = HashMap::default();
+        let mut import_handles = HashMap::default();
+        {
+            let bar_handle = Handle::<Shader>::default();
+            shaders.insert(bar_handle.clone_weak(), Shader::from_wgsl(BAR));
+            import_handles.insert(
+                ShaderImport::Custom("BAR".to_string()),
+                bar_handle.clone_weak(),
+            );
+        }
+        {
+            let baz_handle = HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 1).typed();
+            shaders.insert(baz_handle.clone_weak(), Shader::from_wgsl(BAZ));
+            import_handles.insert(
+                ShaderImport::Custom("BAZ".to_string()),
+                baz_handle.clone_weak(),
+            );
+        }
+        let result = processor
+            .process(
+                &Shader::from_wgsl(INPUT),
+                &["FOO".to_string()],
+                &shaders,
+                &import_handles,
+            )
+            .unwrap();
+        assert_eq!(result.get_wgsl_source().unwrap(), EXPECTED_FOO);
+
+        let result = processor
+            .process(&Shader::from_wgsl(INPUT), &[], &shaders, &import_handles)
+            .unwrap();
+        assert_eq!(result.get_wgsl_source().unwrap(), EXPECTED);
+    }
 }