module comments have specific syntax

gfx-rs · Oct 11, 2024 · 02caf3a · 02caf3a
1 parent ff84ff8
commit 02caf3a
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 22 deletions.
diff --git a/naga/src/front/wgsl/error.rs b/naga/src/front/wgsl/error.rs
@@ -300,6 +300,7 @@ impl<'a> Error<'a> {
                             Token::Unknown(c) => format!("unknown ('{c}')"),
                             Token::Trivia => "trivia".to_string(),
                             Token::Comment(s) => format!("documentation ('{s}')"),
+                            Token::CommentModule(s) => format!("module documentation ('{s}')"),
                             Token::End => "end".to_string(),
                         }
                     }

diff --git a/naga/src/front/wgsl/parse/lexer.rs b/naga/src/front/wgsl/parse/lexer.rs
@@ -22,6 +22,7 @@ pub enum Token<'a> {
     Unknown(char),
     Trivia,
     Comment(&'a str),
+    CommentModule(&'a str),
     End,
 }
 
@@ -88,7 +89,11 @@ fn consume_token(input: &str, generic: bool) -> (Token<'_>, &str) {
                     {
                         let end_position = end_position.0 + 1;
                         return (
-                            Token::Comment(&input[..end_position]),
+                            if chars.next() == Some('!') {
+                                Token::CommentModule(&input[..end_position])
+                            } else {
+                                Token::Comment(&input[..end_position])
+                            },
                             &input[end_position..],
                         );
                     }
@@ -253,7 +258,7 @@ impl<'a> Lexer<'a> {
         loop {
             // Eat all trivia because `next` doesn't eat trailing trivia.
             let (token, rest) = consume_token(self.input, false);
-            if let Token::Trivia | Token::Comment(_) = token {
+            if let Token::Trivia | Token::Comment(_) | Token::CommentModule(_) = token {
                 self.input = rest;
             } else {
                 return self.current_byte_offset();
@@ -287,6 +292,25 @@ impl<'a> Lexer<'a> {
             }
         }
     }
+    pub(in crate::front::wgsl) fn start_byte_offset_and_aggregate_comment_module(
+        &'a mut self,
+        comments: &mut Vec<Span>,
+    ) -> usize {
+        loop {
+            let start = self.current_byte_offset();
+            // Eat all trivia because `next` doesn't eat trailing trivia.
+            let (token, rest) = consume_token(self.input, false);
+            if let Token::CommentModule(_) = token {
+                self.input = rest;
+                let next = self.current_byte_offset();
+                comments.push(Span::new(start as u32, next as u32));
+            } else if let Token::Trivia = token {
+                self.input = rest;
+            } else {
+                return self.current_byte_offset();
+            }
+        }
+    }
 
     pub const fn current_byte_offset(&self) -> usize {
         self.source.len() - self.input.len()

diff --git a/naga/src/front/wgsl/parse/mod.rs b/naga/src/front/wgsl/parse/mod.rs
@@ -2283,9 +2283,22 @@ impl Parser {
             (None, None) => {}
         }
 
+        // read module docs
+        {
+            let mut module_comments = Vec::new();
+            let mut lexer_comment_modules = lexer_comments.clone();
+            lexer_comment_modules
+                .start_byte_offset_and_aggregate_comment_module(&mut module_comments);
+            let mut module_comments: Vec<_> = module_comments
+                .into_iter()
+                .map(|comment_span| lexer.source.index(comment_span))
+                .collect();
+            out.comments.append(&mut module_comments);
+        }
+        let token_span = lexer.next();
+
         // read item
         let start = lexer.start_byte_offset();
-        let token_span = lexer.next();
 
         let kind = match token_span {
             (Token::Separator(';'), _) => None,
@@ -2438,25 +2451,6 @@ impl Parser {
 
         let mut lexer = Lexer::new(source);
         let mut tu = ast::TranslationUnit::default();
-        let mut comments = Vec::new();
-
-        fn peek_any_next<'a>(lexer: &'a Lexer) -> (Token<'a>, Span) {
-            let mut cloned = lexer.clone();
-            let token = cloned.next_until(|_| true, false);
-            token
-        }
-        loop {
-            match peek_any_next(&lexer) {
-                (Token::Comment(_), span) => {
-                    comments.push(lexer.source.index(span));
-                    let _ = lexer.next_until(|_| true, false);
-                }
-                _ => {
-                    break;
-                }
-            }
-        }
-        tu.comments = comments;
 
         loop {
             match self.global_decl(&mut lexer, &mut tu) {