compio_buf/
io_vec_buf.rs

1use std::{iter, mem::MaybeUninit};
2
3use crate::{IntoInner, IoBuf, IoBufMut, 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    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]>;
147
148    /// The total capacity of all buffers.
149    fn total_capacity(&mut self) -> usize {
150        self.iter_uninit_slice().map(|buf| buf.len()).sum()
151    }
152
153    /// Get an owned view of the vectored buffer.
154    ///
155    /// Unlike [`IoVectoredBuf::slice`], the iterator returned by this function
156    /// will skip both initialized and uninitialized bytes.
157    ///
158    /// # Examples
159    ///
160    /// ```
161    /// use compio_buf::{IoBuf, IoVectoredBuf, IoVectoredBufMut, VectoredSlice};
162    ///
163    /// # fn main() {
164    /// /// Create a buffer with given content and capacity.
165    /// fn new_buf(slice: &[u8], cap: usize) -> Vec<u8> {
166    ///     let mut buf = Vec::new();
167    ///     buf.reserve_exact(cap);
168    ///     buf.extend_from_slice(slice);
169    ///     buf
170    /// }
171    ///
172    /// let bufs = [new_buf(b"hello", 10), new_buf(b"world", 10)];
173    /// let vectored_buf = bufs.slice_mut(13);
174    /// let mut iter = vectored_buf.iter_slice();
175    /// let buf1 = iter.next().unwrap();
176    /// assert!(iter.next().is_none());
177    /// assert_eq!(buf1.as_init(), b"ld");
178    /// # }
179    /// ```
180    fn slice_mut(mut self, begin: usize) -> VectoredSlice<Self>
181    where
182        Self: Sized,
183    {
184        let mut offset = begin;
185        let mut idx = 0;
186
187        for b in self.iter_uninit_slice() {
188            let len = b.len();
189            if len > offset {
190                break;
191            }
192            offset -= len;
193            idx += 1;
194        }
195
196        VectoredSlice::new(self, begin, idx, offset)
197    }
198}
199
200impl<T: IoBufMut> IoVectoredBufMut for &'static mut [T] {
201    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
202        self.iter_mut().map(|buf| buf.as_uninit())
203    }
204}
205
206impl<T: IoBufMut, const N: usize> IoVectoredBufMut for [T; N] {
207    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
208        self.iter_mut().map(|buf| buf.as_uninit())
209    }
210}
211
212impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: std::alloc::Allocator + 'static>
213    IoVectoredBufMut for t_alloc!(Vec, T, A)
214{
215    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
216        self.iter_mut().map(|buf| buf.as_uninit())
217    }
218}
219
220#[cfg(feature = "arrayvec")]
221impl<T: IoBufMut, const N: usize> IoVectoredBufMut for arrayvec::ArrayVec<T, N> {
222    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
223        self.iter_mut().map(|buf| buf.as_uninit())
224    }
225}
226
227#[cfg(feature = "smallvec")]
228impl<T: IoBufMut, const N: usize> IoVectoredBufMut for smallvec::SmallVec<[T; N]>
229where
230    [T; N]: smallvec::Array<Item = T>,
231{
232    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
233        self.iter_mut().map(|buf| buf.as_uninit())
234    }
235}
236
237impl<T: IoBufMut, Rest: IoVectoredBufMut> IoVectoredBufMut for (T, Rest) {
238    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
239        let (h, t) = self;
240        iter::once(h.as_uninit()).chain(t.iter_uninit_slice())
241    }
242}
243
244impl<T: IoBufMut> IoVectoredBufMut for (T,) {
245    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
246        iter::once(self.0.as_uninit())
247    }
248}
249
250impl IoVectoredBufMut for () {
251    fn iter_uninit_slice(&mut self) -> impl Iterator<Item = &mut [MaybeUninit<u8>]> {
252        iter::empty()
253    }
254}
255
256impl<T: IoBufMut, Rest: IoVectoredBufMut> SetLen for (T, Rest) {
257    unsafe fn set_len(&mut self, len: usize) {
258        let head_len = std::cmp::min(len, self.0.buf_capacity());
259        let rest_len = len - head_len;
260
261        // SAFETY: head_len <= self.0.buf_capacity()
262        unsafe { self.0.set_len(head_len) };
263        // SAFETY: propagate
264        unsafe { self.1.set_len(rest_len) };
265    }
266}
267
268impl<T: IoBufMut> SetLen for (T,) {
269    unsafe fn set_len(&mut self, len: usize) {
270        unsafe { self.0.set_len(len) };
271    }
272}
273
274impl SetLen for () {
275    unsafe fn set_len(&mut self, len: usize) {
276        assert_eq!(len, 0, "set_len called with non-zero len on empty buffer");
277    }
278}
279
280/// An owned iterator over a vectored buffer.
281///
282/// Normally one would use [`IoVectoredBuf::owned_iter`] to create this
283/// iterator.
284pub struct VectoredBufIter<T> {
285    buf: T,
286    total_filled: usize,
287    index: usize,
288    len: usize,
289    filled: usize,
290}
291
292impl<T> VectoredBufIter<T> {
293    /// Create a new [`VectoredBufIter`] from an indexable container. If the
294    /// container is empty, return the buffer back in `Err(T)`.
295    pub fn next(mut self) -> Result<Self, T> {
296        self.index += 1;
297        if self.index < self.len {
298            self.total_filled += self.filled;
299            self.filled = 0;
300            Ok(self)
301        } else {
302            Err(self.buf)
303        }
304    }
305}
306
307impl<T: IoVectoredBuf> VectoredBufIter<T> {
308    fn new(buf: T) -> Result<Self, T> {
309        let len = buf.iter_slice().count();
310        if len > 0 {
311            Ok(Self {
312                buf,
313                index: 0,
314                len,
315                total_filled: 0,
316                filled: 0,
317            })
318        } else {
319            Err(buf)
320        }
321    }
322}
323
324impl<T> IntoInner for VectoredBufIter<T> {
325    type Inner = T;
326
327    fn into_inner(self) -> Self::Inner {
328        self.buf
329    }
330}
331
332impl<T: IoVectoredBuf> IoBuf for VectoredBufIter<T> {
333    fn as_init(&self) -> &[u8] {
334        let curr = self
335            .buf
336            .iter_slice()
337            .nth(self.index)
338            .expect("`index` should not exceed `len`");
339
340        &curr[self.filled..]
341    }
342}
343
344impl<T: IoVectoredBuf + SetLen> SetLen for VectoredBufIter<T> {
345    unsafe fn set_len(&mut self, len: usize) {
346        self.filled = len;
347
348        unsafe { self.buf.set_len(self.total_filled + self.filled) };
349    }
350}
351
352impl<T: IoVectoredBufMut> IoBufMut for VectoredBufIter<T> {
353    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
354        self.buf
355            .iter_uninit_slice()
356            .nth(self.index)
357            .expect("`index` should not exceed `len`")
358    }
359}
compio_buf/io_vec_buf.rs

compio_buf/
io_vec_buf.rs
