Add `AsyncArgument` to AST.

This commit adds an `AsyncArgument` struct to the AST that contains the
generated argument and statement that will be used in HIR lowering, name
resolution and def collection.

5 files changed, 129 insertions, 26 deletions

diff --git a/src/libsyntax/ast.rs b/src/libsyntax/ast.rs
index 2fe0ebcdd28..e6669e0d6ed 100644
--- a/src/libsyntax/ast.rs
+++ b/src/libsyntax/ast.rs
@@ -1849,18 +1849,35 @@ pub enum Unsafety {
     Normal,
 }
 
-#[derive(Copy, Clone, RustcEncodable, RustcDecodable, Debug)]
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug)]
+pub struct AsyncArgument {
+    /// `__arg0`
+    pub ident: Ident,
+    /// `__arg0: <ty>` argument to replace existing function argument `<pat>: <ty>`.
+    pub arg: Arg,
+    /// `let <pat>: <ty> = __arg0;` statement to be inserted at the start of the block.
+    pub stmt: Stmt,
+}
+
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug)]
 pub enum IsAsync {
     Async {
         closure_id: NodeId,
         return_impl_trait_id: NodeId,
+        /// This field stores the arguments and statements that are used in HIR lowering to
+        /// ensure that `async fn` arguments are dropped at the correct time.
+        ///
+        /// The argument and statements here are generated at parse time as they are required in
+        /// both the hir lowering, def collection and name resolution and this stops them needing
+        /// to be created in each place.
+        arguments: Vec<AsyncArgument>,
     },
     NotAsync,
 }
 
 impl IsAsync {
-    pub fn is_async(self) -> bool {
-        if let IsAsync::Async { .. } = self {
+    pub fn is_async(&self) -> bool {
+        if let IsAsync::Async { .. } = *self {
             true
         } else {
             false
@@ -1868,12 +1885,12 @@ impl IsAsync {
     }
 
     /// In ths case this is an `async` return, the `NodeId` for the generated `impl Trait` item.
-    pub fn opt_return_id(self) -> Option<NodeId> {
+    pub fn opt_return_id(&self) -> Option<NodeId> {
         match self {
             IsAsync::Async {
                 return_impl_trait_id,
                 ..
-            } => Some(return_impl_trait_id),
+            } => Some(*return_impl_trait_id),
             IsAsync::NotAsync => None,
         }
     }
@@ -2213,7 +2230,7 @@ impl Item {
 ///
 /// All the information between the visibility and the name of the function is
 /// included in this struct (e.g., `async unsafe fn` or `const extern "C" fn`).
-#[derive(Clone, Copy, RustcEncodable, RustcDecodable, Debug)]
+#[derive(Clone, RustcEncodable, RustcDecodable, Debug)]
 pub struct FnHeader {
     pub unsafety: Unsafety,
     pub asyncness: Spanned<IsAsync>,
diff --git a/src/libsyntax/ext/placeholders.rs b/src/libsyntax/ext/placeholders.rs
index 3e60dd81a3b..68cd3c28676 100644
--- a/src/libsyntax/ext/placeholders.rs
+++ b/src/libsyntax/ext/placeholders.rs
@@ -102,6 +102,13 @@ impl<'a, 'b> PlaceholderExpander<'a, 'b> {
     fn remove(&mut self, id: ast::NodeId) -> AstFragment {
         self.expanded_fragments.remove(&id).unwrap()
     }
+
+    fn next_id(&mut self, id: &mut ast::NodeId) {
+        if self.monotonic {
+            assert_eq!(*id, ast::DUMMY_NODE_ID);
+            *id = self.cx.resolver.next_node_id()
+        }
+    }
 }
 
 impl<'a, 'b> MutVisitor for PlaceholderExpander<'a, 'b> {
@@ -183,9 +190,16 @@ impl<'a, 'b> MutVisitor for PlaceholderExpander<'a, 'b> {
         noop_visit_block(block, self);
 
         for stmt in block.stmts.iter_mut() {
-            if self.monotonic {
-                assert_eq!(stmt.id, ast::DUMMY_NODE_ID);
-                stmt.id = self.cx.resolver.next_node_id();
+            self.next_id(&mut stmt.id);
+        }
+    }
+
+    fn visit_asyncness(&mut self, a: &mut ast::IsAsync) {
+        noop_visit_asyncness(a, self);
+
+        if let ast::IsAsync::Async { ref mut arguments, .. } = a {
+            for argument in arguments.iter_mut() {
+                self.next_id(&mut argument.stmt.id);
             }
         }
     }
diff --git a/src/libsyntax/mut_visit.rs b/src/libsyntax/mut_visit.rs
index 87826ccc891..bb9116e678e 100644
--- a/src/libsyntax/mut_visit.rs
+++ b/src/libsyntax/mut_visit.rs
@@ -679,9 +679,17 @@ pub fn noop_visit_interpolated<T: MutVisitor>(nt: &mut token::Nonterminal, vis:
 
 pub fn noop_visit_asyncness<T: MutVisitor>(asyncness: &mut IsAsync, vis: &mut T) {
     match asyncness {
-        IsAsync::Async { closure_id, return_impl_trait_id } => {
+        IsAsync::Async { closure_id, return_impl_trait_id, ref mut arguments } => {
             vis.visit_id(closure_id);
             vis.visit_id(return_impl_trait_id);
+            for AsyncArgument { ident, arg, stmt } in arguments.iter_mut() {
+                vis.visit_ident(ident);
+                vis.visit_arg(arg);
+                visit_clobber(stmt, |stmt| {
+                    vis.flat_map_stmt(stmt)
+                        .expect_one("expected visitor to produce exactly one item")
+                });
+            }
         }
         IsAsync::NotAsync => {}
     }
diff --git a/src/libsyntax/parse/parser.rs b/src/libsyntax/parse/parser.rs
index 8620627765f..26ede5a1057 100644
--- a/src/libsyntax/parse/parser.rs
+++ b/src/libsyntax/parse/parser.rs
@@ -1,4 +1,4 @@
-use crate::ast::{AngleBracketedArgs, ParenthesizedArgs, AttrStyle, BareFnTy};
+use crate::ast::{AngleBracketedArgs, AsyncArgument, ParenthesizedArgs, AttrStyle, BareFnTy};
 use crate::ast::{GenericBound, TraitBoundModifier};
 use crate::ast::Unsafety;
 use crate::ast::{Mod, AnonConst, Arg, Arm, Guard, Attribute, BindingMode, TraitItemKind};
@@ -1517,6 +1517,7 @@ impl<'a> Parser<'a> {
             IsAsync::Async {
                 closure_id: ast::DUMMY_NODE_ID,
                 return_impl_trait_id: ast::DUMMY_NODE_ID,
+                arguments: Vec::new(),
             }
         } else {
             IsAsync::NotAsync
@@ -1575,7 +1576,7 @@ impl<'a> Parser<'a> {
             // trait item macro.
             (keywords::Invalid.ident(), ast::TraitItemKind::Macro(mac), ast::Generics::default())
         } else {
-            let (constness, unsafety, asyncness, abi) = self.parse_fn_front_matter()?;
+            let (constness, unsafety, mut asyncness, abi) = self.parse_fn_front_matter()?;
 
             let ident = self.parse_ident()?;
             let mut generics = self.parse_generics()?;
@@ -1589,6 +1590,7 @@ impl<'a> Parser<'a> {
                 p.parse_arg_general(p.span.rust_2018(), true, false)
             })?;
             generics.where_clause = self.parse_where_clause()?;
+            self.construct_async_arguments(&mut asyncness, &d);
 
             let sig = ast::MethodSig {
                 header: FnHeader {
@@ -6567,7 +6569,7 @@ impl<'a> Parser<'a> {
     /// Parses an item-position function declaration.
     fn parse_item_fn(&mut self,
                      unsafety: Unsafety,
-                     asyncness: Spanned<IsAsync>,
+                     mut asyncness: Spanned<IsAsync>,
                      constness: Spanned<Constness>,
                      abi: Abi)
                      -> PResult<'a, ItemInfo> {
@@ -6576,6 +6578,7 @@ impl<'a> Parser<'a> {
         let decl = self.parse_fn_decl(allow_c_variadic)?;
         generics.where_clause = self.parse_where_clause()?;
         let (inner_attrs, body) = self.parse_inner_attrs_and_block()?;
+        self.construct_async_arguments(&mut asyncness, &decl);
         let header = FnHeader { unsafety, asyncness, constness, abi };
         Ok((ident, ItemKind::Fn(decl, header, generics, body), Some(inner_attrs)))
     }
@@ -6751,11 +6754,12 @@ impl<'a> Parser<'a> {
             Ok((keywords::Invalid.ident(), vec![], ast::Generics::default(),
                 ast::ImplItemKind::Macro(mac)))
         } else {
-            let (constness, unsafety, asyncness, abi) = self.parse_fn_front_matter()?;
+            let (constness, unsafety, mut asyncness, abi) = self.parse_fn_front_matter()?;
             let ident = self.parse_ident()?;
             let mut generics = self.parse_generics()?;
             let decl = self.parse_fn_decl_with_self(|p| p.parse_arg())?;
             generics.where_clause = self.parse_where_clause()?;
+            self.construct_async_arguments(&mut asyncness, &decl);
             *at_end = true;
             let (inner_attrs, body) = self.parse_inner_attrs_and_block()?;
             let header = ast::FnHeader { abi, unsafety, constness, asyncness };
@@ -8177,6 +8181,7 @@ impl<'a> Parser<'a> {
                                    respan(async_span, IsAsync::Async {
                                        closure_id: ast::DUMMY_NODE_ID,
                                        return_impl_trait_id: ast::DUMMY_NODE_ID,
+                                       arguments: Vec::new(),
                                    }),
                                    respan(fn_span, Constness::NotConst),
                                    Abi::Rust)?;
@@ -8822,6 +8827,66 @@ impl<'a> Parser<'a> {
             }
         }
     }
+
+    /// When lowering a `async fn` to the HIR, we need to move all of the arguments of the function
+    /// into the generated closure so that they are dropped when the future is polled and not when
+    /// it is created.
+    ///
+    /// The arguments of the function are replaced in HIR lowering with the arguments created by
+    /// this function and the statements created here are inserted at the top of the closure body.
+    fn construct_async_arguments(&mut self, asyncness: &mut Spanned<IsAsync>, decl: &FnDecl) {
+        if let IsAsync::Async { ref mut arguments, .. } = asyncness.node {
+            for (index, input) in decl.inputs.iter().enumerate() {
+                let id = ast::DUMMY_NODE_ID;
+                let span = input.pat.span;
+
+                // Construct a name for our temporary argument.
+                let name = format!("__arg{}", index);
+                let ident = Ident::from_str(&name);
+
+                // Construct an argument representing `__argN: <ty>` to replace the argument of the
+                // async function.
+                let arg = Arg {
+                    ty: input.ty.clone(),
+                    id,
+                    pat: P(Pat {
+                        id,
+                        node: PatKind::Ident(
+                            BindingMode::ByValue(Mutability::Immutable), ident, None,
+                        ),
+                        span,
+                    }),
+                };
+
+                // Construct a `let <pat>: <ty> = __argN;` statement to insert at the top of the
+                // async closure.
+                let local = P(Local {
+                    pat: input.pat.clone(),
+                    ty: Some(P(Ty {
+                        id,
+                        node: input.ty.node.clone(),
+                        span: input.ty.span,
+                    })),
+                    init: Some(P(Expr {
+                        id,
+                        node: ExprKind::Path(None, ast::Path {
+                            span,
+                            segments: vec![PathSegment { ident, id, args: None }],
+                        }),
+                        span,
+                        attrs: ThinVec::new(),
+                    })),
+                    id,
+                    span,
+                    attrs: ThinVec::new(),
+                    source: LocalSource::AsyncFn,
+                });
+                let stmt = Stmt { id, node: StmtKind::Local(local), span, };
+
+                arguments.push(AsyncArgument { ident, arg, stmt });
+            }
+        }
+    }
 }
 
 pub fn emit_unclosed_delims(unclosed_delims: &mut Vec<UnmatchedBrace>, handler: &errors::Handler) {
diff --git a/src/libsyntax/print/pprust.rs b/src/libsyntax/print/pprust.rs
index ca05ff71c94..d440a02f2fd 100644
--- a/src/libsyntax/print/pprust.rs
+++ b/src/libsyntax/print/pprust.rs
@@ -372,7 +372,7 @@ pub fn vis_to_string(v: &ast::Visibility) -> String {
 }
 
 pub fn fun_to_string(decl: &ast::FnDecl,
-                     header: ast::FnHeader,
+                     header: &ast::FnHeader,
                      name: ast::Ident,
                      generics: &ast::Generics)
                      -> String {
@@ -1133,7 +1133,7 @@ impl<'a> State<'a> {
         match item.node {
             ast::ForeignItemKind::Fn(ref decl, ref generics) => {
                 self.head("")?;
-                self.print_fn(decl, ast::FnHeader::default(),
+                self.print_fn(decl, &ast::FnHeader::default(),
                               Some(item.ident),
                               generics, &item.vis)?;
                 self.end()?; // end head-ibox
@@ -1263,7 +1263,7 @@ impl<'a> State<'a> {
                 self.s.word(";")?;
                 self.end()?; // end the outer cbox
             }
-            ast::ItemKind::Fn(ref decl, header, ref param_names, ref body) => {
+            ast::ItemKind::Fn(ref decl, ref header, ref param_names, ref body) => {
                 self.head("")?;
                 self.print_fn(
                     decl,
@@ -1615,7 +1615,7 @@ impl<'a> State<'a> {
                             vis: &ast::Visibility)
                             -> io::Result<()> {
         self.print_fn(&m.decl,
-                      m.header,
+                      &m.header,
                       Some(ident),
                       &generics,
                       vis)
@@ -2213,7 +2213,7 @@ impl<'a> State<'a> {
                 self.bclose_(expr.span, INDENT_UNIT)?;
             }
             ast::ExprKind::Closure(
-                capture_clause, asyncness, movability, ref decl, ref body, _) => {
+                capture_clause, ref asyncness, movability, ref decl, ref body, _) => {
                 self.print_movability(movability)?;
                 self.print_asyncness(asyncness)?;
                 self.print_capture_clause(capture_clause)?;
@@ -2798,7 +2798,7 @@ impl<'a> State<'a> {
 
     pub fn print_fn(&mut self,
                     decl: &ast::FnDecl,
-                    header: ast::FnHeader,
+                    header: &ast::FnHeader,
                     name: Option<ast::Ident>,
                     generics: &ast::Generics,
                     vis: &ast::Visibility) -> io::Result<()> {
@@ -2853,8 +2853,7 @@ impl<'a> State<'a> {
         }
     }
 
-    pub fn print_asyncness(&mut self, asyncness: ast::IsAsync)
-                                -> io::Result<()> {
+    pub fn print_asyncness(&mut self, asyncness: &ast::IsAsync) -> io::Result<()> {
         if asyncness.is_async() {
             self.word_nbsp("async")?;
         }
@@ -3126,7 +3125,7 @@ impl<'a> State<'a> {
             span: syntax_pos::DUMMY_SP,
         };
         self.print_fn(decl,
-                      ast::FnHeader { unsafety, abi, ..ast::FnHeader::default() },
+                      &ast::FnHeader { unsafety, abi, ..ast::FnHeader::default() },
                       name,
                       &generics,
                       &source_map::dummy_spanned(ast::VisibilityKind::Inherited))?;
@@ -3189,7 +3188,7 @@ impl<'a> State<'a> {
     }
 
     pub fn print_fn_header_info(&mut self,
-                                header: ast::FnHeader,
+                                header: &ast::FnHeader,
                                 vis: &ast::Visibility) -> io::Result<()> {
         self.s.word(visibility_qualified(vis, ""))?;
 
@@ -3198,7 +3197,7 @@ impl<'a> State<'a> {
             ast::Constness::Const => self.word_nbsp("const")?
         }
 
-        self.print_asyncness(header.asyncness.node)?;
+        self.print_asyncness(&header.asyncness.node)?;
         self.print_unsafety(header.unsafety)?;
 
         if header.abi != Abi::Rust {
@@ -3247,7 +3246,7 @@ mod tests {
             assert_eq!(
                 fun_to_string(
                     &decl,
-                    ast::FnHeader {
+                    &ast::FnHeader {
                         unsafety: ast::Unsafety::Normal,
                         constness: source_map::dummy_spanned(ast::Constness::NotConst),
                         asyncness: source_map::dummy_spanned(ast::IsAsync::NotAsync),