use rust like module names for shader imports

.gitignore

@@ -1 +1,2 @@
 /target
+out.wgsl

Cargo.lock

@@ -154,6 +154,7 @@ version = "0.1.0"
 dependencies = [
  "pest",
  "pest_derive",
+ "thiserror",
 ]
 
 [[package]]

Cargo.toml

@@ -6,3 +6,4 @@ edition = "2021"
 [dependencies]
 pest = "2.7.11"
 pest_derive = "2.7.11"
+thiserror = "1.0.63"

shaders/base.wgsl

@@ -1,4 +1,5 @@
-#import "simple.wgsl"
+#define_module base
+#import simple
 
 fn main() -> vec4<f32> {
     let a = do_something_cool(10.0);

shaders/inner_include.wgsl

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

shaders/out.wgsl

@@ -1,10 +0,0 @@
-// ==== START OF INCLUDE OF 'shaders/simple.wgsl' ====
-fn do_something_cool(in: f32) -> f32 {
-    return in * 2.0;
-}
-// ==== END OF INCLUDE OF 'shaders/simple.wgsl' ====
-
-fn main() -> vec4<f32> {
-    let a = do_something_cool(10.0);
-    return vec4<f32>(vec3<f32>(a), 1.0);
-}

shaders/simple.wgsl

@@ -1,3 +1,6 @@
+#define_module simple
+#import inner::some_include
+
 fn do_something_cool(in: f32) -> f32 {
-    return in * 2.0;
+    return mult_some_nums(in, 2.0);
 }

src/main.rs

@@ -1,4 +1,6 @@
-use std::{fs, io::Write};
+use std::{
+    collections::HashMap, fmt::Write, fs, path::{Path, PathBuf}
+};
 
 use pest::Parser;
 use pest_derive::Parser;
@@ -7,68 +9,230 @@ use pest_derive::Parser;
 #[grammar = "wgsl.pest"]
 pub struct WgslParser;
 
+#[derive(Debug, thiserror::Error)]
+pub enum PreprocessorError {
+    #[error("{0}")]
+    IoError(#[from] std::io::Error),
+    #[error("error parsing {path}: {err}")]
+    ParserError {
+        path: PathBuf,
+        err: pest::error::Error<Rule>,
+    },
+    #[error("failure formatting preprocessor output to string ({0})")]
+    FormatError(#[from] std::fmt::Error),
+    #[error("unknown module import '{module}', in {from_path}")]
+    UnknownModule { from_path: PathBuf, module: String },
+}
+
 fn main() {
-    /* let unparsed_file = fs::read_to_string("shaders/base.wgsl")
+    /* let mut successful_parse = WgslParser::parse(Rule::command_line, "#define_module inner::some_include").unwrap();
-        .expect("cannot read file");
+    let p = successful_parse.next().unwrap();
-    let mut successful_parse = WgslParser::parse(Rule::file, &unparsed_file).unwrap();
+    println!("test: {}", p.as_str()); */
 
-    let a = successful_parse.next().unwrap();
+    let mut p = Processor::new();
-    println!("got {}", a.as_str()); */
+    let f = p.parse_modules("shaders", ["wgsl"]).unwrap();
+    println!("Parsed {} modules:", f);
 
-    let unparsed_file = fs::read_to_string("shaders/base.wgsl")
+    for mod_name in p.modules.keys() {
-        .expect("cannot read file");
+        println!("  {mod_name}");
-    // add a new line to the end of the input to make the grammar happy
+    }
-    let unparsed_file = format!("{unparsed_file}\n");
 
-    let mut out_file = fs::File::create("shaders/out.wgsl").unwrap();
+    let processed = p.process_file("shaders/base.wgsl").unwrap();
+    fs::write("out.wgsl", processed).unwrap();
+}
 
-    let file = WgslParser::parse(Rule::file, &unparsed_file)
+#[derive(Default)]
-        .expect("unsuccessful parse") // unwrap the parse result
+pub struct Module {
-        .next().unwrap(); // get and unwrap the `file` rule; never fails
+    name: String,
+    path: String,
+}
 
-    for record in file.into_inner() {
+#[derive(Default)]
-        match record.as_rule() {
+pub struct Processor {
-            Rule::command_line => {
+    modules: HashMap<String, Module>,
+}
+
+impl Processor {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Parse a module file to attempt to find the include identifier.
+    ///
+    /// Returns `None` if the module does not define an include identifier.
+    pub fn parse_module<P: AsRef<Path>>(
+        &mut self,
+        path: P,
+    ) -> Result<Option<String>, PreprocessorError> {
+        let unparsed_file = fs::read_to_string(path.as_ref())?;
+
+        // add a new line to the end of the input to make the grammar happy
+        let unparsed_file = format!("{unparsed_file}\n");
+
+        let file = WgslParser::parse(Rule::file, &unparsed_file)
+            .map_err(|e| PreprocessorError::ParserError {
+                path: path.as_ref().to_path_buf(),
+                err: e,
+            })?
+            .next()
+            .unwrap(); // get and unwrap the `file` rule; never fails
+
+        for record in file.into_inner() {
+            match record.as_rule() {
+                Rule::command_line => {
+                    // the parser has found a preprocessor command, figure out what it is
+                    let mut pairs = record.into_inner();
+                    let command_line = pairs.next().unwrap();
+
+                    match command_line.as_rule() {
+                        Rule::import_command => (),
+                        Rule::define_module_command => {
+                            let mut shader_file_pairs = command_line.into_inner();
+                            let shader_file = shader_file_pairs.next().unwrap();
+                            let shader_file = shader_file.as_str().to_string();
+
+                            self.modules.insert(
+                                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::EOI => (),
+                _ => unreachable!(),
+            }
+        }
+
+        Ok(None)
+    }
+
+    /// Find files recursively in `path` with an extension in `extensions`, and parse them.
+    ///
+    /// For each file that's found, [`Processor::parse_module`] is used to parse them.
+    ///
+    /// Parameters:
+    /// * `path` - The path to search for files in.
+    /// * `extensions` - The extensions that the discovered files must have. Make sure they have
+    /// no leading '.'
+    pub fn parse_modules<P: AsRef<Path>, const N: usize>(
+        &mut self,
+        path: P,
+        extensions: [&str; N],
+    ) -> Result<usize, PreprocessorError> {
+        //debug_assert!(!extension.starts_with("."), "remove leading '.' from extension");
+
+        let files = recurse_files(path)?;
+
+        let mut parsed = 0;
+
+        for file in files {
+            if let Some(ext) = file.extension().and_then(|p| p.to_str()) {
+                if extensions.contains(&ext) {
+                    self.parse_module(file)?;
+                    parsed += 1;
+                }
+            }
+        }
+
+        Ok(parsed)
+    }
+
+    pub fn process_file<P: AsRef<Path>>(&mut self, path: P) -> Result<String, PreprocessorError> {
+        let unparsed_file = fs::read_to_string(path.as_ref())?;
+        // add a new line to the end of the input to make the grammar happy
+        let unparsed_file = format!("{unparsed_file}\n");
+
+        let mut out_string = String::new();
+
+        let file = WgslParser::parse(Rule::file, &unparsed_file)
+            .map_err(|e| PreprocessorError::ParserError {
+                path: path.as_ref().to_path_buf(),
+                err: e,
+            })?
+            .next()
+            .unwrap(); // get and unwrap the `file` rule; never fails
+
+        for record in file.into_inner() {
+            match record.as_rule() {
+                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();
 
-                for command_line in record.into_inner() {
                     match command_line.as_rule() {
-                        Rule::preproc_command => {},
                         Rule::import_command => {
                             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];
+                            //let shader_file = &shader_file[1..shader_file.len() - 1];
 
-                            println!("found include for file: {}", shader_file);
+                            println!("found module import: {}", shader_file);
+                            let imported_mod = self.modules.get(shader_file).ok_or_else(|| {
+                                PreprocessorError::UnknownModule {
+                                    from_path: path.as_ref().to_path_buf(),
+                                    module: shader_file.into(),
+                                }
+                            })?;
 
-                            let path = format!("shaders/{}", shader_file);
+                            //let shader_path = parent_dir.join(shader_file);
-                            let included_file = fs::read(&path)
+                            let included_file = self.process_file(imported_mod.path.clone())?;
-                                .expect("cannot read file");
 
-                            let start_header = format!("// ==== START OF INCLUDE OF '{}' ====\n", path);
+                            let start_header =
-                            let end_header = format!("\n// ==== END OF INCLUDE OF '{}' ====\n", path);
+                                format!("// ==== START OF INCLUDE OF '{}' ====\n", shader_file);
-
+                            let end_header =
-                            out_file.write(start_header.as_bytes()).unwrap();
+                                format!("\n// ==== END OF INCLUDE OF '{}' ====\n", shader_file);
-                            out_file.write(&included_file).unwrap();
-                            out_file.write(end_header.as_bytes()).unwrap();
 
+                            out_string.write_str(&start_header)?;
+                            out_string.write_str(&included_file)?;
+                            out_string.write_str(&end_header)?;
                         },
-                        _ => unreachable!()
+                        Rule::define_module_command => (),
+                        _ => unreachable!(),
                     }
                 }
-                println!("found line");
+                Rule::shader_code_line => {
-            },
+                    let input = record.as_str();
-            Rule::shader_code_line => {
+                    // add new line to end of input
-                let input = record.as_str();
+                    let input = format!("{input}\n");
-                println!("in: {}", input);
 
-                let input = format!("{input}\n");
+                    out_string.write_str(&input)?;
-                out_file.write(input.as_bytes()).unwrap();
+                }
-            },
+                Rule::EOI => (),
-            Rule::EOI => (),
+                _ => unreachable!(),
-            Rule::eol => (),
+            }
-            _ => unreachable!(),
         }
+
+        Ok(out_string)
     }
 }
+
+/// Recursively find files in `path`.
+fn recurse_files(path: impl AsRef<Path>) -> std::io::Result<Vec<PathBuf>> {
+    let mut buf = vec![];
+    let entries = fs::read_dir(path)?;
+
+    for entry in entries {
+        let entry = entry?;
+        let meta = entry.metadata()?;
+
+        if meta.is_dir() {
+            let mut subdir = recurse_files(entry.path())?;
+            buf.append(&mut subdir);
+        }
+
+        if meta.is_file() {
+            buf.push(entry.path());
+        }
+    }
+
+    Ok(buf)
+}

src/wgsl.pest

@@ -1,12 +1,14 @@
-shader_file = { "\"" ~ ASCII_ALPHA* ~ "." ~ ASCII_ALPHA* ~ "\"" }
+shader_module_text = { ASCII_ALPHANUMERIC | "_" | "-" }
-import_command = { "import" ~ ws ~ shader_file }
+shader_module = { shader_module_text* ~ ( "::" ~ shader_module_text*)* }
-preproc_prefix = { "#" }
+import_command = { "import" ~ ws ~ shader_module }
+define_module_command = { "define_module" ~ ws ~ shader_module }
+preproc_prefix = _{ "#" }
 
 // a line of preprocessor commands
-command_line = { preproc_prefix ~ import_command }
+command_line = { preproc_prefix ~ (define_module_command | import_command) }
 
 // all characters used by wgsl
-shader_code = { ASCII_ALPHANUMERIC | "{" | "}" | "@" | "-" | "+" | "=" | "(" | ")" | ">" | "<" | ";" | "." | "_" | "," }
+shader_code = { ASCII_ALPHANUMERIC | "{" | "}" | "@" | "-" | "+" | "*" | "/" | "=" | "(" | ")" | ">" | "<" | ";" | ":" | "." | "_" | "," }
 // a line of shader code, including white space
 shader_code_line = { (ws* ~ shader_code ~ ws*)* }