[naga] Comments parsing (naga + wgsl)

CHANGELOG.md

@@ -139,6 +139,7 @@ By @cwfitzgerald in [#7030](https://github.com/gfx-rs/wgpu/pull/7030).
 - Support @must_use attribute on function declarations. By @turbocrime in [#6801](https://github.com/gfx-rs/wgpu/pull/6801).
 - Support for generating the candidate intersections from AABB geometry, and confirming the hits. By @kvark in [#7047](https://github.com/gfx-rs/wgpu/pull/7047).
 - Make naga::back::spv::Function::to_words write the OpFunctionEnd instruction in itself, instead of making another call after it. By @junjunjd in [#7156](https://github.com/gfx-rs/wgpu/pull/7156).
+- Support comments parsing for wgsl through `naga::front::glsl::Frontend::new_with_options`. By @Vrixyz in [#6364](https://github.com/gfx-rs/wgpu/pull/6364).
 
 ### Changes
 

naga/src/compact/mod.rs

@@ -264,6 +264,48 @@ pub fn compact(module: &mut crate::Module) {
             module_map.global_expressions.adjust(init);
         }
     }
+    // Adjust comments
+    if let Some(ref mut comments) = module.comments {
+        let crate::Comments {
+            module: _,
+            types: ref mut comment_types,
+            struct_members: ref mut comment_struct_members,
+            entry_points: _,
+            functions: _,
+            constants: ref mut comment_constants,
+            global_variables: _,
+        } = **comments;
+        log::trace!("adjusting comments for types");
+        for (mut comment_type_handle, comment) in std::mem::take(comment_types) {
+            if !module_map.types.used(comment_type_handle) {
+                continue;
+            }
+            module_map.types.adjust(&mut comment_type_handle);
+            comment_types.insert(comment_type_handle, comment);
+        }
+        log::trace!("adjusting comments for struct members");
+        for (mut comment_struct_member_handle, comment) in std::mem::take(comment_struct_members) {
+            if !module_map.types.used(comment_struct_member_handle.0) {
+                continue;
+            }
+            module_map.types.adjust(&mut comment_struct_member_handle.0);
+            comment_struct_members.insert(
+                (
+                    comment_struct_member_handle.0,
+                    comment_struct_member_handle.1,
+                ),
+                comment,
+            );
+        }
+        log::trace!("adjusting comments for constants");
+        for (mut comment_constant_handle, comment) in std::mem::take(comment_constants) {
+            if !module_map.constants.used(comment_constant_handle) {
+                continue;
+            }
+            module_map.constants.adjust(&mut comment_constant_handle);
+            comment_constants.insert(comment_constant_handle, comment);
+        }
+    }
 
     // Temporary storage to help us reuse allocations of existing
     // named expression tables.

naga/src/front/mod.rs

@@ -1,5 +1,5 @@
 /*!
-Frontend parsers that consume binary and text shaders and load them into [`Module`](super::Module)s.
+Frontend parsers that consume binary and text shaders and load them into [`Module`]s.
 */
 
 mod interpolator;
@@ -328,3 +328,12 @@ impl<Name: fmt::Debug, Var: fmt::Debug> fmt::Debug for SymbolTable<Name, Var> {
             .finish()
     }
 }
+
+use crate::{Comments, Module};
+
+impl Module {
+    pub fn get_comments_or_insert_default(&mut self) -> &mut Box<Comments> {
+        self.comments
+            .get_or_insert_with(|| Box::new(Comments::default()))
+    }
+}

naga/src/front/wgsl/error.rs

@@ -377,6 +377,8 @@ impl<'a> Error<'a> {
                         Token::Arrow => "->".to_string(),
                         Token::Unknown(c) => format!("unknown (`{c}`)"),
                         Token::Trivia => "trivia".to_string(),
+                        Token::CommentDoc(s) => format!("documentation ('{s}')"),
+                        Token::CommentDocModule(s) => format!("module documentation ('{s}')"),
                         Token::End => "end".to_string(),
                     },
                     ExpectedToken::Identifier => "identifier".to_string(),

naga/src/front/wgsl/lower/mod.rs

@@ -957,7 +957,7 @@ enum LoweredGlobalDecl {
     Const(Handle<crate::Constant>),
     Override(Handle<crate::Override>),
     Type(Handle<crate::Type>),
-    EntryPoint,
+    EntryPoint(usize),
 }
 
 enum Texture {
@@ -1057,6 +1057,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             layouter: &mut Layouter::default(),
             global_expression_kind_tracker: &mut crate::proc::ExpressionKindTracker::new(),
         };
+        if !tu.comments.is_empty() {
+            ctx.module.get_comments_or_insert_default().module =
+                tu.comments.iter().map(|s| s.to_string()).collect();
+        }
 
         for decl_handle in self.index.visit_ordered() {
             let span = tu.decls.get_span(decl_handle);
@@ -1065,6 +1069,29 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             match decl.kind {
                 ast::GlobalDeclKind::Fn(ref f) => {
                     let lowered_decl = self.function(f, span, &mut ctx)?;
+                    if !f.comments.is_empty() {
+                        match lowered_decl {
+                            LoweredGlobalDecl::Function { handle, .. } => {
+                                ctx.module
+                                    .get_comments_or_insert_default()
+                                    .functions
+                                    .insert(
+                                        handle,
+                                        f.comments.iter().map(|s| s.to_string()).collect(),
+                                    );
+                            }
+                            LoweredGlobalDecl::EntryPoint(index) => {
+                                ctx.module
+                                    .get_comments_or_insert_default()
+                                    .entry_points
+                                    .insert(
+                                        index,
+                                        f.comments.iter().map(|s| s.to_string()).collect(),
+                                    );
+                            }
+                            _ => {}
+                        }
+                    }
                     ctx.globals.insert(f.name.name, lowered_decl);
                 }
                 ast::GlobalDeclKind::Var(ref v) => {
@@ -1095,6 +1122,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         span,
                     );
 
+                    if !v.comments.is_empty() {
+                        ctx.module
+                            .get_comments_or_insert_default()
+                            .global_variables
+                            .insert(handle, v.comments.iter().map(|s| s.to_string()).collect());
+                    }
                     ctx.globals
                         .insert(v.name.name, LoweredGlobalDecl::Var(handle));
                 }
@@ -1120,6 +1153,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
 
                     ctx.globals
                         .insert(c.name.name, LoweredGlobalDecl::Const(handle));
+                    if !c.comments.is_empty() {
+                        ctx.module
+                            .get_comments_or_insert_default()
+                            .constants
+                            .insert(handle, c.comments.iter().map(|s| s.to_string()).collect());
+                    }
                 }
                 ast::GlobalDeclKind::Override(ref o) => {
                     let explicit_ty =
@@ -1160,6 +1199,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     let handle = self.r#struct(s, span, &mut ctx)?;
                     ctx.globals
                         .insert(s.name.name, LoweredGlobalDecl::Type(handle));
+                    if !s.comments.is_empty() {
+                        ctx.module
+                            .get_comments_or_insert_default()
+                            .types
+                            .insert(handle, s.comments.iter().map(|s| s.to_string()).collect());
+                    }
                 }
                 ast::GlobalDeclKind::Type(ref alias) => {
                     let ty = self.resolve_named_ast_type(
@@ -1375,7 +1420,9 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 workgroup_size_overrides,
                 function,
             });
-            Ok(LoweredGlobalDecl::EntryPoint)
+            Ok(LoweredGlobalDecl::EntryPoint(
+                ctx.module.entry_points.len() - 1,
+            ))
         } else {
             let handle = ctx.module.functions.append(function, span);
             Ok(LoweredGlobalDecl::Function {
@@ -1933,7 +1980,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     }
                     LoweredGlobalDecl::Function { .. }
                     | LoweredGlobalDecl::Type(_)
-                    | LoweredGlobalDecl::EntryPoint => {
+                    | LoweredGlobalDecl::EntryPoint(_) => {
                         return Err(Error::Unexpected(span, ExpectedToken::Variable));
                     }
                 };
@@ -2193,7 +2240,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 | &LoweredGlobalDecl::Override(_)
                 | &LoweredGlobalDecl::Var(_),
             ) => Err(Error::Unexpected(function_span, ExpectedToken::Function)),
-            Some(&LoweredGlobalDecl::EntryPoint) => Err(Error::CalledEntryPoint(function_span)),
+            Some(&LoweredGlobalDecl::EntryPoint(_)) => Err(Error::CalledEntryPoint(function_span)),
             Some(&LoweredGlobalDecl::Function {
                 handle: function,
                 must_use,
@@ -3058,6 +3105,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         let mut struct_alignment = Alignment::ONE;
         let mut members = Vec::with_capacity(s.members.len());
 
+        let mut comments: Vec<Option<Vec<String>>> = Vec::new();
+
         for member in s.members.iter() {
             let ty = self.resolve_ast_type(member.ty, &mut ctx.as_const())?;
 
@@ -3097,6 +3146,11 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             offset = member_alignment.round_up(offset);
             struct_alignment = struct_alignment.max(member_alignment);
 
+            if !member.comments.is_empty() {
+                comments.push(Some(
+                    member.comments.iter().map(|s| s.to_string()).collect(),
+                ));
+            }
             members.push(crate::StructMember {
                 name: Some(member.name.name.to_owned()),
                 ty,
@@ -3120,6 +3174,12 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             },
             span,
         );
+        for (i, c) in comments.drain(..).enumerate() {
+            if let Some(comment) = c {
+                let comments = ctx.module.get_comments_or_insert_default();
+                comments.struct_members.insert((handle, i), comment);
+            }
+        }
         Ok(handle)
     }
 

naga/src/front/wgsl/mod.rs

@@ -19,19 +19,28 @@ use thiserror::Error;
 pub use crate::front::wgsl::error::ParseError;
 use crate::front::wgsl::lower::Lowerer;
 use crate::Scalar;
+pub use parse::Options;
 
 pub use crate::front::wgsl::parse::directive::language_extension::{
     ImplementedLanguageExtension, LanguageExtension, UnimplementedLanguageExtension,
 };
 
 pub struct Frontend {
     parser: Parser,
+    options: Options,
 }
 
 impl Frontend {
     pub const fn new() -> Self {
         Self {
             parser: Parser::new(),
+            options: Options::new(),
+        }
+    }
+    pub const fn new_with_options(options: Options) -> Self {
+        Self {
+            parser: Parser::new(),
+            options,
         }
     }
 
@@ -40,7 +49,7 @@ impl Frontend {
     }
 
     fn inner<'a>(&mut self, source: &'a str) -> Result<crate::Module, Error<'a>> {
-        let tu = self.parser.parse(source)?;
+        let tu = self.parser.parse(source, &self.options)?;
         let index = index::Index::generate(&tu)?;
         let module = Lowerer::new(&index).lower(&tu)?;
 

naga/src/front/wgsl/parse/ast.rs

@@ -38,6 +38,10 @@ pub struct TranslationUnit<'a> {
     /// See [`DiagnosticFilterNode`] for details on how the tree is represented and used in
     /// validation.
     pub diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
+
+    /// Comments appearing first in the file.
+    /// This serves as documentation for the whole TranslationUnit.
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -135,6 +139,7 @@ pub struct Function<'a> {
     pub result: Option<FunctionResult<'a>>,
     pub body: Block<'a>,
     pub diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -161,6 +166,7 @@ pub struct GlobalVariable<'a> {
     pub binding: Option<ResourceBinding<'a>>,
     pub ty: Option<Handle<Type<'a>>>,
     pub init: Option<Handle<Expression<'a>>>,
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -170,12 +176,15 @@ pub struct StructMember<'a> {
     pub binding: Option<Binding<'a>>,
     pub align: Option<Handle<Expression<'a>>>,
     pub size: Option<Handle<Expression<'a>>>,
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
 pub struct Struct<'a> {
     pub name: Ident<'a>,
     pub members: Vec<StructMember<'a>>,
+    // TODO: Make it optional ? Store Span ? Add it to other elements
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -189,6 +198,7 @@ pub struct Const<'a> {
     pub name: Ident<'a>,
     pub ty: Option<Handle<Type<'a>>>,
     pub init: Handle<Expression<'a>>,
+    pub comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]