Auto merge of #64525 - nikomatsakis:issue-64512-drop-order-tail-temp, r=davidtwco

adjust desugaring for async fn to correct drop order

Old desugaring, given a user function body `{ $stmts; $expr }`

```
{
    let $param_pattern0 = $raw_param0;
    ...
    let $param_patternN = $raw_paramN;
    $stmts;
    $expr
}
```

New desugaring:

```
{
    let $param_pattern0 = $raw_param0;
    ...
    let $param_patternN = $raw_paramN;
    drop-temps {
        $stmts;
        $expr
    }
}
```

The drop-temps is an internal bit of HIR that drops temporaries from the resulting expression, but it should be equivalent to `return { $stmts; $expr }`.

Fixes #64512
Fixes #64391

Author: bors

src/librustc/hir/lowering.rs

@@ -2682,12 +2682,8 @@ impl<'a> LoweringContext<'a> {
         bounds.iter().map(|bound| self.lower_param_bound(bound, itctx.reborrow())).collect()
     }
 
-    fn lower_block_with_stmts(
-        &mut self,
-        b: &Block,
-        targeted_by_break: bool,
-        mut stmts: Vec<hir::Stmt>,
-    ) -> P<hir::Block> {
+    fn lower_block(&mut self, b: &Block, targeted_by_break: bool) -> P<hir::Block> {
+        let mut stmts = vec![];
         let mut expr = None;
 
         for (index, stmt) in b.stmts.iter().enumerate() {
@@ -2712,8 +2708,11 @@ impl<'a> LoweringContext<'a> {
         })
     }
 
-    fn lower_block(&mut self, b: &Block, targeted_by_break: bool) -> P<hir::Block> {
-        self.lower_block_with_stmts(b, targeted_by_break, vec![])
+    /// Lowers a block directly to an expression, presuming that it
+    /// has no attributes and is not targeted by a `break`.
+    fn lower_block_expr(&mut self, b: &Block) -> hir::Expr {
+        let block = self.lower_block(b, false);
+        self.expr_block(block, ThinVec::new())
     }
 
     fn lower_pat(&mut self, p: &Pat) -> P<hir::Pat> {

src/librustc/hir/lowering/expr.rs

@@ -90,10 +90,7 @@ impl LoweringContext<'_> {
             ),
             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| {
-                        let block = this.lower_block(block, false);
-                        this.expr_block(block, ThinVec::new())
-                    })
+                    this.with_new_scopes(|this| this.lower_block_expr(block))
                 })
             }
             ExprKind::Await(ref expr) => self.lower_expr_await(e.span, expr),
@@ -284,8 +281,7 @@ impl LoweringContext<'_> {
         let else_arm = self.arm(hir_vec![else_pat], P(else_expr));
 
         // Handle then + scrutinee:
-        let then_blk = self.lower_block(then, false);
-        let then_expr = self.expr_block(then_blk, ThinVec::new());
+        let then_expr = self.lower_block_expr(then);
         let (then_pat, scrutinee, desugar) = match cond.node {
             // `<pat> => <then>`:
             ExprKind::Let(ref pat, ref scrutinee) => {
@@ -335,8 +331,7 @@ impl LoweringContext<'_> {
         };
 
         // Handle then + scrutinee:
-        let then_blk = self.lower_block(body, false);
-        let then_expr = self.expr_block(then_blk, ThinVec::new());
+        let then_expr = self.lower_block_expr(body);
         let (then_pat, scrutinee, desugar, source) = match cond.node {
             ExprKind::Let(ref pat, ref scrutinee) => {
                 // to:
@@ -356,7 +351,7 @@ impl LoweringContext<'_> {
                 //
                 // ```
                 // 'label: loop {
-                //     match DropTemps($cond) {
+                //     match drop-temps { $cond } {
                 //         true => $body,
                 //         _ => break,
                 //     }
@@ -1310,7 +1305,7 @@ impl LoweringContext<'_> {
     /// `{ let _t = $expr; _t }` but should provide better compile-time performance.
     ///
     /// The drop order can be important in e.g. `if expr { .. }`.
-    fn expr_drop_temps(
+    pub(super) fn expr_drop_temps(
         &mut self,
         span: Span,
         expr: P<hir::Expr>,

src/librustc/hir/lowering/item.rs

@@ -1071,10 +1071,7 @@ impl LoweringContext<'_> {
     }
 
     fn lower_fn_body_block(&mut self, decl: &FnDecl, body: &Block) -> hir::BodyId {
-        self.lower_fn_body(decl, |this| {
-            let body = this.lower_block(body, false);
-            this.expr_block(body, ThinVec::new())
-        })
+        self.lower_fn_body(decl, |this| this.lower_block_expr(body))
     }
 
     pub(super) fn lower_const_body(&mut self, expr: &Expr) -> hir::BodyId {
@@ -1102,8 +1099,7 @@ impl LoweringContext<'_> {
             // from:
             //
             //     async fn foo(<pattern>: <ty>, <pattern>: <ty>, <pattern>: <ty>) {
-            //       async move {
-            //       }
+            //         <body>
             //     }
             //
             // into:
@@ -1116,11 +1112,19 @@ impl LoweringContext<'_> {
             //         let <pattern> = __arg1;
             //         let __arg0 = __arg0;
             //         let <pattern> = __arg0;
+            //         drop-temps { <body> } // see comments later in fn for details
             //       }
             //     }
             //
             // If `<pattern>` is a simple ident, then it is lowered to a single
             // `let <pattern> = <pattern>;` statement as an optimization.
+            //
+            // Note that the body is embedded in `drop-temps`; an
+            // equivalent desugaring would be `return { <body>
+            // };`. The key point is that we wish to drop all the
+            // let-bound variables and temporaries created in the body
+            // (and its tail expression!) before we drop the
+            // parameters (c.f. rust-lang/rust#64512).
             for (index, parameter) in decl.inputs.iter().enumerate() {
                 let parameter = this.lower_param(parameter);
                 let span = parameter.pat.span;
@@ -1219,8 +1223,36 @@ impl LoweringContext<'_> {
             let async_expr = this.make_async_expr(
                 CaptureBy::Value, closure_id, None, body.span,
                 |this| {
-                    let body = this.lower_block_with_stmts(body, false, statements);
-                    this.expr_block(body, ThinVec::new())
+                    // Create a block from the user's function body:
+                    let user_body = this.lower_block_expr(body);
+
+                    // Transform into `drop-temps { <user-body> }`, an expression:
+                    let desugared_span = this.mark_span_with_reason(
+                        DesugaringKind::Async,
+                        user_body.span,
+                        None,
+                    );
+                    let user_body = this.expr_drop_temps(
+                        desugared_span,
+                        P(user_body),
+                        ThinVec::new(),
+                    );
+
+                    // As noted above, create the final block like
+                    //
+                    // ```
+                    // {
+                    //   let $param_pattern = $raw_param;
+                    //   ...
+                    //   drop-temps { <user-body> }
+                    // }
+                    // ```
+                    let body = this.block_all(
+                        desugared_span,
+                        statements.into(),
+                        Some(P(user_body)),
+                    );
+                    this.expr_block(P(body), ThinVec::new())
                 });
             (HirVec::from(parameters), this.expr(body.span, async_expr, ThinVec::new()))
         })

src/test/ui/async-await/drop-order/drop-order-for-temporary-in-tail-return-expr.rs

@@ -0,0 +1,95 @@
+// aux-build:arc_wake.rs
+// edition:2018
+// run-pass
+
+#![allow(unused_variables)]
+
+// Test the drop order for parameters relative to local variables and
+// temporaries created in the tail return expression of the function
+// body. In particular, check that this drop order is the same between
+// a `async fn` and an ordinary `fn`. See #64512.
+
+extern crate arc_wake;
+
+use arc_wake::ArcWake;
+use std::cell::RefCell;
+use std::future::Future;
+use std::sync::Arc;
+use std::rc::Rc;
+use std::task::Context;
+
+struct EmptyWaker;
+
+impl ArcWake for EmptyWaker {
+    fn wake(self: Arc<Self>) {}
+}
+
+#[derive(Debug, Eq, PartialEq)]
+enum DropOrder {
+    Function,
+    Val(&'static str),
+}
+
+type DropOrderListPtr = Rc<RefCell<Vec<DropOrder>>>;
+
+struct D(&'static str, DropOrderListPtr);
+
+impl Drop for D {
+    fn drop(&mut self) {
+        self.1.borrow_mut().push(DropOrder::Val(self.0));
+    }
+}
+
+/// Check drop order of temporary "temp" as compared to `x`, `y`, and `z`.
+///
+/// Expected order:
+/// - `z`
+/// - temp
+/// - `y`
+/// - `x`
+async fn foo_async(x: D, _y: D) {
+    let l = x.1.clone();
+    let z = D("z", l.clone());
+    l.borrow_mut().push(DropOrder::Function);
+    helper_async(&D("temp", l)).await
+}
+
+async fn helper_async(v: &D) { }
+
+fn foo_sync(x: D, _y: D) {
+    let l = x.1.clone();
+    let z = D("z", l.clone());
+    l.borrow_mut().push(DropOrder::Function);
+    helper_sync(&D("temp", l))
+}
+
+fn helper_sync(v: &D) { }
+
+fn assert_drop_order_after_poll<Fut: Future<Output = ()>>(
+    f: impl FnOnce(DropOrderListPtr) -> Fut,
+    g: impl FnOnce(DropOrderListPtr),
+) {
+    let empty = Arc::new(EmptyWaker);
+    let waker = ArcWake::into_waker(empty);
+    let mut cx = Context::from_waker(&waker);
+
+    let actual_order = Rc::new(RefCell::new(Vec::new()));
+    let mut fut = Box::pin(f(actual_order.clone()));
+    let r = fut.as_mut().poll(&mut cx);
+
+    assert!(match r {
+        std::task::Poll::Ready(()) => true,
+        _ => false,
+    });
+
+    let expected_order = Rc::new(RefCell::new(Vec::new()));
+    g(expected_order.clone());
+
+    assert_eq!(*actual_order.borrow(), *expected_order.borrow());
+}
+
+fn main() {
+    // Free functions (see doc comment on function for what it tests).
+    assert_drop_order_after_poll(|l| foo_async(D("x", l.clone()), D("_y", l.clone())),
+                                 |l| foo_sync(D("x", l.clone()), D("_y", l.clone())));
+}

src/test/ui/async-await/issue-64391.rs

@@ -0,0 +1,14 @@
+// Regression test for Issue #64391. The goal here is that this
+// function compiles. In the past, due to incorrect drop order for
+// temporaries in the tail expression, we failed to compile this
+// example. The drop order itself is directly tested in
+// `drop-order/drop-order-for-temporary-in-tail-return-expr.rs`.
+//
+// check-pass
+// edition:2018
+
+async fn add(x: u32, y: u32) -> u32 {
+    async { x + y }.await
+}
+
+fn main() { }