1 files changed, 58 insertions, 25 deletions

diff --git a/crates/ide/src/annotations.rs b/crates/ide/src/annotations.rs
index 2e8e82b70..72492f826 100644
--- a/crates/ide/src/annotations.rs
+++ b/crates/ide/src/annotations.rs
@@ -1,12 +1,13 @@
-use hir::Semantics;
+use either::Either;
+use hir::{HasSource, Semantics};
 use ide_db::{
-    base_db::{FileId, FilePosition, FileRange, SourceDatabase},
-    RootDatabase, SymbolKind,
+    base_db::{FileId, FilePosition, FileRange},
+    helpers::visit_file_defs,
+    RootDatabase,
 };
-use syntax::TextRange;
+use syntax::{ast::NameOwner, AstNode, TextRange, TextSize};
 
 use crate::{
-    file_structure::file_structure,
     fn_references::find_all_methods,
     goto_implementation::goto_implementation,
     references::find_all_refs,
@@ -75,39 +76,56 @@ pub(crate) fn annotations(
         }
     }
 
-    file_structure(&db.parse(file_id).tree())
-        .into_iter()
-        .filter(|node| {
-            matches!(
-                node.kind,
-                SymbolKind::Trait
-                    | SymbolKind::Struct
-                    | SymbolKind::Enum
-                    | SymbolKind::Union
-                    | SymbolKind::Const
-            )
-        })
-        .for_each(|node| {
-            if config.annotate_impls && node.kind != SymbolKind::Const {
+    visit_file_defs(&Semantics::new(db), file_id, &mut |def| match def {
+        Either::Left(def) => {
+            let node = match def {
+                hir::ModuleDef::Const(konst) => {
+                    konst.source(db).and_then(|node| range_and_position_of(&node.value))
+                }
+                hir::ModuleDef::Trait(trait_) => {
+                    trait_.source(db).and_then(|node| range_and_position_of(&node.value))
+                }
+                hir::ModuleDef::Adt(hir::Adt::Struct(strukt)) => {
+                    strukt.source(db).and_then(|node| range_and_position_of(&node.value))
+                }
+                hir::ModuleDef::Adt(hir::Adt::Enum(enum_)) => {
+                    enum_.source(db).and_then(|node| range_and_position_of(&node.value))
+                }
+                hir::ModuleDef::Adt(hir::Adt::Union(union)) => {
+                    union.source(db).and_then(|node| range_and_position_of(&node.value))
+                }
+                _ => None,
+            };
+            let (offset, range) = match node {
+                Some(node) => node,
+                None => return,
+            };
+
+            if config.annotate_impls && !matches!(def, hir::ModuleDef::Const(_)) {
                 annotations.push(Annotation {
-                    range: node.node_range,
+                    range,
                     kind: AnnotationKind::HasImpls {
-                        position: FilePosition { file_id, offset: node.navigation_range.start() },
+                        position: FilePosition { file_id, offset },
                         data: None,
                     },
                 });
             }
-
             if config.annotate_references {
                 annotations.push(Annotation {
-                    range: node.node_range,
+                    range,
                     kind: AnnotationKind::HasReferences {
-                        position: FilePosition { file_id, offset: node.navigation_range.start() },
+                        position: FilePosition { file_id, offset },
                         data: None,
                     },
                 });
             }
-        });
+
+            fn range_and_position_of(node: &dyn NameOwner) -> Option<(TextSize, TextRange)> {
+                Some((node.name()?.syntax().text_range().start(), node.syntax().text_range()))
+            }
+        }
+        Either::Right(_) => (),
+    });
 
     if config.annotate_method_references {
         annotations.extend(find_all_methods(db, file_id).into_iter().map(|method| Annotation {
@@ -934,4 +952,19 @@ mod tests {
             "#]],
         );
     }
+
+    #[test]
+    fn test_no_annotations_outside_module_tree() {
+        check(
+            r#"
+//- /foo.rs
+struct Foo;
+//- /lib.rs
+// this file comes last since `check` checks the first file only
+"#,
+            expect![[r#"
+                []
+            "#]],
+        );
+    }
 }