compio_buf/
io_vec_buf.rs

1use std::mem::MaybeUninit;
2
3use crate::{
4    IntoInner, IoBuf, IoBufMut, IoBuffer, IoSlice, IoSliceMut, SetBufInit, VectoredSlice,
5    VectoredSliceMut, t_alloc,
6};
7
8/// A trait for vectored buffers.
9pub trait IoVectoredBuf: 'static {
10    /// An iterator over the [`IoBuffer`]s.
11    ///
12    /// # Safety
13    ///
14    /// The returned slice must not live longer than `self`.
15    /// It is static to provide convenience from writing self-referenced
16    /// structure.
17    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer>;
18
19    /// Collected [`IoSlice`]s of the buffers.
20    ///
21    /// # Safety
22    ///
23    /// The returned slice must not live longer than `self`.
24    /// It is static to provide convenience from writing self-referenced
25    /// structure.
26    unsafe fn io_slices(&self) -> Vec<IoSlice> {
27        unsafe { self.iter_io_slice() }.collect()
28    }
29
30    /// An iterator over the [`IoSlice`]s.
31    ///
32    /// # Safety
33    ///
34    /// The returned slice must not live longer than `self`.
35    /// It is static to provide convenience from writing self-referenced
36    /// structure.
37    unsafe fn iter_io_slice(&self) -> impl Iterator<Item = IoSlice> {
38        unsafe { self.iter_io_buffer() }.map(IoSlice::from)
39    }
40
41    /// An iterator over slices.
42    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
43        unsafe {
44            self.iter_io_slice()
45                .map(|slice| std::slice::from_raw_parts(slice.as_ptr(), slice.len()))
46        }
47    }
48
49    /// The total length of all buffers.
50    fn total_len(&self) -> usize {
51        unsafe { self.iter_io_buffer().map(|buf| buf.len()).sum() }
52    }
53
54    /// The total capacity of all buffers.
55    fn total_capacity(&self) -> usize {
56        unsafe { self.iter_io_buffer().map(|buf| buf.capacity()).sum() }
57    }
58
59    /// Wrap self into an owned iterator.
60    fn owned_iter(self) -> Result<VectoredBufIter<Self>, Self>
61    where
62        Self: Sized,
63    {
64        let len = unsafe { self.iter_io_buffer().count() };
65        if len > 0 {
66            Ok(VectoredBufIter {
67                buf: self,
68                index: 0,
69                len,
70                filled: 0,
71                cur_filled: 0,
72            })
73        } else {
74            Err(self)
75        }
76    }
77
78    /// Get an owned view with a specific offset.
79    fn slice(self, begin: usize) -> VectoredSlice<Self>
80    where
81        Self: Sized,
82    {
83        VectoredSlice::new(self, begin)
84    }
85}
86
87impl<T: IoBuf> IoVectoredBuf for &'static [T] {
88    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
89        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
90    }
91}
92
93impl<T: IoBuf> IoVectoredBuf for &'static mut [T] {
94    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
95        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
96    }
97}
98
99impl<T: IoBuf, const N: usize> IoVectoredBuf for [T; N] {
100    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
101        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
102    }
103}
104
105impl<T: IoBuf, #[cfg(feature = "allocator_api")] A: std::alloc::Allocator + 'static> IoVectoredBuf
106    for t_alloc!(Vec, T, A)
107{
108    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
109        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
110    }
111}
112
113#[cfg(feature = "arrayvec")]
114impl<T: IoBuf, const N: usize> IoVectoredBuf for arrayvec::ArrayVec<T, N> {
115    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
116        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
117    }
118}
119
120#[cfg(feature = "smallvec")]
121impl<T: IoBuf, const N: usize> IoVectoredBuf for smallvec::SmallVec<[T; N]>
122where
123    [T; N]: smallvec::Array<Item = T>,
124{
125    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
126        self.iter().map(|buf| unsafe { buf.as_io_buffer() })
127    }
128}
129
130impl<T: IoBuf, Rest: IoVectoredBuf> IoVectoredBuf for (T, Rest) {
131    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
132        unsafe { std::iter::once(self.0.as_io_buffer()).chain(self.1.iter_io_buffer()) }
133    }
134}
135
136impl<T: IoBuf> IoVectoredBuf for (T,) {
137    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
138        unsafe { std::iter::once(self.0.as_io_buffer()) }
139    }
140}
141
142impl IoVectoredBuf for () {
143    unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer> {
144        std::iter::empty()
145    }
146}
147
148/// A trait for mutable vectored buffers.
149pub trait IoVectoredBufMut: IoVectoredBuf + SetBufInit {
150    /// An iterator for the [`IoSliceMut`]s of the buffers.
151    ///
152    /// # Safety
153    ///
154    /// The returned slice must not live longer than `self`.
155    /// It is static to provide convenience from writing self-referenced
156    /// structure.
157    unsafe fn io_slices_mut(&mut self) -> Vec<IoSliceMut> {
158        unsafe { self.iter_io_slice_mut() }.collect()
159    }
160
161    /// An iterator over the [`IoSliceMut`]s.
162    ///
163    /// # Safety
164    ///
165    /// The returned slice must not live longer than `self`.
166    /// It is static to provide convenience from writing self-referenced
167    /// structure.
168    unsafe fn iter_io_slice_mut(&mut self) -> impl Iterator<Item = IoSliceMut> {
169        unsafe { self.iter_io_buffer() }.map(IoSliceMut::from)
170    }
171
172    /// An iterator over mutable slices.
173    fn iter_slice_mut(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
174        unsafe {
175            self.iter_io_slice_mut()
176                .map(|slice| std::slice::from_raw_parts_mut(slice.as_ptr().cast(), slice.len()))
177        }
178    }
179
180    /// Get an owned mutable view with a specific offset.
181    fn slice_mut(self, begin: usize) -> VectoredSliceMut<Self>
182    where
183        Self: Sized,
184    {
185        VectoredSliceMut::new(self, begin)
186    }
187}
188
189impl<T: IoBufMut> IoVectoredBufMut for &'static mut [T] {}
190
191impl<T: IoBufMut, const N: usize> IoVectoredBufMut for [T; N] {}
192
193impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: std::alloc::Allocator + 'static>
194    IoVectoredBufMut for t_alloc!(Vec, T, A)
195{
196}
197
198#[cfg(feature = "arrayvec")]
199impl<T: IoBufMut, const N: usize> IoVectoredBufMut for arrayvec::ArrayVec<T, N> {}
200
201#[cfg(feature = "smallvec")]
202impl<T: IoBufMut, const N: usize> IoVectoredBufMut for smallvec::SmallVec<[T; N]> where
203    [T; N]: smallvec::Array<Item = T>
204{
205}
206
207impl<T: IoBufMut, Rest: IoVectoredBufMut> IoVectoredBufMut for (T, Rest) {}
208
209impl<T: IoBufMut> IoVectoredBufMut for (T,) {}
210
211impl IoVectoredBufMut for () {}
212
213impl<T: IoBufMut, Rest: IoVectoredBufMut> SetBufInit for (T, Rest) {
214    unsafe fn set_buf_init(&mut self, len: usize) {
215        let buf0_len = std::cmp::min(len, self.0.buf_capacity());
216
217        // SAFETY: buf0_len <= self.0.buf_capacity()
218        unsafe { self.0.set_buf_init(buf0_len) };
219        // SAFETY: propogate
220        unsafe { self.1.set_buf_init(len - buf0_len) };
221    }
222}
223
224impl<T: IoBufMut> SetBufInit for (T,) {
225    unsafe fn set_buf_init(&mut self, len: usize) {
226        unsafe { self.0.set_buf_init(len) };
227    }
228}
229
230impl SetBufInit for () {
231    unsafe fn set_buf_init(&mut self, len: usize) {
232        assert_eq!(
233            len, 0,
234            "set_buf_init called with non-zero len on empty buffer"
235        );
236    }
237}
238
239/// An owned iterator over a vectored buffer.
240pub struct VectoredBufIter<T> {
241    buf: T,
242    index: usize,
243    len: usize,
244    filled: usize,
245    cur_filled: usize,
246}
247
248impl<T> VectoredBufIter<T> {
249    /// Create a new [`VectoredBufIter`] from an indexable container. If the
250    /// container is empty, return the buffer back in `Err(T)`.
251    pub fn next(mut self) -> Result<Self, T> {
252        self.index += 1;
253        if self.index < self.len {
254            self.filled += self.cur_filled;
255            self.cur_filled = 0;
256            Ok(self)
257        } else {
258            Err(self.buf)
259        }
260    }
261}
262
263impl<T> IntoInner for VectoredBufIter<T> {
264    type Inner = T;
265
266    fn into_inner(self) -> Self::Inner {
267        self.buf
268    }
269}
270
271unsafe impl<T: IoVectoredBuf> IoBuf for VectoredBufIter<T> {
272    fn as_buf_ptr(&self) -> *const u8 {
273        unsafe {
274            self.buf
275                .iter_io_buffer()
276                .nth(self.index)
277                .unwrap()
278                .as_ptr()
279                .add(self.cur_filled)
280                .cast_const()
281                .cast()
282        }
283    }
284
285    fn buf_len(&self) -> usize {
286        unsafe { self.buf.iter_io_buffer().nth(self.index).unwrap().len() - self.cur_filled }
287    }
288
289    fn buf_capacity(&self) -> usize {
290        unsafe {
291            self.buf
292                .iter_io_buffer()
293                .nth(self.index)
294                .unwrap()
295                .capacity()
296                - self.cur_filled
297        }
298    }
299}
300
301impl<T: IoVectoredBuf + SetBufInit> SetBufInit for VectoredBufIter<T> {
302    unsafe fn set_buf_init(&mut self, len: usize) {
303        self.cur_filled = len;
304        unsafe { self.buf.set_buf_init(self.filled + self.cur_filled) };
305    }
306}
307
308unsafe impl<T: IoVectoredBufMut> IoBufMut for VectoredBufIter<T> {
309    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
310        unsafe {
311            self.buf
312                .iter_io_buffer()
313                .nth(self.index)
314                .unwrap()
315                .as_ptr()
316                .add(self.cur_filled)
317                .cast()
318        }
319    }
320}
compio_buf/io_vec_buf.rs

compio_buf/
io_vec_buf.rs
