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

@@ -306,7 +306,7 @@ pub struct ImportDefinition {
 
 #[derive(Debug, Clone)]
 pub struct ImportDefWithOffset {
-    definition: ImportDefinition,
+    pub definition: ImportDefinition,
     offset: usize,
 }
 
@@ -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,
@@ -688,7 +690,6 @@ impl Composer {
                     }
                 })?,
         };
-
         Ok(IrBuildResult {
             module,
             start_offset,
@@ -962,7 +963,6 @@ impl Composer {
             module_definition.name,
             source.len()
         );
-
         let IrBuildResult {
             module: mut source_ir,
             start_offset,
@@ -1160,7 +1160,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 +1282,8 @@ impl Composer {
             }
         }
 
+        derived.import_comments(&composable.module_ir.comments);
+
         derived.clear_shader_source();
     }
 
@@ -1835,6 +1838,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 +1852,7 @@ pub fn get_preprocessor_data(
         imports,
         defines,
         ..
-    }) = PREPROCESSOR.get_preprocessor_metadata(source, true)
+    }) = try_get_preprocessor_data(source)
     {
         (
             name,
@@ -1862,3 +1867,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,

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,54 @@ 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,
+            functions,
+            constants,
+            global_variables,
+            struct_members,
+            module,
+        } = comments;
+        self.comments.module = module.clone();
+        for comment in types.iter() {
+            if let Some(new_handle) = self.type_map.get(comment.0) {
+                self.comments.types.insert(*new_handle, comment.1.clone());
+            }
+        }
+        for ((struct_handle, index), comment) in struct_members.iter() {
+            if let Some(new_struct_handle) = self.type_map.get(struct_handle) {
+                self.comments
+                    .struct_members
+                    .insert((*new_struct_handle, *index), comment.clone());
+            }
+        }
+        for function in functions.iter() {
+            self.comments
+                .functions
+                .insert(function.0.to_string(), function.1.clone());
+        }
+        for constant in constants.iter() {
+            if let Some(new_handle) = self.const_map.get(constant.0) {
+                self.comments
+                    .constants
+                    .insert(*new_handle, constant.1.clone());
+            }
+        }
+        for global_variable in global_variables.iter() {
+            if let Some(new_handle) = self.global_map.get(global_variable.0) {
+                self.comments
+                    .global_variables
+                    .insert(*new_handle, global_variable.1.clone());
+            } else {
+                // NOTE: Could not migrate global variable, I'm not sure why this happens.
+                // It's probably because this module is not the owner of the global.
+                // Chances are that comment will be migrated by another module.
+            }
+        }
+    }
+
     pub fn into_module_with_entrypoints(mut self) -> naga::Module {
         let entry_points = self
             .shader
@@ -842,6 +892,7 @@ impl<'a> From<DerivedModule<'a>> for naga::Module {
             special_types: Default::default(),
             entry_points: Default::default(),
             overrides: derived.pipeline_overrides,
+            comments: derived.comments,
         }
     }
 }