compio_buf/
io_vec_buf.rs

1use std::{iter, mem::MaybeUninit};
2
3use crate::{IntoInner, IoBuf, IoBufMut, IoBufMutExt, SetLen, VectoredSlice, t_alloc};
4
5/// A trait for vectored buffers.
6///
7/// # Note for implementors
8///
9/// The iterator must be idemptotent and always yield the same slices in the
10/// exact same orders, i.e., [`Iterator::enumerate`] will mark the same buffer
11/// with same index.
12pub trait IoVectoredBuf: 'static {
13    /// An iterator of initialized slice of the buffers.
14    fn iter_slice(&self) -> impl Iterator<Item = &[u8]>;
15
16    /// The total length of all buffers.
17    fn total_len(&self) -> usize {
18        self.iter_slice().map(|buf| buf.len()).sum()
19    }
20
21    /// Wrap self into an owned iterator.
22    fn owned_iter(self) -> Result<VectoredBufIter<Self>, Self>
23    where
24        Self: Sized,
25    {
26        VectoredBufIter::new(self)
27    }
28
29    /// Get an owned view of the vectored buffer that skips the first
30    /// `begin`-many **initialized** bytes.
31    ///
32    /// # Examples
33    ///
34    /// ```
35    /// use compio_buf::{IoBuf, IoVectoredBuf, VectoredSlice};
36    ///
37    /// # fn main() {
38    /// /// Create a buffer with given content and capacity.
39    /// fn new_buf(slice: &[u8], cap: usize) -> Vec<u8> {
40    ///     let mut buf = Vec::new();
41    ///     buf.reserve_exact(cap);
42    ///     buf.extend_from_slice(slice);
43    ///     buf
44    /// }
45    ///
46    /// let bufs = [new_buf(b"hello", 10), new_buf(b"world", 10)];
47    /// let vectored_buf = bufs.slice(3);
48    /// let mut iter = vectored_buf.iter_slice();
49    /// let buf1 = iter.next().unwrap();
50    /// let buf2 = iter.next().unwrap();
51    /// assert_eq!(&buf1.as_init()[..], b"lo");
52    /// assert_eq!(&buf2.as_init()[..], b"world");
53    ///
54    /// let bufs = [new_buf(b"hello", 10), new_buf(b"world", 10)];
55    /// let vectored_buf = bufs.slice(6);
56    /// let mut iter = vectored_buf.iter_slice();
57    /// let buf1 = iter.next().unwrap();
58    /// assert!(iter.next().is_none());
59    /// assert_eq!(&buf1.as_init()[..], b"orld");
60    /// # }
61    /// ```
62    fn slice(self, begin: usize) -> VectoredSlice<Self>
63    where
64        Self: Sized,
65    {
66        let mut offset = begin;
67        let mut idx = 0;
68
69        for b in self.iter_slice() {
70            let len = b.len();
71            if len > offset {
72                break;
73            }
74            offset -= len;
75            idx += 1;
76        }
77
78        VectoredSlice::new(self, begin, idx, offset)
79    }
80}
81
82impl<T: IoBuf> IoVectoredBuf for &'static [T] {
83    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
84        self.iter().map(|buf| buf.as_init())
85    }
86}
87
88impl<T: IoBuf> IoVectoredBuf for &'static mut [T] {
89    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
90        self.iter().map(|buf| buf.as_init())
91    }
92}
93
94impl<T: IoBuf, const N: usize> IoVectoredBuf for [T; N] {
95    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
96        self.iter().map(|buf| buf.as_init())
97    }
98}
99
100impl<T: IoBuf, #[cfg(feature = "allocator_api")] A: std::alloc::Allocator + 'static> IoVectoredBuf
101    for t_alloc!(Vec, T, A)
102{
103    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
104        self.iter().map(|buf| buf.as_init())
105    }
106}
107
108#[cfg(feature = "arrayvec")]
109impl<T: IoBuf, const N: usize> IoVectoredBuf for arrayvec::ArrayVec<T, N> {
110    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
111        self.iter().map(|buf| buf.as_init())
112    }
113}
114
115#[cfg(feature = "smallvec")]
116impl<T: IoBuf, const N: usize> IoVectoredBuf for smallvec::SmallVec<[T; N]>
117where
118    [T; N]: smallvec::Array<Item = T>,
119{
120    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
121        self.iter().map(|buf| buf.as_init())
122    }
123}
124
125impl<T: IoBuf, Rest: IoVectoredBuf> IoVectoredBuf for (T, Rest) {
126    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
127        std::iter::once(self.0.as_init()).chain(self.1.iter_slice())
128    }
129}
130
131impl<T: IoBuf> IoVectoredBuf for (T,) {
132    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
133        std::iter::once(self.0.as_init())
134    }
135}
136
137impl IoVectoredBuf for () {
138    fn iter_slice(&self) -> impl Iterator<Item = &[u8]> {
139        std::iter::empty()
140    }
141}
142
143/// A trait for mutable vectored buffers.
144pub trait IoVectoredBufMut: IoVectoredBuf + SetLen {
145    /// An iterator of maybe uninitialized slice of the buffers.
146    ///
147    /// # Safety
148    ///
149    /// The caller must not de-initialize any byte that belongs to the
150    /// initialized slices returned by [`IoVectoredBuf::iter_slice`].
151    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]>;
152
153    /// The total capacity of all buffers.
154    fn total_capacity(&mut self) -> usize {
155        // SAFETY: Only the lengths of the slices are read.
156        unsafe { self.iter_uninit_slice() }
157            .map(|buf| buf.len())
158            .sum()
159    }
160
161    /// Get an owned view of the vectored buffer.
162    ///
163    /// Unlike [`IoVectoredBuf::slice`], the iterator returned by this function
164    /// will skip both initialized and uninitialized bytes.
165    ///
166    /// # Examples
167    ///
168    /// ```
169    /// use compio_buf::{IoBuf, IoVectoredBuf, IoVectoredBufMut, VectoredSlice};
170    ///
171    /// # fn main() {
172    /// /// Create a buffer with given content and capacity.
173    /// fn new_buf(slice: &[u8], cap: usize) -> Vec<u8> {
174    ///     let mut buf = Vec::new();
175    ///     buf.reserve_exact(cap);
176    ///     buf.extend_from_slice(slice);
177    ///     buf
178    /// }
179    ///
180    /// let bufs = [new_buf(b"hello", 10), new_buf(b"world", 10)];
181    /// let vectored_buf = bufs.slice_mut(13);
182    /// let mut iter = vectored_buf.iter_slice();
183    /// let buf1 = iter.next().unwrap();
184    /// assert!(iter.next().is_none());
185    /// assert_eq!(buf1.as_init(), b"ld");
186    /// # }
187    /// ```
188    fn slice_mut(mut self, begin: usize) -> VectoredSlice<Self>
189    where
190        Self: Sized,
191    {
192        let mut offset = begin;
193        let mut idx = 0;
194
195        // SAFETY: The iterator is only used to read the lengths of the
196        // buffers.
197        for b in unsafe { self.iter_uninit_slice() } {
198            let len = b.len();
199            if len > offset {
200                break;
201            }
202            offset -= len;
203            idx += 1;
204        }
205
206        VectoredSlice::new(self, begin, idx, offset)
207    }
208}
209
210impl<T: IoBufMut> IoVectoredBufMut for &'static mut [T] {
211    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
212        self.iter_mut().map(|buf| unsafe { buf.as_uninit() })
213    }
214}
215
216impl<T: IoBufMut, const N: usize> IoVectoredBufMut for [T; N] {
217    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
218        self.iter_mut().map(|buf| unsafe { buf.as_uninit() })
219    }
220}
221
222impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: std::alloc::Allocator + 'static>
223    IoVectoredBufMut for t_alloc!(Vec, T, A)
224{
225    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
226        self.iter_mut().map(|buf| unsafe { buf.as_uninit() })
227    }
228}
229
230#[cfg(feature = "arrayvec")]
231impl<T: IoBufMut, const N: usize> IoVectoredBufMut for arrayvec::ArrayVec<T, N> {
232    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
233        self.iter_mut().map(|buf| unsafe { buf.as_uninit() })
234    }
235}
236
237#[cfg(feature = "smallvec")]
238impl<T: IoBufMut, const N: usize> IoVectoredBufMut for smallvec::SmallVec<[T; N]>
239where
240    [T; N]: smallvec::Array<Item = T>,
241{
242    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
243        self.iter_mut().map(|buf| unsafe { buf.as_uninit() })
244    }
245}
246
247impl<T: IoBufMut, Rest: IoVectoredBufMut> IoVectoredBufMut for (T, Rest) {
248    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
249        let (h, t) = self;
250        unsafe { iter::once(h.as_uninit()).chain(t.iter_uninit_slice()) }
251    }
252}
253
254impl<T: IoBufMut> IoVectoredBufMut for (T,) {
255    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
256        iter::once(unsafe { self.0.as_uninit() })
257    }
258}
259
260impl IoVectoredBufMut for () {
261    unsafe fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
262        iter::empty()
263    }
264}
265
266impl<T: IoBufMut, Rest: IoVectoredBufMut> SetLen for (T, Rest) {
267    unsafe fn set_len(&mut self, len: usize) {
268        let head_len = std::cmp::min(len, self.0.buf_capacity());
269        let rest_len = len - head_len;
270
271        // SAFETY: head_len <= self.0.buf_capacity()
272        unsafe { self.0.set_len(head_len) };
273        // SAFETY: propagate
274        unsafe { self.1.set_len(rest_len) };
275    }
276}
277
278impl<T: IoBufMut> SetLen for (T,) {
279    unsafe fn set_len(&mut self, len: usize) {
280        unsafe { self.0.set_len(len) };
281    }
282}
283
284impl SetLen for () {
285    unsafe fn set_len(&mut self, len: usize) {
286        assert_eq!(len, 0, "set_len called with non-zero len on empty buffer");
287    }
288}
289
290/// An owned iterator over a vectored buffer.
291///
292/// Normally one would use [`IoVectoredBuf::owned_iter`] to create this
293/// iterator.
294pub struct VectoredBufIter<T> {
295    buf: T,
296    total_filled: usize,
297    index: usize,
298    len: usize,
299    filled: usize,
300}
301
302impl<T> VectoredBufIter<T> {
303    /// Create a new [`VectoredBufIter`] from an indexable container. If the
304    /// container is empty, return the buffer back in `Err(T)`.
305    pub fn next(mut self) -> Result<Self, T> {
306        self.index += 1;
307        if self.index < self.len {
308            self.total_filled += self.filled;
309            self.filled = 0;
310            Ok(self)
311        } else {
312            Err(self.buf)
313        }
314    }
315}
316
317impl<T: IoVectoredBuf> VectoredBufIter<T> {
318    fn new(buf: T) -> Result<Self, T> {
319        let len = buf.iter_slice().count();
320        if len > 0 {
321            Ok(Self {
322                buf,
323                index: 0,
324                len,
325                total_filled: 0,
326                filled: 0,
327            })
328        } else {
329            Err(buf)
330        }
331    }
332}
333
334impl<T> IntoInner for VectoredBufIter<T> {
335    type Inner = T;
336
337    fn into_inner(self) -> Self::Inner {
338        self.buf
339    }
340}
341
342impl<T: IoVectoredBuf> IoBuf for VectoredBufIter<T> {
343    fn as_init(&self) -> &[u8] {
344        let curr = self
345            .buf
346            .iter_slice()
347            .nth(self.index)
348            .expect("`index` should not exceed `len`");
349
350        &curr[self.filled..]
351    }
352}
353
354impl<T: IoVectoredBuf + SetLen> SetLen for VectoredBufIter<T> {
355    unsafe fn set_len(&mut self, len: usize) {
356        self.filled = len;
357
358        unsafe { self.buf.set_len(self.total_filled + self.filled) };
359    }
360}
361
362impl<T: IoVectoredBufMut> IoBufMut for VectoredBufIter<T> {
363    unsafe fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
364        let filled = self.filled;
365        // SAFETY: The initialized prefix `..filled` of the current buffer is
366        // sliced off below.
367        let curr = unsafe {
368            self.buf
369                .iter_uninit_slice()
370                .nth(self.index)
371                .expect("`index` should not exceed `len`")
372        };
373        &mut curr[filled..]
374    }
375}
376
377#[cfg(test)]
378mod test {
379    use std::mem::MaybeUninit;
380
381    use crate::{IntoInner, IoBuf, IoBufMut, IoVectoredBuf, SetLen};
382
383    #[test]
384    fn vectored_buf_iter_uninit_skips_filled() {
385        let bufs = [Vec::with_capacity(8), Vec::with_capacity(8)];
386        let mut iter = bufs.owned_iter().unwrap();
387
388        // SAFETY: `iter` has not been advanced, so all of its capacity is
389        // exposed, and `0..3` is within it.
390        unsafe {
391            let slice = iter.as_uninit();
392            assert_eq!(slice.len(), 8);
393            slice[..3].fill(MaybeUninit::new(7));
394            // SAFETY: The range `0..3` has just been initialized.
395            iter.set_len(3);
396        }
397
398        // The filled prefix is not exposed by `as_init`.
399        assert!(iter.as_init().is_empty());
400
401        // SAFETY: The remaining tail is still uninitialized, and the filled
402        // prefix is excluded.
403        let tail = unsafe { iter.as_uninit() };
404        assert_eq!(tail.len(), 5);
405        tail[..2].fill(MaybeUninit::new(9));
406
407        let bufs = iter.into_inner();
408        assert_eq!(bufs[0].as_init(), &[7, 7, 7]);
409        assert_eq!(bufs[1].as_init(), &[]);
410    }
411}
compio_buf/io_vec_buf.rs

compio_buf/
io_vec_buf.rs
