fix: Include frontmatter in -Zunpretty output

In the implementation (#140035), this was left as an open question for
the tracking issue (#136889).  My assumption is that this should be
carried over.

Thankfully, either way, `-Zunpretty` is unstable and we can always
change it even if we stabilize frontmatter.

Author: epage

compiler/rustc_ast_pretty/src/pprust/state.rs

@@ -120,7 +120,7 @@ fn gather_comments(sm: &SourceMap, path: FileName, src: String) -> Vec<Comment>
         pos += shebang_len;
     }
 
-    for token in rustc_lexer::tokenize(&text[pos..]) {
+    for token in rustc_lexer::tokenize_document(&text[pos..]) {
         let token_text = &text[pos..pos + token.len as usize];
         match token.kind {
             rustc_lexer::TokenKind::Whitespace => {
@@ -171,6 +171,14 @@ fn gather_comments(sm: &SourceMap, path: FileName, src: String) -> Vec<Comment>
                     })
                 }
             }
+            rustc_lexer::TokenKind::Frontmatter { .. } => {
+                code_to_the_left = false;
+                comments.push(Comment {
+                    style: CommentStyle::Isolated,
+                    lines: vec![token_text.to_string()],
+                    pos: start_bpos + BytePos(pos as u32),
+                });
+            }
             _ => {
                 code_to_the_left = true;
             }

compiler/rustc_lexer/src/lib.rs

@@ -311,6 +311,17 @@ pub fn tokenize(input: &str) -> impl Iterator<Item = Token> {
     })
 }
 
+/// Creates an iterator that produces tokens from the input string.
+///
+/// Be sure to call [`strip_shebang`] first
+pub fn tokenize_document(input: &str) -> impl Iterator<Item = Token> {
+    let mut cursor = Cursor::new(input, FrontmatterAllowed::Yes);
+    std::iter::from_fn(move || {
+        let token = cursor.advance_token();
+        if token.kind != TokenKind::Eof { Some(token) } else { None }
+    })
+}
+
 /// True if `c` is considered a whitespace according to Rust language definition.
 /// See [Rust language reference](https://doc.rust-lang.org/reference/whitespace.html)
 /// for definitions of these classes.

tests/ui/unpretty/frontmatter.stdout

@@ -1,3 +1,5 @@
+---
+---
 
 //@ compile-flags: -Zunpretty=normal
 //@ check-pass