Merge #9375

9375: feat: Highlight exit and yield points r=Veykril a=Veykril

![Code_YBHOCF3DbU](https://user-images.githubusercontent.com/3757771/123128986-e1270a80-d44b-11eb-9854-065459a2dd50.png)
![Code_YyMhqES0LX](https://user-images.githubusercontent.com/3757771/123128988-e1bfa100-d44b-11eb-9c81-6a6031aad740.png)

Fixes #4691
Fixes #9365

Co-authored-by: Lukas Wirth <lukastw97@gmail.com>

crates/hir/src/lib.rs

@@ -2075,10 +2075,15 @@ impl Type {
     pub fn is_unit(&self) -> bool {
         matches!(self.ty.kind(&Interner), TyKind::Tuple(0, ..))
     }
+
     pub fn is_bool(&self) -> bool {
         matches!(self.ty.kind(&Interner), TyKind::Scalar(Scalar::Bool))
     }
 
+    pub fn is_never(&self) -> bool {
+        matches!(self.ty.kind(&Interner), TyKind::Never)
+    }
+
     pub fn is_mutable_reference(&self) -> bool {
         matches!(self.ty.kind(&Interner), TyKind::Ref(hir_ty::Mutability::Mut, ..))
     }

crates/ide/src/highlight_related.rs

@@ -0,0 +1,586 @@
+use hir::Semantics;
+use ide_db::{
+    base_db::FilePosition,
+    defs::Definition,
+    helpers::pick_best_token,
+    search::{FileReference, ReferenceAccess, SearchScope},
+    RootDatabase,
+};
+use syntax::{
+    ast::{self, LoopBodyOwner},
+    match_ast, AstNode, SyntaxNode, SyntaxToken, TextRange, WalkEvent, T,
+};
+
+use crate::{display::TryToNav, references, NavigationTarget};
+
+pub struct HighlightedRange {
+    pub range: TextRange,
+    pub access: Option<ReferenceAccess>,
+}
+
+// Feature: Highlight Related
+//
+// Highlights constructs related to the thing under the cursor:
+// - if on an identifier, highlights all references to that identifier in the current file
+// - if on an `async` or `await token, highlights all yield points for that async context
+// - if on a `return` token, `?` character or `->` return type arrow, highlights all exit points for that context
+pub(crate) fn highlight_related(
+    sema: &Semantics<RootDatabase>,
+    position: FilePosition,
+) -> Option<Vec<HighlightedRange>> {
+    let _p = profile::span("highlight_related");
+    let syntax = sema.parse(position.file_id).syntax().clone();
+
+    let token = pick_best_token(syntax.token_at_offset(position.offset), |kind| match kind {
+        T![?] => 2, // prefer `?` when the cursor is sandwiched like `await$0?`
+        T![await] | T![async] | T![return] | T![->] => 1,
+        _ => 0,
+    })?;
+
+    match token.kind() {
+        T![return] | T![?] | T![->] => highlight_exit_points(sema, token),
+        T![await] | T![async] => highlight_yield_points(token),
+        _ => highlight_references(sema, &syntax, position),
+    }
+}
+
+fn highlight_references(
+    sema: &Semantics<RootDatabase>,
+    syntax: &SyntaxNode,
+    FilePosition { offset, file_id }: FilePosition,
+) -> Option<Vec<HighlightedRange>> {
+    let def = references::find_def(sema, syntax, offset)?;
+    let usages = def.usages(sema).set_scope(Some(SearchScope::single_file(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 == file_id)
+    .and_then(|decl| {
+        let range = decl.focus_range?;
+        let access = references::decl_access(&def, syntax, range);
+        Some(HighlightedRange { range, access })
+    });
+
+    let file_refs = usages.references.get(&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, .. }| HighlightedRange { range, access }),
+    );
+    Some(res)
+}
+
+fn highlight_exit_points(
+    sema: &Semantics<RootDatabase>,
+    token: SyntaxToken,
+) -> Option<Vec<HighlightedRange>> {
+    fn hl(
+        sema: &Semantics<RootDatabase>,
+        body: Option<ast::Expr>,
+    ) -> Option<Vec<HighlightedRange>> {
+        let mut highlights = Vec::new();
+        let body = body?;
+        walk(&body, &mut |expr| match expr {
+            ast::Expr::ReturnExpr(expr) => {
+                if let Some(token) = expr.return_token() {
+                    highlights.push(HighlightedRange { access: None, range: token.text_range() });
+                }
+            }
+            ast::Expr::TryExpr(try_) => {
+                if let Some(token) = try_.question_mark_token() {
+                    highlights.push(HighlightedRange { access: None, range: token.text_range() });
+                }
+            }
+            ast::Expr::MethodCallExpr(_) | ast::Expr::CallExpr(_) | ast::Expr::MacroCall(_) => {
+                if sema.type_of_expr(&expr).map_or(false, |ty| ty.is_never()) {
+                    highlights
+                        .push(HighlightedRange { access: None, range: expr.syntax().text_range() });
+                }
+            }
+            _ => (),
+        });
+        let tail = match body {
+            ast::Expr::BlockExpr(b) => b.tail_expr(),
+            e => Some(e),
+        };
+
+        if let Some(tail) = tail {
+            for_each_inner_tail(&tail, &mut |tail| {
+                highlights
+                    .push(HighlightedRange { access: None, range: tail.syntax().text_range() })
+            });
+        }
+        Some(highlights)
+    }
+    for anc in token.ancestors() {
+        return match_ast! {
+            match anc {
+                ast::Fn(fn_) => hl(sema, fn_.body().map(ast::Expr::BlockExpr)),
+                ast::ClosureExpr(closure) => hl(sema, closure.body()),
+                ast::EffectExpr(effect) => if matches!(effect.effect(), ast::Effect::Async(_) | ast::Effect::Try(_)| ast::Effect::Const(_)) {
+                    hl(sema, effect.block_expr().map(ast::Expr::BlockExpr))
+                } else {
+                    continue;
+                },
+                _ => continue,
+            }
+        };
+    }
+    None
+}
+
+fn highlight_yield_points(token: SyntaxToken) -> Option<Vec<HighlightedRange>> {
+    fn hl(
+        async_token: Option<SyntaxToken>,
+        body: Option<ast::Expr>,
+    ) -> Option<Vec<HighlightedRange>> {
+        let mut highlights = Vec::new();
+        highlights.push(HighlightedRange { access: None, range: async_token?.text_range() });
+        if let Some(body) = body {
+            walk(&body, &mut |expr| {
+                if let ast::Expr::AwaitExpr(expr) = expr {
+                    if let Some(token) = expr.await_token() {
+                        highlights
+                            .push(HighlightedRange { access: None, range: token.text_range() });
+                    }
+                }
+            });
+        }
+        Some(highlights)
+    }
+    for anc in token.ancestors() {
+        return match_ast! {
+            match anc {
+                ast::Fn(fn_) => hl(fn_.async_token(), fn_.body().map(ast::Expr::BlockExpr)),
+                ast::EffectExpr(effect) => hl(effect.async_token(), effect.block_expr().map(ast::Expr::BlockExpr)),
+                ast::ClosureExpr(closure) => hl(closure.async_token(), closure.body()),
+                _ => continue,
+            }
+        };
+    }
+    None
+}
+
+/// Preorder walk all the expression's child expressions
+fn walk(expr: &ast::Expr, cb: &mut dyn FnMut(ast::Expr)) {
+    let mut preorder = expr.syntax().preorder();
+    while let Some(event) = preorder.next() {
+        let node = match event {
+            WalkEvent::Enter(node) => node,
+            WalkEvent::Leave(_) => continue,
+        };
+        match ast::Stmt::cast(node.clone()) {
+            // recursively walk the initializer, skipping potential const pat expressions
+            // lets statements aren't usually nested too deeply so this is fine to recurse on
+            Some(ast::Stmt::LetStmt(l)) => {
+                if let Some(expr) = l.initializer() {
+                    walk(&expr, cb);
+                }
+                preorder.skip_subtree();
+            }
+            // Don't skip subtree since we want to process the expression child next
+            Some(ast::Stmt::ExprStmt(_)) => (),
+            // skip inner items which might have their own expressions
+            Some(ast::Stmt::Item(_)) => preorder.skip_subtree(),
+            None => {
+                if let Some(expr) = ast::Expr::cast(node) {
+                    let is_different_context = match &expr {
+                        ast::Expr::EffectExpr(effect) => {
+                            matches!(
+                                effect.effect(),
+                                ast::Effect::Async(_) | ast::Effect::Try(_) | ast::Effect::Const(_)
+                            )
+                        }
+                        ast::Expr::ClosureExpr(__) => true,
+                        _ => false,
+                    };
+                    cb(expr);
+                    if is_different_context {
+                        preorder.skip_subtree();
+                    }
+                } else {
+                    preorder.skip_subtree();
+                }
+            }
+        }
+    }
+}
+
+// FIXME: doesn't account for labeled breaks in labeled blocks
+fn for_each_inner_tail(expr: &ast::Expr, cb: &mut dyn FnMut(&ast::Expr)) {
+    match expr {
+        ast::Expr::BlockExpr(b) => {
+            if let Some(e) = b.tail_expr() {
+                for_each_inner_tail(&e, cb);
+            }
+        }
+        ast::Expr::EffectExpr(e) => match e.effect() {
+            ast::Effect::Label(_) | ast::Effect::Unsafe(_) => {
+                if let Some(e) = e.block_expr().and_then(|b| b.tail_expr()) {
+                    for_each_inner_tail(&e, cb);
+                }
+            }
+            ast::Effect::Async(_) | ast::Effect::Try(_) | ast::Effect::Const(_) => cb(expr),
+        },
+        ast::Expr::IfExpr(if_) => {
+            if_.blocks().for_each(|block| for_each_inner_tail(&ast::Expr::BlockExpr(block), cb))
+        }
+        ast::Expr::LoopExpr(l) => for_each_break(l, cb),
+        ast::Expr::MatchExpr(m) => {
+            if let Some(arms) = m.match_arm_list() {
+                arms.arms().filter_map(|arm| arm.expr()).for_each(|e| for_each_inner_tail(&e, cb));
+            }
+        }
+        ast::Expr::ArrayExpr(_)
+        | ast::Expr::AwaitExpr(_)
+        | ast::Expr::BinExpr(_)
+        | ast::Expr::BoxExpr(_)
+        | ast::Expr::BreakExpr(_)
+        | ast::Expr::CallExpr(_)
+        | ast::Expr::CastExpr(_)
+        | ast::Expr::ClosureExpr(_)
+        | ast::Expr::ContinueExpr(_)
+        | ast::Expr::FieldExpr(_)
+        | ast::Expr::ForExpr(_)
+        | ast::Expr::IndexExpr(_)
+        | ast::Expr::Literal(_)
+        | ast::Expr::MacroCall(_)
+        | ast::Expr::MacroStmts(_)
+        | ast::Expr::MethodCallExpr(_)
+        | ast::Expr::ParenExpr(_)
+        | ast::Expr::PathExpr(_)
+        | ast::Expr::PrefixExpr(_)
+        | ast::Expr::RangeExpr(_)
+        | ast::Expr::RecordExpr(_)
+        | ast::Expr::RefExpr(_)
+        | ast::Expr::ReturnExpr(_)
+        | ast::Expr::TryExpr(_)
+        | ast::Expr::TupleExpr(_)
+        | ast::Expr::WhileExpr(_)
+        | ast::Expr::YieldExpr(_) => cb(expr),
+    }
+}
+
+fn for_each_break(l: &ast::LoopExpr, cb: &mut dyn FnMut(&ast::Expr)) {
+    let label = l.label().and_then(|lbl| lbl.lifetime());
+    let mut depth = 0;
+    if let Some(b) = l.loop_body() {
+        let preorder = &mut b.syntax().preorder();
+        let ev_as_expr = |ev| match ev {
+            WalkEvent::Enter(it) => Some(WalkEvent::Enter(ast::Expr::cast(it)?)),
+            WalkEvent::Leave(it) => Some(WalkEvent::Leave(ast::Expr::cast(it)?)),
+        };
+        let eq_label = |lt: Option<ast::Lifetime>| {
+            lt.zip(label.as_ref()).map_or(false, |(lt, lbl)| lt.text() == lbl.text())
+        };
+        while let Some(node) = preorder.find_map(ev_as_expr) {
+            match node {
+                WalkEvent::Enter(expr) => match &expr {
+                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_) => {
+                        depth += 1
+                    }
+                    ast::Expr::EffectExpr(e) if e.label().is_some() => depth += 1,
+                    ast::Expr::BreakExpr(b) if depth == 0 || eq_label(b.lifetime()) => {
+                        cb(&expr);
+                    }
+                    _ => (),
+                },
+                WalkEvent::Leave(expr) => match expr {
+                    ast::Expr::LoopExpr(_) | ast::Expr::WhileExpr(_) | ast::Expr::ForExpr(_) => {
+                        depth -= 1
+                    }
+                    ast::Expr::EffectExpr(e) if e.label().is_some() => depth -= 1,
+                    _ => (),
+                },
+            }
+        }
+    }
+}
+
+#[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_related(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#"
+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#"
+fn foo() {
+    let mut bar = 3;
+         // ^^^ write
+    bar$0;
+ // ^^^ read
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_yield_points() {
+        check(
+            r#"
+pub async fn foo() {
+ // ^^^^^
+    let x = foo()
+        .await$0
+      // ^^^^^
+        .await;
+      // ^^^^^
+    || { 0.await };
+    (async { 0.await }).await
+                     // ^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_yield_points2() {
+        check(
+            r#"
+pub async$0 fn foo() {
+ // ^^^^^
+    let x = foo()
+        .await
+      // ^^^^^
+        .await;
+      // ^^^^^
+    || { 0.await };
+    (async { 0.await }).await
+                     // ^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_yield_nested_fn() {
+        check(
+            r#"
+async fn foo() {
+    async fn foo2() {
+ // ^^^^^
+        async fn foo3() {
+            0.await
+        }
+        0.await$0
+       // ^^^^^
+    }
+    0.await
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_yield_nested_async_blocks() {
+        check(
+            r#"
+async fn foo() {
+    (async {
+  // ^^^^^
+        (async {
+           0.await
+        }).await$0 }
+        // ^^^^^
+    ).await;
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_exit_points() {
+        check(
+            r#"
+fn foo() -> u32 {
+    if true {
+        return$0 0;
+     // ^^^^^^
+    }
+
+    0?;
+  // ^
+    0xDEAD_BEEF
+ // ^^^^^^^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_exit_points2() {
+        check(
+            r#"
+fn foo() ->$0 u32 {
+    if true {
+        return 0;
+     // ^^^^^^
+    }
+
+    0?;
+  // ^
+    0xDEAD_BEEF
+ // ^^^^^^^^^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_prefer_ref_over_tail_exit() {
+        check(
+            r#"
+fn foo() -> u32 {
+// ^^^
+    if true {
+        return 0;
+    }
+
+    0?;
+
+    foo$0()
+ // ^^^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_never_call_is_exit_point() {
+        check(
+            r#"
+struct Never;
+impl Never {
+    fn never(self) -> ! { loop {} }
+}
+macro_rules! never {
+    () => { never() }
+}
+fn never() -> ! { loop {} }
+fn foo() ->$0 u32 {
+    never();
+ // ^^^^^^^
+    never!();
+ // FIXME sema doesn't give us types for macrocalls
+
+    Never.never();
+ // ^^^^^^^^^^^^^
+
+    0
+ // ^
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_hl_inner_tail_exit_points() {
+        check(
+            r#"
+fn foo() ->$0 u32 {
+    if true {
+        unsafe {
+            return 5;
+         // ^^^^^^
+            5
+         // ^
+        }
+    } else {
+        match 5 {
+            6 => 100,
+              // ^^^
+            7 => loop {
+                break 5;
+             // ^^^^^^^
+            }
+            8 => 'a: loop {
+                'b: loop {
+                    break 'a 5;
+                 // ^^^^^^^^^^
+                    break 'b 5;
+                    break 5;
+                };
+            }
+            //
+            _ => 500,
+              // ^^^
+        }
+    }
+}
+"#,
+        );
+    }
+}

crates/ide/src/lib.rs

@@ -24,33 +24,34 @@ mod display;
 
 mod annotations;
 mod call_hierarchy;
+mod doc_links;
+mod highlight_related;
 mod expand_macro;
 mod extend_selection;
 mod file_structure;
+mod fn_references;
 mod folding_ranges;
 mod goto_declaration;
 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;
@@ -76,6 +77,7 @@ pub use crate::{
     expand_macro::ExpandedMacro,
     file_structure::{StructureNode, StructureNodeKind},
     folding_ranges::{Fold, FoldKind},
+    highlight_related::HighlightedRange,
     hover::{HoverAction, HoverConfig, HoverDocFormat, HoverGotoTypeData, HoverResult},
     inlay_hints::{InlayHint, InlayHintsConfig, InlayKind},
     markup::Markup,
@@ -492,6 +494,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_related(
+        &self,
+        position: FilePosition,
+    ) -> Cancellable<Option<Vec<HighlightedRange>>> {
+        self.with_db(|db| highlight_related::highlight_related(&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

@@ -20,7 +20,7 @@ use rustc_hash::FxHashMap;
 use syntax::{
     algo::find_node_at_offset,
     ast::{self, NameOwner},
-    match_ast, AstNode, SyntaxNode, TextRange, T,
+    match_ast, AstNode, SyntaxNode, TextRange, TextSize, T,
 };
 
 use crate::{display::TryToNav, FilePosition, NavigationTarget};
@@ -60,7 +60,7 @@ pub(crate) fn find_all_refs(
         if let Some(name) = get_name_of_item_declaration(&syntax, position) {
             (NameClass::classify(sema, &name)?.referenced_or_defined(sema.db), true)
         } else {
-            (find_def(sema, &syntax, position)?, false)
+            (find_def(sema, &syntax, position.offset)?, false)
         };
 
     let mut usages = def.usages(sema).set_scope(search_scope).include_self_refs().all();
@@ -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,12 +110,12 @@ 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,
+    offset: TextSize,
 ) -> Option<Definition> {
-    let def = match sema.find_node_at_offset_with_descend(syntax, position.offset)? {
+    let def = match sema.find_node_at_offset_with_descend(syntax, offset)? {
         ast::NameLike::NameRef(name_ref) => {
             NameRefClass::classify(sema, &name_ref)?.referenced(sema.db)
         }
@@ -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(
@@ -1021,7 +1059,7 @@ impl Foo {
                 actual += "\n";
             }
         }
-        expect.assert_eq(&actual)
+        expect.assert_eq(actual.trim_start())
     }
 
     #[test]

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,
+    HoverAction, HoverGotoTypeData, Query, RangeInfo, Runnable, RunnableKind, SingleResolve,
-    SingleResolve, SourceChange, TextEdit,
+    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,
+    CodeLens, CompletionItem, Diagnostic, DiagnosticTag, DocumentFormattingParams, FoldingRange,
-    DocumentHighlight, FoldingRange, FoldingRangeParams, HoverContents, Location, NumberOrString,
+    FoldingRangeParams, HoverContents, Location, NumberOrString, Position, PrepareRenameResponse,
-    Position, PrepareRenameResponse, Range, RenameParams, SemanticTokensDeltaParams,
+    Range, RenameParams, SemanticTokensDeltaParams, SemanticTokensFullDeltaResult,
-    SemanticTokensFullDeltaResult, SemanticTokensParams, SemanticTokensRangeParams,
+    SemanticTokensParams, SemanticTokensRangeParams, SemanticTokensRangeResult,
-    SemanticTokensRangeResult, SemanticTokensResult, SymbolInformation, SymbolTag,
+    SemanticTokensResult, SymbolInformation, SymbolTag, TextDocumentIdentifier,
-    TextDocumentIdentifier, TextDocumentPositionParams, Url, WorkspaceEdit,
+    TextDocumentPositionParams, Url, WorkspaceEdit,
 };
 use project_model::TargetKind;
 use serde::{Deserialize, Serialize};
@@ -1178,33 +1178,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
+    let refs = match snap.analysis.highlight_related(position)? {
-        .analysis
-        .find_all_refs(position, Some(SearchScope::single_file(position.file_id)))?
-    {
         None => return Ok(None),
         Some(refs) => refs,
     };
-
+    let res = refs
-    let decl = refs.declaration.filter(|decl| decl.nav.file_id == position.file_id).map(|decl| {
+        .into_iter()
-        DocumentHighlight {
+        .map(|ide::HighlightedRange { range, access }| lsp_types::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),
-    }));
+        })
+        .collect();
     Ok(Some(res))
 }