[wgsl-in,ir] Add support for parsing rust-style doc comments (#6364)

* [wgsl-in,ir] add support for parsing rust-style doc comments

* rename relevant items to `doc_comments` (or variations of it)

* address comments

* remove `next_until`

* rename `save_doc_comments` to `ignore_doc_comments`

* expand snapshot test and ignore blankspace when accumulating doc comments

* make tokenizer more straightforward

---------

Co-authored-by: teoxoy <28601907+teoxoy@users.noreply.github.com>

Author: Vrixyz

CHANGELOG.md

@@ -45,6 +45,7 @@ Bottom level categories:
 #### Naga
 
 - Added `no_std` support with default features disabled. By @Bushrat011899 in [#7585](https://github.com/gfx-rs/wgpu/pull/7585).
+- [wgsl-in,ir] Add support for parsing rust-style doc comments via `naga::front::glsl::Frontend::new_with_options`. By @Vrixyz in [#6364](https://github.com/gfx-rs/wgpu/pull/6364).
 
 #### General
 

naga/src/compact/mod.rs

@@ -272,6 +272,42 @@ pub fn compact(module: &mut crate::Module) {
             module_map.global_expressions.adjust(init);
         }
     }
+    // Adjust doc comments
+    if let Some(ref mut doc_comments) = module.doc_comments {
+        let crate::DocComments {
+            module: _,
+            types: ref mut doc_comments_for_types,
+            struct_members: ref mut doc_comments_for_struct_members,
+            entry_points: _,
+            functions: _,
+            constants: ref mut doc_comments_for_constants,
+            global_variables: _,
+        } = **doc_comments;
+        log::trace!("adjusting doc comments for types");
+        for (mut ty, doc_comment) in core::mem::take(doc_comments_for_types) {
+            if !module_map.types.used(ty) {
+                continue;
+            }
+            module_map.types.adjust(&mut ty);
+            doc_comments_for_types.insert(ty, doc_comment);
+        }
+        log::trace!("adjusting doc comments for struct members");
+        for ((mut ty, index), doc_comment) in core::mem::take(doc_comments_for_struct_members) {
+            if !module_map.types.used(ty) {
+                continue;
+            }
+            module_map.types.adjust(&mut ty);
+            doc_comments_for_struct_members.insert((ty, index), doc_comment);
+        }
+        log::trace!("adjusting doc comments for constants");
+        for (mut constant, doc_comment) in core::mem::take(doc_comments_for_constants) {
+            if !module_map.constants.used(constant) {
+                continue;
+            }
+            module_map.constants.adjust(&mut constant);
+            doc_comments_for_constants.insert(constant, doc_comment);
+        }
+    }
 
     // Temporary storage to help us reuse allocations of existing
     // named expression tables.

naga/src/front/mod.rs

@@ -14,7 +14,7 @@ pub mod spv;
 #[cfg(feature = "wgsl-in")]
 pub mod wgsl;
 
-use alloc::{vec, vec::Vec};
+use alloc::{boxed::Box, vec, vec::Vec};
 use core::ops;
 
 use crate::{
@@ -330,3 +330,10 @@ impl<Name: fmt::Debug, Var: fmt::Debug> fmt::Debug for SymbolTable<Name, Var> {
             .finish()
     }
 }
+
+impl crate::Module {
+    pub fn get_or_insert_default_doc_comments(&mut self) -> &mut Box<crate::DocComments> {
+        self.doc_comments
+            .get_or_insert_with(|| Box::new(crate::DocComments::default()))
+    }
+}

naga/src/front/wgsl/error.rs

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

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

@@ -1022,7 +1022,7 @@ enum LoweredGlobalDecl {
     Const(Handle<ir::Constant>),
     Override(Handle<ir::Override>),
     Type(Handle<ir::Type>),
-    EntryPoint,
+    EntryPoint(usize),
 }
 
 enum Texture {
@@ -1130,6 +1130,10 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             layouter: &mut proc::Layouter::default(),
             global_expression_kind_tracker: &mut proc::ExpressionKindTracker::new(),
         };
+        if !tu.doc_comments.is_empty() {
+            ctx.module.get_or_insert_default_doc_comments().module =
+                tu.doc_comments.iter().map(|s| s.to_string()).collect();
+        }
 
         for decl_handle in self.index.visit_ordered() {
             let span = tu.decls.get_span(decl_handle);
@@ -1138,6 +1142,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.doc_comments.is_empty() {
+                        match lowered_decl {
+                            LoweredGlobalDecl::Function { handle, .. } => {
+                                ctx.module
+                                    .get_or_insert_default_doc_comments()
+                                    .functions
+                                    .insert(
+                                        handle,
+                                        f.doc_comments.iter().map(|s| s.to_string()).collect(),
+                                    );
+                            }
+                            LoweredGlobalDecl::EntryPoint(index) => {
+                                ctx.module
+                                    .get_or_insert_default_doc_comments()
+                                    .entry_points
+                                    .insert(
+                                        index,
+                                        f.doc_comments.iter().map(|s| s.to_string()).collect(),
+                                    );
+                            }
+                            _ => {}
+                        }
+                    }
                     ctx.globals.insert(f.name.name, lowered_decl);
                 }
                 ast::GlobalDeclKind::Var(ref v) => {
@@ -1173,6 +1200,15 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                         span,
                     );
 
+                    if !v.doc_comments.is_empty() {
+                        ctx.module
+                            .get_or_insert_default_doc_comments()
+                            .global_variables
+                            .insert(
+                                handle,
+                                v.doc_comments.iter().map(|s| s.to_string()).collect(),
+                            );
+                    }
                     ctx.globals
                         .insert(v.name.name, LoweredGlobalDecl::Var(handle));
                 }
@@ -1203,6 +1239,15 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
 
                     ctx.globals
                         .insert(c.name.name, LoweredGlobalDecl::Const(handle));
+                    if !c.doc_comments.is_empty() {
+                        ctx.module
+                            .get_or_insert_default_doc_comments()
+                            .constants
+                            .insert(
+                                handle,
+                                c.doc_comments.iter().map(|s| s.to_string()).collect(),
+                            );
+                    }
                 }
                 ast::GlobalDeclKind::Override(ref o) => {
                     let explicit_ty =
@@ -1249,6 +1294,15 @@ 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.doc_comments.is_empty() {
+                        ctx.module
+                            .get_or_insert_default_doc_comments()
+                            .types
+                            .insert(
+                                handle,
+                                s.doc_comments.iter().map(|s| s.to_string()).collect(),
+                            );
+                    }
                 }
                 ast::GlobalDeclKind::Type(ref alias) => {
                     let ty = self.resolve_named_ast_type(
@@ -1469,7 +1523,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 {
@@ -2086,7 +2142,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     }
                     LoweredGlobalDecl::Function { .. }
                     | LoweredGlobalDecl::Type(_)
-                    | LoweredGlobalDecl::EntryPoint => {
+                    | LoweredGlobalDecl::EntryPoint(_) => {
                         return Err(Box::new(Error::Unexpected(span, ExpectedToken::Variable)));
                     }
                 };
@@ -2373,7 +2429,7 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                 function_span,
                 ExpectedToken::Function,
             ))),
-            Some(&LoweredGlobalDecl::EntryPoint) => {
+            Some(&LoweredGlobalDecl::EntryPoint(_)) => {
                 Err(Box::new(Error::CalledEntryPoint(function_span)))
             }
             Some(&LoweredGlobalDecl::Function {
@@ -3581,6 +3637,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
         let mut struct_alignment = proc::Alignment::ONE;
         let mut members = Vec::with_capacity(s.members.len());
 
+        let mut doc_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())?;
 
@@ -3623,6 +3681,11 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             offset = member_alignment.round_up(offset);
             struct_alignment = struct_alignment.max(member_alignment);
 
+            if !member.doc_comments.is_empty() {
+                doc_comments.push(Some(
+                    member.doc_comments.iter().map(|s| s.to_string()).collect(),
+                ));
+            }
             members.push(ir::StructMember {
                 name: Some(member.name.name.to_owned()),
                 ty,
@@ -3646,6 +3709,14 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
             },
             span,
         );
+        for (i, c) in doc_comments.drain(..).enumerate() {
+            if let Some(comment) = c {
+                ctx.module
+                    .get_or_insert_default_doc_comments()
+                    .struct_members
+                    .insert((handle, i), comment);
+            }
+        }
         Ok(handle)
     }
 

naga/src/front/wgsl/mod.rs

@@ -15,6 +15,7 @@ pub use crate::front::wgsl::error::ParseError;
 pub use crate::front::wgsl::parse::directive::language_extension::{
     ImplementedLanguageExtension, LanguageExtension, UnimplementedLanguageExtension,
 };
+pub use crate::front::wgsl::parse::Options;
 
 use alloc::boxed::Box;
 use thiserror::Error;
@@ -31,12 +32,20 @@ pub(crate) type Result<'a, T> = core::result::Result<T, Box<Error<'a>>>;
 
 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,
         }
     }
 
@@ -45,7 +54,7 @@ impl Frontend {
     }
 
     fn inner<'a>(&mut self, source: &'a str) -> Result<'a, crate::Module> {
-        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

@@ -40,6 +40,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>>,
+
+    /// Doc comments appearing first in the file.
+    /// This serves as documentation for the whole TranslationUnit.
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -137,6 +141,7 @@ pub struct Function<'a> {
     pub result: Option<FunctionResult<'a>>,
     pub body: Block<'a>,
     pub diagnostic_filter_leaf: Option<Handle<DiagnosticFilterNode>>,
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -163,6 +168,7 @@ pub struct GlobalVariable<'a> {
     pub binding: Option<ResourceBinding<'a>>,
     pub ty: Option<Handle<Type<'a>>>,
     pub init: Option<Handle<Expression<'a>>>,
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -172,12 +178,14 @@ pub struct StructMember<'a> {
     pub binding: Option<Binding<'a>>,
     pub align: Option<Handle<Expression<'a>>>,
     pub size: Option<Handle<Expression<'a>>>,
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
 pub struct Struct<'a> {
     pub name: Ident<'a>,
     pub members: Vec<StructMember<'a>>,
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]
@@ -191,6 +199,7 @@ pub struct Const<'a> {
     pub name: Ident<'a>,
     pub ty: Option<Handle<Type<'a>>>,
     pub init: Handle<Expression<'a>>,
+    pub doc_comments: Vec<&'a str>,
 }
 
 #[derive(Debug)]