feat: Closure capture inlay hints

Author: Veykril

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

@@ -148,7 +148,7 @@ impl HirPlace {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-pub(crate) enum CaptureKind {
+pub enum CaptureKind {
     ByRef(BorrowKind),
     ByValue,
 }
@@ -166,6 +166,10 @@ impl CapturedItem {
         self.place.local
     }
 
+    pub fn kind(&self) -> CaptureKind {
+        self.kind
+    }
+
     pub fn display_kind(&self) -> &'static str {
         match self.kind {
             CaptureKind::ByRef(k) => match k {

crates/hir-ty/src/lib.rs

@@ -61,8 +61,9 @@ pub use autoderef::autoderef;
 pub use builder::{ParamKind, TyBuilder};
 pub use chalk_ext::*;
 pub use infer::{
-    closure::CapturedItem, could_coerce, could_unify, Adjust, Adjustment, AutoBorrow, BindingMode,
-    InferenceDiagnostic, InferenceResult, OverloadedDeref, PointerCast,
+    closure::{CaptureKind, CapturedItem},
+    could_coerce, could_unify, Adjust, Adjustment, AutoBorrow, BindingMode, InferenceDiagnostic,
+    InferenceResult, OverloadedDeref, PointerCast,
 };
 pub use interner::Interner;
 pub use lower::{

crates/hir/src/lib.rs

@@ -2611,6 +2611,10 @@ impl LocalSource {
         self.source.file_id.original_file(db.upcast())
     }
 
+    pub fn file(&self) -> HirFileId {
+        self.source.file_id
+    }
+
     pub fn name(&self) -> Option<ast::Name> {
         self.source.value.name()
     }
@@ -3232,6 +3236,21 @@ impl ClosureCapture {
         Local { parent: self.owner, binding_id: self.capture.local() }
     }
 
+    pub fn kind(&self) -> CaptureKind {
+        match self.capture.kind() {
+            hir_ty::CaptureKind::ByRef(
+                hir_ty::mir::BorrowKind::Shallow | hir_ty::mir::BorrowKind::Shared,
+            ) => CaptureKind::SharedRef,
+            hir_ty::CaptureKind::ByRef(hir_ty::mir::BorrowKind::Unique) => {
+                CaptureKind::UniqueSharedRef
+            }
+            hir_ty::CaptureKind::ByRef(hir_ty::mir::BorrowKind::Mut { .. }) => {
+                CaptureKind::MutableRef
+            }
+            hir_ty::CaptureKind::ByValue => CaptureKind::Move,
+        }
+    }
+
     pub fn display_kind(&self) -> &'static str {
         self.capture.display_kind()
     }
@@ -3241,6 +3260,13 @@ impl ClosureCapture {
     }
 }
 
+pub enum CaptureKind {
+    SharedRef,
+    UniqueSharedRef,
+    MutableRef,
+    Move,
+}
+
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Type {
     env: Arc<TraitEnvironment>,

crates/ide/src/inlay_hints.rs

@@ -20,16 +20,17 @@ use text_edit::TextEdit;
 
 use crate::{navigation_target::TryToNav, FileId};
 
-mod closing_brace;
-mod implicit_static;
-mod fn_lifetime_fn;
-mod closure_ret;
 mod adjustment;
-mod chaining;
-mod param_name;
-mod binding_mode;
 mod bind_pat;
+mod binding_mode;
+mod chaining;
+mod closing_brace;
+mod closure_ret;
+mod closure_captures;
 mod discriminant;
+mod fn_lifetime_fn;
+mod implicit_static;
+mod param_name;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct InlayHintsConfig {
@@ -42,6 +43,7 @@ pub struct InlayHintsConfig {
     pub adjustment_hints_mode: AdjustmentHintsMode,
     pub adjustment_hints_hide_outside_unsafe: bool,
     pub closure_return_type_hints: ClosureReturnTypeHints,
+    pub closure_capture_hints: bool,
     pub binding_mode_hints: bool,
     pub lifetime_elision_hints: LifetimeElisionHints,
     pub param_names_for_lifetime_elision_hints: bool,
@@ -88,6 +90,8 @@ pub enum AdjustmentHintsMode {
     PreferPostfix,
 }
 
+// FIXME: Clean up this mess, the kinds are mainly used for setting different rendering properties in the lsp layer
+// We should probably turns this into such a property holding struct. Or clean this up in some other form.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum InlayKind {
     BindingMode,
@@ -98,6 +102,7 @@ pub enum InlayKind {
     Adjustment,
     AdjustmentPostfix,
     Lifetime,
+    ClosureCapture,
     Parameter,
     Type,
     Discriminant,
@@ -444,10 +449,10 @@ fn hints(
                     ast::Expr::MethodCallExpr(it) => {
                         param_name::hints(hints, sema, config, ast::Expr::from(it))
                     }
-                    ast::Expr::ClosureExpr(it) => closure_ret::hints(hints, famous_defs, config, file_id, it),
-                    // We could show reborrows for all expressions, but usually that is just noise to the user
-                    // and the main point here is to show why "moving" a mutable reference doesn't necessarily move it
-                    // ast::Expr::PathExpr(_) => reborrow_hints(hints, sema, config, &expr),
+                    ast::Expr::ClosureExpr(it) => {
+                        closure_captures::hints(hints, famous_defs, config, file_id, it.clone());
+                        closure_ret::hints(hints, famous_defs, config, file_id, it)
+                    },
                     _ => None,
                 }
             },
@@ -535,6 +540,7 @@ mod tests {
         chaining_hints: false,
         lifetime_elision_hints: LifetimeElisionHints::Never,
         closure_return_type_hints: ClosureReturnTypeHints::Never,
+        closure_capture_hints: false,
         adjustment_hints: AdjustmentHints::Never,
         adjustment_hints_mode: AdjustmentHintsMode::Prefix,
         adjustment_hints_hide_outside_unsafe: false,

crates/ide/src/inlay_hints/closure_captures.rs

@@ -0,0 +1,194 @@
+//! Implementation of "closure return type" inlay hints.
+//!
+//! Tests live in [`bind_pat`][super::bind_pat] module.
+use ide_db::{base_db::FileId, famous_defs::FamousDefs};
+use syntax::ast::{self, AstNode};
+use text_edit::{TextRange, TextSize};
+
+use crate::{InlayHint, InlayHintLabel, InlayHintsConfig, InlayKind};
+
+pub(super) fn hints(
+    acc: &mut Vec<InlayHint>,
+    FamousDefs(sema, _): &FamousDefs<'_, '_>,
+    config: &InlayHintsConfig,
+    _file_id: FileId,
+    closure: ast::ClosureExpr,
+) -> Option<()> {
+    if !config.closure_capture_hints {
+        return None;
+    }
+    let ty = &sema.type_of_expr(&closure.clone().into())?.original;
+    let c = ty.as_closure()?;
+    let captures = c.captured_items(sema.db);
+
+    if captures.is_empty() {
+        return None;
+    }
+
+    let move_kw_range = match closure.move_token() {
+        Some(t) => t.text_range(),
+        None => {
+            let range = closure.syntax().first_token()?.prev_token()?.text_range();
+            let range = TextRange::new(range.end() - TextSize::from(1), range.end());
+            acc.push(InlayHint {
+                range,
+                kind: InlayKind::ClosureCapture,
+                label: InlayHintLabel::simple("move", None, None),
+                text_edit: None,
+            });
+            range
+        }
+    };
+    acc.push(InlayHint {
+        range: move_kw_range,
+        kind: InlayKind::ClosureCapture,
+        label: InlayHintLabel::from("("),
+        text_edit: None,
+    });
+    let last = captures.len() - 1;
+    for (idx, capture) in captures.into_iter().enumerate() {
+        let local = capture.local();
+        let source = local.primary_source(sema.db);
+
+        // force cache the source file, otherwise sema lookup will potentially panic
+        _ = sema.parse_or_expand(source.file());
+
+        acc.push(InlayHint {
+            range: move_kw_range,
+            kind: InlayKind::ClosureCapture,
+            label: InlayHintLabel::simple(
+                format!(
+                    "{}{}",
+                    match capture.kind() {
+                        hir::CaptureKind::SharedRef => "&",
+                        hir::CaptureKind::UniqueSharedRef => "&unique ",
+                        hir::CaptureKind::MutableRef => "&mut ",
+                        hir::CaptureKind::Move => "",
+                    },
+                    local.name(sema.db)
+                ),
+                None,
+                source.name().and_then(|name| sema.original_range_opt(name.syntax())),
+            ),
+            text_edit: None,
+        });
+
+        if idx != last {
+            acc.push(InlayHint {
+                range: move_kw_range,
+                kind: InlayKind::ClosureCapture,
+                label: InlayHintLabel::simple(", ", None, None),
+                text_edit: None,
+            });
+        }
+    }
+    acc.push(InlayHint {
+        range: move_kw_range,
+        kind: InlayKind::ClosureCapture,
+        label: InlayHintLabel::from(")"),
+        text_edit: None,
+    });
+
+    Some(())
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        inlay_hints::tests::{check_with_config, DISABLED_CONFIG},
+        InlayHintsConfig,
+    };
+
+    #[test]
+    fn all_capture_kinds() {
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, derive
+
+
+#[derive(Copy, Clone)]
+struct Copy;
+
+struct NonCopy;
+
+fn main() {
+    let foo = Copy;
+    let bar = NonCopy;
+    let mut baz = NonCopy;
+    let qux = &mut NonCopy;
+    // FIXME: &foo should be foo
+    || {
+// ^ move
+// ^ (
+// ^ &foo
+// ^ ,
+// ^ bar
+// ^ ,
+// ^ baz
+// ^ ,
+// ^ qux
+// ^ )
+        foo;
+        bar;
+        baz;
+        qux;
+    };
+    || {
+// ^ move
+// ^ (
+// ^ &foo
+// ^ ,
+// ^ &bar
+// ^ ,
+// ^ &baz
+// ^ ,
+// ^ &qux
+// ^ )
+        &foo;
+        &bar;
+        &baz;
+        &qux;
+    };
+    || {
+// ^ move
+// ^ (
+// ^ &mut baz
+// ^ )
+        &mut baz;
+    };
+    // FIXME: &mut qux should be &unique qux
+    || {
+// ^ move
+// ^ (
+// ^ &mut baz
+// ^ ,
+// ^ &mut qux
+// ^ )
+        baz = NonCopy;
+        *qux = NonCopy;
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn move_token() {
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, derive
+fn main() {
+    let foo = u32;
+    move || {
+//  ^^^^ (
+//  ^^^^ foo
+//  ^^^^ )
+        foo;
+    };
+}
+"#,
+        );
+    }
+}

crates/ide/src/static_index.rs

@@ -122,6 +122,7 @@ impl StaticIndex<'_> {
                     param_names_for_lifetime_elision_hints: false,
                     binding_mode_hints: false,
                     max_length: Some(25),
+                    closure_capture_hints: false,
                     closing_brace_hints_min_lines: Some(25),
                 },
                 file_id,

crates/rust-analyzer/src/config.rs

@@ -338,6 +338,8 @@ config_data! {
         /// Minimum number of lines required before the `}` until the hint is shown (set to 0 or 1
         /// to always show them).
         inlayHints_closingBraceHints_minLines: usize               = "25",
+        /// Whether to show inlay hints for closure captures.
+        inlayHints_closureCaptureHints_enable: bool                          = "false",
         /// Whether to show inlay type hints for return types of closures.
         inlayHints_closureReturnTypeHints_enable: ClosureReturnTypeHintsDef  = "\"never\"",
         /// Closure notation in type and chaining inlay hints.
@@ -1312,6 +1314,7 @@ impl Config {
                 ClosureStyle::WithId => hir::ClosureStyle::ClosureWithId,
                 ClosureStyle::Hide => hir::ClosureStyle::Hide,
             },
+            closure_capture_hints: self.data.inlayHints_closureCaptureHints_enable,
             adjustment_hints: match self.data.inlayHints_expressionAdjustmentHints_enable {
                 AdjustmentHintsDef::Always => ide::AdjustmentHints::Always,
                 AdjustmentHintsDef::Never => match self.data.inlayHints_reborrowHints_enable {