rustdoc: Recognize more weak keywords when highlighting Rust code

Author: fmease

src/librustdoc/html/highlight.rs

@@ -754,21 +754,6 @@ impl<'a> Iterator for TokenIter<'a> {
     }
 }
 
-/// Classifies into identifier class; returns `None` if this is a non-keyword identifier.
-fn get_real_ident_class(text: &str, allow_path_keywords: bool) -> Option<Class> {
-    let ignore: &[&str] =
-        if allow_path_keywords { &["self", "Self", "super", "crate"] } else { &["self", "Self"] };
-    if ignore.contains(&text) {
-        return None;
-    }
-    Some(match text {
-        "ref" | "mut" => Class::RefKeyWord,
-        "false" | "true" => Class::Bool,
-        _ if Symbol::intern(text).is_reserved(|| Edition::Edition2021) => Class::KeyWord,
-        _ => return None,
-    })
-}
-
 /// This iterator comes from the same idea than "Peekable" except that it allows to "peek" more than
 /// just the next item by using `peek_next`. The `peek` method always returns the next item after
 /// the current one whereas `peek_next` will return the next item after the last one peeked.
@@ -786,16 +771,16 @@ impl<'a> PeekIter<'a> {
         Self { stored: VecDeque::new(), peek_pos: 0, iter }
     }
     /// Returns the next item after the current one. It doesn't interfere with `peek_next` output.
-    fn peek(&mut self) -> Option<&(TokenKind, &'a str)> {
+    fn peek(&mut self) -> Option<(TokenKind, &'a str)> {
         if self.stored.is_empty()
             && let Some(next) = self.iter.next()
         {
             self.stored.push_back(next);
         }
-        self.stored.front()
+        self.stored.front().copied()
     }
     /// Returns the next item after the last one peeked. It doesn't interfere with `peek` output.
-    fn peek_next(&mut self) -> Option<&(TokenKind, &'a str)> {
+    fn peek_next(&mut self) -> Option<(TokenKind, &'a str)> {
         self.peek_pos += 1;
         if self.peek_pos - 1 < self.stored.len() {
             self.stored.get(self.peek_pos - 1)
@@ -805,6 +790,7 @@ impl<'a> PeekIter<'a> {
         } else {
             None
         }
+        .copied()
     }
 
     fn stop_peeking(&mut self) {
@@ -951,15 +937,9 @@ impl<'src> Classifier<'src> {
                 }
             }
 
-            if let Some((None, text)) = self.tokens.peek().map(|(token, text)| {
-                if *token == TokenKind::Ident {
-                    let class = get_real_ident_class(text, true);
-                    (class, text)
-                } else {
-                    // Doesn't matter which Class we put in here...
-                    (Some(Class::Comment), text)
-                }
-            }) {
+            if let Some((TokenKind::Ident, text)) = self.tokens.peek()
+                && !ident_is_keyword(text)
+            {
                 // We only "add" the colon if there is an ident behind.
                 pos += text.len() + nb;
                 has_ident = true;
@@ -1197,34 +1177,15 @@ impl<'src> Classifier<'src> {
             },
             TokenKind::GuardedStrPrefix => return no_highlight(sink),
             TokenKind::Ident | TokenKind::RawIdent
-                if self.peek_non_whitespace() == Some(TokenKind::Bang) =>
+                if let Some((TokenKind::Bang, _)) = self.peek_non_trivia() =>
             {
                 self.in_macro = true;
                 let span = new_span(before, text, file_span);
                 sink(DUMMY_SP, Highlight::EnterSpan { class: Class::Macro(span) });
                 sink(span, Highlight::Token { text, class: None });
                 return;
             }
-            TokenKind::Ident => {
-                match get_real_ident_class(text, false) {
-                    None => match text {
-                        "Option" | "Result" => Class::PreludeTy(new_span(before, text, file_span)),
-                        "Some" | "None" | "Ok" | "Err" => {
-                            Class::PreludeVal(new_span(before, text, file_span))
-                        }
-                        // "union" is a weak keyword and is only considered as a keyword when declaring
-                        // a union type.
-                        "union" if self.check_if_is_union_keyword() => Class::KeyWord,
-                        _ if self.in_macro_nonterminal => {
-                            self.in_macro_nonterminal = false;
-                            Class::MacroNonTerminal
-                        }
-                        "self" | "Self" => Class::Self_(new_span(before, text, file_span)),
-                        _ => Class::Ident(new_span(before, text, file_span)),
-                    },
-                    Some(c) => c,
-                }
-            }
+            TokenKind::Ident => self.classify_ident(before, text),
             TokenKind::RawIdent | TokenKind::UnknownPrefix | TokenKind::InvalidIdent => {
                 Class::Ident(new_span(before, text, file_span))
             }
@@ -1245,25 +1206,62 @@ impl<'src> Classifier<'src> {
         }
     }
 
+    fn classify_ident(&mut self, before: u32, text: &'src str) -> Class {
+        // Macro non-terminals (meta vars) take precedence.
+        if self.in_macro_nonterminal {
+            self.in_macro_nonterminal = false;
+            return Class::MacroNonTerminal;
+        }
+
+        let file_span = self.file_span;
+        let span = || new_span(before, text, file_span);
+
+        match text {
+            "ref" | "mut" => Class::RefKeyWord,
+            "false" | "true" => Class::Bool,
+            "self" | "Self" => Class::Self_(span()),
+            "Option" | "Result" => Class::PreludeTy(span()),
+            "Some" | "None" | "Ok" | "Err" => Class::PreludeVal(span()),
+            _ if ident_is_keyword(text) || self.ident_is_weak_keyword(text) => Class::KeyWord,
+            _ => Class::Ident(span()),
+        }
+    }
+
+    fn ident_is_weak_keyword(&mut self, text: &str) -> bool {
+        let matches = match text {
+            "auto" => |text| text == "trait", // `auto trait Trait {}` (`auto_traits`)
+            "pin" => |text| text == "const" || text == "mut", // `&pin mut Type` (`pin_ergonomics`)
+            "raw" => |text| text == "const" || text == "mut", // `&raw const local`
+            "safe" => |text| text == "fn" || text == "extern", // `unsafe extern { safe fn f(); }`
+            "union" => |_| true,              // `union Untagged { field: () }`
+            _ => return false,
+        };
+        matches!(self.peek_non_trivia(), Some((TokenKind::Ident, text)) if matches(text))
+    }
+
     fn peek(&mut self) -> Option<TokenKind> {
-        self.tokens.peek().map(|(token_kind, _text)| *token_kind)
+        self.tokens.peek().map(|(kind, _)| kind)
     }
 
-    fn peek_non_whitespace(&mut self) -> Option<TokenKind> {
-        while let Some((token_kind, _)) = self.tokens.peek_next() {
-            if *token_kind != TokenKind::Whitespace {
-                let token_kind = *token_kind;
-                self.tokens.stop_peeking();
-                return Some(token_kind);
+    fn peek_non_trivia(&mut self) -> Option<(TokenKind, &str)> {
+        while let Some(token @ (kind, _)) = self.tokens.peek_next() {
+            if let TokenKind::Whitespace
+            | TokenKind::LineComment { doc_style: None }
+            | TokenKind::BlockComment { doc_style: None, .. } = kind
+            {
+                continue;
             }
+            self.tokens.stop_peeking();
+            return Some(token);
         }
         self.tokens.stop_peeking();
         None
     }
+}
 
-    fn check_if_is_union_keyword(&mut self) -> bool {
-        self.peek_non_whitespace().is_some_and(|kind| kind == TokenKind::Ident)
-    }
+fn ident_is_keyword(text: &str) -> bool {
+    // FIXME(#148221): Don't hard-code the edition. The classifier should take it as an argument.
+    Symbol::intern(text).is_reserved(|| Edition::Edition2024)
 }
 
 fn generate_link_to_def(