tokio_io_utility/
reusable_io_slices.rs

Help
1use std::{
2    io::IoSlice,
3    mem::{ManuallyDrop, MaybeUninit},
4    num::NonZeroUsize,
5    ptr::NonNull,
6    slice::{from_raw_parts, from_raw_parts_mut},
7};
8
9/// [`Box`]ed [`IoSlice`] that can be reused for different io_slices
10/// with different lifetime.
11#[derive(Debug)]
12pub struct ReusableIoSlices {
13    ptr: NonNull<()>,
14    cap: NonZeroUsize,
15}
16
17unsafe impl Send for ReusableIoSlices {}
18unsafe impl Sync for ReusableIoSlices {}
19
20impl ReusableIoSlices {
21    /// Create new [`ReusableIoSlices`].
22    pub fn new(cap: NonZeroUsize) -> Self {
23        let mut v = ManuallyDrop::new(Vec::<MaybeUninit<IoSlice<'_>>>::with_capacity(cap.get()));
24
25        debug_assert_eq!(v.capacity(), cap.get());
26        debug_assert_eq!(v.len(), 0);
27
28        let ptr = v.as_mut_ptr();
29
30        // Safety:
31        //
32        //  - ptr is allocated using Vec::with_capacity, with non-zero cap
33        //  - Vec::as_mut_ptr returns a non-null, valid pointer when
34        //    the capacity is non-zero.
35        //  - It is valid after the vec is dropped since it is wrapped
36        //    in ManuallyDrop.
37        let ptr = unsafe { NonNull::new_unchecked(ptr as *mut ()) };
38
39        Self { ptr, cap }
40    }
41
42    /// Return the underlying io_slices
43    pub fn get_mut(&mut self) -> &mut [MaybeUninit<IoSlice<'_>>] {
44        unsafe {
45            from_raw_parts_mut(
46                self.ptr.as_ptr() as *mut MaybeUninit<IoSlice<'_>>,
47                self.cap.get(),
48            )
49        }
50    }
51
52    /// Return the underlying io_slices
53    pub fn get(&self) -> &[MaybeUninit<IoSlice<'_>>] {
54        unsafe {
55            from_raw_parts(
56                self.ptr.as_ptr() as *const MaybeUninit<IoSlice<'_>>,
57                self.cap.get(),
58            )
59        }
60    }
61}
62
63impl Drop for ReusableIoSlices {
64    fn drop(&mut self) {
65        // Cast it back to its original type
66        let ptr = self.ptr.as_ptr() as *mut MaybeUninit<IoSlice<'_>>;
67
68        // Safety:
69        //
70        //  - ptr are obtained using `Vec::as_mut_ptr` from
71        //    a `ManuallyDrop<Vec<_>>` with non-zero cap.
72        //  - `IoSlice` does not have a `Drop` implementation and is
73        //    `Copy`able, so it is ok to pass `0` as length.
74        //  - self.cap.get() == v.capacity()
75        let v: Vec<MaybeUninit<IoSlice<'_>>> =
76            unsafe { Vec::from_raw_parts(ptr, 0, self.cap.get()) };
77
78        drop(v);
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn test_reusable_io_slices() {
88        let io_slice = IoSlice::new(b"123exr3x");
89
90        for size in 1..300 {
91            let cap = NonZeroUsize::new(size).unwrap();
92            let mut reusable_io_slices = ReusableIoSlices::new(cap);
93
94            for uninit_io_slice in reusable_io_slices.get_mut() {
95                uninit_io_slice.write(io_slice);
96            }
97        }
98    }
99}
tokio_io_utility/reusable_io_slices.rs

tokio_io_utility/
reusable_io_slices.rs
