Rollup merge of #147189 - yotamofek:pr/rustdoc/highlight-optimizations-2, r=GuillaumeGomez

Replace `rustc_span::Span` with a stripped down version for librustdoc's highlighter

While profiling rustdoc's syntax highlighter, I noticed a lot of time being spent in the `Span` interner, due to the highlighter creating a lot of (new) spans.
Since the only data from the `Span` that we use is the `hi` and `lo` byte positions - I replaced the regular `Span` with a simple one with two fields, and in my benchmarks it seemed to make a big dent in the highlighter's perf, so thought I would see what the perf runner says.

7 files changed, 80 insertions, 19 deletions

diff --git a/src/bootstrap/src/core/builder/mod.rs b/src/bootstrap/src/core/builder/mod.rs
index 006dea4b98d..fc06db8f80b 100644
--- a/src/bootstrap/src/core/builder/mod.rs
+++ b/src/bootstrap/src/core/builder/mod.rs
@@ -1145,7 +1145,7 @@ impl<'a> Builder<'a> {
                 test::RunMakeCargo,
             ),
             Kind::Miri => describe!(test::Crate),
-            Kind::Bench => describe!(test::Crate, test::CrateLibrustc),
+            Kind::Bench => describe!(test::Crate, test::CrateLibrustc, test::CrateRustdoc),
             Kind::Doc => describe!(
                 doc::UnstableBook,
                 doc::UnstableBookGen,
diff --git a/src/librustdoc/html/highlight.rs b/src/librustdoc/html/highlight.rs
index fad15573cde..1dcb4dcc3ff 100644
--- a/src/librustdoc/html/highlight.rs
+++ b/src/librustdoc/html/highlight.rs
@@ -12,15 +12,16 @@ use std::iter;
 
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_lexer::{Cursor, FrontmatterAllowed, LiteralKind, TokenKind};
+use rustc_span::BytePos;
 use rustc_span::edition::Edition;
 use rustc_span::symbol::Symbol;
-use rustc_span::{BytePos, DUMMY_SP, Span};
 
 use super::format;
 use crate::clean::PrimitiveType;
 use crate::display::Joined as _;
 use crate::html::escape::EscapeBodyText;
 use crate::html::macro_expansion::ExpandedCode;
+use crate::html::render::span_map::{DUMMY_SP, Span};
 use crate::html::render::{Context, LinkFromSrc};
 
 /// This type is needed in case we want to render links on items to allow to go to their definition.
diff --git a/src/librustdoc/html/highlight/tests.rs b/src/librustdoc/html/highlight/tests.rs
index 2603e887bea..4d1bee9b3a1 100644
--- a/src/librustdoc/html/highlight/tests.rs
+++ b/src/librustdoc/html/highlight/tests.rs
@@ -1,6 +1,7 @@
 use expect_test::expect_file;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_span::create_default_session_globals_then;
+use test::Bencher;
 
 use super::{DecorationInfo, write_code};
 
@@ -81,3 +82,16 @@ let a = 4;";
         expect_file!["fixtures/decorations.html"].assert_eq(&html);
     });
 }
+
+#[bench]
+fn bench_html_highlighting(b: &mut Bencher) {
+    let src = include_str!("../../../../compiler/rustc_ast/src/visit.rs");
+
+    create_default_session_globals_then(|| {
+        b.iter(|| {
+            let mut out = String::new();
+            write_code(&mut out, src, None, None, None);
+            out
+        });
+    });
+}
diff --git a/src/librustdoc/html/render/context.rs b/src/librustdoc/html/render/context.rs
index 5f92ab2fada..4c06d0da470 100644
--- a/src/librustdoc/html/render/context.rs
+++ b/src/librustdoc/html/render/context.rs
@@ -30,6 +30,7 @@ use crate::formats::item_type::ItemType;
 use crate::html::escape::Escape;
 use crate::html::macro_expansion::ExpandedCode;
 use crate::html::markdown::{self, ErrorCodes, IdMap, plain_text_summary};
+use crate::html::render::span_map::Span;
 use crate::html::render::write_shared::write_shared;
 use crate::html::url_parts_builder::UrlPartsBuilder;
 use crate::html::{layout, sources, static_files};
@@ -139,7 +140,7 @@ pub(crate) struct SharedContext<'tcx> {
 
     /// Correspondence map used to link types used in the source code pages to allow to click on
     /// links to jump to the type's definition.
-    pub(crate) span_correspondence_map: FxHashMap<rustc_span::Span, LinkFromSrc>,
+    pub(crate) span_correspondence_map: FxHashMap<Span, LinkFromSrc>,
     pub(crate) expanded_codes: FxHashMap<BytePos, Vec<ExpandedCode>>,
     /// The [`Cache`] used during rendering.
     pub(crate) cache: Cache,
diff --git a/src/librustdoc/html/render/mod.rs b/src/librustdoc/html/render/mod.rs
index 97dcaf57cdf..d6371e4dbab 100644
--- a/src/librustdoc/html/render/mod.rs
+++ b/src/librustdoc/html/render/mod.rs
@@ -36,7 +36,7 @@ mod ordered_json;
 mod print_item;
 pub(crate) mod sidebar;
 mod sorted_template;
-mod span_map;
+pub(crate) mod span_map;
 mod type_layout;
 mod write_shared;
 
diff --git a/src/librustdoc/html/render/span_map.rs b/src/librustdoc/html/render/span_map.rs
index ef7ce33298d..bc9417b1bb1 100644
--- a/src/librustdoc/html/render/span_map.rs
+++ b/src/librustdoc/html/render/span_map.rs
@@ -8,11 +8,48 @@ use rustc_hir::{ExprKind, HirId, Item, ItemKind, Mod, Node, QPath};
 use rustc_middle::hir::nested_filter;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::hygiene::MacroKind;
-use rustc_span::{BytePos, ExpnKind, Span};
+use rustc_span::{BytePos, ExpnKind};
 
 use crate::clean::{self, PrimitiveType, rustc_span};
 use crate::html::sources;
 
+/// This is a stripped down version of [`rustc_span::Span`] that only contains the start and end byte positions of the span.
+///
+/// Profiling showed that the `Span` interner was taking up a lot of the run-time when highlighting, and since we
+/// never actually use the context and parent that are stored in a normal `Span`, we can replace its usages with this
+/// one, which is much cheaper to construct.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct Span {
+    lo: BytePos,
+    hi: BytePos,
+}
+
+impl From<rustc_span::Span> for Span {
+    fn from(value: rustc_span::Span) -> Self {
+        Self { lo: value.lo(), hi: value.hi() }
+    }
+}
+
+impl Span {
+    pub(crate) fn lo(self) -> BytePos {
+        self.lo
+    }
+
+    pub(crate) fn hi(self) -> BytePos {
+        self.hi
+    }
+
+    pub(crate) fn with_lo(self, lo: BytePos) -> Self {
+        Self { lo, hi: self.hi() }
+    }
+
+    pub(crate) fn with_hi(self, hi: BytePos) -> Self {
+        Self { lo: self.lo(), hi }
+    }
+}
+
+pub(crate) const DUMMY_SP: Span = Span { lo: BytePos(0), hi: BytePos(0) };
+
 /// This enum allows us to store two different kinds of information:
 ///
 /// In case the `span` definition comes from the same crate, we can simply get the `span` and use
@@ -96,7 +133,7 @@ impl SpanMapVisitor<'_> {
                         })
                         .unwrap_or(path.span)
                 };
-                self.matches.insert(span, link);
+                self.matches.insert(span.into(), link);
             }
             Res::Local(_) if let Some(span) = self.tcx.hir_res_span(path.res) => {
                 let path_span = if only_use_last_segment
@@ -106,11 +143,12 @@ impl SpanMapVisitor<'_> {
                 } else {
                     path.span
                 };
-                self.matches.insert(path_span, LinkFromSrc::Local(clean::Span::new(span)));
+                self.matches.insert(path_span.into(), LinkFromSrc::Local(clean::Span::new(span)));
             }
             Res::PrimTy(p) => {
                 // FIXME: Doesn't handle "path-like" primitives like arrays or tuples.
-                self.matches.insert(path.span, LinkFromSrc::Primitive(PrimitiveType::from(p)));
+                self.matches
+                    .insert(path.span.into(), LinkFromSrc::Primitive(PrimitiveType::from(p)));
             }
             Res::Err => {}
             _ => {}
@@ -127,7 +165,7 @@ impl SpanMapVisitor<'_> {
             if cspan.inner().is_dummy() || cspan.cnum(self.tcx.sess) != LOCAL_CRATE {
                 return;
             }
-            self.matches.insert(span, LinkFromSrc::Doc(item.owner_id.to_def_id()));
+            self.matches.insert(span.into(), LinkFromSrc::Doc(item.owner_id.to_def_id()));
         }
     }
 
@@ -138,7 +176,7 @@ impl SpanMapVisitor<'_> {
     /// so, we loop until we find the macro definition by using `outer_expn_data` in a loop.
     /// Finally, we get the information about the macro itself (`span` if "local", `DefId`
     /// otherwise) and store it inside the span map.
-    fn handle_macro(&mut self, span: Span) -> bool {
+    fn handle_macro(&mut self, span: rustc_span::Span) -> bool {
         if !span.from_expansion() {
             return false;
         }
@@ -176,7 +214,7 @@ impl SpanMapVisitor<'_> {
         // The "call_site" includes the whole macro with its "arguments". We only want
         // the macro name.
         let new_span = new_span.with_hi(new_span.lo() + BytePos(macro_name.len() as u32));
-        self.matches.insert(new_span, link_from_src);
+        self.matches.insert(new_span.into(), link_from_src);
         true
     }
 
@@ -233,7 +271,7 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
         intravisit::walk_path(self, path);
     }
 
-    fn visit_qpath(&mut self, qpath: &QPath<'tcx>, id: HirId, _span: Span) {
+    fn visit_qpath(&mut self, qpath: &QPath<'tcx>, id: HirId, _span: rustc_span::Span) {
         match *qpath {
             QPath::TypeRelative(qself, path) => {
                 if matches!(path.res, Res::Err) {
@@ -249,7 +287,7 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
                         self.handle_path(&path, false);
                     }
                 } else {
-                    self.infer_id(path.hir_id, Some(id), path.ident.span);
+                    self.infer_id(path.hir_id, Some(id), path.ident.span.into());
                 }
 
                 rustc_ast::visit::try_visit!(self.visit_ty_unambig(qself));
@@ -267,7 +305,7 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
         }
     }
 
-    fn visit_mod(&mut self, m: &'tcx Mod<'tcx>, span: Span, id: HirId) {
+    fn visit_mod(&mut self, m: &'tcx Mod<'tcx>, span: rustc_span::Span, id: HirId) {
         // To make the difference between "mod foo {}" and "mod foo;". In case we "import" another
         // file, we want to link to it. Otherwise no need to create a link.
         if !span.overlaps(m.spans.inner_span) {
@@ -275,8 +313,10 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
             // name only and not all the "mod foo;".
             if let Node::Item(item) = self.tcx.hir_node(id) {
                 let (ident, _) = item.expect_mod();
-                self.matches
-                    .insert(ident.span, LinkFromSrc::Local(clean::Span::new(m.spans.inner_span)));
+                self.matches.insert(
+                    ident.span.into(),
+                    LinkFromSrc::Local(clean::Span::new(m.spans.inner_span)),
+                );
             }
         } else {
             // If it's a "mod foo {}", we want to look to its documentation page.
@@ -288,9 +328,9 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
     fn visit_expr(&mut self, expr: &'tcx rustc_hir::Expr<'tcx>) {
         match expr.kind {
             ExprKind::MethodCall(segment, ..) => {
-                self.infer_id(segment.hir_id, Some(expr.hir_id), segment.ident.span)
+                self.infer_id(segment.hir_id, Some(expr.hir_id), segment.ident.span.into())
             }
-            ExprKind::Call(call, ..) => self.infer_id(call.hir_id, None, call.span),
+            ExprKind::Call(call, ..) => self.infer_id(call.hir_id, None, call.span.into()),
             _ => {
                 if self.handle_macro(expr.span) {
                     // We don't want to go deeper into the macro.
diff --git a/src/librustdoc/html/sources.rs b/src/librustdoc/html/sources.rs
index 9c5518a780e..c79f63fbc20 100644
--- a/src/librustdoc/html/sources.rs
+++ b/src/librustdoc/html/sources.rs
@@ -348,7 +348,12 @@ pub(crate) fn print_src(
         highlight::write_code(
             fmt,
             s,
-            Some(highlight::HrefContext { context, file_span, root_path, current_href }),
+            Some(highlight::HrefContext {
+                context,
+                file_span: file_span.into(),
+                root_path,
+                current_href,
+            }),
             Some(decoration_info),
             Some(line_info),
         );