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        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        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| 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| 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| 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| 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| 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| 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        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        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        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        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        self.0.set_buf_init(buf0_len);
218        self.1.set_buf_init(len - buf0_len);
219    }
220}
221
222impl<T: IoBufMut> SetBufInit for (T,) {
223    unsafe fn set_buf_init(&mut self, len: usize) {
224        self.0.set_buf_init(len);
225    }
226}
227
228impl SetBufInit for () {
229    unsafe fn set_buf_init(&mut self, len: usize) {
230        assert_eq!(
231            len, 0,
232            "set_buf_init called with non-zero len on empty buffer"
233        );
234    }
235}
236
237/// An owned iterator over a vectored buffer.
238pub struct VectoredBufIter<T> {
239    buf: T,
240    index: usize,
241    len: usize,
242    filled: usize,
243    cur_filled: usize,
244}
245
246impl<T> VectoredBufIter<T> {
247    /// Create a new [`VectoredBufIter`] from an indexable container. If the
248    /// container is empty, return the buffer back in `Err(T)`.
249    pub fn next(mut self) -> Result<Self, T> {
250        self.index += 1;
251        if self.index < self.len {
252            self.filled += self.cur_filled;
253            self.cur_filled = 0;
254            Ok(self)
255        } else {
256            Err(self.buf)
257        }
258    }
259}
260
261impl<T> IntoInner for VectoredBufIter<T> {
262    type Inner = T;
263
264    fn into_inner(self) -> Self::Inner {
265        self.buf
266    }
267}
268
269unsafe impl<T: IoVectoredBuf> IoBuf for VectoredBufIter<T> {
270    fn as_buf_ptr(&self) -> *const u8 {
271        unsafe {
272            self.buf
273                .iter_io_buffer()
274                .nth(self.index)
275                .unwrap()
276                .as_ptr()
277                .add(self.cur_filled)
278                .cast_const()
279                .cast()
280        }
281    }
282
283    fn buf_len(&self) -> usize {
284        unsafe { self.buf.iter_io_buffer().nth(self.index).unwrap().len() - self.cur_filled }
285    }
286
287    fn buf_capacity(&self) -> usize {
288        unsafe {
289            self.buf
290                .iter_io_buffer()
291                .nth(self.index)
292                .unwrap()
293                .capacity()
294                - self.cur_filled
295        }
296    }
297}
298
299impl<T: IoVectoredBuf + SetBufInit> SetBufInit for VectoredBufIter<T> {
300    unsafe fn set_buf_init(&mut self, len: usize) {
301        self.cur_filled = len;
302        self.buf.set_buf_init(self.filled + self.cur_filled);
303    }
304}
305
306unsafe impl<T: IoVectoredBufMut> IoBufMut for VectoredBufIter<T> {
307    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
308        unsafe {
309            self.buf
310                .iter_io_buffer()
311                .nth(self.index)
312                .unwrap()
313                .as_ptr()
314                .add(self.cur_filled)
315                .cast()
316        }
317    }
318}