Enforce consistent drop order w/ async methods.

This commit extends the previous commit to apply to trait methods as
well as free functions.

2 files changed, 130 insertions, 9 deletions

diff --git a/src/librustc/hir/lowering.rs b/src/librustc/hir/lowering.rs
index 5dabf10cbf8..5afac999fc0 100644
--- a/src/librustc/hir/lowering.rs
+++ b/src/librustc/hir/lowering.rs
@@ -3611,15 +3611,33 @@ impl<'a> LoweringContext<'a> {
                 )
             }
             ImplItemKind::Method(ref sig, ref body) => {
-                let body_id = self.lower_async_body(&sig.decl, &sig.header.asyncness.node, body);
-                let impl_trait_return_allow = !self.is_in_trait_impl;
-                let (generics, sig) = self.lower_method_sig(
-                    &i.generics,
-                    sig,
-                    impl_item_def_id,
-                    impl_trait_return_allow,
-                    sig.header.asyncness.node.opt_return_id(),
-                );
+                let mut lower_method = |sig: &MethodSig| {
+                    let body_id = self.lower_async_body(
+                        &sig.decl, &sig.header.asyncness.node, body
+                    );
+                    let impl_trait_return_allow = !self.is_in_trait_impl;
+                    let (generics, sig) = self.lower_method_sig(
+                        &i.generics,
+                        sig,
+                        impl_item_def_id,
+                        impl_trait_return_allow,
+                        sig.header.asyncness.node.opt_return_id(),
+                    );
+                    (body_id, generics, sig)
+                };
+
+                let (body_id, generics, sig) = if let IsAsync::Async {
+                    ref arguments, ..
+                } = sig.header.asyncness.node {
+                    let mut sig = sig.clone();
+                    // Replace the arguments of this async function with the generated
+                    // arguments that will be moved into the closure.
+                    sig.decl.inputs = arguments.clone().drain(..).map(|a| a.arg).collect();
+                    lower_method(&sig)
+                } else {
+                    lower_method(sig)
+                };
+
                 (generics, hir::ImplItemKind::Method(sig, body_id))
             }
             ImplItemKind::Type(ref ty) => (
diff --git a/src/test/run-pass/issue-54716.rs b/src/test/run-pass/issue-54716.rs
index 7f64c6e120e..d075d2d6191 100644
--- a/src/test/run-pass/issue-54716.rs
+++ b/src/test/run-pass/issue-54716.rs
@@ -10,6 +10,7 @@ extern crate arc_wake;
 use arc_wake::ArcWake;
 use std::cell::RefCell;
 use std::future::Future;
+use std::marker::PhantomData;
 use std::sync::Arc;
 use std::task::Context;
 
@@ -49,6 +50,46 @@ async fn foobar(x: D, (a, _, _c): (D, D, D), _: D, _y: D) {
     x.1.borrow_mut().push(DropOrder::Function);
 }
 
+struct Foo;
+
+impl Foo {
+    async fn foo(x: D, _y: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn bar(x: D, _: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn baz((x, _): (D, D)) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn foobar(x: D, (a, _, _c): (D, D, D), _: D, _y: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+}
+
+struct Bar<'a>(PhantomData<&'a ()>);
+
+impl<'a> Bar<'a> {
+    async fn foo(&'a self, x: D, _y: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn bar(&'a self, x: D, _: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn baz(&'a self, (x, _): (D, D)) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+
+    async fn foobar(&'a self, x: D, (a, _, _c): (D, D, D), _: D, _y: D) {
+        x.1.borrow_mut().push(DropOrder::Function);
+    }
+}
+
 fn main() {
     let empty = Arc::new(EmptyWaker);
     let waker = ArcWake::into_waker(empty);
@@ -60,6 +101,8 @@ fn main() {
     // non-async functions. This is because the drop order of captured variables doesn't match the
     // drop order of arguments in a function.
 
+    // Free functions
+
     let af = Arc::new(RefCell::new(Vec::new()));
     let mut fut = Box::pin(foo(D("x", af.clone()), D("_y", af.clone())));
     let _ = fut.as_mut().poll(&mut cx);
@@ -86,4 +129,64 @@ fn main() {
     assert_eq!(*af.borrow(), &[
        Function, Val("_y"), Val("_c"), Val("a"), Val("x"), Val("_"), Val("_"),
     ]);
+
+    // Methods w/out self
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(Foo::foo(D("x", af.clone()), D("_y", af.clone())));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("_y"), Val("x")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(Foo::bar(D("x", af.clone()), D("_", af.clone())));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("x"), Val("_")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(Foo::baz((D("x", af.clone()), D("_", af.clone()))));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("x"), Val("_")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(Foo::foobar(
+        D("x", af.clone()),
+        (D("a", af.clone()), D("_", af.clone()), D("_c", af.clone())),
+        D("_", af.clone()),
+        D("_y", af.clone()),
+    ));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[
+       Function, Val("_y"), Val("_c"), Val("a"), Val("x"), Val("_"), Val("_"),
+    ]);
+
+    // Methods
+
+    let b = Bar(Default::default());
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(b.foo(D("x", af.clone()), D("_y", af.clone())));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("_y"), Val("x")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(b.bar(D("x", af.clone()), D("_", af.clone())));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("x"), Val("_")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(b.baz((D("x", af.clone()), D("_", af.clone()))));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[Function, Val("x"), Val("_")]);
+
+    let af = Arc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(b.foobar(
+        D("x", af.clone()),
+        (D("a", af.clone()), D("_", af.clone()), D("_c", af.clone())),
+        D("_", af.clone()),
+        D("_y", af.clone()),
+    ));
+    let _ = fut.as_mut().poll(&mut cx);
+    assert_eq!(*af.borrow(), &[
+       Function, Val("_y"), Val("_c"), Val("a"), Val("x"), Val("_"), Val("_"),
+    ]);
 }