Move document highlighting computation from rust-analyzer to ide

crates/ide/src/document_highlight.rs

@@ -0,0 +1,138 @@
+use hir::Semantics;
+use ide_db::{
+    base_db::FilePosition,
+    defs::Definition,
+    search::{FileReference, ReferenceAccess, SearchScope},
+    RootDatabase,
+};
+use syntax::{AstNode, TextRange};
+
+use crate::{display::TryToNav, references, NavigationTarget};
+
+pub struct DocumentHighlight {
+    pub range: TextRange,
+    pub access: Option<ReferenceAccess>,
+}
+
+// Feature: Document highlight
+//
+// Highlights the definition and its all references of the item at the cursor location in the current file.
+pub(crate) fn document_highlight(
+    sema: &Semantics<RootDatabase>,
+    position: FilePosition,
+) -> Option<Vec<DocumentHighlight>> {
+    let _p = profile::span("document_highlight");
+    let syntax = sema.parse(position.file_id).syntax().clone();
+    let def = references::find_def(sema, &syntax, position)?;
+    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(position.file_id))).all();
+
+    let declaration = match def {
+        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
+            Some(NavigationTarget::from_module_to_decl(sema.db, module))
+        }
+        def => def.try_to_nav(sema.db),
+    }
+    .filter(|decl| decl.file_id == position.file_id)
+    .and_then(|decl| {
+        let range = decl.focus_range?;
+        let access = references::decl_access(&def, &syntax, range);
+        Some(DocumentHighlight { range, access })
+    });
+
+    let file_refs = usages.references.get(&position.file_id).map_or(&[][..], Vec::as_slice);
+    let mut res = Vec::with_capacity(file_refs.len() + 1);
+    res.extend(declaration);
+    res.extend(
+        file_refs
+            .iter()
+            .map(|&FileReference { access, range, .. }| DocumentHighlight { range, access }),
+    );
+    Some(res)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::fixture;
+
+    use super::*;
+
+    fn check(ra_fixture: &str) {
+        let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
+        let hls = analysis.highlight_document(pos).unwrap().unwrap();
+
+        let mut expected = annotations
+            .into_iter()
+            .map(|(r, access)| (r.range, (!access.is_empty()).then(|| access)))
+            .collect::<Vec<_>>();
+
+        let mut actual = hls
+            .into_iter()
+            .map(|hl| {
+                (
+                    hl.range,
+                    hl.access.map(|it| {
+                        match it {
+                            ReferenceAccess::Read => "read",
+                            ReferenceAccess::Write => "write",
+                        }
+                        .to_string()
+                    }),
+                )
+            })
+            .collect::<Vec<_>>();
+        actual.sort_by_key(|(range, _)| range.start());
+        expected.sort_by_key(|(range, _)| range.start());
+
+        assert_eq!(expected, actual);
+    }
+
+    #[test]
+    fn test_hl_module() {
+        check(
+            r#"
+//- /lib.rs
+mod foo$0;
+ // ^^^
+//- /foo.rs
+struct Foo;
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_self_in_crate_root() {
+        check(
+            r#"
+//- /lib.rs
+use self$0;
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_self_in_module() {
+        check(
+            r#"
+//- /lib.rs
+mod foo;
+//- /foo.rs
+use self$0;
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_local() {
+        check(
+            r#"
+//- /lib.rs
+fn foo() {
+    let mut bar = 3;
+         // ^^^ write
+    bar$0;
+ // ^^^ read
+}
+"#,
+        );
+    }
+}

crates/ide/src/lib.rs

@@ -24,32 +24,33 @@ mod display;
 
 mod annotations;
 mod call_hierarchy;
+mod doc_links;
+mod document_highlight;
 mod expand_macro;
 mod extend_selection;
 mod file_structure;
+mod fn_references;
 mod folding_ranges;
 mod goto_definition;
 mod goto_implementation;
 mod goto_type_definition;
-mod view_hir;
 mod hover;
 mod inlay_hints;
 mod join_lines;
+mod markdown_remove;
 mod matching_brace;
 mod move_item;
 mod parent_module;
 mod references;
 mod rename;
-mod fn_references;
 mod runnables;
 mod ssr;
 mod status;
 mod syntax_highlighting;
 mod syntax_tree;
 mod typing;
-mod markdown_remove;
-mod doc_links;
 mod view_crate_graph;
+mod view_hir;
 mod view_item_tree;
 
 use std::sync::Arc;
@@ -72,6 +73,7 @@ pub use crate::{
     annotations::{Annotation, AnnotationConfig, AnnotationKind},
     call_hierarchy::CallItem,
     display::navigation_target::NavigationTarget,
+    document_highlight::DocumentHighlight,
     expand_macro::ExpandedMacro,
     file_structure::{StructureNode, StructureNodeKind},
     folding_ranges::{Fold, FoldKind},
@@ -483,6 +485,14 @@ impl Analysis {
         self.with_db(|db| syntax_highlighting::highlight(db, file_id, None, false))
     }
 
+    /// Computes all ranges to highlight for a given item in a file.
+    pub fn highlight_document(
+        &self,
+        position: FilePosition,
+    ) -> Cancellable<Option<Vec<DocumentHighlight>>> {
+        self.with_db(|db| document_highlight::document_highlight(&Semantics::new(db), position))
+    }
+
     /// Computes syntax highlighting for the given file range.
     pub fn highlight_range(&self, frange: FileRange) -> Cancellable<Vec<HlRange>> {
         self.with_db(|db| {

crates/ide/src/references.rs

@@ -90,7 +90,13 @@ pub(crate) fn find_all_refs(
             _ => {}
         }
     }
-    let declaration = def.try_to_nav(sema.db).map(|nav| {
+    let declaration = match def {
+        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
+            Some(NavigationTarget::from_module_to_decl(sema.db, module))
+        }
+        def => def.try_to_nav(sema.db),
+    }
+    .map(|nav| {
         let decl_range = nav.focus_or_full_range();
         Declaration { nav, access: decl_access(&def, &syntax, decl_range) }
     });
@@ -104,7 +110,7 @@ pub(crate) fn find_all_refs(
     Some(ReferenceSearchResult { declaration, references })
 }
 
-fn find_def(
+pub(crate) fn find_def(
     sema: &Semantics<RootDatabase>,
     syntax: &SyntaxNode,
     position: FilePosition,
@@ -126,7 +132,11 @@ fn find_def(
     Some(def)
 }
 
-fn decl_access(def: &Definition, syntax: &SyntaxNode, range: TextRange) -> Option<ReferenceAccess> {
+pub(crate) fn decl_access(
+    def: &Definition,
+    syntax: &SyntaxNode,
+    range: TextRange,
+) -> Option<ReferenceAccess> {
     match def {
         Definition::Local(_) | Definition::Field(_) => {}
         _ => return None,
@@ -658,9 +668,6 @@ fn f() {
         );
     }
 
-    // `mod foo;` is not in the results because `foo` is an `ast::Name`.
-    // So, there are two references: the first one is a definition of the `foo` module,
-    // which is the whole `foo.rs`, and the second one is in `use foo::Foo`.
     #[test]
     fn test_find_all_refs_decl_module() {
         check(
@@ -680,13 +687,44 @@ pub struct Foo {
 }
 "#,
             expect![[r#"
-                foo Module FileId(1) 0..35
+                foo Module FileId(0) 0..8 4..7
 
                 FileId(0) 14..17
             "#]],
         );
     }
 
+    #[test]
+    fn test_find_all_refs_decl_module_on_self() {
+        check(
+            r#"
+//- /lib.rs
+mod foo;
+
+//- /foo.rs
+use self$0;
+"#,
+            expect![[r#"
+                foo Module FileId(0) 0..8 4..7
+
+            "#]],
+        );
+    }
+
+    #[test]
+    fn test_find_all_refs_decl_module_on_self_crate_root() {
+        check(
+            r#"
+//- /lib.rs
+use self$0;
+"#,
+            expect![[r#"
+                Module FileId(0) 0..10
+
+            "#]],
+        );
+    }
+
     #[test]
     fn test_find_all_refs_super_mod_vis() {
         check(

crates/rust-analyzer/src/handlers.rs

@@ -9,8 +9,8 @@ use std::{
 
 use ide::{
     AnnotationConfig, AssistKind, AssistResolveStrategy, FileId, FilePosition, FileRange,
-    HoverAction, HoverGotoTypeData, Query, RangeInfo, Runnable, RunnableKind, SearchScope,
-    SingleResolve, SourceChange, TextEdit,
+    HoverAction, HoverGotoTypeData, Query, RangeInfo, Runnable, RunnableKind, SingleResolve,
+    SourceChange, TextEdit,
 };
 use ide_db::SymbolKind;
 use itertools::Itertools;
@@ -18,12 +18,12 @@ use lsp_server::ErrorCode;
 use lsp_types::{
     CallHierarchyIncomingCall, CallHierarchyIncomingCallsParams, CallHierarchyItem,
     CallHierarchyOutgoingCall, CallHierarchyOutgoingCallsParams, CallHierarchyPrepareParams,
-    CodeLens, CompletionItem, Diagnostic, DiagnosticTag, DocumentFormattingParams,
-    DocumentHighlight, FoldingRange, FoldingRangeParams, HoverContents, Location, NumberOrString,
-    Position, PrepareRenameResponse, Range, RenameParams, SemanticTokensDeltaParams,
-    SemanticTokensFullDeltaResult, SemanticTokensParams, SemanticTokensRangeParams,
-    SemanticTokensRangeResult, SemanticTokensResult, SymbolInformation, SymbolTag,
-    TextDocumentIdentifier, TextDocumentPositionParams, Url, WorkspaceEdit,
+    CodeLens, CompletionItem, Diagnostic, DiagnosticTag, DocumentFormattingParams, FoldingRange,
+    FoldingRangeParams, HoverContents, Location, NumberOrString, Position, PrepareRenameResponse,
+    Range, RenameParams, SemanticTokensDeltaParams, SemanticTokensFullDeltaResult,
+    SemanticTokensParams, SemanticTokensRangeParams, SemanticTokensRangeResult,
+    SemanticTokensResult, SymbolInformation, SymbolTag, TextDocumentIdentifier,
+    TextDocumentPositionParams, Url, WorkspaceEdit,
 };
 use project_model::TargetKind;
 use serde::{Deserialize, Serialize};
@@ -1163,33 +1163,22 @@ pub(crate) fn handle_code_lens_resolve(
 pub(crate) fn handle_document_highlight(
     snap: GlobalStateSnapshot,
     params: lsp_types::DocumentHighlightParams,
-) -> Result<Option<Vec<DocumentHighlight>>> {
+) -> Result<Option<Vec<lsp_types::DocumentHighlight>>> {
     let _p = profile::span("handle_document_highlight");
     let position = from_proto::file_position(&snap, params.text_document_position_params)?;
     let line_index = snap.file_line_index(position.file_id)?;
 
-    let refs = match snap
-        .analysis
-        .find_all_refs(position, Some(SearchScope::single_file(position.file_id)))?
-    {
+    let refs = match snap.analysis.highlight_document(position)? {
         None => return Ok(None),
         Some(refs) => refs,
     };
-
-    let decl = refs.declaration.filter(|decl| decl.nav.file_id == position.file_id).map(|decl| {
-        DocumentHighlight {
-            range: to_proto::range(&line_index, decl.nav.focus_or_full_range()),
-            kind: decl.access.map(to_proto::document_highlight_kind),
-        }
-    });
-
-    let file_refs = refs.references.get(&position.file_id).map_or(&[][..], Vec::as_slice);
-    let mut res = Vec::with_capacity(file_refs.len() + 1);
-    res.extend(decl);
-    res.extend(file_refs.iter().map(|&(range, access)| DocumentHighlight {
-        range: to_proto::range(&line_index, range),
-        kind: access.map(to_proto::document_highlight_kind),
-    }));
+    let res = refs
+        .into_iter()
+        .map(|ide::DocumentHighlight { range, access }| lsp_types::DocumentHighlight {
+            range: to_proto::range(&line_index, range),
+            kind: access.map(to_proto::document_highlight_kind),
+        })
+        .collect();
     Ok(Some(res))
 }