Suggest `.await` when type impls IntoFuture

2 files changed, 42 insertions, 7 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 800ea58ba29..7f16634afe1 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2780,19 +2780,28 @@ impl Type {
     /// Checks that particular type `ty` implements `std::future::Future`.
     /// This function is used in `.await` syntax completion.
     pub fn impls_future(&self, db: &dyn HirDatabase) -> bool {
-        // FIXME: This should be checking for IntoFuture trait, but I don't know how to find the
-        // right TraitId in this crate.
-        let std_future_trait = db
-            .lang_item(self.env.krate, SmolStr::new_inline("future_trait"))
-            .and_then(|it| it.as_trait());
-        let std_future_trait = match std_future_trait {
+        let trait_ = db
+            .lang_item(self.env.krate, SmolStr::new_inline("into_future"))
+            .and_then(|it| {
+                let into_future_fn = it.as_function()?;
+                let assoc_item = as_assoc_item(db, AssocItem::Function, into_future_fn)?;
+                let into_future_trait = assoc_item.containing_trait_or_trait_impl(db)?;
+                Some(into_future_trait.id)
+            })
+            .or_else(|| {
+                let future_trait =
+                    db.lang_item(self.env.krate, SmolStr::new_inline("future_trait"))?;
+                future_trait.as_trait()
+            });
+
+        let trait_ = match trait_ {
             Some(it) => it,
             None => return false,
         };
 
         let canonical_ty =
             Canonical { value: self.ty.clone(), binders: CanonicalVarKinds::empty(Interner) };
-        method_resolution::implements_trait(&canonical_ty, db, self.env.clone(), std_future_trait)
+        method_resolution::implements_trait(&canonical_ty, db, self.env.clone(), trait_)
     }
 
     /// Checks that particular type `ty` implements `std::ops::FnOnce`.
diff --git a/crates/ide-completion/src/completions/keyword.rs b/crates/ide-completion/src/completions/keyword.rs
index 032b23725f8..1d03c8cc5ca 100644
--- a/crates/ide-completion/src/completions/keyword.rs
+++ b/crates/ide-completion/src/completions/keyword.rs
@@ -115,6 +115,32 @@ fn foo() {
     }
 
     #[test]
+    fn test_completion_await_impls_into_future() {
+        check(
+            r#"
+//- minicore: future
+use core::future::*;
+struct A {}
+impl IntoFuture for A {}
+fn foo(a: A) { a.$0 }
+"#,
+            expect![[r#"
+                kw await                  expr.await
+                me into_future() (as IntoFuture) fn(self) -> <Self as IntoFuture>::IntoFuture
+                sn box                    Box::new(expr)
+                sn call                   function(expr)
+                sn dbg                    dbg!(expr)
+                sn dbgr                   dbg!(&expr)
+                sn let                    let
+                sn letm                   let mut
+                sn match                  match expr {}
+                sn ref                    &expr
+                sn refm                   &mut expr
+            "#]],
+        );
+    }
+
+    #[test]
     fn let_semi() {
         cov_mark::check!(let_semi);
         check_edit(