Rollup merge of #135771 - GuillaumeGomez:jump-to-def-perf, r=fmease

[rustdoc] Add support for associated items in "jump to def" feature

Fixes https://github.com/rust-lang/rust/issues/135485.

r? ``@fmease``

2 files changed, 152 insertions, 92 deletions

diff --git a/src/librustdoc/html/highlight.rs b/src/librustdoc/html/highlight.rs
index 0e06361024b..8037654e037 100644
--- a/src/librustdoc/html/highlight.rs
+++ b/src/librustdoc/html/highlight.rs
@@ -1262,6 +1262,64 @@ fn string<W: Write>(
     }
 }
 
+fn generate_link_to_def(
+    out: &mut impl Write,
+    text_s: &str,
+    klass: Class,
+    href_context: &Option<HrefContext<'_, '_>>,
+    def_span: Span,
+    open_tag: bool,
+) -> bool {
+    if let Some(href_context) = href_context
+        && let Some(href) =
+            href_context.context.shared.span_correspondence_map.get(&def_span).and_then(|href| {
+                let context = href_context.context;
+                // FIXME: later on, it'd be nice to provide two links (if possible) for all items:
+                // one to the documentation page and one to the source definition.
+                // FIXME: currently, external items only generate a link to their documentation,
+                // a link to their definition can be generated using this:
+                // https://github.com/rust-lang/rust/blob/60f1a2fc4b535ead9c85ce085fdce49b1b097531/src/librustdoc/html/render/context.rs#L315-L338
+                match href {
+                    LinkFromSrc::Local(span) => {
+                        context.href_from_span_relative(*span, &href_context.current_href)
+                    }
+                    LinkFromSrc::External(def_id) => {
+                        format::href_with_root_path(*def_id, context, Some(href_context.root_path))
+                            .ok()
+                            .map(|(url, _, _)| url)
+                    }
+                    LinkFromSrc::Primitive(prim) => format::href_with_root_path(
+                        PrimitiveType::primitive_locations(context.tcx())[prim],
+                        context,
+                        Some(href_context.root_path),
+                    )
+                    .ok()
+                    .map(|(url, _, _)| url),
+                    LinkFromSrc::Doc(def_id) => {
+                        format::href_with_root_path(*def_id, context, Some(href_context.root_path))
+                            .ok()
+                            .map(|(doc_link, _, _)| doc_link)
+                    }
+                }
+            })
+    {
+        if !open_tag {
+            // We're already inside an element which has the same klass, no need to give it
+            // again.
+            write!(out, "<a href=\"{href}\">{text_s}").unwrap();
+        } else {
+            let klass_s = klass.as_html();
+            if klass_s.is_empty() {
+                write!(out, "<a href=\"{href}\">{text_s}").unwrap();
+            } else {
+                write!(out, "<a class=\"{klass_s}\" href=\"{href}\">{text_s}").unwrap();
+            }
+        }
+        return true;
+    }
+    false
+}
+
 /// This function writes `text` into `out` with some modifications depending on `klass`:
 ///
 /// * If `klass` is `None`, `text` is written into `out` with no modification.
@@ -1291,10 +1349,14 @@ fn string_without_closing_tag<T: Display>(
         return Some("</span>");
     };
 
+    let mut added_links = false;
     let mut text_s = text.to_string();
     if text_s.contains("::") {
+        let mut span = def_span.with_hi(def_span.lo());
         text_s = text_s.split("::").intersperse("::").fold(String::new(), |mut path, t| {
+            span = span.with_hi(span.hi() + BytePos(t.len() as _));
             match t {
+                "::" => write!(&mut path, "::"),
                 "self" | "Self" => write!(
                     &mut path,
                     "<span class=\"{klass}\">{t}</span>",
@@ -1307,58 +1369,24 @@ fn string_without_closing_tag<T: Display>(
                         klass = Class::KeyWord.as_html(),
                     )
                 }
-                t => write!(&mut path, "{t}"),
+                t => {
+                    if !t.is_empty()
+                        && generate_link_to_def(&mut path, t, klass, href_context, span, open_tag)
+                    {
+                        added_links = true;
+                        write!(&mut path, "</a>")
+                    } else {
+                        write!(&mut path, "{t}")
+                    }
+                }
             }
             .expect("Failed to build source HTML path");
+            span = span.with_lo(span.lo() + BytePos(t.len() as _));
             path
         });
     }
 
-    if let Some(href_context) = href_context
-        && let Some(href) = href_context.context.shared.span_correspondence_map.get(&def_span)
-        && let Some(href) = {
-            let context = href_context.context;
-            // FIXME: later on, it'd be nice to provide two links (if possible) for all items:
-            // one to the documentation page and one to the source definition.
-            // FIXME: currently, external items only generate a link to their documentation,
-            // a link to their definition can be generated using this:
-            // https://github.com/rust-lang/rust/blob/60f1a2fc4b535ead9c85ce085fdce49b1b097531/src/librustdoc/html/render/context.rs#L315-L338
-            match href {
-                LinkFromSrc::Local(span) => {
-                    context.href_from_span_relative(*span, &href_context.current_href)
-                }
-                LinkFromSrc::External(def_id) => {
-                    format::href_with_root_path(*def_id, context, Some(href_context.root_path))
-                        .ok()
-                        .map(|(url, _, _)| url)
-                }
-                LinkFromSrc::Primitive(prim) => format::href_with_root_path(
-                    PrimitiveType::primitive_locations(context.tcx())[prim],
-                    context,
-                    Some(href_context.root_path),
-                )
-                .ok()
-                .map(|(url, _, _)| url),
-                LinkFromSrc::Doc(def_id) => {
-                    format::href_with_root_path(*def_id, context, Some(href_context.root_path))
-                        .ok()
-                        .map(|(doc_link, _, _)| doc_link)
-                }
-            }
-        }
-    {
-        if !open_tag {
-            // We're already inside an element which has the same klass, no need to give it
-            // again.
-            write!(out, "<a href=\"{href}\">{text_s}").unwrap();
-        } else {
-            let klass_s = klass.as_html();
-            if klass_s.is_empty() {
-                write!(out, "<a href=\"{href}\">{text_s}").unwrap();
-            } else {
-                write!(out, "<a class=\"{klass_s}\" href=\"{href}\">{text_s}").unwrap();
-            }
-        }
+    if !added_links && generate_link_to_def(out, &text_s, klass, href_context, def_span, open_tag) {
         return Some("</a>");
     }
     if !open_tag {
diff --git a/src/librustdoc/html/render/span_map.rs b/src/librustdoc/html/render/span_map.rs
index 8bc2e0bd957..ef7ce33298d 100644
--- a/src/librustdoc/html/render/span_map.rs
+++ b/src/librustdoc/html/render/span_map.rs
@@ -2,11 +2,9 @@ use std::path::{Path, PathBuf};
 
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_hir::def::{DefKind, Res};
-use rustc_hir::def_id::{DefId, LOCAL_CRATE};
-use rustc_hir::intravisit::{self, Visitor};
-use rustc_hir::{
-    ExprKind, HirId, Item, ItemKind, Mod, Node, Pat, PatExpr, PatExprKind, PatKind, QPath,
-};
+use rustc_hir::def_id::{DefId, LOCAL_CRATE, LocalDefId};
+use rustc_hir::intravisit::{self, Visitor, VisitorExt};
+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;
@@ -67,7 +65,7 @@ struct SpanMapVisitor<'tcx> {
 
 impl SpanMapVisitor<'_> {
     /// This function is where we handle `hir::Path` elements and add them into the "span map".
-    fn handle_path(&mut self, path: &rustc_hir::Path<'_>) {
+    fn handle_path(&mut self, path: &rustc_hir::Path<'_>, only_use_last_segment: bool) {
         match path.res {
             // FIXME: For now, we handle `DefKind` if it's not a `DefKind::TyParam`.
             // Would be nice to support them too alongside the other `DefKind`
@@ -79,24 +77,36 @@ impl SpanMapVisitor<'_> {
                     LinkFromSrc::External(def_id)
                 };
                 // In case the path ends with generics, we remove them from the span.
-                let span = path
-                    .segments
-                    .last()
-                    .map(|last| {
-                        // In `use` statements, the included item is not in the path segments.
-                        // However, it doesn't matter because you can't have generics on `use`
-                        // statements.
-                        if path.span.contains(last.ident.span) {
-                            path.span.with_hi(last.ident.span.hi())
-                        } else {
-                            path.span
-                        }
-                    })
-                    .unwrap_or(path.span);
+                let span = if only_use_last_segment
+                    && let Some(path_span) = path.segments.last().map(|segment| segment.ident.span)
+                {
+                    path_span
+                } else {
+                    path.segments
+                        .last()
+                        .map(|last| {
+                            // In `use` statements, the included item is not in the path segments.
+                            // However, it doesn't matter because you can't have generics on `use`
+                            // statements.
+                            if path.span.contains(last.ident.span) {
+                                path.span.with_hi(last.ident.span.hi())
+                            } else {
+                                path.span
+                            }
+                        })
+                        .unwrap_or(path.span)
+                };
                 self.matches.insert(span, link);
             }
             Res::Local(_) if let Some(span) = self.tcx.hir_res_span(path.res) => {
-                self.matches.insert(path.span, LinkFromSrc::Local(clean::Span::new(span)));
+                let path_span = if only_use_last_segment
+                    && let Some(path_span) = path.segments.last().map(|segment| segment.ident.span)
+                {
+                    path_span
+                } else {
+                    path.span
+                };
+                self.matches.insert(path_span, LinkFromSrc::Local(clean::Span::new(span)));
             }
             Res::PrimTy(p) => {
                 // FIXME: Doesn't handle "path-like" primitives like arrays or tuples.
@@ -189,31 +199,23 @@ impl SpanMapVisitor<'_> {
             self.matches.insert(span, link);
         }
     }
+}
 
-    fn handle_pat(&mut self, p: &Pat<'_>) {
-        let mut check_qpath = |qpath, hir_id| match qpath {
-            QPath::TypeRelative(_, path) if matches!(path.res, Res::Err) => {
-                self.infer_id(path.hir_id, Some(hir_id), qpath.span());
-            }
-            QPath::Resolved(_, path) => self.handle_path(path),
-            _ => {}
-        };
-        match p.kind {
-            PatKind::Binding(_, _, _, Some(p)) => self.handle_pat(p),
-            PatKind::Struct(qpath, _, _) | PatKind::TupleStruct(qpath, _, _) => {
-                check_qpath(qpath, p.hir_id)
-            }
-            PatKind::Expr(PatExpr { kind: PatExprKind::Path(qpath), hir_id, .. }) => {
-                check_qpath(*qpath, *hir_id)
-            }
-            PatKind::Or(pats) => {
-                for pat in pats {
-                    self.handle_pat(pat);
-                }
-            }
-            _ => {}
+// This is a reimplementation of `hir_enclosing_body_owner` which allows to fail without
+// panicking.
+fn hir_enclosing_body_owner(tcx: TyCtxt<'_>, hir_id: HirId) -> Option<LocalDefId> {
+    for (_, node) in tcx.hir_parent_iter(hir_id) {
+        // FIXME: associated type impl items don't have an associated body, so we don't handle
+        // them currently.
+        if let Node::ImplItem(impl_item) = node
+            && matches!(impl_item.kind, rustc_hir::ImplItemKind::Type(_))
+        {
+            return None;
+        } else if let Some((def_id, _)) = node.associated_body() {
+            return Some(def_id);
         }
     }
+    None
 }
 
 impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
@@ -227,12 +229,42 @@ impl<'tcx> Visitor<'tcx> for SpanMapVisitor<'tcx> {
         if self.handle_macro(path.span) {
             return;
         }
-        self.handle_path(path);
+        self.handle_path(path, false);
         intravisit::walk_path(self, path);
     }
 
-    fn visit_pat(&mut self, p: &Pat<'tcx>) {
-        self.handle_pat(p);
+    fn visit_qpath(&mut self, qpath: &QPath<'tcx>, id: HirId, _span: Span) {
+        match *qpath {
+            QPath::TypeRelative(qself, path) => {
+                if matches!(path.res, Res::Err) {
+                    let tcx = self.tcx;
+                    if let Some(body_id) = hir_enclosing_body_owner(tcx, id) {
+                        let typeck_results = tcx.typeck_body(tcx.hir_body_owned_by(body_id).id());
+                        let path = rustc_hir::Path {
+                            // We change the span to not include parens.
+                            span: path.ident.span,
+                            res: typeck_results.qpath_res(qpath, id),
+                            segments: &[],
+                        };
+                        self.handle_path(&path, false);
+                    }
+                } else {
+                    self.infer_id(path.hir_id, Some(id), path.ident.span);
+                }
+
+                rustc_ast::visit::try_visit!(self.visit_ty_unambig(qself));
+                self.visit_path_segment(path);
+            }
+            QPath::Resolved(maybe_qself, path) => {
+                self.handle_path(path, true);
+
+                rustc_ast::visit::visit_opt!(self, visit_ty_unambig, maybe_qself);
+                if !self.handle_macro(path.span) {
+                    intravisit::walk_path(self, path);
+                }
+            }
+            _ => {}
+        }
     }
 
     fn visit_mod(&mut self, m: &'tcx Mod<'tcx>, span: Span, id: HirId) {