Auto merge of #12982 - jridgewell:into_future, r=Veykril

Implement IntoFuture type inference

One of my projects is using [IntoFuture](https://doc.rust-lang.org/std/future/trait.IntoFuture.html) to make our async code a little less verbose. However, rust-analyzer can't infer the output type of an await expression if the value uses `IntoFuture` to convert into another type. So we're getting `{unknown}` types everywhere since switching.

`foo.await` itself [desugars](https://github.com/rust-lang/rust/blob/e4417cf020fbcd6182c11637bc6b8694434bd81a/compiler/rustc_ast_lowering/src/expr.rs#L644-L658) into a `match into_future(foo) {}`, with every `Future` impl getting a [default](https://github.com/rust-lang/rust/blob/e4417cf020fbcd6182c11637bc6b8694434bd81a/library/core/src/future/into_future.rs#L131-L139) `IntoFuture` implementation. I'm not sure if we want to disable the old `future_trait` paths, since this only recently [stabilize](https://github.com/rust-lang/rust/pull/98718).

11 files changed, 178 insertions, 35 deletions

diff --git a/crates/hir-expand/src/mod_path.rs b/crates/hir-expand/src/mod_path.rs
index d0f73ec8208..d7586d129b7 100644
--- a/crates/hir-expand/src/mod_path.rs
+++ b/crates/hir-expand/src/mod_path.rs
@@ -257,6 +257,7 @@ macro_rules! __known_path {
     (core::ops::RangeToInclusive) => {};
     (core::ops::RangeInclusive) => {};
     (core::future::Future) => {};
+    (core::future::IntoFuture) => {};
     (core::ops::Try) => {};
     ($path:path) => {
         compile_error!("Please register your known path in the path module")
diff --git a/crates/hir-expand/src/name.rs b/crates/hir-expand/src/name.rs
index a1bb353374b..2b859f77509 100644
--- a/crates/hir-expand/src/name.rs
+++ b/crates/hir-expand/src/name.rs
@@ -266,6 +266,7 @@ pub mod known {
         Try,
         Ok,
         Future,
+        IntoFuture,
         Result,
         Option,
         Output,
@@ -399,6 +400,7 @@ pub mod known {
         future_trait,
         index,
         index_mut,
+        into_future,
         mul_assign,
         mul,
         neg,
diff --git a/crates/hir-ty/src/infer.rs b/crates/hir-ty/src/infer.rs
index 46eeea0e6fc..95a7229e879 100644
--- a/crates/hir-ty/src/infer.rs
+++ b/crates/hir-ty/src/infer.rs
@@ -875,7 +875,10 @@ impl<'a> InferenceContext<'a> {
     }
 
     fn resolve_future_future_output(&self) -> Option<TypeAliasId> {
-        let trait_ = self.resolve_lang_item(name![future_trait])?.as_trait()?;
+        let trait_ = self
+            .resolver
+            .resolve_known_trait(self.db.upcast(), &path![core::future::IntoFuture])
+            .or_else(|| self.resolve_lang_item(name![future_trait])?.as_trait())?;
         self.db.trait_data(trait_).associated_type_by_name(&name![Output])
     }
 
diff --git a/crates/hir-ty/src/tests/traits.rs b/crates/hir-ty/src/tests/traits.rs
index c128a051f72..0f37970e2b3 100644
--- a/crates/hir-ty/src/tests/traits.rs
+++ b/crates/hir-ty/src/tests/traits.rs
@@ -138,6 +138,31 @@ fn not_send() -> Box<dyn Future<Output = ()> + 'static> {
 }
 
 #[test]
+fn into_future_trait() {
+    check_types(
+        r#"
+//- minicore: future
+struct Futurable;
+impl core::future::IntoFuture for Futurable {
+    type Output = u64;
+    type IntoFuture = IntFuture;
+}
+
+struct IntFuture;
+impl core::future::Future for IntFuture {
+    type Output = u64;
+}
+
+fn test() {
+    let r = Futurable;
+    let v = r.await;
+    v;
+} //^ u64
+"#,
+    );
+}
+
+#[test]
 fn infer_try() {
     check_types(
         r#"
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 3561bdeba7f..aa019ca4838 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2778,20 +2778,32 @@ impl Type {
         self.ty.is_unknown()
     }
 
-    /// Checks that particular type `ty` implements `std::future::Future`.
+    /// Checks that particular type `ty` implements `std::future::IntoFuture` or
+    /// `std::future::Future`.
     /// This function is used in `.await` syntax completion.
-    pub fn impls_future(&self, db: &dyn HirDatabase) -> bool {
-        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 {
+    pub fn impls_into_future(&self, db: &dyn HirDatabase) -> bool {
+        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/hir/src/source_analyzer.rs b/crates/hir/src/source_analyzer.rs
index ae2896e1932..bd35af06e23 100644
--- a/crates/hir/src/source_analyzer.rs
+++ b/crates/hir/src/source_analyzer.rs
@@ -27,6 +27,7 @@ use hir_def::{
 use hir_expand::{
     builtin_fn_macro::BuiltinFnLikeExpander,
     hygiene::Hygiene,
+    mod_path::path,
     name,
     name::{AsName, Name},
     HirFileId, InFile,
@@ -269,14 +270,35 @@ impl SourceAnalyzer {
         db: &dyn HirDatabase,
         await_expr: &ast::AwaitExpr,
     ) -> Option<FunctionId> {
-        let ty = self.ty_of_expr(db, &await_expr.expr()?.into())?;
+        let mut ty = self.ty_of_expr(db, &await_expr.expr()?.into())?.clone();
+
+        let into_future_trait = self
+            .resolver
+            .resolve_known_trait(db.upcast(), &path![core::future::IntoFuture])
+            .map(Trait::from);
+
+        if let Some(into_future_trait) = into_future_trait {
+            let type_ = Type::new_with_resolver(db, &self.resolver, ty.clone());
+            if type_.impls_trait(db, into_future_trait, &[]) {
+                let items = into_future_trait.items(db);
+                let into_future_type = items.into_iter().find_map(|item| match item {
+                    AssocItem::TypeAlias(alias)
+                        if alias.name(db) == hir_expand::name![IntoFuture] =>
+                    {
+                        Some(alias)
+                    }
+                    _ => None,
+                })?;
+                let future_trait = type_.normalize_trait_assoc_type(db, &[], into_future_type)?;
+                ty = future_trait.ty;
+            }
+        }
 
-        let op_fn = db
+        let poll_fn = db
             .lang_item(self.resolver.krate(), hir_expand::name![poll].to_smol_str())?
             .as_function()?;
-        let substs = hir_ty::TyBuilder::subst_for_def(db, op_fn).push(ty.clone()).build();
-
-        Some(self.resolve_impl_method_or_trait_def(db, op_fn, &substs))
+        let substs = hir_ty::TyBuilder::subst_for_def(db, poll_fn).push(ty.clone()).build();
+        Some(self.resolve_impl_method_or_trait_def(db, poll_fn, &substs))
     }
 
     pub(crate) fn resolve_prefix_expr(
diff --git a/crates/ide-assists/src/utils/suggest_name.rs b/crates/ide-assists/src/utils/suggest_name.rs
index 662ddd063f3..c521a10fccf 100644
--- a/crates/ide-assists/src/utils/suggest_name.rs
+++ b/crates/ide-assists/src/utils/suggest_name.rs
@@ -55,6 +55,7 @@ const USELESS_METHODS: &[&str] = &[
     "iter",
     "into_iter",
     "iter_mut",
+    "into_future",
 ];
 
 pub(crate) fn for_generic_parameter(ty: &ast::ImplTraitType) -> SmolStr {
diff --git a/crates/ide-completion/src/completions/dot.rs b/crates/ide-completion/src/completions/dot.rs
index cf40ca489c0..02004ff7b68 100644
--- a/crates/ide-completion/src/completions/dot.rs
+++ b/crates/ide-completion/src/completions/dot.rs
@@ -19,7 +19,7 @@ pub(crate) fn complete_dot(
     };
 
     // Suggest .await syntax for types that implement Future trait
-    if receiver_ty.impls_future(ctx.db) {
+    if receiver_ty.impls_into_future(ctx.db) {
         let mut item =
             CompletionItem::new(CompletionItemKind::Keyword, ctx.source_range(), "await");
         item.detail("expr.await");
diff --git a/crates/ide-completion/src/completions/keyword.rs b/crates/ide-completion/src/completions/keyword.rs
index 3989a451bde..1d03c8cc5ca 100644
--- a/crates/ide-completion/src/completions/keyword.rs
+++ b/crates/ide-completion/src/completions/keyword.rs
@@ -75,16 +75,17 @@ impl Future for A {}
 fn foo(a: A) { a.$0 }
 "#,
             expect![[r#"
-                kw await expr.await
-                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
+                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
             "#]],
         );
 
@@ -98,18 +99,45 @@ fn foo() {
 }
 "#,
             expect![[r#"
-                kw await expr.await
-                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
+                kw await                  expr.await
+                me into_future() (use core::future::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 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]
diff --git a/crates/ide/src/goto_definition.rs b/crates/ide/src/goto_definition.rs
index d6cd5783f05..36a648fe4a8 100644
--- a/crates/ide/src/goto_definition.rs
+++ b/crates/ide/src/goto_definition.rs
@@ -1665,6 +1665,40 @@ fn f() {
     }
 
     #[test]
+    fn goto_await_into_future_poll() {
+        check(
+            r#"
+//- minicore: future
+
+struct Futurable;
+
+impl core::future::IntoFuture for Futurable {
+    type IntoFuture = MyFut;
+}
+
+struct MyFut;
+
+impl core::future::Future for MyFut {
+    type Output = ();
+
+    fn poll(
+     //^^^^
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>
+    ) -> std::task::Poll<Self::Output>
+    {
+        ()
+    }
+}
+
+fn f() {
+    Futurable.await$0;
+}
+"#,
+        );
+    }
+
+    #[test]
     fn goto_try_op() {
         check(
             r#"
diff --git a/crates/test-utils/src/minicore.rs b/crates/test-utils/src/minicore.rs
index f48d1ec66aa..6df29db4745 100644
--- a/crates/test-utils/src/minicore.rs
+++ b/crates/test-utils/src/minicore.rs
@@ -471,6 +471,21 @@ pub mod future {
         #[lang = "poll"]
         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>;
     }
+
+    pub trait IntoFuture {
+        type Output;
+        type IntoFuture: Future<Output = Self::Output>;
+        #[lang = "into_future"]
+        fn into_future(self) -> Self::IntoFuture;
+    }
+
+    impl<F: Future> IntoFuture for F {
+        type Output = F::Output;
+        type IntoFuture = F;
+        fn into_future(self) -> F {
+            self
+        }
+    }
 }
 pub mod task {
     pub enum Poll<T> {