Find IntoFuture::IntoFuture's poll method

4 files changed, 65 insertions, 10 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 7f16634afe1..46d24b9f146 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2777,9 +2777,10 @@ 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 {
+    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| {
diff --git a/crates/hir/src/source_analyzer.rs b/crates/hir/src/source_analyzer.rs
index 756772cf84e..9418afa91da 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,16 +270,35 @@ impl SourceAnalyzer {
         db: &dyn HirDatabase,
         await_expr: &ast::AwaitExpr,
     ) -> Option<FunctionId> {
-        // FIXME This should be pointing to the poll of IntoFuture::Output's Future impl, but I
-        // don't know how to resolve the Output type so that we can query for its poll method.
-        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-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/src/goto_definition.rs b/crates/ide/src/goto_definition.rs
index b2123b9a879..c7a922e6cc6 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#"