1use std::mem::MaybeUninit;
2
3use crate::{
4 IntoInner, IoBuf, IoBufMut, IoBuffer, IoSlice, IoSliceMut, SetBufInit, VectoredSlice,
5 VectoredSliceMut, t_alloc,
6};
7
8pub trait IoVectoredBuf: 'static {
10 unsafe fn iter_io_buffer(&self) -> impl Iterator<Item = IoBuffer>;
18
19 unsafe fn io_slices(&self) -> Vec<IoSlice> {
27 self.iter_io_slice().collect()
28 }
29
30 unsafe fn iter_io_slice(&self) -> impl Iterator<Item = IoSlice> {
38 self.iter_io_buffer().map(IoSlice::from)
39 }
40
41 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 fn total_len(&self) -> usize {
51 unsafe { self.iter_io_buffer().map(|buf| buf.len()).sum() }
52 }
53
54 fn total_capacity(&self) -> usize {
56 unsafe { self.iter_io_buffer().map(|buf| buf.capacity()).sum() }
57 }
58
59 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 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
148pub trait IoVectoredBufMut: IoVectoredBuf + SetBufInit {
150 unsafe fn io_slices_mut(&mut self) -> Vec<IoSliceMut> {
158 self.iter_io_slice_mut().collect()
159 }
160
161 unsafe fn iter_io_slice_mut(&mut self) -> impl Iterator<Item = IoSliceMut> {
169 self.iter_io_buffer().map(IoSliceMut::from)
170 }
171
172 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 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
237pub 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 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}