Lower the return type of async functions properly.

Author: Without Boats

src/librustc/hir/lowering.rs

@@ -296,6 +296,14 @@ enum AnonymousLifetimeMode {
     PassThrough,
 }
 
+trait FnRetTransform {
+    fn lower(&self, ctx: &mut LoweringContext, output: &FunctionRetTy, impl_trait_context: ImplTraitContext) -> hir::FunctionRetTy;
+}
+
+struct NonAsyncRet;
+struct AsyncClosureRet;
+struct AsyncFunctionRet;
+
 impl<'a> LoweringContext<'a> {
     fn lower_crate(mut self, c: &Crate) -> hir::Crate {
         /// Full-crate AST visitor that inserts into a fresh
@@ -1050,7 +1058,7 @@ impl<'a> LoweringContext<'a> {
                         ),
                         unsafety: this.lower_unsafety(f.unsafety),
                         abi: f.abi,
-                        decl: this.lower_fn_decl(&f.decl, None, false),
+                        decl: this.lower_fn_decl(&f.decl, None, false, &NonAsyncRet),
                         arg_names: this.lower_fn_args_to_names(&f.decl),
                     }))
                 },
@@ -1729,7 +1737,121 @@ impl<'a> LoweringContext<'a> {
         decl: &FnDecl,
         fn_def_id: Option<DefId>,
         impl_trait_return_allow: bool,
+        fn_ret_transform: &FnRetTransform,
     ) -> P<hir::FnDecl> {
+        // For non async functions, the return type is lowered in a straightforward way
+        impl FnRetTransform for NonAsyncRet {
+            fn lower(&self, ctx: &mut LoweringContext, output: &FunctionRetTy, impl_trait_context: ImplTraitContext) -> hir::FunctionRetTy {
+                match output {
+                    FunctionRetTy::Ty(ty) => hir::Return(ctx.lower_ty(ty, impl_trait_context)),
+                    FunctionRetTy::Default(span) => hir::DefaultReturn(*span),
+                }
+            }
+        }
+
+        // For async closures:
+        //  - If the return type is omitted, a straightforward lowering is performed
+        //  - If the return type is explicit, given a return type T, it is lowered to
+        //          ::std::raw::GenFuture<_, T>
+        impl FnRetTransform for AsyncClosureRet {
+            fn lower(&self, ctx: &mut LoweringContext, output: &FunctionRetTy, impl_trait_context: ImplTraitContext) -> hir::FunctionRetTy {
+                match output {
+                    FunctionRetTy::Ty(ty) => {
+                        let inner_ty = ctx.lower_ty(ty, impl_trait_context);
+                        let span = inner_ty.span;
+                        let hir::Path { def, segments, .. } = ctx.std_path(span, &["raw", "GenFuture"], false);
+                        let LoweredNodeId { node_id, hir_id } = ctx.next_id();
+                        let gen_future_path = hir::Path {
+                            segments: segments.map_slice(|mut v| {
+                                v.last_mut().unwrap().parameters = Some(P(hir::PathParameters {
+                                    lifetimes: hir_vec![],
+                                    bindings: hir_vec![],
+                                    types: hir_vec![
+                                        P(hir::Ty {
+                                            id: node_id,
+                                            node: hir::TyInfer,
+                                            hir_id, span,
+                                        }),
+                                        inner_ty,
+                                    ],
+                                    parenthesized: false,
+                                }));
+                                v
+                            }),
+                            def, span
+                        };
+
+                        let LoweredNodeId { node_id, hir_id } = ctx.next_id();
+                        hir::Return(P(hir::Ty {
+                            id: node_id,
+                            node: hir::TyPath(hir::QPath::Resolved(None, P(gen_future_path))),
+                            hir_id, span,
+                        }))
+                    }
+                    FunctionRetTy::Default(span) => hir::DefaultReturn(*span),
+                }
+            }
+        }
+
+        // For async functions the return type is lowered to impl ::std::future::Future<Output = T>
+        impl FnRetTransform for AsyncFunctionRet {
+            fn lower(&self, ctx: &mut LoweringContext, output: &FunctionRetTy, impl_trait_context: ImplTraitContext) -> hir::FunctionRetTy {
+                let inner_ty = match output {
+                    FunctionRetTy::Ty(ty) => ctx.lower_ty(ty, impl_trait_context),
+                    FunctionRetTy::Default(span) => {
+                        let LoweredNodeId { node_id, hir_id } = ctx.next_id();
+                        P(hir::Ty {
+                            id: node_id,
+                            hir_id: hir_id,
+                            node: hir::TyTup(hir_vec![]),
+                            span: *span,
+                        })
+                    }
+                };
+                let span = inner_ty.span;
+                let hir::Path { def, segments, .. }  = ctx.std_path(span, &["future", "Future"], false);
+                let future_path = hir::Path {
+                    segments: segments.map_slice(|mut v| {
+                        v.last_mut().unwrap().parameters = Some(P(hir::PathParameters {
+                            lifetimes: hir_vec![],
+                            types: hir_vec![],
+                            bindings: hir_vec![hir::TypeBinding {
+                                name: ctx.str_to_ident("Output"),
+                                ty: inner_ty,
+                                id: ctx.next_id().node_id,
+                                span,
+                            }],
+                            parenthesized: false,
+                        }));
+                        v
+                    }),
+                    def, span
+                };
+                let LoweredNodeId { hir_id, node_id } = ctx.next_id();
+                hir::Return(P(hir::Ty {
+                    id: node_id,
+                    node: hir::TyImplTraitExistential(hir::ExistTy {
+                        generics: hir::Generics::empty(),
+                        bounds: hir_vec![hir::TyParamBound::TraitTyParamBound(hir::PolyTraitRef {
+                            trait_ref: hir::TraitRef {
+                                path: future_path,
+                                ref_id: ctx.next_id().node_id,
+                            },
+                            bound_generic_params: hir_vec![],
+                            span,
+                        }, hir::TraitBoundModifier::None)],
+                    }, hir_vec![]),
+                    hir_id, span,
+                }))
+            }
+        }
+
+        let output_impl_trait_context = if fn_def_id.is_some() && impl_trait_return_allow {
+            ImplTraitContext::Existential
+        } else {
+            ImplTraitContext::Disallowed
+        };
+
         // NOTE: The two last parameters here have to do with impl Trait. If fn_def_id is Some,
         //       then impl Trait arguments are lowered into generic parameters on the given
         //       fn_def_id, otherwise impl Trait is disallowed. (for now)
@@ -1748,15 +1870,7 @@ impl<'a> LoweringContext<'a> {
                     }
                 })
                 .collect(),
-            output: match decl.output {
-                FunctionRetTy::Ty(ref ty) => match fn_def_id {
-                    Some(_) if impl_trait_return_allow => {
-                        hir::Return(self.lower_ty(ty, ImplTraitContext::Existential))
-                    }
-                    _ => hir::Return(self.lower_ty(ty, ImplTraitContext::Disallowed)),
-                },
-                FunctionRetTy::Default(span) => hir::DefaultReturn(span),
-            },
+            output: fn_ret_transform.lower(self, &decl.output, output_impl_trait_context),
             variadic: decl.variadic,
             has_implicit_self: decl.inputs.get(0).map_or(false, |arg| match arg.ty.node {
                 TyKind::ImplicitSelf => true,
@@ -2201,11 +2315,15 @@ impl<'a> LoweringContext<'a> {
                         let body = this.lower_block(body, false);
                         this.expr_block(body, ThinVec::new())
                     });
+                    let fn_ret_transform: &FnRetTransform = match header.asyncness {
+                        IsAsync::Async => &AsyncFunctionRet,
+                        IsAsync::NotAsync => &NonAsyncRet,
+                    };
                     let (generics, fn_decl) = this.add_in_band_defs(
                         generics,
                         fn_def_id,
                         AnonymousLifetimeMode::PassThrough,
-                        |this| this.lower_fn_decl(decl, Some(fn_def_id), true),
+                        |this| this.lower_fn_decl(decl, Some(fn_def_id), true, fn_ret_transform),
                     );
 
                     hir::ItemFn(
@@ -2721,7 +2839,7 @@ impl<'a> LoweringContext<'a> {
                             |this| {
                                 (
                                     // Disallow impl Trait in foreign items
-                                    this.lower_fn_decl(fdec, None, false),
+                                    this.lower_fn_decl(fdec, None, false, &NonAsyncRet),
                                     this.lower_fn_args_to_names(fdec),
                                 )
                             },
@@ -2747,9 +2865,13 @@ impl<'a> LoweringContext<'a> {
         fn_def_id: DefId,
         impl_trait_return_allow: bool,
     ) -> hir::MethodSig {
+        let fn_ret_transform: &FnRetTransform = match sig.header.asyncness {
+            IsAsync::Async => &AsyncFunctionRet,
+            IsAsync::NotAsync => &NonAsyncRet,
+        };
         hir::MethodSig {
             header: self.lower_fn_header(sig.header),
-            decl: self.lower_fn_decl(&sig.decl, Some(fn_def_id), impl_trait_return_allow),
+            decl: self.lower_fn_decl(&sig.decl, Some(fn_def_id), impl_trait_return_allow, fn_ret_transform),
         }
     }
 
@@ -3074,7 +3196,7 @@ impl<'a> LoweringContext<'a> {
                 arms.iter().map(|x| self.lower_arm(x)).collect(),
                 hir::MatchSource::Normal,
             ),
-            ExprKind::Closure(capture_clause, movability, ref decl, ref body, fn_decl_span) => {
+            ExprKind::Closure(capture_clause, asyncness, movability, ref decl, ref body, fn_decl_span) => {
                 self.with_new_scopes(|this| {
                     this.with_parent_def(e.id, |this| {
                         let mut is_generator = false;
@@ -3108,9 +3230,13 @@ impl<'a> LoweringContext<'a> {
                             }
                             None
                         };
+                        let fn_ret_transform: &FnRetTransform = match asyncness {
+                            IsAsync::Async => &AsyncClosureRet,
+                            IsAsync::NotAsync => &NonAsyncRet,
+                        };
                         hir::ExprClosure(
                             this.lower_capture_clause(capture_clause),
-                            this.lower_fn_decl(decl, None, false),
+                            this.lower_fn_decl(decl, None, false, fn_ret_transform),
                             body_id,
                             fn_decl_span,
                             generator_option,

src/librustc_resolve/lib.rs

@@ -790,7 +790,7 @@ impl<'a, 'tcx> Visitor<'tcx> for Resolver<'a> {
             FnKind::Method(_, _, _, _) => {
                 TraitOrImplItemRibKind
             }
-            FnKind::Closure(_) => ClosureRibKind(node_id),
+            FnKind::Closure(..) => ClosureRibKind(node_id),
         };
 
         // Create a value rib for the function.
@@ -816,7 +816,7 @@ impl<'a, 'tcx> Visitor<'tcx> for Resolver<'a> {
             FnKind::Method(.., body) => {
                 self.visit_block(body);
             }
-            FnKind::Closure(body) => {
+            FnKind::Closure(_, body) => {
                 self.visit_expr(body);
             }
         };

src/librustc_save_analysis/dump_visitor.rs

@@ -1549,7 +1549,7 @@ impl<'l, 'tcx: 'l, 'll, O: DumpOutput + 'll> Visitor<'l> for DumpVisitor<'l, 'tc
                     }
                 }
             }
-            ast::ExprKind::Closure(_, _, ref decl, ref body, _fn_decl_span) => {
+            ast::ExprKind::Closure(_, _, _, ref decl, ref body, _fn_decl_span) => {
                 let mut id = String::from("$");
                 id.push_str(&ex.id.to_string());
 

src/libsyntax/ast.rs

@@ -1133,7 +1133,7 @@ pub enum ExprKind {
     /// A closure (for example, `move |a, b, c| a + b + c`)
     ///
     /// The final span is the span of the argument block `|...|`
-    Closure(CaptureBy, Movability, P<FnDecl>, P<Expr>, Span),
+    Closure(CaptureBy, IsAsync, Movability, P<FnDecl>, P<Expr>, Span),
     /// A block (`'label: { ... }`)
     Block(P<Block>, Option<Label>),
     /// A catch block (`catch { ... }`)

src/libsyntax/ext/build.rs

@@ -914,6 +914,7 @@ impl<'a> AstBuilder for ExtCtxt<'a> {
                       fn_decl_span: Span) // span of the `|...|` part
                       -> P<ast::Expr> {
         self.expr(span, ast::ExprKind::Closure(ast::CaptureBy::Ref,
+                                               ast::IsAsync::NotAsync,
                                                ast::Movability::Movable,
                                                fn_decl,
                                                body,
@@ -934,6 +935,7 @@ impl<'a> AstBuilder for ExtCtxt<'a> {
         // the entire lambda body. Probably we should extend the API
         // here, but that's not entirely clear.
         self.expr(span, ast::ExprKind::Closure(ast::CaptureBy::Ref,
+                                               ast::IsAsync::NotAsync,
                                                ast::Movability::Movable,
                                                fn_decl,
                                                body,

src/libsyntax/fold.rs

@@ -1265,8 +1265,9 @@ pub fn noop_fold_expr<T: Folder>(Expr {id, node, span, attrs}: Expr, folder: &mu
                 ExprKind::Match(folder.fold_expr(expr),
                           arms.move_map(|x| folder.fold_arm(x)))
             }
-            ExprKind::Closure(capture_clause, movability, decl, body, span) => {
+            ExprKind::Closure(capture_clause, asyncness, movability, decl, body, span) => {
                 ExprKind::Closure(capture_clause,
+                                  asyncness,
                                   movability,
                                   folder.fold_fn_decl(decl),
                                   folder.fold_expr(body),

src/libsyntax/parse/parser.rs

@@ -3312,7 +3312,7 @@ impl<'a> Parser<'a> {
 
         Ok(self.mk_expr(
             lo.to(body.span),
-            ExprKind::Closure(capture_clause, movability, decl, body, lo.to(decl_hi)),
+            ExprKind::Closure(capture_clause, IsAsync::NotAsync, movability, decl, body, lo.to(decl_hi)),
             attrs))
     }
 

src/libsyntax/print/pprust.rs

@@ -2158,8 +2158,9 @@ impl<'a> State<'a> {
                 }
                 self.bclose_(expr.span, INDENT_UNIT)?;
             }
-            ast::ExprKind::Closure(capture_clause, movability, ref decl, ref body, _) => {
+            ast::ExprKind::Closure(capture_clause, asyncness, movability, ref decl, ref body, _) => {
                 self.print_movability(movability)?;
+                self.print_asyncness(asyncness)?;
                 self.print_capture_clause(capture_clause)?;
 
                 self.print_fn_block_args(decl)?;
@@ -2793,6 +2794,14 @@ impl<'a> State<'a> {
         }
     }
 
+    pub fn print_asyncness(&mut self, asyncness: ast::IsAsync)
+                                -> io::Result<()> {
+        match asyncness {
+            ast::IsAsync::Async => self.word_space("async"),
+            ast::IsAsync::NotAsync => Ok(()),
+        }
+    }
+
     pub fn print_capture_clause(&mut self, capture_clause: ast::CaptureBy)
                                 -> io::Result<()> {
         match capture_clause {

src/libsyntax/ptr.rs

@@ -95,6 +95,16 @@ impl<T: 'static> P<T> {
     }
 }
 
+impl<T: 'static> P<[T]> {
+    pub fn map_slice<F>(self, f: F) -> P<[T]> where
+        F: FnOnce(Vec<T>) -> Vec<T>
+    {
+        P {
+            ptr: f(self.ptr.into()).into(),
+        }
+    }
+}
+
 impl<T: ?Sized> Deref for P<T> {
     type Target = T;
 

src/libsyntax/visit.rs

@@ -37,7 +37,7 @@ pub enum FnKind<'a> {
     Method(Ident, &'a MethodSig, Option<&'a Visibility>, &'a Block),
 
     /// |x, y| body
-    Closure(&'a Expr),
+    Closure(IsAsync, &'a Expr),
 }
 
 /// Each method of the Visitor trait is a hook to be potentially
@@ -552,7 +552,7 @@ pub fn walk_fn<'a, V>(visitor: &mut V, kind: FnKind<'a>, declaration: &'a FnDecl
             walk_fn_decl(visitor, declaration);
             visitor.visit_block(body);
         }
-        FnKind::Closure(body) => {
+        FnKind::Closure(_, body) => {
             walk_fn_decl(visitor, declaration);
             visitor.visit_expr(body);
         }
@@ -733,8 +733,8 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             visitor.visit_expr(subexpression);
             walk_list!(visitor, visit_arm, arms);
         }
-        ExprKind::Closure(_, _, ref function_declaration, ref body, _decl_span) => {
-            visitor.visit_fn(FnKind::Closure(body),
+        ExprKind::Closure(_, is_async, _, ref function_declaration, ref body, _decl_span) => {
+            visitor.visit_fn(FnKind::Closure(is_async, body),
                              function_declaration,
                              expression.span,
                              expression.id)