Add support for comments

Cargo.toml

@@ -34,3 +34,7 @@ indexmap = "2"
 wgpu = { version = "22", features = ["naga-ir"] }
 futures-lite = "1"
 tracing-subscriber = { version = "0.3", features = ["std", "fmt"] }
+
+[patch.crates-io]
+# naga = { path = "../wgpu/naga" }
+naga = { git = "https://github.com/Vrixyz/wgpu.git", branch = "token-comment" }

src/compose/mod.rs

@@ -570,7 +570,9 @@ impl Composer {
         }
     }
 
-    // build naga module for a given shader_def configuration. builds a minimal self-contained module built against headers for imports
+    /// Build naga module for a given shader_def configuration.
+    ///
+    /// Builds a minimal self-contained module built against headers for imports.
     fn create_module_ir(
         &self,
         name: &str,
@@ -1160,6 +1162,8 @@ impl Composer {
                 // todo figure out how to get span info for entrypoints
             }
         }
+        // Copy comments, should be done after each commented item has been imported.
+        module_builder.import_comments(&source_ir.comments);
 
         let module_ir = module_builder.into_module_with_entrypoints();
         let mut header_ir: naga::Module = header_builder.into();
@@ -1281,6 +1285,8 @@ impl Composer {
             }
         }
 
+        derived.import_comments(&composable.module_ir.comments);
+
         derived.clear_shader_source();
     }
 
@@ -1835,6 +1841,8 @@ static PREPROCESSOR: once_cell::sync::Lazy<Preprocessor> =
     once_cell::sync::Lazy::new(Preprocessor::default);
 
 /// Get module name and all required imports (ignoring shader_defs) from a shader string
+///
+/// Returns nothing in case of error. The actual error will be displayed when the caller attempts to use the shader.
 pub fn get_preprocessor_data(
     source: &str,
 ) -> (
@@ -1847,7 +1855,7 @@ pub fn get_preprocessor_data(
         imports,
         defines,
         ..
-    }) = PREPROCESSOR.get_preprocessor_metadata(source, true)
+    }) = try_get_preprocessor_data(source)
     {
         (
             name,
@@ -1862,3 +1870,8 @@ pub fn get_preprocessor_data(
         Default::default()
     }
 }
+
+/// Get module name and all required imports (ignoring shader_defs) from a shader string
+pub fn try_get_preprocessor_data(source: &str) -> Result<PreprocessorMetaData, ComposerErrorInner> {
+    PREPROCESSOR.get_preprocessor_metadata(source, true)
+}

src/compose/preprocess.rs

@@ -225,10 +225,11 @@ impl Preprocessor {
         Ok((false, None))
     }
 
-    // process #if[(n)?def]? / #else / #endif preprocessor directives,
-    // strip module name and imports
-    // also strip "#version xxx"
-    // replace items with resolved decorated names
+    /// - strip comments
+    /// - Process `#if[(n)?def]?` / `#else` / `#endif` preprocessor directives
+    /// - strip module name and imports
+    /// - strip `#version xxx`
+    /// - replace items with resolved decorated names
     pub fn preprocess(
         &self,
         shader_str: &str,
@@ -242,7 +243,9 @@ impl Preprocessor {
 
         // this code broadly stolen from bevy_render::ShaderProcessor
         let mut lines = shader_str.lines();
-        let mut lines = lines.replace_comments().zip(shader_str.lines()).peekable();
+        let mut lines = lines //.replace_comments()
+            .zip(shader_str.lines())
+            .peekable();
 
         while let Some((mut line, original_line)) = lines.next() {
             let mut output = false;

src/derive.rs

@@ -1,8 +1,9 @@
 use indexmap::IndexMap;
 use naga::{
-    Arena, AtomicFunction, Block, Constant, EntryPoint, Expression, Function, FunctionArgument,
-    FunctionResult, GatherMode, GlobalVariable, Handle, ImageQuery, LocalVariable, Module,
-    Override, SampleLevel, Span, Statement, StructMember, SwitchCase, Type, TypeInner, UniqueArena,
+    Arena, AtomicFunction, Block, Comments, Constant, EntryPoint, Expression, Function,
+    FunctionArgument, FunctionResult, GatherMode, GlobalVariable, Handle, ImageQuery,
+    LocalVariable, Module, Override, SampleLevel, Span, Statement, StructMember, SwitchCase, Type,
+    TypeInner, UniqueArena,
 };
 use std::{cell::RefCell, rc::Rc};
 
@@ -30,6 +31,7 @@ pub struct DerivedModule<'a> {
     globals: Arena<GlobalVariable>,
     functions: Arena<Function>,
     pipeline_overrides: Arena<Override>,
+    comments: Comments,
 }
 
 impl<'a> DerivedModule<'a> {
@@ -807,6 +809,16 @@ impl<'a> DerivedModule<'a> {
         self.import_function(func, span)
     }
 
+    pub(crate) fn import_comments(&mut self, comments: &Comments) {
+        // Deconstructing to not miss a new property to map if we add more.
+        let Comments { types } = comments;
+        for comment in types.iter() {
+            self.comments
+                .types
+                .insert(*self.type_map.get(comment.0).unwrap(), comment.1.clone());
+        }
+    }
+
     pub fn into_module_with_entrypoints(mut self) -> naga::Module {
         let entry_points = self
             .shader
@@ -842,6 +854,7 @@ impl<'a> From<DerivedModule<'a>> for naga::Module {
             special_types: Default::default(),
             entry_points: Default::default(),
             overrides: derived.pipeline_overrides,
+            comments: derived.comments,
         }
     }
 }