Make the computation of coroutine_captures_by_ref_ty more sophisticated

compiler/rustc_hir_typeck/src/upvar.rs

@@ -367,37 +367,48 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 ty::INNERMOST,
                 ty::BoundRegion { var: ty::BoundVar::ZERO, kind: ty::BoundRegionKind::BrEnv },
             );
+
+            let num_args = args
+                .as_coroutine_closure()
+                .coroutine_closure_sig()
+                .skip_binder()
+                .tupled_inputs_ty
+                .tuple_fields()
+                .len();
+            let typeck_results = self.typeck_results.borrow();
+
             let tupled_upvars_ty_for_borrow = Ty::new_tup_from_iter(
                 self.tcx,
-                self.typeck_results
-                    .borrow()
-                    .closure_min_captures_flattened(
-                        self.tcx.coroutine_for_closure(closure_def_id).expect_local(),
-                    )
-                    // Skip the captures that are just moving the closure's args
-                    // into the coroutine. These are always by move, and we append
-                    // those later in the `CoroutineClosureSignature` helper functions.
-                    .skip(
-                        args.as_coroutine_closure()
-                            .coroutine_closure_sig()
-                            .skip_binder()
-                            .tupled_inputs_ty
-                            .tuple_fields()
-                            .len(),
-                    )
-                    .map(|captured_place| {
-                        let upvar_ty = captured_place.place.ty();
-                        let capture = captured_place.info.capture_kind;
+                ty::analyze_coroutine_closure_captures(
+                    typeck_results.closure_min_captures_flattened(closure_def_id),
+                    typeck_results
+                        .closure_min_captures_flattened(
+                            self.tcx.coroutine_for_closure(closure_def_id).expect_local(),
+                        )
+                        // Skip the captures that are just moving the closure's args
+                        // into the coroutine. These are always by move, and we append
+                        // those later in the `CoroutineClosureSignature` helper functions.
+                        .skip(num_args),
+                    |(_, parent_capture), (_, child_capture)| {
+                        // This is subtle. See documentation on function.
+                        let needs_ref = should_reborrow_from_env_of_parent_coroutine_closure(
+                            parent_capture,
+                            child_capture,
+                        );
+
+                        let upvar_ty = child_capture.place.ty();
+                        let capture = child_capture.info.capture_kind;
                         // Not all upvars are captured by ref, so use
                         // `apply_capture_kind_on_capture_ty` to ensure that we
                         // compute the right captured type.
-                        apply_capture_kind_on_capture_ty(
+                        return apply_capture_kind_on_capture_ty(
                             self.tcx,
                             upvar_ty,
                             capture,
-                            Some(closure_env_region),
-                        )
-                    }),
+                            if needs_ref { Some(closure_env_region) } else { child_capture.region },
+                        );
+                    },
+                ),
             );
             let coroutine_captures_by_ref_ty = Ty::new_fn_ptr(
                 self.tcx,
@@ -1761,6 +1772,61 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     }
 }
 
+/// Determines whether a child capture that is derived from a parent capture
+/// should be borrowed with the lifetime of the parent coroutine-closure's env.
+///
+/// There are two cases when this needs to happen:
+///
+/// (1.) Are we borrowing data owned by the parent closure? We can determine if
+/// that is the case by checking if the parent capture is by move, EXCEPT if we
+/// apply a deref projection, which means we're reborrowing a reference that we
+/// captured by move.
+///
+/// ```rust
+/// #![feature(async_closure)]
+/// let x = &i32; // Let's call this lifetime `'1`.
+/// let c = async move || {
+///     println!("{:?}", *x);
+///     // Even though the inner coroutine borrows by ref, we're only capturing `*x`,
+///     // not `x`, so the inner closure is allowed to reborrow the data for `'1`.
+/// };
+/// ```
+///
+/// (2.) If a coroutine is mutably borrowing from a parent capture, then that
+/// mutable borrow cannot live for longer than either the parent *or* the borrow
+/// that we have on the original upvar. Therefore we always need to borrow the
+/// child capture with the lifetime of the parent coroutine-closure's env.
+///
+/// ```rust
+/// #![feature(async_closure)]
+/// let mut x = 1i32;
+/// let c = async || {
+///     x = 1;
+///     // The parent borrows `x` for some `&'1 mut i32`.
+///     // However, when we call `c()`, we implicitly autoref for the signature of
+///     // `AsyncFnMut::async_call_mut`. Let's call that lifetime `'call`. Since
+///     // the maximum that `&'call mut &'1 mut i32` can be reborrowed is `&'call mut i32`,
+///     // the inner coroutine should capture w/ the lifetime of the coroutine-closure.
+/// };
+/// ```
+///
+/// If either of these cases apply, then we should capture the borrow with the
+/// lifetime of the parent coroutine-closure's env. Luckily, if this function is
+/// not correct, then the program is not unsound, since we still borrowck and validate
+/// the choices made from this function -- the only side-effect is that the user
+/// may receive unnecessary borrowck errors.
+fn should_reborrow_from_env_of_parent_coroutine_closure<'tcx>(
+    parent_capture: &ty::CapturedPlace<'tcx>,
+    child_capture: &ty::CapturedPlace<'tcx>,
+) -> bool {
+    (!parent_capture.is_by_ref()
+        && !matches!(
+            child_capture.place.projections.get(parent_capture.place.projections.len()),
+            Some(Projection { kind: ProjectionKind::Deref, .. })
+        ))
+        || matches!(child_capture.info.capture_kind, UpvarCapture::ByRef(ty::BorrowKind::MutBorrow))
+}
+
 /// Truncate the capture so that the place being borrowed is in accordance with RFC 1240,
 /// which states that it's unsafe to take a reference into a struct marked `repr(packed)`.
 fn restrict_repr_packed_field_ref_capture<'tcx>(

tests/ui/async-await/async-borrowck-escaping-closure-error.rs

@@ -5,7 +5,6 @@ fn foo() -> Box<dyn std::future::Future<Output = u32>> {
     let x = 0u32;
     Box::new((async || x)())
     //~^ ERROR cannot return value referencing local variable `x`
-    //~| ERROR cannot return value referencing temporary value
 }
 
 fn main() {

tests/ui/async-await/async-borrowck-escaping-closure-error.stderr

@@ -7,15 +7,6 @@ LL |     Box::new((async || x)())
    |     |        `x` is borrowed here
    |     returns a value referencing data owned by the current function
 
-error[E0515]: cannot return value referencing temporary value
-  --> $DIR/async-borrowck-escaping-closure-error.rs:6:5
-   |
-LL |     Box::new((async || x)())
-   |     ^^^^^^^^^------------^^^
-   |     |        |
-   |     |        temporary value created here
-   |     returns a value referencing data owned by the current function
-
-error: aborting due to 2 previous errors
+error: aborting due to 1 previous error
 
 For more information about this error, try `rustc --explain E0515`.

tests/ui/async-await/async-closures/moro-example.rs

@@ -0,0 +1,43 @@
+//@ check-pass
+//@ edition: 2021
+
+#![feature(async_closure)]
+
+use std::future::Future;
+use std::pin::Pin;
+use std::{marker::PhantomData, sync::Mutex};
+
+type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
+
+pub struct Scope<'scope, 'env: 'scope> {
+    enqueued: Mutex<Vec<BoxFuture<'scope, ()>>>,
+    phantom: PhantomData<&'env ()>,
+}
+
+impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
+    pub fn spawn(&'scope self, future: impl Future<Output = ()> + Send + 'scope) {
+        self.enqueued.lock().unwrap().push(Box::pin(future));
+    }
+}
+
+fn scope_with_closure<'env, B>(_body: B) -> BoxFuture<'env, ()>
+where
+    for<'scope> B: async FnOnce(&'scope Scope<'scope, 'env>),
+{
+    todo!()
+}
+
+type ScopeRef<'scope, 'env> = &'scope Scope<'scope, 'env>;
+
+async fn go<'a>(value: &'a i32) {
+    let closure = async |scope: ScopeRef<'_, 'a>| {
+        let _future1 = scope.spawn(async {
+            // Make sure that `*value` is immutably borrowed with lifetime of
+            // `'a` and not with the lifetime of the containing coroutine-closure.
+            let _v = *value;
+        });
+    };
+    scope_with_closure(closure).await;
+}
+
+fn main() {}

tests/ui/async-await/async-closures/no-borrow-from-env.rs

@@ -0,0 +1,44 @@
+//@ edition: 2021
+//@ check-pass
+
+#![feature(async_closure)]
+
+fn outlives<'a>(_: impl Sized + 'a) {}
+
+async fn call_once(f: impl async FnOnce()) {
+    f().await;
+}
+
+fn simple<'a>(x: &'a i32) {
+    let c = async || { println!("{}", *x); };
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+
+    let c = async move || { println!("{}", *x); };
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+}
+
+struct S<'a>(&'a i32);
+
+fn through_field<'a>(x: S<'a>) {
+    let c = async || { println!("{}", *x.0); };
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+
+    let c = async move || { println!("{}", *x.0); };
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+}
+
+fn through_field_and_ref<'a>(x: &S<'a>) {
+    let c = async || { println!("{}", *x.0); };
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+
+    let c = async move || { println!("{}", *x.0); };
+    outlives::<'a>(c());
+    // outlives::<'a>(call_once(c)); // FIXME(async_closures): Figure out why this fails
+}
+
+fn main() {}

tests/ui/async-await/async-closures/without-precise-captures-we-are-powerless.rs

@@ -0,0 +1,47 @@
+//@ edition: 2018
+
+// This is `no-borrow-from-env.rs`, but under edition 2018 we still want to make
+// sure that we don't ICE or anything, even if precise closure captures means
+// that we can't actually borrowck successfully.
+
+#![feature(async_closure)]
+
+fn outlives<'a>(_: impl Sized + 'a) {}
+
+async fn call_once(f: impl async FnOnce()) {
+    f().await;
+}
+
+fn simple<'a>(x: &'a i32) {
+    let c = async || { println!("{}", *x); }; //~ ERROR `x` does not live long enough
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+
+    let c = async move || { println!("{}", *x); };
+    outlives::<'a>(c()); //~ ERROR `c` does not live long enough
+    outlives::<'a>(call_once(c)); //~ ERROR cannot move out of `c`
+}
+
+struct S<'a>(&'a i32);
+
+fn through_field<'a>(x: S<'a>) {
+    let c = async || { println!("{}", *x.0); }; //~ ERROR `x` does not live long enough
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c));
+
+    let c = async move || { println!("{}", *x.0); }; //~ ERROR cannot move out of `x`
+    outlives::<'a>(c()); //~ ERROR `c` does not live long enough
+    outlives::<'a>(call_once(c)); //~ ERROR cannot move out of `c`
+}
+
+fn through_field_and_ref<'a>(x: &S<'a>) {
+    let c = async || { println!("{}", *x.0); }; //~ ERROR `x` does not live long enough
+    outlives::<'a>(c());
+    outlives::<'a>(call_once(c)); //~ ERROR explicit lifetime required in the type of `x`
+
+    let c = async move || { println!("{}", *x.0); };
+    outlives::<'a>(c()); //~ ERROR `c` does not live long enough
+    // outlives::<'a>(call_once(c)); // FIXME(async_closures): Figure out why this fails
+}
+
+fn main() {}