Auto merge of #64999 - nikomatsakis:issue-60424-async-return-inference, r=cramertj

extract expected return type for async fn generators

Fixes #60424

cc @Centril, I know you've been eager to see this fixed.

r? @cramertj

Author: bors

src/librustc/hir/lowering/expr.rs

@@ -89,9 +89,14 @@ impl LoweringContext<'_> {
                 hir::MatchSource::Normal,
             ),
             ExprKind::Async(capture_clause, closure_node_id, ref block) => {
-                self.make_async_expr(capture_clause, closure_node_id, None, block.span, |this| {
-                    this.with_new_scopes(|this| this.lower_block_expr(block))
-                })
+                self.make_async_expr(
+                    capture_clause,
+                    closure_node_id,
+                    None,
+                    block.span,
+                    hir::AsyncGeneratorKind::Block,
+                    |this| this.with_new_scopes(|this| this.lower_block_expr(block)),
+                )
             }
             ExprKind::Await(ref expr) => self.lower_expr_await(e.span, expr),
             ExprKind::Closure(
@@ -457,6 +462,7 @@ impl LoweringContext<'_> {
         closure_node_id: NodeId,
         ret_ty: Option<AstP<Ty>>,
         span: Span,
+        async_gen_kind: hir::AsyncGeneratorKind,
         body: impl FnOnce(&mut LoweringContext<'_>) -> hir::Expr,
     ) -> hir::ExprKind {
         let capture_clause = self.lower_capture_clause(capture_clause);
@@ -470,7 +476,7 @@ impl LoweringContext<'_> {
         };
         let decl = self.lower_fn_decl(&ast_decl, None, /* impl trait allowed */ false, None);
         let body_id = self.lower_fn_body(&ast_decl, |this| {
-            this.generator_kind = Some(hir::GeneratorKind::Async);
+            this.generator_kind = Some(hir::GeneratorKind::Async(async_gen_kind));
             body(this)
         });
 
@@ -522,7 +528,7 @@ impl LoweringContext<'_> {
     /// ```
     fn lower_expr_await(&mut self, await_span: Span, expr: &Expr) -> hir::ExprKind {
         match self.generator_kind {
-            Some(hir::GeneratorKind::Async) => {},
+            Some(hir::GeneratorKind::Async(_)) => {},
             Some(hir::GeneratorKind::Gen) |
             None => {
                 let mut err = struct_span_err!(
@@ -727,7 +733,7 @@ impl LoweringContext<'_> {
                     Movability::Static => hir::GeneratorMovability::Static,
                 })
             },
-            Some(hir::GeneratorKind::Async) => {
+            Some(hir::GeneratorKind::Async(_)) => {
                 bug!("non-`async` closure body turned `async` during lowering");
             },
             None => {
@@ -786,10 +792,12 @@ impl LoweringContext<'_> {
                     None
                 };
                 let async_body = this.make_async_expr(
-                    capture_clause, closure_id, async_ret_ty, body.span,
-                    |this| {
-                        this.with_new_scopes(|this| this.lower_expr(body))
-                    }
+                    capture_clause,
+                    closure_id,
+                    async_ret_ty,
+                    body.span,
+                    hir::AsyncGeneratorKind::Closure,
+                    |this| this.with_new_scopes(|this| this.lower_expr(body)),
                 );
                 this.expr(fn_decl_span, async_body, ThinVec::new())
             });
@@ -1005,7 +1013,7 @@ impl LoweringContext<'_> {
     fn lower_expr_yield(&mut self, span: Span, opt_expr: Option<&Expr>) -> hir::ExprKind {
         match self.generator_kind {
             Some(hir::GeneratorKind::Gen) => {},
-            Some(hir::GeneratorKind::Async) => {
+            Some(hir::GeneratorKind::Async(_)) => {
                 span_err!(
                     self.sess,
                     span,

src/librustc/hir/lowering/item.rs

@@ -1222,7 +1222,11 @@ impl LoweringContext<'_> {
             }
 
             let async_expr = this.make_async_expr(
-                CaptureBy::Value, closure_id, None, body.span,
+                CaptureBy::Value,
+                closure_id,
+                None,
+                body.span,
+                hir::AsyncGeneratorKind::Fn,
                 |this| {
                     // Create a block from the user's function body:
                     let user_body = this.lower_block_expr(body);

src/librustc/hir/mod.rs

@@ -1362,21 +1362,49 @@ impl Body {
 }
 
 /// The type of source expression that caused this generator to be created.
-// Not `IsAsync` because we want to eventually add support for `AsyncGen`
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, HashStable,
          RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
 pub enum GeneratorKind {
-    /// An `async` block or function.
-    Async,
+    /// An explicit `async` block or the body of an async function.
+    Async(AsyncGeneratorKind),
+
     /// A generator literal created via a `yield` inside a closure.
     Gen,
 }
 
 impl fmt::Display for GeneratorKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            GeneratorKind::Async(k) => fmt::Display::fmt(k, f),
+            GeneratorKind::Gen => f.write_str("generator"),
+        }
+    }
+}
+
+/// In the case of a generator created as part of an async construct,
+/// which kind of async construct caused it to be created?
+///
+/// This helps error messages but is also used to drive coercions in
+/// type-checking (see #60424).
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, HashStable,
+         RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
+pub enum AsyncGeneratorKind {
+    /// An explicit `async` block written by the user.
+    Block,
+
+    /// An explicit `async` block written by the user.
+    Closure,
+
+    /// The `async` block generated as the body of an async function.
+    Fn,
+}
+
+impl fmt::Display for AsyncGeneratorKind {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.write_str(match self {
-            GeneratorKind::Async => "`async` object",
-            GeneratorKind::Gen => "generator",
+            AsyncGeneratorKind::Block => "`async` block",
+            AsyncGeneratorKind::Closure => "`async` closure body",
+            AsyncGeneratorKind::Fn => "`async fn` body",
         })
     }
 }
@@ -1758,6 +1786,7 @@ pub struct Destination {
 pub enum GeneratorMovability {
     /// May contain self-references, `!Unpin`.
     Static,
+
     /// Must not contain self-references, `Unpin`.
     Movable,
 }

src/librustc/infer/mod.rs

@@ -1319,6 +1319,14 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
         }
     }
 
+    /// Resolve any type variables found in `value` -- but only one
+    /// level.  So, if the variable `?X` is bound to some type
+    /// `Foo<?Y>`, then this would return `Foo<?Y>` (but `?Y` may
+    /// itself be bound to a type).
+    ///
+    /// Useful when you only need to inspect the outermost level of
+    /// the type and don't care about nested types (or perhaps you
+    /// will be resolving them as well, e.g. in a loop).
     pub fn shallow_resolve<T>(&self, value: T) -> T
     where
         T: TypeFoldable<'tcx>,
@@ -1579,6 +1587,9 @@ impl<'a, 'tcx> ShallowResolver<'a, 'tcx> {
         ShallowResolver { infcx }
     }
 
+    /// If `typ` is a type variable of some kind, resolve it one level
+    /// (but do not resolve types found in the result). If `typ` is
+    /// not a type variable, just return it unmodified.
     pub fn shallow_resolve(&mut self, typ: Ty<'tcx>) -> Ty<'tcx> {
         match typ.kind {
             ty::Infer(ty::TyVar(v)) => {

src/librustc_typeck/check/closure.rs

@@ -337,7 +337,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     ) -> ClosureSignatures<'tcx> {
         debug!("sig_of_closure_no_expectation()");
 
-        let bound_sig = self.supplied_sig_of_closure(expr_def_id, decl);
+        let bound_sig = self.supplied_sig_of_closure(expr_def_id, decl, body);
 
         self.closure_sigs(expr_def_id, body, bound_sig)
     }
@@ -490,7 +490,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         //
         // (See comment on `sig_of_closure_with_expectation` for the
         // meaning of these letters.)
-        let supplied_sig = self.supplied_sig_of_closure(expr_def_id, decl);
+        let supplied_sig = self.supplied_sig_of_closure(expr_def_id, decl, body);
 
         debug!(
             "check_supplied_sig_against_expectation: supplied_sig={:?}",
@@ -591,14 +591,31 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         &self,
         expr_def_id: DefId,
         decl: &hir::FnDecl,
+        body: &hir::Body,
     ) -> ty::PolyFnSig<'tcx> {
         let astconv: &dyn AstConv<'_> = self;
 
+        debug!(
+            "supplied_sig_of_closure(decl={:?}, body.generator_kind={:?})",
+            decl,
+            body.generator_kind,
+        );
+
         // First, convert the types that the user supplied (if any).
         let supplied_arguments = decl.inputs.iter().map(|a| astconv.ast_ty_to_ty(a));
         let supplied_return = match decl.output {
             hir::Return(ref output) => astconv.ast_ty_to_ty(&output),
-            hir::DefaultReturn(_) => astconv.ty_infer(None, decl.output.span()),
+            hir::DefaultReturn(_) => match body.generator_kind {
+                // In the case of the async block that we create for a function body,
+                // we expect the return type of the block to match that of the enclosing
+                // function.
+                Some(hir::GeneratorKind::Async(hir::AsyncGeneratorKind::Fn)) => {
+                    debug!("supplied_sig_of_closure: closure is async fn body");
+                    self.deduce_future_output_from_obligations(expr_def_id)
+                }
+
+                _ => astconv.ty_infer(None, decl.output.span()),
+            }
         };
 
         let result = ty::Binder::bind(self.tcx.mk_fn_sig(
@@ -620,6 +637,117 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         result
     }
 
+    /// Invoked when we are translating the generator that results
+    /// from desugaring an `async fn`. Returns the "sugared" return
+    /// type of the `async fn` -- that is, the return type that the
+    /// user specified. The "desugared" return type is a `impl
+    /// Future<Output = T>`, so we do this by searching through the
+    /// obligations to extract the `T`.
+    fn deduce_future_output_from_obligations(
+        &self,
+        expr_def_id: DefId,
+    ) -> Ty<'tcx> {
+        debug!("deduce_future_output_from_obligations(expr_def_id={:?})", expr_def_id);
+
+        let ret_coercion =
+            self.ret_coercion
+            .as_ref()
+            .unwrap_or_else(|| span_bug!(
+                self.tcx.def_span(expr_def_id),
+                "async fn generator outside of a fn"
+            ));
+
+        // In practice, the return type of the surrounding function is
+        // always a (not yet resolved) inference variable, because it
+        // is the hidden type for an `impl Trait` that we are going to
+        // be inferring.
+        let ret_ty = ret_coercion.borrow().expected_ty();
+        let ret_ty = self.inh.infcx.shallow_resolve(ret_ty);
+        let ret_vid = match ret_ty.kind {
+            ty::Infer(ty::TyVar(ret_vid)) => ret_vid,
+            _ => {
+                span_bug!(
+                    self.tcx.def_span(expr_def_id),
+                    "async fn generator return type not an inference variable"
+                )
+            }
+        };
+
+        // Search for a pending obligation like
+        //
+        // `<R as Future>::Output = T`
+        //
+        // where R is the return type we are expecting. This type `T`
+        // will be our output.
+        let output_ty = self.obligations_for_self_ty(ret_vid)
+            .find_map(|(_, obligation)| {
+                if let ty::Predicate::Projection(ref proj_predicate) = obligation.predicate {
+                    self.deduce_future_output_from_projection(
+                        obligation.cause.span,
+                        proj_predicate
+                    )
+                } else {
+                    None
+                }
+            })
+            .unwrap();
+
+        debug!("deduce_future_output_from_obligations: output_ty={:?}", output_ty);
+        output_ty
+    }
+
+    /// Given a projection like
+    ///
+    /// `<X as Future>::Output = T`
+    ///
+    /// where `X` is some type that has no late-bound regions, returns
+    /// `Some(T)`. If the projection is for some other trait, returns
+    /// `None`.
+    fn deduce_future_output_from_projection(
+        &self,
+        cause_span: Span,
+        predicate: &ty::PolyProjectionPredicate<'tcx>,
+    ) -> Option<Ty<'tcx>> {
+        debug!("deduce_future_output_from_projection(predicate={:?})", predicate);
+
+        // We do not expect any bound regions in our predicate, so
+        // skip past the bound vars.
+        let predicate = match predicate.no_bound_vars() {
+            Some(p) => p,
+            None => {
+                debug!("deduce_future_output_from_projection: has late-bound regions");
+                return None;
+            }
+        };
+
+        // Check that this is a projection from the `Future` trait.
+        let trait_ref = predicate.projection_ty.trait_ref(self.tcx);
+        let future_trait = self.tcx.lang_items().future_trait().unwrap();
+        if trait_ref.def_id != future_trait {
+            debug!("deduce_future_output_from_projection: not a future");
+            return None;
+        }
+
+        // The `Future` trait has only one associted item, `Output`,
+        // so check that this is what we see.
+        let output_assoc_item = self.tcx.associated_items(future_trait).nth(0).unwrap().def_id;
+        if output_assoc_item != predicate.projection_ty.item_def_id {
+            span_bug!(
+                cause_span,
+                "projecting associated item `{:?}` from future, which is not Output `{:?}`",
+                predicate.projection_ty.item_def_id,
+                output_assoc_item,
+            );
+        }
+
+        // Extract the type from the projection. Note that there can
+        // be no bound variables in this type because the "self type"
+        // does not have any regions in it.
+        let output_ty = self.resolve_vars_if_possible(&predicate.ty);
+        debug!("deduce_future_output_from_projection: output_ty={:?}", output_ty);
+        Some(output_ty)
+    }
+
     /// Converts the types that the user supplied, in case that doing
     /// so should yield an error, but returns back a signature where
     /// all parameters are of type `TyErr`.

src/librustc_typeck/check/generator_interior.rs

@@ -55,7 +55,7 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
             expr_and_pat_count: 0,
             source: match self.kind { // Guess based on the kind of the current generator.
                 hir::GeneratorKind::Gen => hir::YieldSource::Yield,
-                hir::GeneratorKind::Async => hir::YieldSource::Await,
+                hir::GeneratorKind::Async(_) => hir::YieldSource::Await,
             },
         }));
 

src/librustc_typeck/check/mod.rs

@@ -562,7 +562,19 @@ pub struct FnCtxt<'a, 'tcx> {
     // if type checking is run in parallel.
     err_count_on_creation: usize,
 
+    /// If `Some`, this stores coercion information for returned
+    /// expressions. If `None`, this is in a context where return is
+    /// inappropriate, such as a const expression.
+    ///
+    /// This is a `RefCell<DynamicCoerceMany>`, which means that we
+    /// can track all the return expressions and then use them to
+    /// compute a useful coercion from the set, similar to a match
+    /// expression or other branching context. You can use methods
+    /// like `expected_ty` to access the declared return type (if
+    /// any).
     ret_coercion: Option<RefCell<DynamicCoerceMany<'tcx>>>,
+
+    /// First span of a return site that we find. Used in error messages.
     ret_coercion_span: RefCell<Option<Span>>,
 
     yield_ty: Option<Ty<'tcx>>,
@@ -4534,7 +4546,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         let item_id = self.tcx().hir().get_parent_node(self.body_id);
         if let Some(body_id) = self.tcx().hir().maybe_body_owned_by(item_id) {
             let body = self.tcx().hir().body(body_id);
-            if let Some(hir::GeneratorKind::Async) = body.generator_kind {
+            if let Some(hir::GeneratorKind::Async(_)) = body.generator_kind {
                 let sp = expr.span;
                 // Check for `Future` implementations by constructing a predicate to
                 // prove: `<T as Future>::Output == U`

src/test/ui/async-await/async-block-control-flow-static-semantics.rs

@@ -20,7 +20,7 @@ fn return_targets_async_block_not_fn() -> u8 {
 }
 
 async fn return_targets_async_block_not_async_fn() -> u8 {
-    //~^ ERROR type mismatch resolving
+    //~^ ERROR mismatched types
     let block = async {
         return 0u8;
     };