Rollup merge of #73807 - euclio:rustdoc-highlighting, r=ollie27,GuillaumeGomez

rustdoc: glue tokens before highlighting

Fixes #72684.

This commit also modifies the signature of `Classifier::new` to avoid
copying the source being highlighted.

5 files changed, 133 insertions, 22 deletions

diff --git a/src/librustdoc/html/highlight.rs b/src/librustdoc/html/highlight.rs
index c4bc73770a7..d4302d0cb54 100644
--- a/src/librustdoc/html/highlight.rs
+++ b/src/librustdoc/html/highlight.rs
@@ -12,15 +12,17 @@ use std::io;
 use std::io::prelude::*;
 
 use rustc_ast::token::{self, Token};
+use rustc_data_structures::sync::Lrc;
 use rustc_parse::lexer;
 use rustc_session::parse::ParseSess;
+use rustc_span::hygiene::SyntaxContext;
 use rustc_span::source_map::SourceMap;
 use rustc_span::symbol::{kw, sym};
-use rustc_span::{FileName, Span};
+use rustc_span::{BytePos, FileName, SourceFile, Span};
 
 /// Highlights `src`, returning the HTML output.
 pub fn render_with_highlighting(
-    src: &str,
+    src: String,
     class: Option<&str>,
     playground_button: Option<&str>,
     tooltip: Option<(&str, &str)>,
@@ -38,12 +40,13 @@ pub fn render_with_highlighting(
     }
 
     let sess = ParseSess::with_silent_emitter();
-    let sf = sess
+    let source_file = sess
         .source_map()
-        .new_source_file(FileName::Custom(String::from("rustdoc-highlighting")), src.to_owned());
+        .new_source_file(FileName::Custom(String::from("rustdoc-highlighting")), src);
+
+    let classifier_source_file = Lrc::clone(&source_file);
     let highlight_result = rustc_driver::catch_fatal_errors(|| {
-        let lexer = lexer::StringReader::new(&sess, sf, None);
-        let mut classifier = Classifier::new(lexer, sess.source_map());
+        let mut classifier = Classifier::new(&sess, classifier_source_file);
 
         let mut highlighted_source = vec![];
         if classifier.write_source(&mut highlighted_source).is_err() {
@@ -61,9 +64,17 @@ pub fn render_with_highlighting(
             write_footer(&mut out, playground_button).unwrap();
         }
         Err(()) => {
+            // Get the source back out of the source map to avoid a copy in the happy path.
+            let span =
+                Span::new(BytePos(0), BytePos(source_file.byte_length()), SyntaxContext::root());
+            let src = sess
+                .source_map()
+                .span_to_snippet(span)
+                .expect("could not retrieve snippet from artificial source file");
+
             // If errors are encountered while trying to highlight, just emit
             // the unhighlighted source.
-            write!(out, "<pre><code>{}</code></pre>", Escape(src)).unwrap();
+            write!(out, "<pre><code>{}</code></pre>", Escape(&src)).unwrap();
         }
     }
 
@@ -73,10 +84,10 @@ pub fn render_with_highlighting(
 /// Processes a program (nested in the internal `lexer`), classifying strings of
 /// text by highlighting category (`Class`). Calls out to a `Writer` to write
 /// each span of text in sequence.
-struct Classifier<'a> {
-    lexer: lexer::StringReader<'a>,
+struct Classifier<'sess> {
+    lexer: lexer::StringReader<'sess>,
     peek_token: Option<Token>,
-    source_map: &'a SourceMap,
+    source_map: &'sess SourceMap,
 
     // State of the classifier.
     in_attribute: bool,
@@ -154,6 +165,7 @@ impl<U: Write> Writer for U {
     }
 }
 
+#[derive(Debug)]
 enum HighlightError {
     LexError,
     IoError(io::Error),
@@ -165,12 +177,14 @@ impl From<io::Error> for HighlightError {
     }
 }
 
-impl<'a> Classifier<'a> {
-    fn new(lexer: lexer::StringReader<'a>, source_map: &'a SourceMap) -> Classifier<'a> {
+impl<'sess> Classifier<'sess> {
+    fn new(sess: &ParseSess, source_file: Lrc<SourceFile>) -> Classifier<'_> {
+        let lexer = lexer::StringReader::new(sess, source_file, None);
+
         Classifier {
             lexer,
             peek_token: None,
-            source_map,
+            source_map: sess.source_map(),
             in_attribute: false,
             in_macro: false,
             in_macro_nonterminal: false,
@@ -209,11 +223,17 @@ impl<'a> Classifier<'a> {
     /// source.
     fn write_source<W: Writer>(&mut self, out: &mut W) -> Result<(), HighlightError> {
         loop {
-            let next = self.try_next_token()?;
+            let mut next = self.try_next_token()?;
             if next == token::Eof {
                 break;
             }
 
+            // Glue any tokens that need to be glued.
+            if let Some(joint) = next.glue(self.peek()?) {
+                next = joint;
+                let _ = self.try_next_token()?;
+            }
+
             self.write_token(out, next)?;
         }
 
@@ -429,3 +449,6 @@ fn write_header(class: Option<&str>, out: &mut dyn Write) -> io::Result<()> {
 fn write_footer(out: &mut dyn Write, playground_button: Option<&str>) -> io::Result<()> {
     write!(out, "</pre>{}</div>\n", if let Some(button) = playground_button { button } else { "" })
 }
+
+#[cfg(test)]
+mod tests;
diff --git a/src/librustdoc/html/highlight/tests.rs b/src/librustdoc/html/highlight/tests.rs
new file mode 100644
index 00000000000..01b25fd6be4
--- /dev/null
+++ b/src/librustdoc/html/highlight/tests.rs
@@ -0,0 +1,82 @@
+use rustc_ast::attr::with_session_globals;
+use rustc_session::parse::ParseSess;
+use rustc_span::edition::Edition;
+use rustc_span::FileName;
+
+use super::Classifier;
+
+fn highlight(src: &str) -> String {
+    let mut out = vec![];
+
+    with_session_globals(Edition::Edition2018, || {
+        let sess = ParseSess::with_silent_emitter();
+        let source_file = sess.source_map().new_source_file(
+            FileName::Custom(String::from("rustdoc-highlighting")),
+            src.to_owned(),
+        );
+
+        let mut classifier = Classifier::new(&sess, source_file);
+        classifier.write_source(&mut out).unwrap();
+    });
+
+    String::from_utf8(out).unwrap()
+}
+
+#[test]
+fn function() {
+    assert_eq!(
+        highlight("fn main() {}"),
+        r#"<span class="kw">fn</span> <span class="ident">main</span>() {}"#,
+    );
+}
+
+#[test]
+fn statement() {
+    assert_eq!(
+        highlight("let foo = true;"),
+        concat!(
+            r#"<span class="kw">let</span> <span class="ident">foo</span> "#,
+            r#"<span class="op">=</span> <span class="bool-val">true</span>;"#,
+        ),
+    );
+}
+
+#[test]
+fn inner_attr() {
+    assert_eq!(
+        highlight(r##"#![crate_type = "lib"]"##),
+        concat!(
+            r##"<span class="attribute">#![<span class="ident">crate_type</span> "##,
+            r##"<span class="op">=</span> <span class="string">&quot;lib&quot;</span>]</span>"##,
+        ),
+    );
+}
+
+#[test]
+fn outer_attr() {
+    assert_eq!(
+        highlight(r##"#[cfg(target_os = "linux")]"##),
+        concat!(
+            r##"<span class="attribute">#[<span class="ident">cfg</span>("##,
+            r##"<span class="ident">target_os</span> <span class="op">=</span> "##,
+            r##"<span class="string">&quot;linux&quot;</span>)]</span>"##,
+        ),
+    );
+}
+
+#[test]
+fn mac() {
+    assert_eq!(
+        highlight("mac!(foo bar)"),
+        concat!(
+            r#"<span class="macro">mac</span><span class="macro">!</span>("#,
+            r#"<span class="ident">foo</span> <span class="ident">bar</span>)"#,
+        ),
+    );
+}
+
+// Regression test for #72684
+#[test]
+fn andand() {
+    assert_eq!(highlight("&&"), r#"<span class="op">&amp;&amp;</span>"#);
+}
diff --git a/src/librustdoc/html/markdown.rs b/src/librustdoc/html/markdown.rs
index a0f8eb04e2e..d09fe454e13 100644
--- a/src/librustdoc/html/markdown.rs
+++ b/src/librustdoc/html/markdown.rs
@@ -292,7 +292,7 @@ impl<'a, I: Iterator<Item = Event<'a>>> Iterator for CodeBlocks<'_, 'a, I> {
 
         if let Some((s1, s2)) = tooltip {
             s.push_str(&highlight::render_with_highlighting(
-                &text,
+                text,
                 Some(&format!(
                     "rust-example-rendered{}",
                     if ignore != Ignore::None {
@@ -313,7 +313,7 @@ impl<'a, I: Iterator<Item = Event<'a>>> Iterator for CodeBlocks<'_, 'a, I> {
             Some(Event::Html(s.into()))
         } else {
             s.push_str(&highlight::render_with_highlighting(
-                &text,
+                text,
                 Some(&format!(
                     "rust-example-rendered{}",
                     if ignore != Ignore::None {
diff --git a/src/librustdoc/html/render.rs b/src/librustdoc/html/render.rs
index 9fa3f6cc396..81e3c490aaf 100644
--- a/src/librustdoc/html/render.rs
+++ b/src/librustdoc/html/render.rs
@@ -4526,7 +4526,12 @@ fn sidebar_foreign_type(buf: &mut Buffer, it: &clean::Item) {
 
 fn item_macro(w: &mut Buffer, cx: &Context, it: &clean::Item, t: &clean::Macro) {
     wrap_into_docblock(w, |w| {
-        w.write_str(&highlight::render_with_highlighting(&t.source, Some("macro"), None, None))
+        w.write_str(&highlight::render_with_highlighting(
+            t.source.clone(),
+            Some("macro"),
+            None,
+            None,
+        ))
     });
     document(w, cx, it)
 }
diff --git a/src/librustdoc/html/sources.rs b/src/librustdoc/html/sources.rs
index 03f79b93186..e3215921f12 100644
--- a/src/librustdoc/html/sources.rs
+++ b/src/librustdoc/html/sources.rs
@@ -75,7 +75,7 @@ impl<'a> SourceCollector<'a> {
             return Ok(());
         }
 
-        let contents = match fs::read_to_string(&p) {
+        let mut contents = match fs::read_to_string(&p) {
             Ok(contents) => contents,
             Err(e) => {
                 return Err(Error::new(e, &p));
@@ -83,8 +83,9 @@ impl<'a> SourceCollector<'a> {
         };
 
         // Remove the utf-8 BOM if any
-        let contents =
-            if contents.starts_with("\u{feff}") { &contents[3..] } else { &contents[..] };
+        if contents.starts_with("\u{feff}") {
+            contents.drain(..3);
+        }
 
         // Create the intermediate directories
         let mut cur = self.dst.clone();
@@ -122,7 +123,7 @@ impl<'a> SourceCollector<'a> {
             &self.scx.layout,
             &page,
             "",
-            |buf: &mut _| print_src(buf, &contents),
+            |buf: &mut _| print_src(buf, contents),
             &self.scx.style_files,
         );
         self.scx.fs.write(&cur, v.as_bytes())?;
@@ -160,7 +161,7 @@ where
 
 /// Wrapper struct to render the source code of a file. This will do things like
 /// adding line numbers to the left-hand side.
-fn print_src(buf: &mut Buffer, s: &str) {
+fn print_src(buf: &mut Buffer, s: String) {
     let lines = s.lines().count();
     let mut cols = 0;
     let mut tmp = lines;