Highlight closure captures when cursor is on pipe

crates/ide/src/highlight_related.rs

@@ -1,6 +1,6 @@
 use hir::Semantics;
 use ide_db::{
-    base_db::{FileId, FilePosition},
+    base_db::{FileId, FilePosition, FileRange},
     defs::{Definition, IdentClass},
     helpers::pick_best_token,
     search::{FileReference, ReferenceCategory, SearchScope},
@@ -30,6 +30,7 @@ pub struct HighlightRelatedConfig {
     pub references: bool,
     pub exit_points: bool,
     pub break_points: bool,
+    pub closure_captures: bool,
     pub yield_points: bool,
 }
 
@@ -53,11 +54,12 @@ pub(crate) fn highlight_related(
 
     let token = pick_best_token(syntax.token_at_offset(offset), |kind| match kind {
         T![?] => 4, // prefer `?` when the cursor is sandwiched like in `await$0?`
-        T![->] => 3,
+        T![->] | T![|] => 3,
         kind if kind.is_keyword() => 2,
         IDENT | INT_NUMBER => 1,
         _ => 0,
     })?;
+    // most if not all of these should be re-implemented with information seeded from hir
     match token.kind() {
         T![?] if config.exit_points && token.parent().and_then(ast::TryExpr::cast).is_some() => {
             highlight_exit_points(sema, token)
@@ -70,11 +72,64 @@ pub(crate) fn highlight_related(
         T![break] | T![loop] | T![while] | T![continue] if config.break_points => {
             highlight_break_points(token)
         }
+        T![|] if config.closure_captures => highlight_closure_captures(
+            sema,
+            token.parent_ancestors().nth(1).and_then(ast::ClosureExpr::cast)?,
+            file_id,
+        ),
+        T![move] if config.closure_captures => highlight_closure_captures(
+            sema,
+            token.parent().and_then(ast::ClosureExpr::cast)?,
+            file_id,
+        ),
         _ if config.references => highlight_references(sema, &syntax, token, file_id),
         _ => None,
     }
 }
 
+fn highlight_closure_captures(
+    sema: &Semantics<'_, RootDatabase>,
+    node: ast::ClosureExpr,
+    file_id: FileId,
+) -> Option<Vec<HighlightedRange>> {
+    let search_range = node.body()?.syntax().text_range();
+    let ty = &sema.type_of_expr(&node.into())?.original;
+    let c = ty.as_closure()?;
+    Some(
+        c.captured_items(sema.db)
+            .into_iter()
+            .map(|capture| capture.local())
+            .flat_map(|local| {
+                let usages = Definition::Local(local)
+                    .usages(sema)
+                    .set_scope(Some(SearchScope::file_range(FileRange {
+                        file_id,
+                        range: search_range,
+                    })))
+                    .include_self_refs()
+                    .all()
+                    .references
+                    .remove(&file_id)
+                    .into_iter()
+                    .flatten()
+                    .map(|FileReference { category, range, .. }| HighlightedRange {
+                        range,
+                        category,
+                    });
+                let category = local.is_mut(sema.db).then_some(ReferenceCategory::Write);
+                local
+                    .sources(sema.db)
+                    .into_iter()
+                    .map(|x| x.to_nav(sema.db))
+                    .filter(|decl| decl.file_id == file_id)
+                    .filter_map(|decl| decl.focus_range)
+                    .map(move |range| HighlightedRange { range, category })
+                    .chain(usages)
+            })
+            .collect(),
+    )
+}
+
 fn highlight_references(
     sema: &Semantics<'_, RootDatabase>,
     node: &SyntaxNode,
@@ -93,10 +148,7 @@ fn highlight_references(
                 .remove(&file_id)
         })
         .flatten()
-        .map(|FileReference { category: access, range, .. }| HighlightedRange {
-            range,
-            category: access,
-        });
+        .map(|FileReference { category, range, .. }| HighlightedRange { range, category });
     let mut res = FxHashSet::default();
     for &def in &defs {
         match def {
@@ -352,16 +404,17 @@ mod tests {
 
     use super::*;
 
+    const ENABLED_CONFIG: HighlightRelatedConfig = HighlightRelatedConfig {
+        break_points: true,
+        exit_points: true,
+        references: true,
+        closure_captures: true,
+        yield_points: true,
+    };
+
     #[track_caller]
     fn check(ra_fixture: &str) {
-        let config = HighlightRelatedConfig {
-            break_points: true,
-            exit_points: true,
-            references: true,
-            yield_points: true,
-        };
-
-        check_with_config(ra_fixture, config);
+        check_with_config(ra_fixture, ENABLED_CONFIG);
     }
 
     #[track_caller]
@@ -1086,12 +1139,7 @@ fn function(field: u32) {
 
     #[test]
     fn test_hl_disabled_ref_local() {
-        let config = HighlightRelatedConfig {
-            references: false,
-            break_points: true,
-            exit_points: true,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { references: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1106,12 +1154,7 @@ fn foo() {
 
     #[test]
     fn test_hl_disabled_ref_local_preserved_break() {
-        let config = HighlightRelatedConfig {
-            references: false,
-            break_points: true,
-            exit_points: true,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { references: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1146,12 +1189,7 @@ fn foo() {
 
     #[test]
     fn test_hl_disabled_ref_local_preserved_yield() {
-        let config = HighlightRelatedConfig {
-            references: false,
-            break_points: true,
-            exit_points: true,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { references: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1182,12 +1220,7 @@ async fn foo() {
 
     #[test]
     fn test_hl_disabled_ref_local_preserved_exit() {
-        let config = HighlightRelatedConfig {
-            references: false,
-            break_points: true,
-            exit_points: true,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { references: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1225,12 +1258,7 @@ fn foo() ->$0 i32 {
 
     #[test]
     fn test_hl_disabled_break() {
-        let config = HighlightRelatedConfig {
-            references: true,
-            break_points: false,
-            exit_points: true,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { break_points: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1246,12 +1274,7 @@ fn foo() {
 
     #[test]
     fn test_hl_disabled_yield() {
-        let config = HighlightRelatedConfig {
-            references: true,
-            break_points: true,
-            exit_points: true,
-            yield_points: false,
-        };
+        let config = HighlightRelatedConfig { yield_points: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1265,12 +1288,7 @@ async$0 fn foo() {
 
     #[test]
     fn test_hl_disabled_exit() {
-        let config = HighlightRelatedConfig {
-            references: true,
-            break_points: true,
-            exit_points: false,
-            yield_points: true,
-        };
+        let config = HighlightRelatedConfig { exit_points: false, ..ENABLED_CONFIG };
 
         check_with_config(
             r#"
@@ -1411,6 +1429,34 @@ impl Trait for () {
     type Output$0 = ();
       // ^^^^^^
 }
+"#,
+        );
+    }
+
+    #[test]
+    fn test_closure_capture_pipe() {
+        check(
+            r#"
+fn f() {
+    let x = 1;
+    //  ^
+    let c = $0|y| x + y;
+    //          ^ read
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn test_closure_capture_move() {
+        check(
+            r#"
+fn f() {
+    let x = 1;
+    //  ^
+    let c = move$0 |y| x + y;
+    //               ^ read
+}
 "#,
         );
     }