From fabc7b32305e1ebdf916237c3755d2359c009759 Mon Sep 17 00:00:00 2001
From: robtfm <50659922+robtfm@users.noreply.github.com>
Date: Sun, 16 Jul 2023 12:42:07 +0100
Subject: [PATCH] limit item substitution

---
 src/compose/mod.rs                            | 11 ++++-
 src/compose/test.rs                           | 40 +++++++++++++++++++
 src/compose/tests/expected/item_sub_point.txt | 17 ++++++++
 src/compose/tests/item_sub_point/mod.wgsl     |  9 +++++
 src/compose/tests/item_sub_point/top.wgsl     |  8 ++++
 5 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 src/compose/tests/expected/item_sub_point.txt
 create mode 100644 src/compose/tests/item_sub_point/mod.wgsl
 create mode 100644 src/compose/tests/item_sub_point/top.wgsl

diff --git a/src/compose/mod.rs b/src/compose/mod.rs
index 4343cfc..666ead7 100644
--- a/src/compose/mod.rs
+++ b/src/compose/mod.rs
@@ -462,6 +462,7 @@ impl Composer {
         let mut item_substituted_source = String::default();
         let mut current_word = String::default();
         let mut line_is_directive = None;
+        let mut is_valid_import_substitution_point = true;
 
         for char in substituted_source.chars() {
             if !current_word.is_empty() {
@@ -473,7 +474,7 @@ impl Composer {
                 let mut output = &current_word;
 
                 // substitute current word if we are not writing a directive (e.g. `#import xyz`)
-                if line_is_directive != Some(true) {
+                if is_valid_import_substitution_point {
                     if let Some(replacement) = imported_items.get(&current_word) {
                         output = replacement;
                     }
@@ -502,6 +503,14 @@ impl Composer {
                 current_word.push(char);
             } else {
                 item_substituted_source.push(char);
+                // we should only substitute global names
+                // '.' -> avoid substituting members with name == import item
+                // '@' -> avoid substituting annotations
+                if char == '.' || char == '@' {
+                    is_valid_import_substitution_point = false;
+                } else {
+                    is_valid_import_substitution_point = !line_is_directive.unwrap_or(false);
+                }
             }
         }
         substituted_source = item_substituted_source;
diff --git a/src/compose/test.rs b/src/compose/test.rs
index 06dabb2..5555ac1 100644
--- a/src/compose/test.rs
+++ b/src/compose/test.rs
@@ -968,6 +968,46 @@ mod test {
         output_eq!(wgsl, "tests/expected/dup_struct_import.txt");
     }
 
+    #[test]
+    fn item_sub_point() {
+        let mut composer = Composer::default();
+
+        composer
+            .add_composable_module(ComposableModuleDescriptor {
+                source: include_str!("tests/item_sub_point/mod.wgsl"),
+                file_path: "tests/item_sub_point/mod.wgsl",
+                ..Default::default()
+            })
+            .unwrap();
+
+        let module = composer
+            .make_naga_module(NagaModuleDescriptor {
+                source: include_str!("tests/item_sub_point/top.wgsl"),
+                file_path: "tests/item_sub_point/top.wgsl",
+                ..Default::default()
+            })
+            .unwrap();
+
+        let info = naga::valid::Validator::new(
+            naga::valid::ValidationFlags::all(),
+            naga::valid::Capabilities::default(),
+        )
+        .validate(&module)
+        .unwrap();
+        let wgsl = naga::back::wgsl::write_string(
+            &module,
+            &info,
+            naga::back::wgsl::WriterFlags::EXPLICIT_TYPES,
+        )
+        .unwrap();
+
+        // let mut f = std::fs::File::create("item_sub_point.txt").unwrap();
+        // f.write_all(wgsl.as_bytes()).unwrap();
+        // drop(f);
+
+        output_eq!(wgsl, "tests/expected/item_sub_point.txt");
+    }
+
     // actually run a shader and extract the result
     // needs the composer to contain a module called "test_module", with a function called "entry_point" returning an f32.
     fn test_shader(composer: &mut Composer) -> f32 {
diff --git a/src/compose/tests/expected/item_sub_point.txt b/src/compose/tests/expected/item_sub_point.txt
new file mode 100644
index 0000000..8b89742
--- /dev/null
+++ b/src/compose/tests/expected/item_sub_point.txt
@@ -0,0 +1,17 @@
+struct _naga_oil_mod_NVXWI_memberFrag {
+    fragment: f32,
+}
+
+fn _naga_oil_mod_NVXWI_memberfragment(f_1: _naga_oil_mod_NVXWI_memberFrag) -> f32 {
+    return (f_1.fragment * 2.0);
+}
+
+@fragment 
+fn main() -> @location(0) f32 {
+    var f: _naga_oil_mod_NVXWI_memberFrag;
+
+    f.fragment = 3.0;
+    let _e3: _naga_oil_mod_NVXWI_memberFrag = f;
+    let _e4: f32 = _naga_oil_mod_NVXWI_memberfragment(_e3);
+    return _e4;
+}
diff --git a/src/compose/tests/item_sub_point/mod.wgsl b/src/compose/tests/item_sub_point/mod.wgsl
new file mode 100644
index 0000000..c9e5f25
--- /dev/null
+++ b/src/compose/tests/item_sub_point/mod.wgsl
@@ -0,0 +1,9 @@
+#define_import_path mod
+
+struct Frag {
+    fragment: f32,
+}
+
+fn fragment(f: Frag) -> f32 {
+    return f.fragment * 2.0;
+}
\ No newline at end of file
diff --git a/src/compose/tests/item_sub_point/top.wgsl b/src/compose/tests/item_sub_point/top.wgsl
new file mode 100644
index 0000000..a81e1f0
--- /dev/null
+++ b/src/compose/tests/item_sub_point/top.wgsl
@@ -0,0 +1,8 @@
+#import mod Frag, fragment
+
+@fragment
+fn main() -> @location(0) f32 {
+    var f: Frag;
+    f.fragment = 3.0;
+    return fragment(f);
+}
\ No newline at end of file