Trick Cargo into adding shader files as dependencies

Add the context of the shader sources with include_bytes!() so that
Cargo will recompile the rust source if the shader source changes.

Cargo.toml

@@ -9,3 +9,4 @@ proc-macro = true
 [dependencies]
 quote = "1.0"
 syn = "2.0.90"
+naga = { version = "23.0.0", features = ["wgsl-in"] }

src/lib.rs

@@ -1,25 +1,84 @@
 #![feature(proc_macro_span)]
 
 use {
+    naga::front::wgsl,
     proc_macro::TokenStream,
-    quote::quote,
-    std::{fs, path::PathBuf},
+    quote::{format_ident, quote},
     syn::{parse, LitStr},
 };
 
-fn emit_create_shader_module(file_contents: &str) -> TokenStream {
-    quote! {
-        wgpu::ShaderModuleDescriptor {
-            label: None,
-            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(#file_contents)),
-        }
+mod source_reader;
+
+struct ShaderModule {
+    source_file: source_reader::SourceFile,
+    naga_module: naga::Module,
+}
+
+impl ShaderModule {
+    fn try_new(
+        source_dir: &std::path::Path,
+        source_filename: &std::path::Path,
+    ) -> Result<Self, Error> {
+        let source_file = source_reader::read_source_file(source_dir, source_filename)?;
+
+        let naga_module = wgsl::parse_str(&source_file.full_text()).map_err(|err| {
+            if let Some(location) = err.location(&source_file.full_text()) {
+                let (inner_source_filename, inner_line_num) =
+                    source_file.map_source_line_num(location.line_number as usize);
+                Error::Message(format!(
+                    "Error in {} at line {}, column {}:\n{}",
+                    &inner_source_filename.as_os_str().to_string_lossy(),
+                    inner_line_num,
+                    location.line_position,
+                    err.message()
+                ))
+            } else {
+                Error::Message(format!(
+                    "Error in {}:\n{}",
+                    &source_file.root_filename().as_os_str().to_string_lossy(),
+                    err.message()
+                ))
+            }
+        })?;
+        Ok(Self {
+            source_file,
+            naga_module,
+        })
+    }
+
+    fn emit_rust_code(&self) -> TokenStream {
+        // We include the text of all the source files into the generated code
+        // with include_bytes!() so that Cargo recognizes them as
+        // dependencies. It's not pretty but it works. Since the file contents
+        // are just stored in unused local variables, they should be removed by
+        // the optimizer and not affect the compiled code.
+        let source_text = &self.source_file.full_text();
+        let dependencies: Vec<_> = self
+            .source_file
+            .all_filenames()
+            .iter()
+            .map(|path| path.as_os_str().to_string_lossy())
+            .collect();
+        let dummy_var: Vec<_> = (0u32..)
+            .map(|n| format_ident!("_{}", n))
+            .take(dependencies.len())
+            .collect();
+        quote! {
+            {
+                #(let #dummy_var = include_bytes!(#dependencies);)*
+                wgpu::ShaderModuleDescriptor {
+                    label: None,
+                    source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(#source_text)),
+                }
+            }
+        }
+        .into()
     }
-    .into()
 }
 
 struct WgslModuleMacroInput {
-    source_dir: PathBuf,
-    filename: String,
+    source_dir: std::path::PathBuf,
+    filename: std::path::PathBuf,
 }
 
 enum Error {
@@ -48,7 +107,7 @@ fn parse_input(token_stream: TokenStream) -> Result<WgslModuleMacroInput, Error>
 
     let source_dir = token.span().source_file().path().parent().unwrap().into();
     let literal: LitStr = parse(token_stream)?;
-    let filename = literal.value();
+    let filename = std::path::Path::new(&literal.value()).into();
     Ok(WgslModuleMacroInput {
         source_dir,
         filename,
@@ -57,21 +116,17 @@ fn parse_input(token_stream: TokenStream) -> Result<WgslModuleMacroInput, Error>
 
 fn wgsl_module_inner(input_token_stream: TokenStream) -> Result<TokenStream, Error> {
     let input = parse_input(input_token_stream)?;
-    let full_filename = input.source_dir.join(&input.filename);
-    let contents = fs::read_to_string(&full_filename).map_err(|err| {
-        Error::from_message(format!(
-            "Could not open \"{}\": {}",
-            &full_filename.as_os_str().to_string_lossy(),
-            err.to_string()
-        ))
-    })?;
-    Ok(emit_create_shader_module(&contents))
+    let module = ShaderModule::try_new(&input.source_dir, &input.filename)?;
+    Ok(module.emit_rust_code())
 }
 
 #[proc_macro]
 pub fn wgsl_module(input: TokenStream) -> TokenStream {
-    wgsl_module_inner(input).unwrap_or_else(|err| match err {
-        Error::Message(message) => quote! { compile_error!(#message) },
-        Error::SynError(syn_err) => syn_err.into_compile_error(),
-    }.into())
+    wgsl_module_inner(input).unwrap_or_else(|err| {
+        match err {
+            Error::Message(message) => quote! { compile_error!(#message) },
+            Error::SynError(syn_err) => syn_err.into_compile_error(),
+        }
+        .into()
+    })
 }

src/source_reader.rs

@@ -0,0 +1,55 @@
+use {
+    crate::Error,
+    std::{
+        fs,
+        path::{Path, PathBuf},
+    },
+};
+
+struct SourceMap {
+    filename: PathBuf,
+    offset: usize,
+    total_lines: usize,
+    insertions: Vec<SourceMap>,
+}
+
+pub struct SourceFile {
+    source_text: String,
+    source_map: SourceMap,
+}
+
+impl SourceFile {
+    pub fn root_filename(&self) -> &Path {
+        &self.source_map.filename
+    }
+    pub fn full_text(&self) -> &str {
+        &self.source_text
+    }
+
+    pub fn map_source_line_num(&self, line_num: usize) -> (&Path, usize) {
+        (&self.source_map.filename, line_num)
+    }
+
+    pub fn all_filenames(&self) -> Vec<&Path> {
+        vec![&self.source_map.filename]
+    }
+}
+
+pub fn read_source_file(source_dir: &Path, filename: &Path) -> Result<SourceFile, Error> {
+    let source_text = fs::read_to_string(&source_dir.join(filename)).map_err(|err| {
+        Error::from_message(format!(
+            "Could not open \"{}\": {}",
+            &filename.as_os_str().to_string_lossy(),
+            err
+        ))
+    })?;
+    Ok(SourceFile {
+        source_text,
+        source_map: SourceMap {
+            filename: filename.into(),
+            offset: 0,
+            total_lines: 0,
+            insertions: vec![],
+        },
+    })
+}