Do not consider async task_context to be alive across yields

compiler/rustc_hir_typeck/src/generator_interior/mod.rs

@@ -33,6 +33,8 @@ struct InteriorVisitor<'a, 'tcx> {
     prev_unresolved_span: Option<Span>,
     linted_values: HirIdSet,
     drop_ranges: DropRanges,
+    task_context_hir_id: Option<HirId>,
+    in_lowered_await: bool,
 }
 
 impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
@@ -53,7 +55,7 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
             ty, hir_id, scope, expr, source_span, self.expr_count,
         );
 
-        let live_across_yield = scope
+        let mut live_across_yield = scope
             .map(|s| {
                 self.region_scope_tree.yield_in_scope(s).and_then(|yield_data| {
                     // If we are recording an expression that is the last yield
@@ -92,6 +94,34 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
                 Some(YieldData { span: DUMMY_SP, expr_and_pat_count: 0, source: self.kind.into() })
             });
 
+        // If this is the `&mut Context<'_>` async resume argument, or we are
+        // just visiting a `_task_context = yield ()` expression from async
+        // lowering, we do *not* consider this type to be live across yields.
+        if Some(hir_id) == self.task_context_hir_id || self.in_lowered_await {
+            #[cfg(debug_assertions)]
+            {
+                // As `record` is being invoked for multiple parts / types of the
+                // lowered `.await`, the `ty` has to either be the `()` going *into*
+                // the `yield`, or a `&mut Context<'_>` coming *out* of it.
+                let tcx = self.fcx.tcx;
+                if ty == tcx.types.unit {
+                    // all good
+                } else if let ty::Ref(_, adt, hir::Mutability::Mut) = ty.kind() && let ty::Adt(adt, _) = adt.kind()
+                {
+                    let context_def_id = tcx.lang_items().context();
+                    assert_eq!(
+                        Some(adt.did()),
+                        context_def_id,
+                        "expected `&mut Context<'_>, found `{:?}` instead`",
+                        ty
+                    );
+                } else {
+                    panic!("expected `()` or `&mut Context<'_>`, found `{:?}` instead", ty);
+                }
+            }
+            live_across_yield = None;
+        }
+
         if let Some(yield_data) = live_across_yield {
             debug!(
                 "type in expr = {:?}, scope = {:?}, type = {:?}, count = {}, yield_span = {:?}",
@@ -183,6 +213,17 @@ pub fn resolve_interior<'a, 'tcx>(
     kind: hir::GeneratorKind,
 ) {
     let body = fcx.tcx.hir().body(body_id);
+
+    // In case we are in an async block, this is the Param/Pat HirId of the
+    // `&mut Context<'_>` resume type. We can use this to explicitly prevent it
+    // from being considered as `live_across_yield`, which it is not, but the
+    // simple scope-based analysis can't tell.
+    let task_context_hir_id = if matches!(kind, hir::GeneratorKind::Async(_)) {
+        Some(body.params[0].pat.hir_id)
+    } else {
+        None
+    };
+
     let typeck_results = fcx.inh.typeck_results.borrow();
     let mut visitor = InteriorVisitor {
         fcx,
@@ -194,6 +235,8 @@ pub fn resolve_interior<'a, 'tcx>(
         prev_unresolved_span: None,
         linted_values: <_>::default(),
         drop_ranges: drop_ranges::compute_drop_ranges(fcx, def_id, body),
+        task_context_hir_id,
+        in_lowered_await: false,
     };
     intravisit::walk_body(&mut visitor, body);
 
@@ -426,6 +469,22 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
                 }
                 _ => intravisit::walk_expr(self, expr),
             },
+            ExprKind::Assign(_, rhs, _) => {
+                // An `.await` expression will be lowered to `_task_context = yield ()`.
+                // In that case, other forms of `yield` are considered errors in lowering.
+                if self.task_context_hir_id.is_some()
+                    && matches!(rhs.kind, hir::ExprKind::Yield(..))
+                {
+                    assert!(!self.in_lowered_await);
+                    self.in_lowered_await = true;
+                }
+                // We are still walking the whole expression including its types.
+                // First, we need to keep `expr_count` in sync as it is asserted
+                // at the very end, and to keep all the other computations in
+                // place just in case they are causing other side effects.
+                intravisit::walk_expr(self, expr);
+                self.in_lowered_await = false;
+            }
             _ => intravisit::walk_expr(self, expr),
         }
 

src/test/ui/async-await/async-await-let-else.drop-tracking.stderr

@@ -40,7 +40,7 @@ LL |   async fn bar2<T>(_: T) -> ! {
 LL | |     panic!()
 LL | | }
    | |_^
-   = note: required because it captures the following types: `&mut Context<'_>`, `Option<bool>`, `impl Future<Output = !>`, `()`
+   = note: required because it captures the following types: `Option<bool>`, `impl Future<Output = !>`, `()`
 note: required because it's used within this `async fn` body
   --> $DIR/async-await-let-else.rs:21:32
    |

src/test/ui/async-await/async-is-unwindsafe.rs

@@ -0,0 +1,31 @@
+// edition:2018
+
+fn is_unwindsafe(_: impl std::panic::UnwindSafe) {}
+
+fn main() {
+    // a normal async block that uses `&mut Context<'_>` implicitly via async lowering is fine
+    // as we should not consider that to be alive across an await point
+    is_unwindsafe(async {
+        async {}.await; // this needs an inner await point
+    });
+
+    is_unwindsafe(async {
+        //~^ ERROR the type `&mut Context<'_>` may not be safely transferred across an unwind boundary
+        use std::ptr::null;
+        use std::task::{Context, RawWaker, RawWakerVTable, Waker};
+        let waker = unsafe {
+            Waker::from_raw(RawWaker::new(
+                null(),
+                &RawWakerVTable::new(|_| todo!(), |_| todo!(), |_| todo!(), |_| todo!()),
+            ))
+        };
+        let mut cx = Context::from_waker(&waker);
+        let cx_ref = &mut cx;
+
+        async {}.await; // this needs an inner await point
+
+        // in this case, `&mut Context<'_>` is *truely* alive across an await point
+        drop(cx_ref);
+        drop(cx);
+    });
+}

src/test/ui/async-await/async-is-unwindsafe.stderr

@@ -0,0 +1,38 @@
+error[E0277]: the type `&mut Context<'_>` may not be safely transferred across an unwind boundary
+  --> $DIR/async-is-unwindsafe.rs:12:19
+   |
+LL |       is_unwindsafe(async {
+   |  ___________________^
+LL | |
+LL | |         use std::ptr::null;
+LL | |         use std::task::{Context, RawWaker, RawWakerVTable, Waker};
+...  |
+LL | |         drop(cx);
+LL | |     });
+   | |     ^
+   | |     |
+   | |_____`&mut Context<'_>` may not be safely transferred across an unwind boundary
+   |       within this `[async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6]`
+   |
+   = help: within `[async block@$DIR/async-is-unwindsafe.rs:12:19: 30:6]`, the trait `UnwindSafe` is not implemented for `&mut Context<'_>`
+   = note: `UnwindSafe` is implemented for `&std::task::Context<'_>`, but not for `&mut std::task::Context<'_>`
+note: future does not implement `UnwindSafe` as this value is used across an await
+  --> $DIR/async-is-unwindsafe.rs:25:17
+   |
+LL |         let cx_ref = &mut cx;
+   |             ------ has type `&mut Context<'_>` which does not implement `UnwindSafe`
+LL |
+LL |         async {}.await; // this needs an inner await point
+   |                 ^^^^^^ await occurs here, with `cx_ref` maybe used later
+...
+LL |     });
+   |     - `cx_ref` is later dropped here
+note: required by a bound in `is_unwindsafe`
+  --> $DIR/async-is-unwindsafe.rs:3:26
+   |
+LL | fn is_unwindsafe(_: impl std::panic::UnwindSafe) {}
+   |                          ^^^^^^^^^^^^^^^^^^^^^^ required by this bound in `is_unwindsafe`
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0277`.

src/test/ui/async-await/issue-68112.drop_tracking.stderr

@@ -57,7 +57,7 @@ note: required because it appears within the type `impl Future<Output = Arc<RefC
    |
 LL | fn make_non_send_future2() -> impl Future<Output = Arc<RefCell<i32>>> {
    |                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-   = note: required because it captures the following types: `&mut Context<'_>`, `impl Future<Output = Arc<RefCell<i32>>>`, `()`, `Ready<i32>`
+   = note: required because it captures the following types: `impl Future<Output = Arc<RefCell<i32>>>`, `()`, `Ready<i32>`
 note: required because it's used within this `async` block
   --> $DIR/issue-68112.rs:60:20
    |

src/test/ui/async-await/issue-68112.no_drop_tracking.stderr

@@ -57,7 +57,7 @@ note: required because it appears within the type `impl Future<Output = Arc<RefC
    |
 LL | fn make_non_send_future2() -> impl Future<Output = Arc<RefCell<i32>>> {
    |                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-   = note: required because it captures the following types: `&mut Context<'_>`, `impl Future<Output = Arc<RefCell<i32>>>`, `()`, `i32`, `Ready<i32>`
+   = note: required because it captures the following types: `impl Future<Output = Arc<RefCell<i32>>>`, `i32`, `Ready<i32>`
 note: required because it's used within this `async` block
   --> $DIR/issue-68112.rs:60:20
    |

src/test/ui/async-await/issue-70935-complex-spans.drop_tracking.stderr

@@ -18,7 +18,7 @@ LL |   async fn baz<T>(_c: impl FnMut() -> T) where T: Future<Output=()> {
    |  ___________________________________________________________________^
 LL | | }
    | |_^
-   = note: required because it captures the following types: `&mut Context<'_>`, `impl Future<Output = ()>`, `()`
+   = note: required because it captures the following types: `impl Future<Output = ()>`, `()`
 note: required because it's used within this `async` block
   --> $DIR/issue-70935-complex-spans.rs:16:5
    |

src/test/ui/async-await/partial-drop-partial-reinit.drop_tracking.stderr

@@ -11,7 +11,7 @@ LL | async fn foo() {
    |
    = help: within `impl Future<Output = ()>`, the trait `Send` is not implemented for `NotSend`
    = note: required because it appears within the type `(NotSend,)`
-   = note: required because it captures the following types: `&mut Context<'_>`, `(NotSend,)`, `()`, `impl Future<Output = ()>`
+   = note: required because it captures the following types: `(NotSend,)`, `()`, `impl Future<Output = ()>`
 note: required because it's used within this `async fn` body
   --> $DIR/partial-drop-partial-reinit.rs:31:16
    |

src/test/ui/async-await/partial-drop-partial-reinit.no_drop_tracking.stderr

@@ -11,7 +11,7 @@ LL | async fn foo() {
    |
    = help: within `impl Future<Output = ()>`, the trait `Send` is not implemented for `NotSend`
    = note: required because it appears within the type `(NotSend,)`
-   = note: required because it captures the following types: `&mut Context<'_>`, `(NotSend,)`, `impl Future<Output = ()>`, `()`
+   = note: required because it captures the following types: `(NotSend,)`, `impl Future<Output = ()>`
 note: required because it's used within this `async fn` body
   --> $DIR/partial-drop-partial-reinit.rs:31:16
    |