Auto merge of #129313 - RalfJung:coroutine-niches, r=compiler-errors

Supress niches in coroutines to avoid aliasing violations

As mentioned [here](https://github.com/rust-lang/rust/issues/63818#issuecomment-2264915918), using niches in fields of coroutines that are referenced by other fields is unsound: the discriminant accesses violate the aliasing requirements of the reference pointing to the relevant field. This issue causes [Miri errors in practice](https://github.com/rust-lang/miri/issues/3780).

The "obvious" fix for this is to suppress niches in coroutines. That's what this PR does. However, we have several tests explicitly ensuring that we *do* use niches in coroutines. So I see two options:
- We guard this behavior behind a `-Z` flag (that Miri will set by default). There is no known case of these aliasing violations causing miscompilations. But absence of evidence is not evidence of absence...
- (What this PR does right now.) We temporarily adjust the coroutine layout logic and the associated tests until the proper fix lands. The "proper fix" here is to wrap fields that other fields can point to in [`UnsafePinned`](https://github.com/rust-lang/rust/issues/125735) and make `UnsafePinned` suppress niches; that would then still permit using niches of *other* fields (those that never get borrowed). However, I know that coroutine sizes are already a problem, so I am not sure if this temporary size regression is acceptable.

`@compiler-errors` any opinion? Also who else should be Cc'd here?

6 files changed, 102 insertions, 27 deletions

diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs
index 234e1a6d55e..c23a7fec5a5 100644
--- a/compiler/rustc_ty_utils/src/layout.rs
+++ b/compiler/rustc_ty_utils/src/layout.rs
@@ -1001,7 +1001,13 @@ fn coroutine_layout<'tcx>(
         },
         fields: outer_fields,
         abi,
-        largest_niche: prefix.largest_niche,
+        // Suppress niches inside coroutines. If the niche is inside a field that is aliased (due to
+        // self-referentiality), getting the discriminant can cause aliasing violations.
+        // `UnsafeCell` blocks niches for the same reason, but we don't yet have `UnsafePinned` that
+        // would do the same for us here.
+        // See <https://github.com/rust-lang/rust/issues/63818>, <https://github.com/rust-lang/miri/issues/3780>.
+        // FIXME: Remove when <https://github.com/rust-lang/rust/issues/125735> is implemented and aliased coroutine fields are wrapped in `UnsafePinned`.
+        largest_niche: None,
         size,
         align,
         max_repr_align: None,
diff --git a/src/tools/miri/tests/pass/async-niche-aliasing.rs b/src/tools/miri/tests/pass/async-niche-aliasing.rs
new file mode 100644
index 00000000000..7f19afb33e4
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-niche-aliasing.rs
@@ -0,0 +1,66 @@
+//@revisions: stack tree
+//@[tree]compile-flags: -Zmiri-tree-borrows
+
+use std::{
+    future::Future,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll, Wake},
+    mem::MaybeUninit,
+};
+
+struct ThingAdder<'a> {
+    // Using `MaybeUninit` to ensure there are no niches here.
+    thing: MaybeUninit<&'a mut String>,
+}
+
+impl Future for ThingAdder<'_> {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
+        unsafe {
+            **self.get_unchecked_mut().thing.assume_init_mut() += ", world";
+        }
+        Poll::Pending
+    }
+}
+
+fn main() {
+    let mut thing = "hello".to_owned();
+    // This future has (at least) two fields, a String (`thing`) and a ThingAdder pointing to that string.
+    let fut = async move { ThingAdder { thing: MaybeUninit::new(&mut thing) }.await };
+
+    let mut fut = MaybeDone::Future(fut);
+    let mut fut = unsafe { Pin::new_unchecked(&mut fut) };
+
+    let waker = Arc::new(DummyWaker).into();
+    let mut ctx = Context::from_waker(&waker);
+    // This ends up reading the discriminant of the `MaybeDone`. If that is stored inside the
+    // `thing: String` as a niche optimization, that causes aliasing conflicts with the reference
+    // stored in `ThingAdder`.
+    assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending);
+    assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending);
+}
+
+struct DummyWaker;
+
+impl Wake for DummyWaker {
+    fn wake(self: Arc<Self>) {}
+}
+
+pub enum MaybeDone<F: Future> {
+    Future(F),
+    Done,
+}
+impl<F: Future<Output = ()>> Future for MaybeDone<F> {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        unsafe {
+            match *self.as_mut().get_unchecked_mut() {
+                MaybeDone::Future(ref mut f) => Pin::new_unchecked(f).poll(cx),
+                MaybeDone::Done => unreachable!(),
+            }
+        }
+    }
+}
diff --git a/tests/ui/async-await/async-drop.rs b/tests/ui/async-await/async-drop.rs
index 12f120a0b12..4e60598661f 100644
--- a/tests/ui/async-await/async-drop.rs
+++ b/tests/ui/async-await/async-drop.rs
@@ -53,18 +53,20 @@ fn main() {
     let i = 13;
     let fut = pin!(async {
         test_async_drop(Int(0), 0).await;
-        test_async_drop(AsyncInt(0), 104).await;
-        test_async_drop([AsyncInt(1), AsyncInt(2)], 152).await;
-        test_async_drop((AsyncInt(3), AsyncInt(4)), 488).await;
+        // FIXME(#63818): niches in coroutines are disabled.
+        // Some of these sizes should be smaller, as indicated in comments.
+        test_async_drop(AsyncInt(0), /*104*/ 112).await;
+        test_async_drop([AsyncInt(1), AsyncInt(2)], /*152*/ 168).await;
+        test_async_drop((AsyncInt(3), AsyncInt(4)), /*488*/ 528).await;
         test_async_drop(5, 0).await;
         let j = 42;
         test_async_drop(&i, 0).await;
         test_async_drop(&j, 0).await;
-        test_async_drop(AsyncStruct { b: AsyncInt(8), a: AsyncInt(7), i: 6 }, 1688).await;
+        test_async_drop(AsyncStruct { b: AsyncInt(8), a: AsyncInt(7), i: 6 }, /*1688*/ 1792).await;
         test_async_drop(ManuallyDrop::new(AsyncInt(9)), 0).await;
 
         let foo = AsyncInt(10);
-        test_async_drop(AsyncReference { foo: &foo }, 104).await;
+        test_async_drop(AsyncReference { foo: &foo }, /*104*/ 112).await;
 
         let foo = AsyncInt(11);
         test_async_drop(
@@ -73,17 +75,17 @@ fn main() {
                 let foo = AsyncInt(10);
                 foo
             },
-            120,
+            /*120*/ 136,
         )
         .await;
 
-        test_async_drop(AsyncEnum::A(AsyncInt(12)), 680).await;
-        test_async_drop(AsyncEnum::B(SyncInt(13)), 680).await;
+        test_async_drop(AsyncEnum::A(AsyncInt(12)), /*680*/ 736).await;
+        test_async_drop(AsyncEnum::B(SyncInt(13)), /*680*/ 736).await;
 
-        test_async_drop(SyncInt(14), 16).await;
+        test_async_drop(SyncInt(14), /*16*/ 24).await;
         test_async_drop(
             SyncThenAsync { i: 15, a: AsyncInt(16), b: SyncInt(17), c: AsyncInt(18) },
-            3064,
+            /*3064*/ 3296,
         )
         .await;
 
@@ -99,11 +101,11 @@ fn main() {
                 black_box(core::future::ready(())).await;
                 foo
             },
-            120,
+            /*120*/ 136,
         )
         .await;
 
-        test_async_drop(AsyncUnion { signed: 21 }, 32).await;
+        test_async_drop(AsyncUnion { signed: 21 }, /*32*/ 40).await;
     });
     let res = fut.poll(&mut cx);
     assert_eq!(res, Poll::Ready(()));
diff --git a/tests/ui/async-await/future-sizes/async-awaiting-fut.stdout b/tests/ui/async-await/future-sizes/async-awaiting-fut.stdout
index def967ba195..642e27b2a57 100644
--- a/tests/ui/async-await/future-sizes/async-awaiting-fut.stdout
+++ b/tests/ui/async-await/future-sizes/async-awaiting-fut.stdout
@@ -14,28 +14,26 @@ print-type-size         field `.value`: 3077 bytes
 print-type-size type: `{async fn body of calls_fut<{async fn body of big_fut()}>()}`: 3077 bytes, alignment: 1 bytes
 print-type-size     discriminant: 1 bytes
 print-type-size     variant `Unresumed`: 1025 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
+print-type-size         upvar `.fut`: 1025 bytes
 print-type-size     variant `Suspend0`: 2052 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
-print-type-size         padding: 1 bytes
-print-type-size         local `.fut`: 1025 bytes, alignment: 1 bytes
+print-type-size         upvar `.fut`: 1025 bytes
+print-type-size         local `.fut`: 1025 bytes
 print-type-size         local `..coroutine_field4`: 1 bytes, type: bool
 print-type-size         local `.__awaitee`: 1 bytes, type: {async fn body of wait()}
 print-type-size     variant `Suspend1`: 3076 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
-print-type-size         padding: 1026 bytes
+print-type-size         upvar `.fut`: 1025 bytes
+print-type-size         padding: 1025 bytes
 print-type-size         local `..coroutine_field4`: 1 bytes, alignment: 1 bytes, type: bool
 print-type-size         local `.__awaitee`: 1025 bytes, type: {async fn body of big_fut()}
 print-type-size     variant `Suspend2`: 2052 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
-print-type-size         padding: 1 bytes
-print-type-size         local `.fut`: 1025 bytes, alignment: 1 bytes
+print-type-size         upvar `.fut`: 1025 bytes
+print-type-size         local `.fut`: 1025 bytes
 print-type-size         local `..coroutine_field4`: 1 bytes, type: bool
 print-type-size         local `.__awaitee`: 1 bytes, type: {async fn body of wait()}
 print-type-size     variant `Returned`: 1025 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
+print-type-size         upvar `.fut`: 1025 bytes
 print-type-size     variant `Panicked`: 1025 bytes
-print-type-size         upvar `.fut`: 1025 bytes, offset: 0 bytes, alignment: 1 bytes
+print-type-size         upvar `.fut`: 1025 bytes
 print-type-size type: `std::mem::ManuallyDrop<{async fn body of big_fut()}>`: 1025 bytes, alignment: 1 bytes
 print-type-size     field `.value`: 1025 bytes
 print-type-size type: `std::mem::MaybeUninit<{async fn body of big_fut()}>`: 1025 bytes, alignment: 1 bytes
diff --git a/tests/ui/coroutine/discriminant.rs b/tests/ui/coroutine/discriminant.rs
index d6879e21825..ca4fcedd7e5 100644
--- a/tests/ui/coroutine/discriminant.rs
+++ b/tests/ui/coroutine/discriminant.rs
@@ -124,12 +124,14 @@ fn main() {
     };
 
     assert_eq!(size_of_val(&gen_u8_tiny_niche()), 1);
-    assert_eq!(size_of_val(&Some(gen_u8_tiny_niche())), 1); // uses niche
+    // FIXME(#63818): niches in coroutines are disabled.
+    // assert_eq!(size_of_val(&Some(gen_u8_tiny_niche())), 1); // uses niche
     assert_eq!(size_of_val(&Some(Some(gen_u8_tiny_niche()))), 2); // cannot use niche anymore
     assert_eq!(size_of_val(&gen_u8_full()), 1);
     assert_eq!(size_of_val(&Some(gen_u8_full())), 2); // cannot use niche
     assert_eq!(size_of_val(&gen_u16()), 2);
-    assert_eq!(size_of_val(&Some(gen_u16())), 2); // uses niche
+    // FIXME(#63818): niches in coroutines are disabled.
+    // assert_eq!(size_of_val(&Some(gen_u16())), 2); // uses niche
 
     cycle(gen_u8_tiny_niche(), 254);
     cycle(gen_u8_full(), 255);
diff --git a/tests/ui/coroutine/niche-in-coroutine.rs b/tests/ui/coroutine/niche-in-coroutine.rs
index 117ee9e6f03..f268ef09f89 100644
--- a/tests/ui/coroutine/niche-in-coroutine.rs
+++ b/tests/ui/coroutine/niche-in-coroutine.rs
@@ -15,5 +15,6 @@ fn main() {
         take(x);
     };
 
-    assert_eq!(size_of_val(&gen1), size_of_val(&Some(gen1)));
+    // FIXME(#63818): niches in coroutines are disabled. Should be `assert_eq`.
+    assert_ne!(size_of_val(&gen1), size_of_val(&Some(gen1)));
 }