@@ -15,7 +15,10 @@ use memchr::memmem::Finder;
1515use nohash_hasher:: IntMap ;
1616use once_cell:: unsync:: Lazy ;
1717use parser:: SyntaxKind ;
18- use syntax:: { ast, match_ast, AstNode , AstToken , SyntaxElement , TextRange , TextSize } ;
18+ use syntax:: {
19+ ast:: { self , HasAttrs as _} ,
20+ match_ast, AstNode , AstToken , SyntaxElement , TextRange , TextSize ,
21+ } ;
1922use triomphe:: Arc ;
2023
2124use crate :: {
@@ -134,6 +137,7 @@ pub enum ReferenceCategory {
134137 // FIXME: Some day should be able to search in doc comments. Would probably
135138 // need to switch from enum to bitflags then?
136139 // DocComment
140+ Test ,
137141}
138142
139143/// Generally, `search_scope` returns files that might contain references for the element.
@@ -872,6 +876,10 @@ fn def_to_ty(sema: &Semantics<'_, RootDatabase>, def: &Definition) -> Option<hir
872876
873877impl ReferenceCategory {
874878 fn new ( def : & Definition , r : & ast:: NameRef ) -> Option < ReferenceCategory > {
879+ if is_name_ref_in_test ( r) {
880+ return Some ( ReferenceCategory :: Test ) ;
881+ }
882+
875883 // Only Locals and Fields have accesses for now.
876884 if !matches ! ( def, Definition :: Local ( _) | Definition :: Field ( _) ) {
877885 return is_name_ref_in_import ( r) . then_some ( ReferenceCategory :: Import ) ;
@@ -910,3 +918,30 @@ fn is_name_ref_in_import(name_ref: &ast::NameRef) -> bool {
910918 . and_then ( |it| it. parent_path ( ) . top_path ( ) . syntax ( ) . parent ( ) )
911919 . map_or ( false , |it| it. kind ( ) == SyntaxKind :: USE_TREE )
912920}
921+
922+ fn is_name_ref_in_test ( name_ref : & ast:: NameRef ) -> bool {
923+ let mode = name_ref. syntax ( ) . ancestors ( ) . find_map ( |node| {
924+ match_ast ! {
925+ match node {
926+ ast:: Fn ( f) => {
927+ let attrs = f. attrs( ) ;
928+ let mut is_test = false ;
929+ for attr in attrs {
930+ if attr. to_string( ) == "#[test]" {
931+ is_test = true ;
932+ break ;
933+ }
934+ }
935+ if is_test {
936+ Some ( ReferenceCategory :: Test )
937+ }
938+ else {
939+ None
940+ }
941+ } ,
942+ _ => None
943+ }
944+ }
945+ } ) ;
946+ mode. is_some ( )
947+ }
0 commit comments