reusable_box/
lib.rs

1use std::alloc::Layout;
2use std::future::Future;
3use std::mem::size_of;
4use std::pin::Pin;
5use std::ptr::{drop_in_place, NonNull};
6use std::task::{Context, Poll};
7
8#[derive(Default)]
9pub struct ReusableBox {
10    // using vec as convenient memory allocator
11    buffer: Vec<u8>,
12}
13
14impl ReusableBox {
15    pub const fn new() -> Self {
16        Self { buffer: Vec::new() }
17    }
18
19    pub fn store_future<'a, F, O>(&'a mut self, f: F) -> ReusedBoxFuture<'a, O>
20    where
21        F: Future<Output = O> + Send + 'a,
22    {
23        const USIZE_SIZE: usize = size_of::<usize>();
24
25        let layout = Layout::new::<F>();
26
27        // Make sure the buffer has the required size (+ size of usize for potential alignment)
28        self.buffer.reserve(layout.size() + USIZE_SIZE);
29
30        let align_offset = self.buffer.as_ptr().align_offset(layout.align());
31
32        assert!(
33            align_offset <= USIZE_SIZE,
34            "Didn't expect the offset to be larger than {USIZE_SIZE} (is {align_offset})"
35        );
36
37        unsafe {
38            let ptr = self.buffer.as_mut_ptr().add(align_offset).cast::<F>();
39
40            ptr.write(f);
41
42            // Cast ptr to dyn Future which can be used later to access and drop the future without any generic parameters
43            let ptr = NonNull::new_unchecked(ptr as *mut (dyn Future<Output = O> + Send + 'a));
44
45            ReusedBoxFuture {
46                ptr_into_buffer: ptr,
47            }
48        }
49    }
50}
51
52pub struct ReusedBoxFuture<'a, O> {
53    ptr_into_buffer: NonNull<(dyn Future<Output = O> + Send + 'a)>,
54}
55
56// SAFETY:
57//
58// The future stored must be Send
59unsafe impl<O: Send> Send for ReusedBoxFuture<'_, O> {}
60
61impl<'a, O> ReusedBoxFuture<'a, O> {
62    fn future(&mut self) -> Pin<&mut (dyn Future<Output = O> + Send + 'a)> {
63        // SAFETY:
64        // self.ptr_into_buffer must always point into a space allocated by a vec
65        // Neither the pointer nor the vec which allocated the memory cannot be modified
66        // while `ReusedBoxFuture` exists.
67        unsafe { Pin::new_unchecked(self.ptr_into_buffer.as_mut()) }
68    }
69}
70
71impl<O> Future for ReusedBoxFuture<'_, O> {
72    type Output = O;
73
74    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
75        self.future().poll(cx)
76    }
77}
78
79impl<O> Drop for ReusedBoxFuture<'_, O> {
80    fn drop(&mut self) {
81        // SAFETY:
82        // ReusedBoxFuture's contract for creation requires the pointer to be valid
83        unsafe {
84            drop_in_place(self.ptr_into_buffer.as_ptr());
85        }
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use std::task::{Context, Poll};
93
94    #[tokio::test]
95    async fn set() {
96        let mut holder = ReusableBox::new();
97
98        let mut x = 0;
99
100        for i in 0..10 {
101            let g = holder.store_future(async {
102                x += i;
103            });
104
105            g.await
106        }
107
108        assert_eq!(x, 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9);
109    }
110
111    struct OnDrop<F: FnOnce()>(Option<F>);
112
113    impl<F: FnOnce()> OnDrop<F> {
114        fn new(f: F) -> Self {
115            Self(Some(f))
116        }
117    }
118
119    impl<F: FnOnce()> Drop for OnDrop<F> {
120        fn drop(&mut self) {
121            self.0.take().unwrap()();
122        }
123    }
124
125    impl<F: FnOnce()> Future for OnDrop<F> {
126        type Output = ();
127
128        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
129            Poll::Ready(())
130        }
131    }
132
133    #[tokio::test]
134    async fn drop_future() {
135        let mut holder = ReusableBox::new();
136
137        let mut x = 0;
138
139        for i in 0..10 {
140            let g = holder.store_future(OnDrop::new(|| x += i));
141
142            g.await
143        }
144
145        assert_eq!(x, 45);
146    }
147
148    #[tokio::test]
149    async fn drop_on_set() {
150        let mut holder = ReusableBox::new();
151
152        let mut x = 0;
153
154        holder.store_future(OnDrop::new(|| x += 1));
155        let g = holder.store_future(OnDrop::new(|| ()));
156
157        drop(g);
158        drop(holder);
159
160        assert_eq!(x, 1);
161    }
162
163    #[tokio::test]
164    async fn return_value() {
165        let mut holder = ReusableBox::new();
166
167        holder.store_future(async { 1 });
168        let g = holder.store_future(async { 2 });
169
170        let v = g.await;
171
172        assert_eq!(v, 2);
173    }
174
175    trait MyAsyncTrait {
176        fn test(&mut self) -> impl Future<Output = u32> + Send;
177    }
178
179    impl MyAsyncTrait for u32 {
180        async fn test(&mut self) -> u32 {
181            *self + 2
182        }
183    }
184
185    trait DynMyAsyncTrait: MyAsyncTrait {
186        fn dyn_test<'this>(
187            &'this mut self,
188            bx: &'this mut ReusableBox,
189        ) -> ReusedBoxFuture<'this, u32> {
190            bx.store_future(MyAsyncTrait::test(self))
191        }
192    }
193
194    impl DynMyAsyncTrait for u32 {}
195
196    #[tokio::test]
197    async fn async_trait_without_realloc() {
198        let mut holder = ReusableBox::new();
199
200        let mut xy = 1u32;
201        let mut g = xy.dyn_test(&mut holder);
202
203        let v = g.future().await;
204
205        assert_eq!(v, 3);
206    }
207}
reusable_box/lib.rs

reusable_box/
lib.rs
