1#[cfg(feature = "allocator_api")]
2use std::alloc::Allocator;
3use std::{mem::MaybeUninit, rc::Rc, sync::Arc};
4
5use crate::*;
6
7pub unsafe trait IoBuf: 'static {
17 fn as_buf_ptr(&self) -> *const u8;
22
23 fn buf_len(&self) -> usize;
30
31 fn buf_capacity(&self) -> usize;
38
39 fn as_slice(&self) -> &[u8] {
41 unsafe { std::slice::from_raw_parts(self.as_buf_ptr(), self.buf_len()) }
42 }
43
44 unsafe fn as_io_slice(&self) -> IoSlice {
52 IoSlice::from_slice(self.as_slice())
53 }
54
55 unsafe fn as_io_buffer(&self) -> IoBuffer {
63 IoBuffer::new(
64 self.as_buf_ptr().cast_mut().cast(),
65 self.buf_len(),
66 self.buf_capacity(),
67 )
68 }
69
70 fn slice(self, range: impl std::ops::RangeBounds<usize>) -> Slice<Self>
84 where
85 Self: Sized,
86 {
87 use std::ops::Bound;
88
89 let begin = match range.start_bound() {
90 Bound::Included(&n) => n,
91 Bound::Excluded(&n) => n + 1,
92 Bound::Unbounded => 0,
93 };
94
95 assert!(begin <= self.buf_capacity());
96
97 let end = match range.end_bound() {
98 Bound::Included(&n) => n.checked_add(1).expect("out of range"),
99 Bound::Excluded(&n) => n,
100 Bound::Unbounded => self.buf_capacity(),
101 };
102
103 assert!(end <= self.buf_capacity());
104 assert!(begin <= self.buf_len());
105
106 Slice::new(self, begin, end)
107 }
108
109 fn uninit(self) -> Uninit<Self>
131 where
132 Self: Sized,
133 {
134 Uninit::new(self)
135 }
136
137 fn filled(&self) -> bool {
139 self.buf_len() == self.buf_capacity()
140 }
141}
142
143unsafe impl<B: IoBuf + ?Sized> IoBuf for &'static B {
144 fn as_buf_ptr(&self) -> *const u8 {
145 (**self).as_buf_ptr()
146 }
147
148 fn buf_len(&self) -> usize {
149 (**self).buf_len()
150 }
151
152 fn buf_capacity(&self) -> usize {
153 (**self).buf_capacity()
154 }
155}
156
157unsafe impl<B: IoBuf + ?Sized> IoBuf for &'static mut B {
158 fn as_buf_ptr(&self) -> *const u8 {
159 (**self).as_buf_ptr()
160 }
161
162 fn buf_len(&self) -> usize {
163 (**self).buf_len()
164 }
165
166 fn buf_capacity(&self) -> usize {
167 (**self).buf_capacity()
168 }
169}
170
171unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
172 for t_alloc!(Box, B, A)
173{
174 fn as_buf_ptr(&self) -> *const u8 {
175 (**self).as_buf_ptr()
176 }
177
178 fn buf_len(&self) -> usize {
179 (**self).buf_len()
180 }
181
182 fn buf_capacity(&self) -> usize {
183 (**self).buf_capacity()
184 }
185}
186
187unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
188 for t_alloc!(Rc, B, A)
189{
190 fn as_buf_ptr(&self) -> *const u8 {
191 (**self).as_buf_ptr()
192 }
193
194 fn buf_len(&self) -> usize {
195 (**self).buf_len()
196 }
197
198 fn buf_capacity(&self) -> usize {
199 (**self).buf_capacity()
200 }
201}
202
203unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
204 for t_alloc!(Arc, B, A)
205{
206 fn as_buf_ptr(&self) -> *const u8 {
207 (**self).as_buf_ptr()
208 }
209
210 fn buf_len(&self) -> usize {
211 (**self).buf_len()
212 }
213
214 fn buf_capacity(&self) -> usize {
215 (**self).buf_capacity()
216 }
217}
218
219unsafe impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
220 for t_alloc!(Vec, u8, A)
221{
222 fn as_buf_ptr(&self) -> *const u8 {
223 self.as_ptr()
224 }
225
226 fn buf_len(&self) -> usize {
227 self.len()
228 }
229
230 fn buf_capacity(&self) -> usize {
231 self.capacity()
232 }
233}
234
235unsafe impl IoBuf for String {
236 fn as_buf_ptr(&self) -> *const u8 {
237 self.as_ptr()
238 }
239
240 fn buf_len(&self) -> usize {
241 self.len()
242 }
243
244 fn buf_capacity(&self) -> usize {
245 self.capacity()
246 }
247}
248
249unsafe impl IoBuf for str {
250 fn as_buf_ptr(&self) -> *const u8 {
251 self.as_ptr()
252 }
253
254 fn buf_len(&self) -> usize {
255 self.len()
256 }
257
258 fn buf_capacity(&self) -> usize {
259 self.len()
260 }
261}
262
263unsafe impl IoBuf for [u8] {
264 fn as_buf_ptr(&self) -> *const u8 {
265 self.as_ptr()
266 }
267
268 fn buf_len(&self) -> usize {
269 self.len()
270 }
271
272 fn buf_capacity(&self) -> usize {
273 self.len()
274 }
275}
276
277unsafe impl<const N: usize> IoBuf for [u8; N] {
278 fn as_buf_ptr(&self) -> *const u8 {
279 self.as_ptr()
280 }
281
282 fn buf_len(&self) -> usize {
283 N
284 }
285
286 fn buf_capacity(&self) -> usize {
287 N
288 }
289}
290
291#[cfg(feature = "bytes")]
292unsafe impl IoBuf for bytes::Bytes {
293 fn as_buf_ptr(&self) -> *const u8 {
294 self.as_ptr()
295 }
296
297 fn buf_len(&self) -> usize {
298 self.len()
299 }
300
301 fn buf_capacity(&self) -> usize {
302 self.len()
303 }
304}
305
306#[cfg(feature = "bytes")]
307unsafe impl IoBuf for bytes::BytesMut {
308 fn as_buf_ptr(&self) -> *const u8 {
309 self.as_ptr()
310 }
311
312 fn buf_len(&self) -> usize {
313 self.len()
314 }
315
316 fn buf_capacity(&self) -> usize {
317 self.capacity()
318 }
319}
320
321#[cfg(feature = "read_buf")]
322unsafe impl IoBuf for std::io::BorrowedBuf<'static> {
323 fn as_buf_ptr(&self) -> *const u8 {
324 self.filled().as_ptr()
325 }
326
327 fn buf_len(&self) -> usize {
328 self.len()
329 }
330
331 fn buf_capacity(&self) -> usize {
332 self.capacity()
333 }
334}
335
336#[cfg(feature = "arrayvec")]
337unsafe impl<const N: usize> IoBuf for arrayvec::ArrayVec<u8, N> {
338 fn as_buf_ptr(&self) -> *const u8 {
339 self.as_ptr()
340 }
341
342 fn buf_len(&self) -> usize {
343 self.len()
344 }
345
346 fn buf_capacity(&self) -> usize {
347 self.capacity()
348 }
349}
350
351#[cfg(feature = "smallvec")]
352unsafe impl<const N: usize> IoBuf for smallvec::SmallVec<[u8; N]>
353where
354 [u8; N]: smallvec::Array<Item = u8>,
355{
356 fn as_buf_ptr(&self) -> *const u8 {
357 self.as_ptr()
358 }
359
360 fn buf_len(&self) -> usize {
361 self.len()
362 }
363
364 fn buf_capacity(&self) -> usize {
365 self.capacity()
366 }
367}
368
369pub unsafe trait IoBufMut: IoBuf + SetBufInit {
380 fn as_buf_mut_ptr(&mut self) -> *mut u8;
385
386 fn as_mut_slice(&mut self) -> &mut [MaybeUninit<u8>] {
388 unsafe {
389 std::slice::from_raw_parts_mut(self.as_buf_mut_ptr().cast(), (*self).buf_capacity())
390 }
391 }
392
393 unsafe fn as_io_slice_mut(&mut self) -> IoSliceMut {
401 IoSliceMut::from_uninit(self.as_mut_slice())
402 }
403}
404
405unsafe impl<B: IoBufMut + ?Sized> IoBufMut for &'static mut B {
406 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
407 (**self).as_buf_mut_ptr()
408 }
409}
410
411unsafe impl<B: IoBufMut + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
412 for t_alloc!(Box, B, A)
413{
414 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
415 (**self).as_buf_mut_ptr()
416 }
417}
418
419unsafe impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
420 for t_alloc!(Vec, u8, A)
421{
422 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
423 self.as_mut_ptr()
424 }
425}
426
427unsafe impl IoBufMut for [u8] {
428 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
429 self.as_mut_ptr()
430 }
431}
432
433unsafe impl<const N: usize> IoBufMut for [u8; N] {
434 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
435 self.as_mut_ptr()
436 }
437}
438
439#[cfg(feature = "bytes")]
440unsafe impl IoBufMut for bytes::BytesMut {
441 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
442 self.as_mut_ptr()
443 }
444}
445
446#[cfg(feature = "read_buf")]
447unsafe impl IoBufMut for std::io::BorrowedBuf<'static> {
448 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
449 (*self).filled().as_ptr() as _
450 }
451}
452
453#[cfg(feature = "arrayvec")]
454unsafe impl<const N: usize> IoBufMut for arrayvec::ArrayVec<u8, N> {
455 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
456 self.as_mut_ptr()
457 }
458}
459
460#[cfg(feature = "smallvec")]
461unsafe impl<const N: usize> IoBufMut for smallvec::SmallVec<[u8; N]>
462where
463 [u8; N]: smallvec::Array<Item = u8>,
464{
465 fn as_buf_mut_ptr(&mut self) -> *mut u8 {
466 self.as_mut_ptr()
467 }
468}
469
470pub trait SetBufInit {
472 unsafe fn set_buf_init(&mut self, len: usize);
479}
480
481impl<B: SetBufInit + ?Sized> SetBufInit for &'static mut B {
482 unsafe fn set_buf_init(&mut self, len: usize) {
483 (**self).set_buf_init(len)
484 }
485}
486
487impl<B: SetBufInit + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
488 for t_alloc!(Box, B, A)
489{
490 unsafe fn set_buf_init(&mut self, len: usize) {
491 (**self).set_buf_init(len)
492 }
493}
494
495impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit for t_alloc!(Vec, u8, A) {
496 unsafe fn set_buf_init(&mut self, len: usize) {
497 if (**self).buf_len() < len {
498 self.set_len(len);
499 }
500 }
501}
502
503impl SetBufInit for [u8] {
504 unsafe fn set_buf_init(&mut self, len: usize) {
505 debug_assert!(len <= self.len());
506 }
507}
508
509impl<const N: usize> SetBufInit for [u8; N] {
510 unsafe fn set_buf_init(&mut self, len: usize) {
511 debug_assert!(len <= N);
512 }
513}
514
515#[cfg(feature = "bytes")]
516impl SetBufInit for bytes::BytesMut {
517 unsafe fn set_buf_init(&mut self, len: usize) {
518 if (**self).buf_len() < len {
519 self.set_len(len);
520 }
521 }
522}
523
524#[cfg(feature = "read_buf")]
525impl SetBufInit for std::io::BorrowedBuf<'static> {
526 unsafe fn set_buf_init(&mut self, len: usize) {
527 let current_len = (*self).buf_len();
528 if current_len < len {
529 self.unfilled().advance(len - current_len);
530 }
531 }
532}
533
534#[cfg(feature = "arrayvec")]
535impl<const N: usize> SetBufInit for arrayvec::ArrayVec<u8, N> {
536 unsafe fn set_buf_init(&mut self, len: usize) {
537 if (**self).buf_len() < len {
538 self.set_len(len);
539 }
540 }
541}
542
543#[cfg(feature = "smallvec")]
544impl<const N: usize> SetBufInit for smallvec::SmallVec<[u8; N]>
545where
546 [u8; N]: smallvec::Array<Item = u8>,
547{
548 unsafe fn set_buf_init(&mut self, len: usize) {
549 if (**self).buf_len() < len {
550 self.set_len(len);
551 }
552 }
553}
554
555impl<T: IoBufMut> SetBufInit for [T] {
556 unsafe fn set_buf_init(&mut self, len: usize) {
557 default_set_buf_init(self.iter_mut(), len)
558 }
559}
560
561impl<T: IoBufMut, const N: usize> SetBufInit for [T; N] {
562 unsafe fn set_buf_init(&mut self, len: usize) {
563 default_set_buf_init(self.iter_mut(), len)
564 }
565}
566
567impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
568 for t_alloc!(Vec, T, A)
569{
570 unsafe fn set_buf_init(&mut self, len: usize) {
571 default_set_buf_init(self.iter_mut(), len)
572 }
573}
574
575#[cfg(feature = "arrayvec")]
576impl<T: IoBufMut, const N: usize> SetBufInit for arrayvec::ArrayVec<T, N> {
577 unsafe fn set_buf_init(&mut self, len: usize) {
578 default_set_buf_init(self.iter_mut(), len)
579 }
580}
581
582#[cfg(feature = "smallvec")]
583impl<T: IoBufMut, const N: usize> SetBufInit for smallvec::SmallVec<[T; N]>
584where
585 [T; N]: smallvec::Array<Item = T>,
586{
587 unsafe fn set_buf_init(&mut self, len: usize) {
588 default_set_buf_init(self.iter_mut(), len)
589 }
590}
591
592unsafe fn default_set_buf_init<'a, B: IoBufMut>(
593 iter: impl IntoIterator<Item = &'a mut B>,
594 mut len: usize,
595) {
596 for buf in iter {
597 let capacity = (*buf).buf_capacity();
598 if len >= capacity {
599 buf.set_buf_init(capacity);
600 len -= capacity;
601 } else {
602 buf.set_buf_init(len);
603 len = 0;
604 }
605 }
606}