5 files changed, 216 insertions, 155 deletions

diff --git a/crates/hir/src/attrs.rs b/crates/hir/src/attrs.rs
index dab8da7bb..4a11622fc 100644
--- a/crates/hir/src/attrs.rs
+++ b/crates/hir/src/attrs.rs
@@ -1,6 +1,6 @@
 //! Attributes & documentation for hir types.
 use hir_def::{
-    attr::{Attrs, Documentation},
+    attr::{AttrsWithOwner, Documentation},
     path::ModPath,
     per_ns::PerNs,
     resolver::HasResolver,
@@ -16,7 +16,7 @@ use crate::{
 };
 
 pub trait HasAttrs {
-    fn attrs(self, db: &dyn HirDatabase) -> Attrs;
+    fn attrs(self, db: &dyn HirDatabase) -> AttrsWithOwner;
     fn docs(self, db: &dyn HirDatabase) -> Option<Documentation>;
     fn resolve_doc_path(
         self,
@@ -36,7 +36,7 @@ pub enum Namespace {
 macro_rules! impl_has_attrs {
     ($(($def:ident, $def_id:ident),)*) => {$(
         impl HasAttrs for $def {
-            fn attrs(self, db: &dyn HirDatabase) -> Attrs {
+            fn attrs(self, db: &dyn HirDatabase) -> AttrsWithOwner {
                 let def = AttrDefId::$def_id(self.into());
                 db.attrs(def)
             }
@@ -70,7 +70,7 @@ impl_has_attrs![
 macro_rules! impl_has_attrs_enum {
     ($($variant:ident),* for $enum:ident) => {$(
         impl HasAttrs for $variant {
-            fn attrs(self, db: &dyn HirDatabase) -> Attrs {
+            fn attrs(self, db: &dyn HirDatabase) -> AttrsWithOwner {
                 $enum::$variant(self).attrs(db)
             }
             fn docs(self, db: &dyn HirDatabase) -> Option<Documentation> {
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 5ebd0a3b8..300087f1f 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -89,7 +89,7 @@ pub use crate::{
 pub use {
     hir_def::{
         adt::StructKind,
-        attr::{Attr, Attrs, Documentation},
+        attr::{Attr, Attrs, AttrsWithOwner, Documentation},
         body::scope::ExprScopes,
         find_path::PrefixKind,
         import_map,
diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index beeaaf117..2c10f46d8 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -21,7 +21,7 @@ use crate::{
     item_tree::{ItemTreeId, ItemTreeNode},
     nameres::ModuleSource,
     path::{ModPath, PathKind},
-    src::HasChildSource,
+    src::{HasChildSource, HasSource},
     AdtId, AttrDefId, EnumId, GenericParamId, HasModule, LocalEnumVariantId, LocalFieldId, Lookup,
     VariantId,
 };
@@ -51,6 +51,12 @@ pub(crate) struct RawAttrs {
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 pub struct Attrs(RawAttrs);
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AttrsWithOwner {
+    attrs: Attrs,
+    owner: AttrDefId,
+}
+
 impl ops::Deref for RawAttrs {
     type Target = [Attr];
 
@@ -73,6 +79,14 @@ impl ops::Deref for Attrs {
     }
 }
 
+impl ops::Deref for AttrsWithOwner {
+    type Target = Attrs;
+
+    fn deref(&self) -> &Attrs {
+        &self.attrs
+    }
+}
+
 impl RawAttrs {
     pub(crate) const EMPTY: Self = Self { entries: None };
 
@@ -169,78 +183,6 @@ impl RawAttrs {
 impl Attrs {
     pub const EMPTY: Self = Self(RawAttrs::EMPTY);
 
-    pub(crate) fn attrs_query(db: &dyn DefDatabase, def: AttrDefId) -> Attrs {
-        let raw_attrs = match def {
-            AttrDefId::ModuleId(module) => {
-                let def_map = module.def_map(db);
-                let mod_data = &def_map[module.local_id];
-                match mod_data.declaration_source(db) {
-                    Some(it) => {
-                        let raw_attrs = RawAttrs::from_attrs_owner(
-                            db,
-                            it.as_ref().map(|it| it as &dyn ast::AttrsOwner),
-                        );
-                        match mod_data.definition_source(db) {
-                            InFile { file_id, value: ModuleSource::SourceFile(file) } => raw_attrs
-                                .merge(RawAttrs::from_attrs_owner(db, InFile::new(file_id, &file))),
-                            _ => raw_attrs,
-                        }
-                    }
-                    None => RawAttrs::from_attrs_owner(
-                        db,
-                        mod_data.definition_source(db).as_ref().map(|src| match src {
-                            ModuleSource::SourceFile(file) => file as &dyn ast::AttrsOwner,
-                            ModuleSource::Module(module) => module as &dyn ast::AttrsOwner,
-                            ModuleSource::BlockExpr(block) => block as &dyn ast::AttrsOwner,
-                        }),
-                    ),
-                }
-            }
-            AttrDefId::FieldId(it) => {
-                return db.fields_attrs(it.parent)[it.local_id].clone();
-            }
-            AttrDefId::EnumVariantId(it) => {
-                return db.variants_attrs(it.parent)[it.local_id].clone();
-            }
-            AttrDefId::AdtId(it) => match it {
-                AdtId::StructId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-                AdtId::EnumId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-                AdtId::UnionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            },
-            AttrDefId::TraitId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::MacroDefId(it) => it
-                .ast_id()
-                .left()
-                .map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db)),
-            AttrDefId::ImplId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::ConstId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::StaticId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::FunctionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::TypeAliasId(it) => attrs_from_item_tree(it.lookup(db).id, db),
-            AttrDefId::GenericParamId(it) => match it {
-                GenericParamId::TypeParamId(it) => {
-                    let src = it.parent.child_source(db);
-                    RawAttrs::from_attrs_owner(
-                        db,
-                        src.with_value(
-                            src.value[it.local_id].as_ref().either(|it| it as _, |it| it as _),
-                        ),
-                    )
-                }
-                GenericParamId::LifetimeParamId(it) => {
-                    let src = it.parent.child_source(db);
-                    RawAttrs::from_attrs_owner(db, src.with_value(&src.value[it.local_id]))
-                }
-                GenericParamId::ConstParamId(it) => {
-                    let src = it.parent.child_source(db);
-                    RawAttrs::from_attrs_owner(db, src.with_value(&src.value[it.local_id]))
-                }
-            },
-        };
-
-        raw_attrs.filter(db, def.krate(db))
-    }
-
     pub(crate) fn variants_attrs_query(
         db: &dyn DefDatabase,
         e: EnumId,
@@ -281,56 +223,6 @@ 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.
-    // FIXME: figure out a better api that doesnt require the for_module hack
-    pub fn source_map(&self, owner: InFile<&dyn ast::AttrsOwner>) -> AttrSourceMap {
-        // FIXME: This doesn't work correctly for modules, as the attributes there can have up to
-        // two different owners
-        AttrSourceMap {
-            attrs: collect_attrs(owner.value)
-                .map(|attr| InFile::new(owner.file_id, attr))
-                .collect(),
-        }
-    }
-
-    pub fn source_map_for_module(
-        &self,
-        db: &dyn DefDatabase,
-        module: crate::ModuleId,
-    ) -> AttrSourceMap {
-        let def_map = module.def_map(db);
-        let mod_data = &def_map[module.local_id];
-        let attrs = match mod_data.declaration_source(db) {
-            Some(it) => {
-                let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
-                    .map(|attr| InFile::new(it.file_id, attr))
-                    .collect();
-                if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
-                    mod_data.definition_source(db)
-                {
-                    attrs.extend(
-                        collect_attrs(&file as &dyn ast::AttrsOwner)
-                            .map(|attr| InFile::new(file_id, attr)),
-                    )
-                }
-                attrs
-            }
-            None => {
-                let InFile { file_id, value } = mod_data.definition_source(db);
-                match &value {
-                    ModuleSource::SourceFile(file) => collect_attrs(file as &dyn ast::AttrsOwner),
-                    ModuleSource::Module(module) => collect_attrs(module as &dyn ast::AttrsOwner),
-                    ModuleSource::BlockExpr(block) => collect_attrs(block as &dyn ast::AttrsOwner),
-                }
-                .map(|attr| InFile::new(file_id, attr))
-                .collect()
-            }
-        };
-        AttrSourceMap { attrs }
-    }
-
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
         AttrQuery { attrs: self, key }
     }
@@ -387,6 +279,180 @@ impl Attrs {
     }
 }
 
+impl AttrsWithOwner {
+    pub(crate) fn attrs_query(db: &dyn DefDatabase, def: AttrDefId) -> Self {
+        // FIXME: this should use `Trace` to avoid duplication in `source_map` below
+        let raw_attrs = match def {
+            AttrDefId::ModuleId(module) => {
+                let def_map = module.def_map(db);
+                let mod_data = &def_map[module.local_id];
+                match mod_data.declaration_source(db) {
+                    Some(it) => {
+                        let raw_attrs = RawAttrs::from_attrs_owner(
+                            db,
+                            it.as_ref().map(|it| it as &dyn ast::AttrsOwner),
+                        );
+                        match mod_data.definition_source(db) {
+                            InFile { file_id, value: ModuleSource::SourceFile(file) } => raw_attrs
+                                .merge(RawAttrs::from_attrs_owner(db, InFile::new(file_id, &file))),
+                            _ => raw_attrs,
+                        }
+                    }
+                    None => RawAttrs::from_attrs_owner(
+                        db,
+                        mod_data.definition_source(db).as_ref().map(|src| match src {
+                            ModuleSource::SourceFile(file) => file as &dyn ast::AttrsOwner,
+                            ModuleSource::Module(module) => module as &dyn ast::AttrsOwner,
+                            ModuleSource::BlockExpr(block) => block as &dyn ast::AttrsOwner,
+                        }),
+                    ),
+                }
+            }
+            AttrDefId::FieldId(it) => {
+                return Self { attrs: db.fields_attrs(it.parent)[it.local_id].clone(), owner: def };
+            }
+            AttrDefId::EnumVariantId(it) => {
+                return Self {
+                    attrs: db.variants_attrs(it.parent)[it.local_id].clone(),
+                    owner: def,
+                };
+            }
+            AttrDefId::AdtId(it) => match it {
+                AdtId::StructId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+                AdtId::EnumId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+                AdtId::UnionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            },
+            AttrDefId::TraitId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::MacroDefId(it) => it
+                .ast_id()
+                .left()
+                .map_or_else(Default::default, |ast_id| attrs_from_ast(ast_id, db)),
+            AttrDefId::ImplId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::ConstId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::StaticId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::FunctionId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::TypeAliasId(it) => attrs_from_item_tree(it.lookup(db).id, db),
+            AttrDefId::GenericParamId(it) => match it {
+                GenericParamId::TypeParamId(it) => {
+                    let src = it.parent.child_source(db);
+                    RawAttrs::from_attrs_owner(
+                        db,
+                        src.with_value(
+                            src.value[it.local_id].as_ref().either(|it| it as _, |it| it as _),
+                        ),
+                    )
+                }
+                GenericParamId::LifetimeParamId(it) => {
+                    let src = it.parent.child_source(db);
+                    RawAttrs::from_attrs_owner(db, src.with_value(&src.value[it.local_id]))
+                }
+                GenericParamId::ConstParamId(it) => {
+                    let src = it.parent.child_source(db);
+                    RawAttrs::from_attrs_owner(db, src.with_value(&src.value[it.local_id]))
+                }
+            },
+        };
+
+        let attrs = raw_attrs.filter(db, def.krate(db));
+        Self { attrs, owner: def }
+    }
+
+    pub fn source_map(&self, db: &dyn DefDatabase) -> AttrSourceMap {
+        let owner = match self.owner {
+            AttrDefId::ModuleId(module) => {
+                // Modules can have 2 attribute owners (the `mod x;` item, and the module file itself).
+
+                let def_map = module.def_map(db);
+                let mod_data = &def_map[module.local_id];
+                let attrs = match mod_data.declaration_source(db) {
+                    Some(it) => {
+                        let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
+                            .map(|attr| InFile::new(it.file_id, attr))
+                            .collect();
+                        if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
+                            mod_data.definition_source(db)
+                        {
+                            attrs.extend(
+                                collect_attrs(&file as &dyn ast::AttrsOwner)
+                                    .map(|attr| InFile::new(file_id, attr)),
+                            )
+                        }
+                        attrs
+                    }
+                    None => {
+                        let InFile { file_id, value } = mod_data.definition_source(db);
+                        match &value {
+                            ModuleSource::SourceFile(file) => {
+                                collect_attrs(file as &dyn ast::AttrsOwner)
+                            }
+                            ModuleSource::Module(module) => {
+                                collect_attrs(module as &dyn ast::AttrsOwner)
+                            }
+                            ModuleSource::BlockExpr(block) => {
+                                collect_attrs(block as &dyn ast::AttrsOwner)
+                            }
+                        }
+                        .map(|attr| InFile::new(file_id, attr))
+                        .collect()
+                    }
+                };
+                return AttrSourceMap { attrs };
+            }
+            AttrDefId::FieldId(id) => {
+                id.parent.child_source(db).map(|source| match &source[id.local_id] {
+                    Either::Left(field) => ast::AttrsOwnerNode::new(field.clone()),
+                    Either::Right(field) => ast::AttrsOwnerNode::new(field.clone()),
+                })
+            }
+            AttrDefId::AdtId(adt) => match adt {
+                AdtId::StructId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+                AdtId::UnionId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+                AdtId::EnumId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            },
+            AttrDefId::FunctionId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::EnumVariantId(id) => id
+                .parent
+                .child_source(db)
+                .map(|source| ast::AttrsOwnerNode::new(source[id.local_id].clone())),
+            AttrDefId::StaticId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::ConstId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::TraitId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::TypeAliasId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::MacroDefId(id) => match id.ast_id() {
+                Either::Left(it) => {
+                    it.with_value(ast::AttrsOwnerNode::new(it.to_node(db.upcast())))
+                }
+                Either::Right(it) => {
+                    it.with_value(ast::AttrsOwnerNode::new(it.to_node(db.upcast())))
+                }
+            },
+            AttrDefId::ImplId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
+            AttrDefId::GenericParamId(id) => match id {
+                GenericParamId::TypeParamId(id) => {
+                    id.parent.child_source(db).map(|source| match &source[id.local_id] {
+                        Either::Left(id) => ast::AttrsOwnerNode::new(id.clone()),
+                        Either::Right(id) => ast::AttrsOwnerNode::new(id.clone()),
+                    })
+                }
+                GenericParamId::LifetimeParamId(id) => id
+                    .parent
+                    .child_source(db)
+                    .map(|source| ast::AttrsOwnerNode::new(source[id.local_id].clone())),
+                GenericParamId::ConstParamId(id) => id
+                    .parent
+                    .child_source(db)
+                    .map(|source| ast::AttrsOwnerNode::new(source[id.local_id].clone())),
+            },
+        };
+
+        AttrSourceMap {
+            attrs: collect_attrs(&owner.value)
+                .map(|attr| InFile::new(owner.file_id, attr))
+                .collect(),
+        }
+    }
+}
+
 fn inner_attributes(
     syntax: &SyntaxNode,
 ) -> Option<(impl Iterator<Item = ast::Attr>, impl Iterator<Item = ast::Comment>)> {
diff --git a/crates/hir_def/src/db.rs b/crates/hir_def/src/db.rs
index 276caf5b3..53df85089 100644
--- a/crates/hir_def/src/db.rs
+++ b/crates/hir_def/src/db.rs
@@ -8,7 +8,7 @@ use syntax::SmolStr;
 
 use crate::{
     adt::{EnumData, StructData},
-    attr::Attrs,
+    attr::{Attrs, AttrsWithOwner},
     body::{scope::ExprScopes, Body, BodySourceMap},
     data::{ConstData, FunctionData, ImplData, StaticData, TraitData, TypeAliasData},
     generics::GenericParams,
@@ -120,8 +120,8 @@ pub trait DefDatabase: InternDatabase + AstDatabase + Upcast<dyn AstDatabase> {
     #[salsa::invoke(Attrs::fields_attrs_query)]
     fn fields_attrs(&self, def: VariantId) -> Arc<ArenaMap<LocalFieldId, Attrs>>;
 
-    #[salsa::invoke(Attrs::attrs_query)]
-    fn attrs(&self, def: AttrDefId) -> Attrs;
+    #[salsa::invoke(AttrsWithOwner::attrs_query)]
+    fn attrs(&self, def: AttrDefId) -> AttrsWithOwner;
 
     #[salsa::invoke(LangItems::crate_lang_items_query)]
     fn crate_lang_items(&self, krate: CrateId) -> Arc<LangItems>;
diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index d2e09dbb6..00493a6b5 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -6,7 +6,7 @@ use either::Either;
 use hir::{HasAttrs, InFile, Semantics};
 use ide_db::{call_info::ActiveParameter, defs::Definition, SymbolKind};
 use syntax::{
-    ast::{self, AstNode, AttrsOwner, AttrsOwnerNode, DocCommentsOwner},
+    ast::{self, AstNode},
     match_ast, AstToken, NodeOrToken, SyntaxNode, SyntaxToken, TextRange, TextSize,
 };
 
@@ -92,24 +92,24 @@ const RUSTDOC_FENCE_TOKENS: &[&'static str] = &[
 fn doc_attributes<'node>(
     sema: &Semantics<RootDatabase>,
     node: &'node SyntaxNode,
-) -> Option<(AttrsOwnerNode, hir::Attrs, Definition)> {
+) -> Option<(hir::AttrsWithOwner, Definition)> {
     match_ast! {
         match node {
-            ast::SourceFile(it)  => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Module(def)))),
-            ast::Module(it)      => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Module(def)))),
-            ast::Fn(it)          => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Function(def)))),
-            ast::Struct(it)      => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Struct(def))))),
-            ast::Union(it)       => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Union(def))))),
-            ast::Enum(it)        => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Enum(def))))),
-            ast::Variant(it)     => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Variant(def)))),
-            ast::Trait(it)       => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Trait(def)))),
-            ast::Static(it)      => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Static(def)))),
-            ast::Const(it)       => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Const(def)))),
-            ast::TypeAlias(it)   => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::TypeAlias(def)))),
-            ast::Impl(it)        => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::SelfType(def))),
-            ast::RecordField(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::Field(def))),
-            ast::TupleField(it)  => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::Field(def))),
-            ast::MacroRules(it)  => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db), Definition::Macro(def))),
+            ast::SourceFile(it)  => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Module(def)))),
+            ast::Module(it)      => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Module(def)))),
+            ast::Fn(it)          => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Function(def)))),
+            ast::Struct(it)      => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Struct(def))))),
+            ast::Union(it)       => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Union(def))))),
+            ast::Enum(it)        => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Adt(hir::Adt::Enum(def))))),
+            ast::Variant(it)     => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Variant(def)))),
+            ast::Trait(it)       => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Trait(def)))),
+            ast::Static(it)      => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Static(def)))),
+            ast::Const(it)       => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::Const(def)))),
+            ast::TypeAlias(it)   => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::ModuleDef(hir::ModuleDef::TypeAlias(def)))),
+            ast::Impl(it)        => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::SelfType(def))),
+            ast::RecordField(it) => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::Field(def))),
+            ast::TupleField(it)  => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::Field(def))),
+            ast::MacroRules(it)  => sema.to_def(&it).map(|def| (def.attrs(sema.db), Definition::Macro(def))),
             // ast::MacroDef(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
             // ast::Use(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
             _ => return None
@@ -123,7 +123,7 @@ pub(super) fn doc_comment(
     sema: &Semantics<RootDatabase>,
     node: InFile<&SyntaxNode>,
 ) {
-    let (owner, attributes, def) = match doc_attributes(sema, node.value) {
+    let (attributes, def) = match doc_attributes(sema, node.value) {
         Some(it) => it,
         None => return,
     };
@@ -131,12 +131,7 @@ pub(super) fn doc_comment(
     let mut inj = Injector::default();
     inj.add_unmapped("fn doctest() {\n");
 
-    let attrs_source_map = match def {
-        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
-            attributes.source_map_for_module(sema.db, module.into())
-        }
-        _ => attributes.source_map(node.with_value(&owner)),
-    };
+    let attrs_source_map = attributes.source_map(sema.db);
 
     let mut is_codeblock = false;
     let mut is_doctest = false;