Fix drop-tracking ICE when a struct containing a field with a Drop impl is used across an await

Previously, drop-tracking would incorrectly assume the struct would be dropped immediately, which
was not true: when the field had a type with a manual `Drop` impl, the drop becomes observable and
has to be dropped after the await instead.

For reasons I don't understand, this also fixes another error crater popped up related to type parameters.

 rust-lang#98476

Author: jyn514

compiler/rustc_typeck/src/check/generator_interior.rs

@@ -14,7 +14,7 @@ use rustc_hir::hir_id::HirIdSet;
 use rustc_hir::intravisit::{self, Visitor};
 use rustc_hir::{Arm, Expr, ExprKind, Guard, HirId, Pat, PatKind};
 use rustc_middle::middle::region::{self, Scope, ScopeData, YieldData};
-use rustc_middle::ty::{self, RvalueScopes, Ty, TyCtxt};
+use rustc_middle::ty::{self, RvalueScopes, Ty, TyCtxt, TypeVisitable};
 use rustc_span::symbol::sym;
 use rustc_span::Span;
 use tracing::debug;
@@ -376,6 +376,17 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
 
         debug!("is_borrowed_temporary: {:?}", self.drop_ranges.is_borrowed_temporary(expr));
 
+        let ty = self.fcx.typeck_results.borrow().expr_ty_adjusted_opt(expr);
+        let may_need_drop = |ty: Ty<'tcx>| {
+            // Avoid ICEs in needs_drop.
+            let ty = self.fcx.resolve_vars_if_possible(ty);
+            let ty = self.fcx.tcx.erase_regions(ty);
+            if ty.needs_infer() {
+                return true;
+            }
+            ty.needs_drop(self.fcx.tcx, self.fcx.param_env)
+        };
+
         // Typically, the value produced by an expression is consumed by its parent in some way,
         // so we only have to check if the parent contains a yield (note that the parent may, for
         // example, store the value into a local variable, but then we already consider local
@@ -384,7 +395,18 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
         // However, in the case of temporary values, we are going to store the value into a
         // temporary on the stack that is live for the current temporary scope and then return a
         // reference to it. That value may be live across the entire temporary scope.
-        let scope = if self.drop_ranges.is_borrowed_temporary(expr) {
+        //
+        // There's another subtlety: if the type has an observable drop, it must be dropped after
+        // the yield, even if it's not borrowed or referenced after the yield. Ideally this would
+        // *only* happen for types with observable drop, not all types which wrap them, but that
+        // doesn't match the behavior of MIR borrowck and causes ICEs. See the FIXME comment in
+        // src/test/ui/generator/drop-tracking-parent-expression.rs.
+        let scope = if self.drop_ranges.is_borrowed_temporary(expr)
+            || ty.map_or(true, |ty| {
+                let needs_drop = may_need_drop(ty);
+                debug!(?needs_drop, ?ty);
+                needs_drop
+            }) {
             self.rvalue_scopes.temporary_scope(self.region_scope_tree, expr.hir_id.local_id)
         } else {
             debug!("parent_node: {:?}", self.fcx.tcx.hir().find_parent_node(expr.hir_id));
@@ -398,7 +420,7 @@ impl<'a, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'tcx> {
 
         // If there are adjustments, then record the final type --
         // this is the actual value that is being produced.
-        if let Some(adjusted_ty) = self.fcx.typeck_results.borrow().expr_ty_adjusted_opt(expr) {
+        if let Some(adjusted_ty) = ty {
             self.record(adjusted_ty, expr.hir_id, scope, Some(expr), expr.span);
         }
 

src/test/ui/async-await/default-struct-update.rs

@@ -0,0 +1,22 @@
+// build-pass
+// edition:2018
+// compile-flags: -Zdrop-tracking=y
+
+fn main() {
+    let _ = foo();
+}
+
+async fn from_config(_: Config) {}
+
+async fn foo() {
+    from_config(Config {
+        nickname: None,
+        ..Default::default()
+    })
+    .await;
+}
+
+#[derive(Default)]
+struct Config {
+    nickname: Option<Box<u8>>,
+}

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

@@ -1,13 +1,13 @@
 error[E0277]: `Sender<i32>` cannot be shared between threads safely
-  --> $DIR/issue-70935-complex-spans.rs:13:45
+  --> $DIR/issue-70935-complex-spans.rs:12:45
    |
 LL | fn foo(tx: std::sync::mpsc::Sender<i32>) -> impl Future + Send {
    |                                             ^^^^^^^^^^^^^^^^^^ `Sender<i32>` cannot be shared between threads safely
    |
    = help: the trait `Sync` is not implemented for `Sender<i32>`
    = note: required because of the requirements on the impl of `Send` for `&Sender<i32>`
 note: required because it's used within this closure
-  --> $DIR/issue-70935-complex-spans.rs:25:13
+  --> $DIR/issue-70935-complex-spans.rs:16:13
    |
 LL |         baz(|| async{
    |             ^^
@@ -16,16 +16,14 @@ note: required because it's used within this `async fn` body
    |
 LL |   async fn baz<T>(_c: impl FnMut() -> T) where T: Future<Output=()> {
    |  ___________________________________________________________________^
-LL | |
 LL | | }
    | |_^
    = note: required because it captures the following types: `ResumeTy`, `impl for<'r, 's, 't0> Future<Output = ()>`, `()`
 note: required because it's used within this `async` block
-  --> $DIR/issue-70935-complex-spans.rs:23:16
+  --> $DIR/issue-70935-complex-spans.rs:15:16
    |
 LL |       async move {
    |  ________________^
-LL | |
 LL | |         baz(|| async{
 LL | |             foo(tx.clone());
 LL | |         }).await;

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

@@ -1,22 +1,22 @@
 error: future cannot be sent between threads safely
-  --> $DIR/issue-70935-complex-spans.rs:13:45
+  --> $DIR/issue-70935-complex-spans.rs:12:45
    |
 LL | fn foo(tx: std::sync::mpsc::Sender<i32>) -> impl Future + Send {
    |                                             ^^^^^^^^^^^^^^^^^^ future created by async block is not `Send`
    |
    = help: the trait `Sync` is not implemented for `Sender<i32>`
 note: future is not `Send` as this value is used across an await
-  --> $DIR/issue-70935-complex-spans.rs:27:11
+  --> $DIR/issue-70935-complex-spans.rs:18:11
    |
 LL |           baz(|| async{
    |  _____________-
 LL | |             foo(tx.clone());
 LL | |         }).await;
    | |         - ^^^^^^ await occurs here, with the value maybe used later
    | |_________|
-   |           has type `[closure@$DIR/issue-70935-complex-spans.rs:25:13: 25:15]` which is not `Send`
+   |           has type `[closure@$DIR/issue-70935-complex-spans.rs:16:13: 16:15]` which is not `Send`
 note: the value is later dropped here
-  --> $DIR/issue-70935-complex-spans.rs:27:17
+  --> $DIR/issue-70935-complex-spans.rs:18:17
    |
 LL |         }).await;
    |                 ^

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

@@ -7,22 +7,13 @@
 use std::future::Future;
 
 async fn baz<T>(_c: impl FnMut() -> T) where T: Future<Output=()> {
-//[drop_tracking]~^ within this `async fn` body
 }
 
 fn foo(tx: std::sync::mpsc::Sender<i32>) -> impl Future + Send {
-    //[normal]~^ ERROR: future cannot be sent between threads safely
-    //[drop_tracking]~^^ ERROR: `Sender<i32>` cannot be shared
-    //[drop_tracking]~|  NOTE: cannot be shared
-    //[drop_tracking]~|  NOTE: requirements on the impl of `Send`
-    //[drop_tracking]~|  NOTE: captures the following types
-    //[drop_tracking]~|  NOTE: in this expansion
-    //[drop_tracking]~|  NOTE: in this expansion
-    //[drop_tracking]~|  NOTE: in this expansion
-    //[drop_tracking]~|  NOTE: in this expansion
+    //[normal]~^ ERROR future cannot be sent between threads safely
+    //[drop_tracking]~^^ ERROR `Sender<i32>` cannot be shared between threads
     async move {
-    //[drop_tracking]~^ within this `async` block
-        baz(|| async{ //[drop_tracking]~ NOTE: used within this closure
+        baz(|| async{
             foo(tx.clone());
         }).await;
     }

src/test/ui/async-await/non-trivial-drop.rs

@@ -0,0 +1,36 @@
+// build-pass
+// edition:2018
+// compile-flags: -Zdrop-tracking=y
+
+#![feature(generators)]
+
+fn main() {
+    let _ = foo();
+}
+
+fn foo() {
+    || {
+        yield drop(Config {
+            nickname: NonCopy,
+            b: NonCopy2,
+        }.nickname);
+    };
+}
+
+#[derive(Default)]
+struct NonCopy;
+impl Drop for NonCopy {
+    fn drop(&mut self) {}
+}
+
+#[derive(Default)]
+struct NonCopy2;
+impl Drop for NonCopy2 {
+    fn drop(&mut self) {}
+}
+
+#[derive(Default)]
+struct Config {
+    nickname: NonCopy,
+    b: NonCopy2,
+}

src/test/ui/async-await/type-parameter-send.rs

@@ -0,0 +1,18 @@
+// check-pass
+// compile-flags: --crate-type lib
+// edition:2018
+
+fn assert_send<F: Send>(_: F) {}
+
+async fn __post<T>() -> T {
+    if false {
+        todo!()
+    } else {
+        async {}.await;
+        todo!()
+    }
+}
+
+fn foo<T>() {
+    assert_send(__post::<T>());
+}

src/test/ui/generator/derived-drop-parent-expr.rs

@@ -0,0 +1,17 @@
+// build-pass
+// compile-flags:-Zdrop-tracking
+
+//! Like drop-tracking-parent-expression, but also tests that this doesn't ICE when building MIR
+#![feature(generators)]
+
+fn assert_send<T: Send>(_thing: T) {}
+
+#[derive(Default)]
+pub struct Client { pub nickname: String }
+
+fn main() {
+    let g = move || match drop(Client { ..Client::default() }) {
+        _status => yield,
+    };
+    assert_send(g);
+}

src/test/ui/generator/drop-tracking-parent-expression.rs

@@ -0,0 +1,69 @@
+// compile-flags: -Zdrop-tracking
+#![feature(generators, negative_impls, rustc_attrs)]
+
+macro_rules! type_combinations {
+    (
+        $( $name:ident => { $( $tt:tt )* } );* $(;)?
+    ) => { $(
+        mod $name {
+            $( $tt )*
+
+            impl !Sync for Client {}
+            impl !Send for Client {}
+        }
+
+        // Struct update syntax. This fails because the Client used in the update is considered
+        // dropped *after* the yield.
+        {
+            let g = move || match drop($name::Client { ..$name::Client::default() }) {
+            //~^ `significant_drop::Client` which is not `Send`
+            //~| `insignificant_dtor::Client` which is not `Send`
+            //~| `derived_drop::Client` which is not `Send`
+                _ => yield,
+            };
+            assert_send(g);
+            //~^ ERROR cannot be sent between threads
+            //~| ERROR cannot be sent between threads
+            //~| ERROR cannot be sent between threads
+        }
+
+        // Simple owned value. This works because the Client is considered moved into `drop`,
+        // even though the temporary expression doesn't end until after the yield.
+        {
+            let g = move || match drop($name::Client::default()) {
+                _ => yield,
+            };
+            assert_send(g);
+        }
+    )* }
+}
+
+fn assert_send<T: Send>(_thing: T) {}
+
+fn main() {
+    type_combinations!(
+        // OK
+        copy => { #[derive(Copy, Clone, Default)] pub struct Client; };
+        // NOT OK: MIR borrowck thinks that this is used after the yield, even though
+        // this has no `Drop` impl and only the drops of the fields are observable.
+        // FIXME: this should compile.
+        derived_drop => { #[derive(Default)] pub struct Client { pub nickname: String } };
+        // NOT OK
+        significant_drop => {
+            #[derive(Default)]
+            pub struct Client;
+            impl Drop for Client {
+                fn drop(&mut self) {}
+            }
+        };
+        // NOT OK (we need to agree with MIR borrowck)
+        insignificant_dtor => {
+            #[derive(Default)]
+            #[rustc_insignificant_dtor]
+            pub struct Client;
+            impl Drop for Client {
+                fn drop(&mut self) {}
+            }
+        };
+    );
+}