parse external usages of variables and functions

.vscode/launch.json

@@ -0,0 +1,26 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "type": "lldb",
+            "request": "launch",
+            "name": "Debug shader_prepoc",
+            "cargo": {
+                "args": [
+                    "build",
+                    //"--manifest-path", "${workspaceFolder}/examples/testbed/Cargo.toml"
+                    "--bin=shader_prepoc",
+                ],
+                "filter": {
+                    "name": "shader_prepoc",
+                    "kind": "bin"
+                }
+            },
+            "args": [],
+            "cwd": "${workspaceFolder}"
+        }
+    ]
+}

Cargo.lock

@@ -46,6 +46,12 @@ dependencies = [
  "crypto-common",
 ]
 
+[[package]]
+name = "either"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0"
+
 [[package]]
 name = "generic-array"
 version = "0.14.7"
@@ -56,6 +62,15 @@ dependencies = [
  "version_check",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "libc"
 version = "0.2.155"
@@ -152,6 +167,7 @@ dependencies = [
 name = "shader_prepoc"
 version = "0.1.0"
 dependencies = [
+ "itertools",
  "pest",
  "pest_derive",
  "thiserror",

Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+itertools = "0.13.0"
 pest = "2.7.11"
 pest_derive = "2.7.11"
 thiserror = "1.0.63"

shaders/inner_include.wgsl

@@ -1,5 +1,7 @@
 #define_module inner::some_include
 
+const scalar: f32 = 5.0;
+
 fn mult_some_nums(a: f32, b: f32) -> f32 {
     return a * b;
 }

shaders/simple.wgsl

@@ -2,5 +2,5 @@
 #import inner::some_include
 
 fn do_something_cool(in: f32) -> f32 {
-    return mult_some_nums(in, 2.0);
+    return inner::some_include::scalar * mult_some_nums(in, 2.0);
 }

src/main.rs

@@ -1,7 +1,11 @@
 use std::{
-    collections::HashMap, fmt::Write, fs, path::{Path, PathBuf}
+    collections::{HashMap, HashSet},
+    fmt::Write,
+    fs,
+    path::{Path, PathBuf},
 };
 
+use itertools::Itertools;
 use pest::Parser;
 use pest_derive::Parser;
 
@@ -22,6 +26,8 @@ pub enum PreprocessorError {
     FormatError(#[from] std::fmt::Error),
     #[error("unknown module import '{module}', in {from_path}")]
     UnknownModule { from_path: PathBuf, module: String },
+    #[error("import usage from `{from_module}` conflicts with local variable/function: `{name}`")]
+    ConflictingImport { from_module: String, name: String },
 }
 
 fn main() {
@@ -30,21 +36,77 @@ fn main() {
     println!("test: {}", p.as_str()); */
 
     let mut p = Processor::new();
-    let f = p.parse_modules("shaders", ["wgsl"]).unwrap();
+    //let f = p.parse_modules("shaders", ["wgsl"]).unwrap();
-    println!("Parsed {} modules:", f);
+    //println!("Parsed {} modules:", f);
+    p.parse_module("shaders/inner_include.wgsl")
+        .unwrap()
+        .expect("failed to find module");
 
-    for mod_name in p.modules.keys() {
+    p.parse_module("shaders/simple.wgsl")
-        println!("  {mod_name}");
+        .unwrap()
+        .expect("failed to find module");
+
+    for (name, module) in &p.modules {
+        println!("  {name}:");
+
+        if !module.constants.is_empty() || !module.functions.is_empty() {
+            println!("    defines:");
+        }
+        
+        for con in &module.constants {
+            println!("      const {con}");
+        }
+
+        for func in &module.functions {
+            println!("      fn {func}");
+        }
+
+        if !module.import_usages.is_empty() {
+            println!("    usages:");
+        }
+
+        for (module, usages) in &module.import_usages {
+            println!("      {}:", module);
+
+            for import in &usages.imports {
+                println!(
+                    "        {:?} `{}` at L{} C{}:",
+                    import.ty, import.name, import.line, import.col
+                );
+            }
+        }
     }
 
-    let processed = p.process_file("shaders/base.wgsl").unwrap();
+    /* let processed = p.process_file("shaders/base.wgsl").unwrap();
-    fs::write("out.wgsl", processed).unwrap();
+    fs::write("out.wgsl", processed).unwrap(); */
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum ExternalUsageType {
+    Variable,
+    Function,
+}
+
+pub struct ExternalUsage {
+    name: String,
+    ty: ExternalUsageType,
+    line: usize,
+    col: usize,
+}
+
+pub struct ImportUsage {
+    module: String,
+    imports: Vec<ExternalUsage>,
 }
 
 #[derive(Default)]
 pub struct Module {
     name: String,
     path: String,
+    constants: HashSet<String>,
+    functions: HashSet<String>,
+    //imports: HashSet<String>,
+    import_usages: HashMap<String, ImportUsage>,
 }
 
 #[derive(Default)]
@@ -77,6 +139,9 @@ impl Processor {
             .next()
             .unwrap(); // get and unwrap the `file` rule; never fails
 
+        let mut module = Module::default();
+        module.path = path.as_ref().to_str().unwrap().into();
+
         for record in file.into_inner() {
             match record.as_rule() {
                 Rule::command_line => {
@@ -91,26 +156,118 @@ impl Processor {
                             let shader_file = shader_file_pairs.next().unwrap();
                             let shader_file = shader_file.as_str().to_string();
 
-                            self.modules.insert(
+                            module.name = shader_file;
-                                shader_file.clone(),
-                                Module {
-                                    name: shader_file.clone(),
-                                    path: path.as_ref().to_str().unwrap().into(),
-                                },
-                            );
-
-                            return Ok(Some(shader_file.clone()));
                         }
                         _ => unreachable!(),
                     }
                 }
-                Rule::shader_code_line => (),
+                Rule::shader_code_line => {
+                    for line in record.into_inner() {
+                        let (pos_line, pos_col) = line.line_col();
+
+                        match line.as_rule() {
+                            Rule::shader_fn_def => {
+                                let mut pairs = line.into_inner();
+                                // shader_ident is the only pair for this rule
+                                let fn_name = pairs.next().unwrap().as_str().to_string();
+
+                                module.functions.insert(fn_name);
+                            }
+                            Rule::shader_const_def => {
+                                let mut pairs = line.into_inner();
+                                // shader_ident is the only pair for this rule
+                                let const_name = pairs.next().unwrap().as_str().to_string();
+
+                                module.constants.insert(const_name);
+                            }
+                            Rule::shader_external_fn => {
+                                let mut pairs = line.into_inner();
+                                // shader_external_variable is the only pair for this rule
+                                let ident_name = pairs.next().unwrap().as_str().to_string();
+
+                                if let Some((module_name, ident)) = ident_name.rsplit_once("::") {
+                                    if module.functions.contains(ident) {
+                                        // TODO: find a way to avoid conflicting imports.
+                                        // maybe this could be done by renaming local variables/functions?
+                                        return Err(PreprocessorError::ConflictingImport {
+                                            from_module: module_name.into(),
+                                            name: ident.into()
+                                        });
+                                    }
+                                    
+                                    let usage = ExternalUsage {
+                                        name: ident.into(),
+                                        ty: ExternalUsageType::Function,
+                                        line: pos_line,
+                                        col: pos_col
+                                    };
+
+                                    module.import_usages.entry(module_name.into())
+                                        .or_insert_with(|| ImportUsage {
+                                            module: module_name.into(),
+                                            imports: vec![],
+                                        })
+                                        .imports.push(usage);
+                                } else {
+                                    // TODO: not really sure how this would get triggered
+                                    unimplemented!("this function is actually not external, i think");
+                                }
+                            }
+                            Rule::shader_external_variable => {
+                                let pairs = line.into_inner();
+                                // shader_external_variable is the only pair for this rule
+                                let ident_name = pairs.as_str();
+
+                                if let Some((module_name, ident)) = ident_name.rsplit_once("::") {
+                                    if module.constants.contains(ident) {
+                                        // TODO: find a way to avoid conflicting imports.
+                                        // maybe this could be done by renaming local variables/functions?
+                                        return Err(PreprocessorError::ConflictingImport {
+                                            from_module: module_name.into(),
+                                            name: ident.into()
+                                        });
+                                    }
+
+                                    let usage = ExternalUsage {
+                                        name: ident.into(),
+                                        ty: ExternalUsageType::Variable,
+                                        line: pos_line,
+                                        col: pos_col
+                                    };
+
+                                    module.import_usages.entry(module_name.into())
+                                        .or_insert_with(|| ImportUsage {
+                                            module: module_name.into(),
+                                            imports: vec![],
+                                        })
+                                        .imports.push(usage);
+                                } else {
+                                    // TODO: not really sure how this would get triggered
+                                    unimplemented!("this function is actually not external, i think");
+                                }
+                            },
+                            Rule::shader_code => (), /* {
+                                let chunk = line.as_str();
+                                println!("shader_code: {chunk}");
+                            }, */
+                            //Rule::shader_unknown_code => (),
+                            _ => unreachable!(),
+                        }
+                    }
+                }
                 Rule::EOI => (),
                 _ => unreachable!(),
             }
         }
 
-        Ok(None)
+        if module.name.is_empty() {
+            Ok(None)
+        } else {
+            let name = module.name.clone();
+            self.modules.insert(name.clone(), module);
+
+            Ok(Some(name))
+        }
     }
 
     /// Find files recursively in `path` with an extension in `extensions`, and parse them.
@@ -164,7 +321,6 @@ impl Processor {
                 Rule::command_line => {
                     // the parser has found a preprocessor command, figure out what it is
                     let mut pairs = record.into_inner();
-                    //assert_eq!(pairs.next().unwrap().as_rule(), Rule::preproc_prefix); // skip preproc_prefix
                     let command_line = pairs.next().unwrap();
 
                     match command_line.as_rule() {
@@ -172,8 +328,6 @@ impl Processor {
                             let mut shader_file_pairs = command_line.into_inner();
                             let shader_file = shader_file_pairs.next().unwrap();
                             let shader_file = shader_file.as_str();
-                            // remove surrounding quotes
-                            //let shader_file = &shader_file[1..shader_file.len() - 1];
 
                             println!("found module import: {}", shader_file);
                             let imported_mod = self.modules.get(shader_file).ok_or_else(|| {
@@ -194,7 +348,7 @@ impl Processor {
                             out_string.write_str(&start_header)?;
                             out_string.write_str(&included_file)?;
                             out_string.write_str(&end_header)?;
-                        },
+                        }
                         Rule::define_module_command => (),
                         _ => unreachable!(),
                     }

src/wgsl.pest

@@ -1,5 +1,5 @@
-shader_module_text = { ASCII_ALPHANUMERIC | "_" | "-" }
+shader_ident = { (ASCII_ALPHANUMERIC | "_")* }
-shader_module = { shader_module_text* ~ ( "::" ~ shader_module_text*)* }
+shader_module = { shader_ident ~ ( "::" ~ shader_ident)* }
 import_command = { "import" ~ ws ~ shader_module }
 define_module_command = { "define_module" ~ ws ~ shader_module }
 preproc_prefix = _{ "#" }
@@ -8,9 +8,37 @@ preproc_prefix = _{ "#" }
 command_line = { preproc_prefix ~ (define_module_command | import_command) }
 
 // all characters used by wgsl
-shader_code = { ASCII_ALPHANUMERIC | "{" | "}" | "@" | "-" | "+" | "*" | "/" | "=" | "(" | ")" | ">" | "<" | ";" | ":" | "." | "_" | "," }
+shader_code_char = { "{" | "}" | "@" | "-" | "+" | "*" | "/" | "=" | "(" | ")" | ">" | "<" | ";" | ":" | "." | "_" | "," }
+shader_code = { shader_code_char | ASCII_ALPHANUMERIC+ }
+// a catch all for some 
+//shader_unknown_code = { (ws* ~ (shader_code ~ ws*)*) }
+
+shader_value_num = { ASCII_DIGIT* ~ ( "." ~ ASCII_DIGIT* )? }
+shader_value_bool = { "true" | "false" }
+shader_value = { shader_value_bool | shader_value_num  }
+
+shader_var_type = { shader_ident ~ ":" ~ ws* ~ shader_ident }
+shader_const_def = { "const" ~ ws ~ shader_ident ~ (ws* ~ shader_var_type)? ~ ws* ~ "=" ~ ws* ~ shader_value ~ ";" }
+// defines type of something i.e., `: f32`, `: u32`, etc.
+shader_fn_def = {
+    "fn" ~ ws ~ shader_ident ~
+      "(" ~ shader_var_type ~ (ws* ~ "," ~ ws* ~ shader_var_type)* ~ ")" ~ ws* ~
+      ("->" ~ ws* ~ shader_code ~ ws*)? ~ "{"?
+}
+//shader_fn_def = { "fn" ~ ws ~ shader_ident ~ "(a: f32, b: f32) -> f32 {" }
+
+// usages of code from another module
+shader_external_variable = { shader_ident ~ ( "::" ~ shader_ident)+ }
+shader_external_fn = { shader_external_variable ~ "(" ~ ANY* ~ ")" }
+shader_external_code = _{ shader_external_fn | shader_external_variable }
+
 // a line of shader code, including white space
-shader_code_line = { (ws* ~ shader_code ~ ws*)* }
+shader_code_line = {
+    shader_fn_def |
+    shader_const_def |
+    (ws* ~ ( (shader_external_code | shader_code) ~ ws*)*)
+}
+//shader_code_line = { shader_fn_def | shader_const_def | (ws* ~ (shader_external_code | shader_code)* ~ ws*) }
 
 file = { SOI ~ ( (command_line | shader_code_line) ~ NEWLINE)* ~ EOI }