Remove quadratic attr source lookup

2 files changed, 27 insertions, 1 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index e7019e0c9..7a6a41dc2 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -294,6 +294,13 @@ impl Attrs {
         Arc::new(res)
     }
 
+    /// Constructs a map that maps the lowered `Attr`s in this `Attrs` back to its original syntax nodes.
+    ///
+    /// `owner` must be the original owner of the attributes.
+    pub fn source_map(&self, owner: &dyn AttrsOwner) -> AttrSourceMap {
+        AttrSourceMap { attrs: collect_attrs(owner).collect() }
+    }
+
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
         AttrQuery { attrs: self, key }
     }
@@ -366,6 +373,24 @@ fn inner_attributes(
     Some((attrs, docs))
 }
 
+pub struct AttrSourceMap {
+    attrs: Vec<Either<ast::Attr, ast::Comment>>,
+}
+
+impl AttrSourceMap {
+    /// Maps the lowered `Attr` back to its original syntax node.
+    ///
+    /// `attr` must come from the `owner` used for AttrSourceMap
+    ///
+    /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
+    /// the attribute represented by `Attr`.
+    pub fn source_of(&self, attr: &Attr) -> &Either<ast::Attr, ast::Comment> {
+        self.attrs
+            .get(attr.index as usize)
+            .unwrap_or_else(|| panic!("cannot find `Attr` at index {}", attr.index))
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
     index: u32,
diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index 0f1de4fb8..d4c367f66 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -153,6 +153,7 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
     if attributes.docs().map_or(true, |docs| !String::from(docs).contains(RUSTDOC_FENCE)) {
         return;
     }
+    let attrs_source_map = attributes.source_map(&owner);
 
     let mut inj = Injector::default();
     inj.add_unmapped("fn doctest() {\n");
@@ -165,7 +166,7 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
     let mut new_comments = Vec::new();
     let mut string;
     for attr in attributes.by_key("doc").attrs() {
-        let src = attr.to_src(&owner);
+        let src = attrs_source_map.source_of(&attr);
         let (line, range, prefix) = match &src {
             Either::Left(it) => {
                 string = match find_doc_string_in_attr(attr, it) {