swap_buffer_queue/write_vectored/
vec.rs

1use std::{io::IoSlice, mem, mem::MaybeUninit, ops::Range};
2
3use crate::{
4    buffer::{Buffer, CellBuffer, Drain, Resize},
5    loom::{
6        cell::Cell,
7        sync::atomic::{AtomicUsize, Ordering},
8    },
9    write_vectored::{VectoredSlice, EMPTY_SLICE},
10};
11
12/// A buffer of [`IoSlice`]
13pub struct WriteVectoredVecBuffer<T> {
14    owned: Box<[Cell<MaybeUninit<T>>]>,
15    slices: Box<[Cell<IoSlice<'static>>]>,
16    total_size: AtomicUsize,
17}
18
19impl<T> Default for WriteVectoredVecBuffer<T> {
20    fn default() -> Self {
21        Self {
22            owned: Default::default(),
23            slices: Default::default(),
24            total_size: Default::default(),
25        }
26    }
27}
28
29// SAFETY: `WriteVectoredVecBuffer::clear` does clear the inserted range from the buffer
30unsafe impl<T> Buffer for WriteVectoredVecBuffer<T>
31where
32    T: AsRef<[u8]>,
33{
34    type Slice<'a> = VectoredSlice<'a>
35    where
36        T: 'a;
37
38    #[inline]
39    fn capacity(&self) -> usize {
40        self.owned.len()
41    }
42
43    #[inline]
44    unsafe fn slice(&mut self, range: Range<usize>) -> Self::Slice<'_> {
45        // SAFETY: [Cell<IoSlice>] has the same layout as [IoSlice]
46        // and function contract guarantees that the range is initialized
47        let slices = unsafe {
48            &mut *(&mut self.slices[range.start..range.end + 2] as *mut _
49                as *mut [IoSlice<'static>])
50        };
51        // SAFETY: slices are never read and live along their owner in the buffer, as they are
52        // inserted and removed together
53        unsafe { VectoredSlice::new(slices, self.total_size.load(Ordering::Acquire)) }
54    }
55
56    #[inline]
57    unsafe fn clear(&mut self, range: Range<usize>) {
58        *self.total_size.get_mut() = 0;
59        for index in range {
60            // SAFETY: function contract guarantees that the range is initialized
61            unsafe { self.remove(index) };
62        }
63    }
64}
65
66// SAFETY: `insert` does initialize the index in the buffer
67unsafe impl<T> CellBuffer<T> for WriteVectoredVecBuffer<T>
68where
69    T: AsRef<[u8]>,
70{
71    unsafe fn insert(&self, index: usize, value: T) {
72        // SAFETY: slice is never read with static lifetime, it will only be used as a reference
73        // with the same lifetime than the slice owner
74        let slice = unsafe { mem::transmute::<IoSlice, IoSlice>(IoSlice::new(value.as_ref())) };
75        self.slices[index + 1].set(slice);
76        self.owned[index].set(MaybeUninit::new(value));
77        self.total_size.fetch_add(slice.len(), Ordering::AcqRel);
78    }
79}
80
81impl<T> Resize for WriteVectoredVecBuffer<T>
82where
83    T: AsRef<[u8]>,
84{
85    fn resize(&mut self, capacity: usize) {
86        self.owned = (0..capacity)
87            .map(|_| Cell::new(MaybeUninit::uninit()))
88            .collect();
89        self.slices = (0..capacity + 2)
90            .map(|_| Cell::new(IoSlice::new(EMPTY_SLICE)))
91            .collect();
92    }
93}
94
95// SAFETY: `WriteVectoredVecBuffer::remove` does remove the index from the buffer
96unsafe impl<T> Drain for WriteVectoredVecBuffer<T>
97where
98    T: AsRef<[u8]>,
99{
100    type Value = T;
101
102    #[inline]
103    unsafe fn remove(&mut self, index: usize) -> Self::Value {
104        // SAFETY: function contract guarantees that the index has been inserted and is then initialized
105        let value = unsafe {
106            self.owned[index]
107                .replace(MaybeUninit::uninit())
108                .assume_init()
109        };
110        self.total_size
111            .fetch_sub(value.as_ref().len(), Ordering::Release);
112        value
113    }
114}