From 781455c5eff6c82b72dd1913999078ca609c1536 Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Fri, 12 Aug 2022 12:19:59 -0700
Subject: [PATCH 1/2] Add regression test for #9

---
 src/lib.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 0e58ad7..436cb5b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -241,12 +241,17 @@ impl<'a, T, const STACK_SIZE: usize> Drop for StackFuture<'a, T, { STACK_SIZE }>
 mod tests {
     use crate::StackFuture;
     use core::task::Poll;
+    use futures::channel::mpsc;
     use futures::executor::block_on;
     use futures::pin_mut;
     use futures::Future;
+    use futures::SinkExt;
+    use futures::Stream;
+    use futures::StreamExt;
     use std::sync::Arc;
     use std::task::Context;
     use std::task::Wake;
+    use std::thread;
 
     #[test]
     fn create_and_run() {
@@ -340,4 +345,31 @@ mod tests {
     fn is_aligned<T>(ptr: *mut T, alignment: usize) -> bool {
         (ptr as usize) & (alignment - 1) == 0
     }
+
+    /// Regression test for #9
+    #[test]
+    fn stress_drop_sender() {
+        const ITER: usize = if cfg!(miri) { 100 } else { 10000 };
+
+        fn list() -> impl Stream<Item = i32> {
+            let (tx, rx) = mpsc::channel(1);
+            thread::spawn(move || {
+                block_on(send_one_two_three(tx));
+            });
+            rx
+        }
+
+        for _ in 0..ITER {
+            let v: Vec<_> = block_on(list().collect());
+            assert_eq!(v, vec![1, 2, 3]);
+        }
+    }
+
+    fn send_one_two_three(mut tx: mpsc::Sender<i32>) -> StackFuture<'static, (), 512> {
+        StackFuture::from(async move {
+            for i in 1..=3 {
+                tx.send(i).await.unwrap();
+            }
+        })
+    }
 }

From 02e53d177e6e809faee41acafa1b6d34b6040ac7 Mon Sep 17 00:00:00 2001
From: Eric Holk <ericholk@microsoft.com>
Date: Fri, 12 Aug 2022 15:07:37 -0700
Subject: [PATCH 2/2] Respect stacked borrows when polling

Polling the wrapped future previously had several reborrow operations,
such as in the use of `map_unchecked_mut` and in the implementation of
`as_mut_ptr`. This caused issues when running under miri in the
`stress_drop_sender` test.

This change re-implements these problematic cases to avoid reborrows,
either through using `Pin::get_unchecked_mut` or transmuting pointers
rather than borrowing and casting.

Fixes #9
---
 src/lib.rs | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 436cb5b..5e80076 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -190,7 +190,10 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
     /// enough to hold F and any required padding.
     fn as_mut_ptr<F>(&mut self) -> *mut F {
         assert!(Self::has_space_for::<F>());
-        self.data.as_mut_ptr().cast()
+        // SAFETY: Self is laid out so that the space for the future comes at offset 0.
+        // This is checked by an assertion in Self::from. Thus it's safe to cast a pointer
+        // to Self into a pointer to the wrapped future.
+        unsafe { mem::transmute(self) }
     }
 
     /// Returns a pinned mutable reference to a type F stored in self.data
@@ -226,8 +229,17 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
 impl<'a, T, const STACK_SIZE: usize> Future for StackFuture<'a, T, { STACK_SIZE }> {
     type Output = T;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        (self.as_mut().poll_fn)(self, cx)
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // SAFETY: This is doing pin projection. We unpin self so we can
+        // access self.poll_fn, and then re-pin self to pass it into poll_in.
+        // The part of the struct that needs to be pinned is data, since it
+        // contains a potentially self-referential future object, but since we
+        // do not touch that while self is unpinned and we do not move self
+        // while unpinned we are okay.
+        unsafe {
+            let this = self.get_unchecked_mut();
+            (this.poll_fn)(Pin::new_unchecked(this), cx)
+        }
     }
 }
 
@@ -346,10 +358,11 @@ mod tests {
         (ptr as usize) & (alignment - 1) == 0
     }
 
-    /// Regression test for #9
     #[test]
     fn stress_drop_sender() {
-        const ITER: usize = if cfg!(miri) { 100 } else { 10000 };
+        // Regression test for #9
+
+        const ITER: usize = if cfg!(miri) { 10 } else { 10000 };
 
         fn list() -> impl Stream<Item = i32> {
             let (tx, rx) = mpsc::channel(1);