Highlight closure captures when cursor is on pipe

Author: Veykril

crates/hir-ty/src/infer/closure.rs

@@ -115,9 +115,10 @@ impl InferenceContext<'_> {
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub(crate) struct HirPlace {
-    pub(crate) local: BindingId,
+    pub local: BindingId,
     pub(crate) projections: Vec<ProjectionElem<Infallible, Ty>>,
 }
+
 impl HirPlace {
     fn ty(&self, ctx: &mut InferenceContext<'_>) -> Ty {
         let mut ty = ctx.table.resolve_completely(ctx.result[self.local].clone());
@@ -161,6 +162,10 @@ pub struct CapturedItem {
 }
 
 impl CapturedItem {
+    pub fn local(&self) -> BindingId {
+        self.place.local
+    }
+
     pub fn display_kind(&self) -> &'static str {
         match self.kind {
             CaptureKind::ByRef(k) => match k {

crates/hir/src/lib.rs

@@ -3209,11 +3209,11 @@ impl Closure {
         self.clone().as_ty().display(db).with_closure_style(ClosureStyle::ImplFn).to_string()
     }
 
-    pub fn captured_items(&self, db: &dyn HirDatabase) -> Vec<hir_ty::CapturedItem> {
+    pub fn captured_items(&self, db: &dyn HirDatabase) -> Vec<ClosureCapture> {
         let owner = db.lookup_intern_closure((self.id).into()).0;
         let infer = &db.infer(owner);
         let info = infer.closure_info(&self.id);
-        info.0.clone()
+        info.0.iter().cloned().map(|capture| ClosureCapture { owner, capture }).collect()
     }
 
     pub fn fn_trait(&self, db: &dyn HirDatabase) -> FnTrait {
@@ -3224,6 +3224,26 @@ impl Closure {
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ClosureCapture {
+    owner: DefWithBodyId,
+    capture: hir_ty::CapturedItem,
+}
+
+impl ClosureCapture {
+    pub fn local(&self) -> Local {
+        Local { parent: self.owner, binding_id: self.capture.local() }
+    }
+
+    pub fn display_kind(&self) -> &'static str {
+        self.capture.display_kind()
+    }
+
+    pub fn display_place(&self, owner: ClosureId, db: &dyn HirDatabase) -> String {
+        self.capture.display_place(owner, db)
+    }
+}
+
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Type {
     env: Arc<TraitEnvironment>,

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
+}
 "#,
         );
     }

crates/rust-analyzer/src/config.rs

@@ -281,6 +281,8 @@ config_data! {
 
         /// Enables highlighting of related references while the cursor is on `break`, `loop`, `while`, or `for` keywords.
         highlightRelated_breakPoints_enable: bool = "true",
+        /// Enables highlighting of all captures of a closure while the cursor is on the `|` or move keyword of a closure.
+        highlightRelated_closureCaptures_enable: bool = "true",
         /// Enables highlighting of all exit points while the cursor is on any `return`, `?`, `fn`, or return type arrow (`->`).
         highlightRelated_exitPoints_enable: bool = "true",
         /// Enables highlighting of related references while the cursor is on any identifier.
@@ -1554,6 +1556,7 @@ impl Config {
             break_points: self.data.highlightRelated_breakPoints_enable,
             exit_points: self.data.highlightRelated_exitPoints_enable,
             yield_points: self.data.highlightRelated_yieldPoints_enable,
+            closure_captures: self.data.highlightRelated_closureCaptures_enable,
         }
     }
 

crates/rust-analyzer/src/lsp_ext.rs

@@ -434,7 +434,7 @@ pub enum HoverRequest {}
 impl Request for HoverRequest {
     type Params = HoverParams;
     type Result = Option<Hover>;
-    const METHOD: &'static str = "textDocument/hover";
+    const METHOD: &'static str = lsp_types::request::HoverRequest::METHOD;
 }
 
 #[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]

crates/test-utils/src/lib.rs

@@ -95,7 +95,7 @@ fn try_extract_range(text: &str) -> Option<(TextRange, String)> {
     Some((TextRange::new(start, end), text))
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 pub enum RangeOrOffset {
     Range(TextRange),
     Offset(TextSize),

docs/dev/lsp-extensions.md

@@ -1,5 +1,5 @@
 <!---
-lsp_ext.rs hash: 37ac44a0f507e05a
+lsp_ext.rs hash: 31ca513a249753ab
 
 If you need to change the above hash to make the test pass, please check if you
 need to adjust this doc as well and ping this issue:

docs/user/generated_config.adoc

@@ -352,6 +352,11 @@ Controls file watching implementation.
 --
 Enables highlighting of related references while the cursor is on `break`, `loop`, `while`, or `for` keywords.
 --
+[[rust-analyzer.highlightRelated.closureCaptures.enable]]rust-analyzer.highlightRelated.closureCaptures.enable (default: `true`)::
++
+--
+Enables highlighting of all captures of a closure while the cursor is on the `|` or move keyword of a closure.
+--
 [[rust-analyzer.highlightRelated.exitPoints.enable]]rust-analyzer.highlightRelated.exitPoints.enable (default: `true`)::
 +
 --

editors/code/package.json

@@ -886,6 +886,11 @@
                     "default": true,
                     "type": "boolean"
                 },
+                "rust-analyzer.highlightRelated.closureCaptures.enable": {
+                    "markdownDescription": "Enables highlighting of all captures of a closure while the cursor is on the `|` or move keyword of a closure.",
+                    "default": true,
+                    "type": "boolean"
+                },
                 "rust-analyzer.highlightRelated.exitPoints.enable": {
                     "markdownDescription": "Enables highlighting of all exit points while the cursor is on any `return`, `?`, `fn`, or return type arrow (`->`).",
                     "default": true,