Merge pull request #22811 from Veykril/lukaswirth/push-vntrlsumxwqz

feat: Add capture hints to coroutines

Author: Veykril

crates/hir/src/lib.rs

@@ -87,7 +87,10 @@ use hir_ty::{
     GenericPredicates, InferBodyId, InferenceResult, ParamEnvAndCrate, TyDefId,
     TyLoweringDiagnostic, ValueTyDefId, all_super_traits, autoderef, check_orphan_rules,
     consteval::try_const_usize,
-    db::{AnonConstId, InternedClosure, InternedClosureId, InternedCoroutineClosureId},
+    db::{
+        AnonConstId, InternedClosure, InternedClosureId, InternedCoroutineClosureId,
+        InternedCoroutineId,
+    },
     diagnostics::BodyValidationDiagnostic,
     direct_super_traits, known_const_to_ast,
     layout::{Layout as TyLayout, RustcEnumVariantIdx, RustcFieldIdx, TagEncoding},
@@ -4937,15 +4940,7 @@ impl<'db> Closure<'db> {
             AnyClosureId::ClosureId(it) => it.loc(db),
             AnyClosureId::CoroutineClosureId(it) => it.loc(db),
         };
-        let InternedClosure { owner: infer_owner, expr: closure, .. } = closure;
-        let infer = InferenceResult::of(db, infer_owner);
-        let owner = infer_owner.expression_store_owner(db);
-        infer.closures_data[&closure]
-            .min_captures
-            .values()
-            .flatten()
-            .map(|capture| ClosureCapture { owner, infer_owner, closure, capture })
-            .collect()
+        captured_items(db, closure)
     }
 
     pub fn fn_trait(&self, _db: &dyn HirDatabase) -> FnTrait {
@@ -4964,6 +4959,34 @@ impl<'db> Closure<'db> {
     }
 }
 
+/// A coroutine expression, including async, generator, and async-generator coroutines.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct Coroutine {
+    id: InternedCoroutineId,
+}
+
+impl Coroutine {
+    /// Returns the values captured by this coroutine.
+    pub fn captured_items<'db>(&self, db: &'db dyn HirDatabase) -> Vec<ClosureCapture<'db>> {
+        captured_items(db, self.id.loc(db))
+    }
+}
+
+fn captured_items<'db>(
+    db: &'db dyn HirDatabase,
+    closure: InternedClosure,
+) -> Vec<ClosureCapture<'db>> {
+    let InternedClosure { owner: infer_owner, expr: closure, .. } = closure;
+    let infer = InferenceResult::of(db, infer_owner);
+    let owner = infer_owner.expression_store_owner(db);
+    infer.closures_data[&closure]
+        .min_captures
+        .values()
+        .flatten()
+        .map(|capture| ClosureCapture { owner, infer_owner, closure, capture })
+        .collect()
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum FnTrait {
     FnOnce,
@@ -5952,6 +5975,14 @@ impl<'db> Type<'db> {
         }
     }
 
+    /// Returns this type as a coroutine.
+    pub fn as_coroutine(&self) -> Option<Coroutine> {
+        match self.ty.skip_binder().kind() {
+            TyKind::Coroutine(id, _) => Some(Coroutine { id: id.0 }),
+            _ => None,
+        }
+    }
+
     pub fn is_fn(&self) -> bool {
         matches!(self.ty.skip_binder().kind(), TyKind::FnDef(..) | TyKind::FnPtr { .. })
     }

crates/ide/src/inlay_hints.rs

@@ -235,9 +235,22 @@ fn hints(
                         param_name::hints(hints, famous_defs, config, file_id, ast::Expr::from(it))
                     }
                     ast::Expr::ClosureExpr(it) => {
-                        closure_captures::hints(hints, famous_defs, config, it.clone(), file_id.edition(sema.db));
+                        closure_captures::hints(
+                            hints,
+                            famous_defs,
+                            config,
+                            Either::Left(it.clone()),
+                            file_id.edition(sema.db),
+                        );
                         closure_ret::hints(hints, famous_defs, config, display_target, it)
                     },
+                    ast::Expr::BlockExpr(it) => closure_captures::hints(
+                        hints,
+                        famous_defs,
+                        config,
+                        Either::Right(it),
+                        file_id.edition(sema.db),
+                    ),
                     ast::Expr::RangeExpr(it) => range_exclusive::hints(hints, famous_defs, config, it),
                     ast::Expr::Literal(it) => ra_fixture::hints(hints, famous_defs.0, file_id, config, it),
                     _ => Some(()),

crates/ide/src/inlay_hints/closure_captures.rs

@@ -1,6 +1,7 @@
 //! Implementation of "closure captures" inlay hints.
 //!
 //! Tests live in [`bind_pat`][super::bind_pat] module.
+use either::Either;
 use ide_db::famous_defs::FamousDefs;
 use span::Edition;
 use stdx::{TupleExt, never};
@@ -14,26 +15,56 @@ pub(super) fn hints(
     acc: &mut Vec<InlayHint>,
     FamousDefs(sema, _): &FamousDefs<'_, '_>,
     config: &InlayHintsConfig<'_>,
-    closure: ast::ClosureExpr,
+    expr: Either<ast::ClosureExpr, ast::BlockExpr>,
     edition: Edition,
 ) -> 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);
+
+    let (expr, move_token, capture_anchor) = match expr {
+        Either::Left(closure) => {
+            let move_token = closure.move_token();
+            let capture_anchor = closure.param_list()?.pipe_token()?;
+            (closure.into(), move_token, capture_anchor)
+        }
+        Either::Right(block) => {
+            let modifier = block.modifier()?;
+            match modifier {
+                ast::BlockModifier::Async(_)
+                | ast::BlockModifier::Gen(_)
+                | ast::BlockModifier::AsyncGen(_) => (),
+                ast::BlockModifier::Unsafe(_)
+                | ast::BlockModifier::Try { .. }
+                | ast::BlockModifier::Const(_)
+                | ast::BlockModifier::Label(_) => return None,
+            }
+            let move_token = block.move_token();
+            let capture_anchor = block.stmt_list()?.l_curly_token()?;
+            (block.into(), move_token, capture_anchor)
+        }
+    };
+
+    let ty = &sema.type_of_expr(&expr)?.original;
+    let captures = match ty.as_closure() {
+        Some(closure) => closure.captured_items(sema.db),
+        None => ty.as_coroutine()?.captured_items(sema.db),
+    };
 
     if captures.is_empty() {
         return None;
     }
 
-    let (range, label, position, pad_right) = match closure.move_token() {
-        Some(t) => (t.text_range(), InlayHintLabel::default(), InlayHintPosition::After, false),
-        None => {
-            let l_pipe = closure.param_list()?.pipe_token()?.text_range();
-            (l_pipe, InlayHintLabel::from("move"), InlayHintPosition::Before, true)
+    let (range, label, position, pad_right) = match move_token {
+        Some(token) => {
+            (token.text_range(), InlayHintLabel::default(), InlayHintPosition::After, false)
         }
+        None => (
+            capture_anchor.text_range(),
+            InlayHintLabel::from("move"),
+            InlayHintPosition::Before,
+            true,
+        ),
     };
     let mut hint = InlayHint {
         range,
@@ -43,7 +74,7 @@ pub(super) fn hints(
         position,
         pad_left: false,
         pad_right,
-        resolve_parent: Some(closure.syntax().text_range()),
+        resolve_parent: Some(expr.syntax().text_range()),
     };
     hint.label.append_str("(");
     let last = captures.len() - 1;
@@ -186,6 +217,92 @@ fn main() {
     };
 }
 
+"#,
+        );
+    }
+
+    #[test]
+    fn all_capture_kinds_async_block() {
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, derive, future
+
+#[derive(Copy, Clone)]
+struct Copy;
+
+struct NonCopy;
+
+fn main() {
+    let foo = Copy;
+    let bar = NonCopy;
+    let mut baz = NonCopy;
+    let qux = &mut NonCopy;
+    async {
+       // ^ move(&foo, bar, baz, qux)
+        foo;
+        bar;
+        baz;
+        qux;
+    };
+    async {
+       // ^ move(&foo, &bar, &baz, &qux)
+        &foo;
+        &bar;
+        &baz;
+        &qux;
+    };
+    async {
+       // ^ move(&mut baz)
+        &mut baz;
+    };
+    async {
+       // ^ move(&mut baz, &mut *qux)
+        baz = NonCopy;
+        *qux = NonCopy;
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn coroutine_blocks() {
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, future
+fn main() {
+    let foo = 0;
+    gen {
+     // ^ move(&foo)
+        foo;
+        yield ();
+    };
+    async gen {
+           // ^ move(&foo)
+        foo;
+        yield ();
+    };
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn legacy_coroutine() {
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, coroutine
+fn main() {
+    let foo = 0;
+    let coroutine = #[coroutine] || {
+                              // ^ move(&foo)
+        foo;
+        yield ();
+    };
+}
 "#,
         );
     }
@@ -216,6 +333,19 @@ fn main() {
         foo;
     };
 }
+"#,
+        );
+        check_with_config(
+            InlayHintsConfig { closure_capture_hints: true, ..DISABLED_CONFIG },
+            r#"
+//- minicore: copy, future
+fn main() {
+    let foo = 0;
+    async move {
+      //  ^^^^ (foo)
+        foo;
+    };
+}
 "#,
         );
     }

crates/rust-analyzer/src/config.rs

@@ -216,7 +216,7 @@ config_data! {
         /// to always show them).
         inlayHints_closingBraceHints_minLines: usize = 25,
 
-        /// Show inlay hints for closure captures.
+        /// Show inlay hints for closure and coroutine captures.
         inlayHints_closureCaptureHints_enable: bool = false,
 
         /// Show inlay type hints for return types of closures.

docs/book/src/configuration_generated.md

@@ -991,7 +991,7 @@ to always show them).
 
 Default: `false`
 
-Show inlay hints for closure captures.
+Show inlay hints for closure and coroutine captures.
 
 
 ## rust-analyzer.inlayHints.closureReturnTypeHints.enable {#inlayHints.closureReturnTypeHints.enable}

editors/code/package.json

@@ -2237,7 +2237,7 @@
                 "title": "Inlay Hints",
                 "properties": {
                     "rust-analyzer.inlayHints.closureCaptureHints.enable": {
-                        "markdownDescription": "Show inlay hints for closure captures.",
+                        "markdownDescription": "Show inlay hints for closure and coroutine captures.",
                         "default": false,
                         "type": "boolean"
                     }