From 53698118e5b0a79765fbe5659379cfb578863277 Mon Sep 17 00:00:00 2001
From: Bryan Garza <1396101+bryangarza@users.noreply.github.com>
Date: Sat, 11 Feb 2023 03:05:05 +0000
Subject: [PATCH] Add tests for moving data across await point

This patch adds a few tests to assert the current behavior when passing
data across an await point. This will help to test out an upcoming fix
for the issue of arguments in async functions growing in size because of
the generator upvar that is generated when we desugar the async
function.

See https://github.com/rust-lang/rust/issues/62958

Also relates to https://github.com/rust-lang/rust/pull/107500
---
 tests/pass/move-data-across-await-point.rs | 78 ++++++++++++++++++++++
 1 file changed, 78 insertions(+)
 create mode 100644 tests/pass/move-data-across-await-point.rs

diff --git a/tests/pass/move-data-across-await-point.rs b/tests/pass/move-data-across-await-point.rs
new file mode 100644
index 0000000000..1e3e9c06ce
--- /dev/null
+++ b/tests/pass/move-data-across-await-point.rs
@@ -0,0 +1,78 @@
+use std::future::Future;
+
+// This test:
+// - Compares addresses of non-Copy data before and after moving it
+// - Writes to the pointer after it has moved across the await point
+async fn data_moved_async() {
+    // Vec<T> is not Copy
+    let mut x: Vec<u8> = vec![2];
+    let raw_pointer = &mut x as *mut Vec<u8>;
+    helper_async(x, raw_pointer).await;
+    unsafe {
+        assert_eq!(*raw_pointer, vec![3]);
+        // Drop to prevent leak.
+        std::ptr::drop_in_place(raw_pointer);
+    }
+}
+
+async fn helper_async(mut data: Vec<u8>, raw_pointer: *mut Vec<u8>) {
+    let raw_pointer2 = &mut data as *mut Vec<u8>;
+    // Addresses are different because the generator upvar for `data is a copy.
+    // To be more precise, there is a `move` happening in the MIR of the
+    // generator closure, which copies the pointer.
+    //
+    // When copy propagation is enabled for generator upvars, the pointer addresses
+    // here should be equal.
+    assert_ne!(raw_pointer, raw_pointer2);
+    unsafe {
+        std::ptr::write(raw_pointer, vec![3]);
+    }
+}
+
+// Same thing as above, but non-async.
+fn data_moved() {
+    let mut x: Vec<u8> = vec![2];
+    let raw_pointer = &mut x as *mut Vec<u8>;
+    helper(x, raw_pointer);
+    unsafe {
+        assert_eq!(*raw_pointer, vec![3]);
+        std::ptr::drop_in_place(raw_pointer);
+    }
+}
+
+#[inline(always)]
+fn helper(mut data: Vec<u8>, raw_pointer: *mut Vec<u8>) {
+    let raw_pointer2 = &mut data as *mut Vec<u8>;
+    assert_ne!(raw_pointer, raw_pointer2);
+    unsafe {
+        std::ptr::write(raw_pointer, vec![3]);
+    }
+}
+
+fn run_fut<T>(fut: impl Future<Output = T>) -> T {
+    use std::sync::Arc;
+    use std::task::{Context, Poll, Wake, Waker};
+
+    struct MyWaker;
+    impl Wake for MyWaker {
+        fn wake(self: Arc<Self>) {
+            unimplemented!()
+        }
+    }
+
+    let waker = Waker::from(Arc::new(MyWaker));
+    let mut context = Context::from_waker(&waker);
+
+    let mut pinned = Box::pin(fut);
+    loop {
+        match pinned.as_mut().poll(&mut context) {
+            Poll::Pending => continue,
+            Poll::Ready(v) => return v,
+        }
+    }
+}
+
+fn main() {
+    run_fut(data_moved_async());
+    data_moved();
+}