4 files changed, 240 insertions, 146 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 0360fb627..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,77 +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().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,
@@ -280,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 }
     }
@@ -386,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/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 4bcdc0fc8..46a3c60cd 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -1294,29 +1294,37 @@ impl ModCollector<'_, '_> {
                 let db = self.def_collector.db;
                 match self.mod_dir.resolve_declaration(db, self.file_id, &module.name, path_attr) {
                     Ok((file_id, is_mod_rs, mod_dir)) => {
-                        let module_id = self.push_child_module(
-                            module.name.clone(),
-                            ast_id,
-                            Some((file_id, is_mod_rs)),
-                            &self.item_tree[module.visibility],
-                        );
                         let item_tree = db.file_item_tree(file_id.into());
-                        ModCollector {
-                            def_collector: &mut *self.def_collector,
-                            macro_depth: self.macro_depth,
-                            module_id,
-                            file_id: file_id.into(),
-                            item_tree: &item_tree,
-                            mod_dir,
-                        }
-                        .collect(item_tree.top_level_items());
-                        if is_macro_use
-                            || item_tree
-                                .top_level_attrs(db, self.def_collector.def_map.krate)
-                                .by_key("macro_use")
-                                .exists()
+                        if item_tree
+                            .top_level_attrs(db, self.def_collector.def_map.krate)
+                            .cfg()
+                            .map_or(true, |cfg| {
+                                self.def_collector.cfg_options.check(&cfg) != Some(false)
+                            })
                         {
-                            self.import_all_legacy_macros(module_id);
+                            let module_id = self.push_child_module(
+                                module.name.clone(),
+                                ast_id,
+                                Some((file_id, is_mod_rs)),
+                                &self.item_tree[module.visibility],
+                            );
+                            ModCollector {
+                                def_collector: &mut *self.def_collector,
+                                macro_depth: self.macro_depth,
+                                module_id,
+                                file_id: file_id.into(),
+                                item_tree: &item_tree,
+                                mod_dir,
+                            }
+                            .collect(item_tree.top_level_items());
+                            if is_macro_use
+                                || item_tree
+                                    .top_level_attrs(db, self.def_collector.def_map.krate)
+                                    .by_key("macro_use")
+                                    .exists()
+                            {
+                                self.import_all_legacy_macros(module_id);
+                            }
                         }
                     }
                     Err(candidate) => {
diff --git a/crates/hir_def/src/nameres/tests/mod_resolution.rs b/crates/hir_def/src/nameres/tests/mod_resolution.rs
index dfbbad1f9..16a2cd27a 100644
--- a/crates/hir_def/src/nameres/tests/mod_resolution.rs
+++ b/crates/hir_def/src/nameres/tests/mod_resolution.rs
@@ -819,3 +819,22 @@ pub mod hash { pub trait Hash {} }
         "#]],
     );
 }
+
+#[test]
+fn cfg_in_module_file() {
+    // Inner `#![cfg]` in a module file makes the whole module disappear.
+    check(
+        r#"
+//- /main.rs
+mod module;
+
+//- /module.rs
+#![cfg(NEVER)]
+
+struct AlsoShoulntAppear;
+        "#,
+        expect![[r#"
+            crate
+        "#]],
+    )
+}