compio_buf/
io_buf.rs

1#[cfg(feature = "allocator_api")]
2use std::alloc::Allocator;
3use std::{error::Error, fmt::Display, mem::MaybeUninit, ops::RangeBounds, rc::Rc, sync::Arc};
4
5use crate::*;
6
7/// A trait for immutable buffers.
8///
9/// The `IoBuf` trait is implemented by buffer types that can be passed to
10/// immutable completion-based IO operations, like writing its content to a
11/// file. This trait will only take initialized bytes of a buffer into account.
12pub trait IoBuf: 'static {
13    /// Get the slice of initialized bytes.
14    fn as_init(&self) -> &[u8];
15}
16
17/// A static assertion that [`IoBuf`] is dyn-compatible (object-safe).
18const _: [&dyn IoBuf; 0] = [];
19
20/// Extension trait for immutable buffers.
21pub trait IoBufExt: IoBuf {
22    /// Length of initialized bytes in the buffer.
23    fn buf_len(&self) -> usize {
24        self.as_init().len()
25    }
26
27    /// Raw pointer to the buffer.
28    fn buf_ptr(&self) -> *const u8 {
29        self.as_init().as_ptr()
30    }
31
32    /// Check if the buffer is empty.
33    fn is_empty(&self) -> bool {
34        self.buf_len() == 0
35    }
36
37    /// Returns a view of the buffer with the specified range.
38    ///
39    /// This method is similar to Rust's slicing (`&buf[..]`), but takes
40    /// ownership of the buffer.
41    ///
42    /// # Examples
43    ///
44    /// ```
45    /// use compio_buf::{IoBuf, IoBufExt};
46    ///
47    /// let buf = b"hello world";
48    /// assert_eq!(buf.slice(6..).as_init(), b"world");
49    /// ```
50    ///
51    /// # Panics
52    /// Panics if:
53    /// * begin > buf_len()
54    /// * end < begin
55    fn slice(self, range: impl std::ops::RangeBounds<usize>) -> Slice<Self>
56    where
57        Self: Sized,
58    {
59        use std::ops::Bound;
60
61        let begin = match range.start_bound() {
62            Bound::Included(&n) => n,
63            Bound::Excluded(&n) => n + 1,
64            Bound::Unbounded => 0,
65        };
66
67        let end = match range.end_bound() {
68            Bound::Included(&n) => Some(n.checked_add(1).expect("out of range")),
69            Bound::Excluded(&n) => Some(n),
70            Bound::Unbounded => None,
71        };
72
73        assert!(begin <= self.buf_len());
74
75        if let Some(end) = end {
76            assert!(begin <= end);
77        }
78
79        // SAFETY: begin <= self.buf_len()
80        unsafe { Slice::new(self, begin, end) }
81    }
82
83    /// Create a [`Reader`] from this buffer, which implements
84    /// [`std::io::Read`].
85    fn into_reader(self) -> Reader<Self>
86    where
87        Self: Sized,
88    {
89        Reader::new(self)
90    }
91
92    /// Create a [`ReaderRef`] from a reference of the buffer, which
93    /// implements [`std::io::Read`].
94    fn as_reader(&self) -> ReaderRef<'_, Self> {
95        ReaderRef::new(self)
96    }
97}
98
99impl<B: IoBuf + ?Sized> IoBufExt for B {}
100
101impl<B: IoBuf + ?Sized> IoBuf for &'static B {
102    fn as_init(&self) -> &[u8] {
103        (**self).as_init()
104    }
105}
106
107impl<B: IoBuf + ?Sized> IoBuf for &'static mut B {
108    fn as_init(&self) -> &[u8] {
109        (**self).as_init()
110    }
111}
112
113impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
114    for t_alloc!(Box, B, A)
115{
116    fn as_init(&self) -> &[u8] {
117        (**self).as_init()
118    }
119}
120
121impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
122    for t_alloc!(Rc, B, A)
123{
124    fn as_init(&self) -> &[u8] {
125        (**self).as_init()
126    }
127}
128
129impl IoBuf for [u8] {
130    fn as_init(&self) -> &[u8] {
131        self
132    }
133}
134
135impl<const N: usize> IoBuf for [u8; N] {
136    fn as_init(&self) -> &[u8] {
137        self
138    }
139}
140
141impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf for t_alloc!(Vec, u8, A) {
142    fn as_init(&self) -> &[u8] {
143        self
144    }
145}
146
147impl IoBuf for str {
148    fn as_init(&self) -> &[u8] {
149        self.as_bytes()
150    }
151}
152
153impl IoBuf for String {
154    fn as_init(&self) -> &[u8] {
155        self.as_bytes()
156    }
157}
158
159impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
160    for t_alloc!(Arc, B, A)
161{
162    fn as_init(&self) -> &[u8] {
163        (**self).as_init()
164    }
165}
166
167#[cfg(feature = "bytes")]
168impl IoBuf for bytes::Bytes {
169    fn as_init(&self) -> &[u8] {
170        self
171    }
172}
173
174#[cfg(feature = "bytes")]
175impl IoBuf for bytes::BytesMut {
176    fn as_init(&self) -> &[u8] {
177        self
178    }
179}
180
181#[cfg(feature = "read_buf")]
182impl IoBuf for std::io::BorrowedBuf<'static, u8> {
183    fn as_init(&self) -> &[u8] {
184        self.filled()
185    }
186}
187
188#[cfg(feature = "arrayvec")]
189impl<const N: usize> IoBuf for arrayvec::ArrayVec<u8, N> {
190    fn as_init(&self) -> &[u8] {
191        self
192    }
193}
194
195#[cfg(feature = "smallvec")]
196impl<const N: usize> IoBuf for smallvec::SmallVec<[u8; N]>
197where
198    [u8; N]: smallvec::Array<Item = u8>,
199{
200    fn as_init(&self) -> &[u8] {
201        self
202    }
203}
204
205#[cfg(feature = "memmap2")]
206impl IoBuf for memmap2::Mmap {
207    fn as_init(&self) -> &[u8] {
208        self
209    }
210}
211
212#[cfg(feature = "memmap2")]
213impl IoBuf for memmap2::MmapMut {
214    fn as_init(&self) -> &[u8] {
215        self
216    }
217}
218
219/// An error indicating that reserving capacity for a buffer failed.
220#[must_use]
221#[derive(Debug)]
222pub enum ReserveError {
223    /// Reservation is not supported.
224    NotSupported,
225
226    /// Reservation failed.
227    ///
228    /// This is usually caused by out-of-memory.
229    ReserveFailed(Box<dyn Error + Send + Sync>),
230}
231
232impl ReserveError {
233    /// Check if the error is [`NotSupported`].
234    ///
235    /// [`NotSupported`]: ReserveError::NotSupported
236    pub fn is_not_supported(&self) -> bool {
237        matches!(self, ReserveError::NotSupported)
238    }
239}
240
241impl Display for ReserveError {
242    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
243        match self {
244            ReserveError::NotSupported => write!(f, "reservation is not supported"),
245            ReserveError::ReserveFailed(src) => write!(f, "reservation failed: {src}"),
246        }
247    }
248}
249
250impl Error for ReserveError {
251    fn source(&self) -> Option<&(dyn Error + 'static)> {
252        match self {
253            ReserveError::ReserveFailed(src) => Some(src.as_ref()),
254            _ => None,
255        }
256    }
257}
258
259impl From<ReserveError> for std::io::Error {
260    fn from(value: ReserveError) -> Self {
261        match value {
262            ReserveError::NotSupported => {
263                std::io::Error::new(std::io::ErrorKind::Unsupported, "reservation not supported")
264            }
265            ReserveError::ReserveFailed(src) => {
266                std::io::Error::new(std::io::ErrorKind::OutOfMemory, src)
267            }
268        }
269    }
270}
271
272/// An error indicating that reserving exact capacity for a buffer failed.
273#[must_use]
274#[derive(Debug)]
275pub enum ReserveExactError {
276    /// Reservation is not supported.
277    NotSupported,
278
279    /// Reservation failed.
280    ///
281    /// This is usually caused by out-of-memory.
282    ReserveFailed(Box<dyn Error + Send + Sync>),
283
284    /// Reserved size does not match the expected size.
285    ExactSizeMismatch {
286        /// Expected size to reserve
287        expected: usize,
288
289        /// Actual size reserved
290        reserved: usize,
291    },
292}
293
294impl ReserveExactError {
295    /// Check if the error is [`NotSupported`]
296    ///
297    /// [`NotSupported`]: ReserveExactError::NotSupported
298    pub fn is_not_supported(&self) -> bool {
299        matches!(self, ReserveExactError::NotSupported)
300    }
301}
302
303impl Display for ReserveExactError {
304    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305        match self {
306            ReserveExactError::NotSupported => write!(f, "reservation is not supported"),
307            ReserveExactError::ReserveFailed(src) => write!(f, "reservation failed: {src}"),
308            ReserveExactError::ExactSizeMismatch { reserved, expected } => {
309                write!(
310                    f,
311                    "reserved size mismatch: expected {}, reserved {}",
312                    expected, reserved
313                )
314            }
315        }
316    }
317}
318
319impl From<ReserveError> for ReserveExactError {
320    fn from(err: ReserveError) -> Self {
321        match err {
322            ReserveError::NotSupported => ReserveExactError::NotSupported,
323            ReserveError::ReserveFailed(src) => ReserveExactError::ReserveFailed(src),
324        }
325    }
326}
327
328impl Error for ReserveExactError {
329    fn source(&self) -> Option<&(dyn Error + 'static)> {
330        match self {
331            ReserveExactError::ReserveFailed(src) => Some(src.as_ref()),
332            _ => None,
333        }
334    }
335}
336
337impl From<ReserveExactError> for std::io::Error {
338    fn from(value: ReserveExactError) -> Self {
339        match value {
340            ReserveExactError::NotSupported => {
341                std::io::Error::new(std::io::ErrorKind::Unsupported, "reservation not supported")
342            }
343            ReserveExactError::ReserveFailed(src) => {
344                std::io::Error::new(std::io::ErrorKind::OutOfMemory, src)
345            }
346            ReserveExactError::ExactSizeMismatch { expected, reserved } => std::io::Error::other(
347                format!("reserved size mismatch: expected {expected}, reserved {reserved}",),
348            ),
349        }
350    }
351}
352
353#[cfg(feature = "smallvec")]
354mod smallvec_err {
355    use std::{error::Error, fmt::Display};
356
357    use smallvec::CollectionAllocErr;
358
359    #[derive(Debug)]
360    pub(super) struct SmallVecErr(pub CollectionAllocErr);
361
362    impl Display for SmallVecErr {
363        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
364            write!(f, "SmallVec allocation error: {}", self.0)
365        }
366    }
367
368    impl Error for SmallVecErr {}
369}
370
371/// A trait for mutable buffers.
372///
373/// The `IoBufMut` trait is implemented by buffer types that can be passed to
374/// mutable completion-based IO operations, like reading content from a file and
375/// write to the buffer. This trait will take all space of a buffer into
376/// account, including uninitialized bytes.
377pub trait IoBufMut: IoBuf + SetLen {
378    /// Get the full mutable slice of the buffer, including both initialized
379    /// and uninitialized bytes.
380    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>];
381
382    /// Reserve additional capacity for the buffer.
383    ///
384    /// By default, this checks if the spare capacity is enough to fit in
385    /// `len`-bytes. If it does, returns `Ok(())`, and otherwise returns
386    /// [`Err(ReserveError::NotSupported)`]. Types that support dynamic
387    /// resizing (like `Vec<u8>`) will override this method to actually
388    /// reserve capacity. The return value indicates whether the reservation
389    /// succeeded. See [`ReserveError`] for details.
390    ///
391    /// Notice that this may move the memory of the buffer, so it's UB to
392    /// call this after the buffer is being pinned.
393    ///
394    /// [`Err(ReserveError::NotSupported)`]: ReserveError::NotSupported
395    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
396        let init = (*self).buf_len();
397        if len <= self.buf_capacity() - init {
398            return Ok(());
399        }
400        Err(ReserveError::NotSupported)
401    }
402
403    /// Reserve exactly `len` additional capacity for the buffer.
404    ///
405    /// By default this falls back to [`IoBufMut::reserve`]. Types that support
406    /// dynamic resizing (like `Vec<u8>`) will override this method to
407    /// actually reserve capacity. The return value indicates whether the
408    /// exact reservation succeeded. See [`ReserveExactError`] for details.
409    ///
410    /// Notice that this may move the memory of the buffer, so it's UB to
411    /// call this after the buffer is being pinned.
412    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
413        self.reserve(len)?;
414        Ok(())
415    }
416}
417
418/// A static assertion that [`IoBufMut`] is dyn-compatible (object-safe).
419const _: [&dyn IoBufMut; 0] = [];
420
421/// Extension trait for mutable buffers.
422pub trait IoBufMutExt: IoBufMut {
423    /// Initialize all bytes in the buffer and return them.
424    ///
425    /// Bytes in the already-initialized prefix (`0..buf_len()`) are preserved.
426    /// Only the uninitialized tail (`buf_len()..buf_capacity()`) is
427    /// zero-initialized.
428    fn ensure_init(&mut self) -> &mut [u8] {
429        let len = (*self).buf_len();
430        let slice = self.as_uninit();
431        slice[len..].fill(MaybeUninit::new(0));
432        unsafe { slice.assume_init_mut() }
433    }
434
435    /// Total capacity of the buffer, including both initialized and
436    /// uninitialized bytes.
437    fn buf_capacity(&mut self) -> usize {
438        self.as_uninit().len()
439    }
440
441    /// Get the raw mutable pointer to the buffer.
442    fn buf_mut_ptr(&mut self) -> *mut MaybeUninit<u8> {
443        self.as_uninit().as_mut_ptr()
444    }
445
446    /// Get the mutable slice of initialized bytes. The content is the same as
447    /// [`IoBuf::as_init`], but mutable.
448    fn as_mut_slice(&mut self) -> &mut [u8] {
449        let len = (*self).buf_len();
450        let ptr = (*self).buf_mut_ptr();
451        // SAFETY:
452        // - lifetime of the returned slice is bounded by &mut self
453        // - bytes within `len` are guaranteed to be initialized
454        // - the pointer is derived from
455        unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, len) }
456    }
457
458    /// Extend the buffer by copying bytes from `src`.
459    ///
460    /// The buffer will reserve additional capacity if necessary, and return an
461    /// error when reservation failed.
462    ///
463    /// Notice that this may move the memory of the buffer, so it's UB to
464    /// call this after the buffer is being pinned.
465    // FIXME: Change to `slice::write_copy_of_slice` when stabilized
466    fn extend_from_slice(&mut self, src: &[u8]) -> Result<(), ReserveError> {
467        let len = src.len();
468        let init = (*self).buf_len();
469        self.reserve(len)?;
470        let ptr = self.buf_mut_ptr().wrapping_add(init);
471
472        unsafe {
473            // SAFETY:
474            // - we have reserved enough capacity so the ptr and len stays in one allocation
475            // - src is valid for len bytes
476            // - ptr is valid for len bytes
477            // - &mut self guarantees that src cannot overlap with dst
478            std::ptr::copy_nonoverlapping(src.as_ptr() as _, ptr, len);
479
480            // SAFETY: the bytes in range [init, init + len) are initialized now
481            self.advance_to(init + len);
482        }
483
484        Ok(())
485    }
486
487    /// Like [`slice::copy_within`], copy a range of bytes within the buffer to
488    /// another location in the same buffer. This will count in both initialized
489    /// and uninitialized bytes.
490    ///
491    /// # Panics
492    ///
493    /// This method will panic if the source or destination range is out of
494    /// bounds.
495    ///
496    /// [`slice::copy_within`]: https://doc.rust-lang.org/std/primitive.slice.html#method.copy_within
497    fn copy_within<R>(&mut self, src: R, dest: usize)
498    where
499        R: RangeBounds<usize>,
500    {
501        self.as_uninit().copy_within(src, dest);
502    }
503
504    /// Returns an [`Uninit`], which is a [`Slice`] that only exposes
505    /// uninitialized bytes.
506    ///
507    /// It will always point to the uninitialized area of a [`IoBufMut`] even
508    /// after reading in some bytes, which is done by [`SetLen`]. This
509    /// is useful for writing data into buffer without overwriting any
510    /// existing bytes.
511    ///
512    /// # Examples
513    ///
514    /// ```
515    /// use compio_buf::{IoBuf, IoBufMut, IoBufMutExt};
516    ///
517    /// let mut buf = Vec::from(b"hello world");
518    /// buf.reserve_exact(10);
519    /// let mut slice = buf.uninit();
520    ///
521    /// assert_eq!(slice.as_init(), b"");
522    /// assert_eq!(slice.buf_capacity(), 10);
523    /// ```
524    fn uninit(self) -> Uninit<Self>
525    where
526        Self: Sized,
527    {
528        Uninit::new(self)
529    }
530
531    /// Create a [`Writer`] from this buffer, which implements
532    /// [`std::io::Write`].
533    fn into_writer(self) -> Writer<Self>
534    where
535        Self: Sized,
536    {
537        Writer::new(self)
538    }
539
540    /// Create a [`Writer`] from a mutable reference of the buffer, which
541    /// implements [`std::io::Write`].
542    fn as_writer(&mut self) -> WriterRef<'_, Self> {
543        WriterRef::new(self)
544    }
545
546    /// Indicate whether the buffer has been filled (uninit portion is empty)
547    fn is_filled(&mut self) -> bool {
548        let len = (*self).as_init().len();
549        let cap = (*self).buf_capacity();
550        len == cap
551    }
552}
553
554impl<B: IoBufMut + ?Sized> IoBufMutExt for B {}
555
556impl<B: IoBufMut + ?Sized> IoBufMut for &'static mut B {
557    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
558        (**self).as_uninit()
559    }
560
561    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
562        (**self).reserve(len)
563    }
564
565    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
566        (**self).reserve_exact(len)
567    }
568}
569
570impl<B: IoBufMut + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
571    for t_alloc!(Box, B, A)
572{
573    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
574        (**self).as_uninit()
575    }
576
577    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
578        (**self).reserve(len)
579    }
580
581    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
582        (**self).reserve_exact(len)
583    }
584}
585
586impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut for t_alloc!(Vec, u8, A) {
587    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
588        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
589        let cap = self.capacity();
590        // SAFETY: Vec guarantees that the pointer is valid for `capacity` bytes
591        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
592    }
593
594    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
595        if let Err(e) = Vec::try_reserve(self, len) {
596            return Err(ReserveError::ReserveFailed(Box::new(e)));
597        }
598
599        Ok(())
600    }
601
602    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
603        if self.capacity() - self.len() >= len {
604            return Ok(());
605        }
606
607        if let Err(e) = Vec::try_reserve_exact(self, len) {
608            return Err(ReserveExactError::ReserveFailed(Box::new(e)));
609        }
610
611        if self.capacity() - self.len() != len {
612            return Err(ReserveExactError::ExactSizeMismatch {
613                reserved: self.capacity() - self.len(),
614                expected: len,
615            });
616        }
617        Ok(())
618    }
619}
620
621impl IoBufMut for [u8] {
622    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
623        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
624        let len = self.len();
625        // SAFETY: slice is fully initialized, so treating it as MaybeUninit is safe
626        unsafe { std::slice::from_raw_parts_mut(ptr, len) }
627    }
628}
629
630impl<const N: usize> IoBufMut for [u8; N] {
631    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
632        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
633        // SAFETY: array is fully initialized, so treating it as MaybeUninit is safe
634        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
635    }
636}
637
638#[cfg(feature = "bytes")]
639impl IoBufMut for bytes::BytesMut {
640    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
641        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
642        let cap = self.capacity();
643        // SAFETY: BytesMut guarantees that the pointer is valid for `capacity` bytes
644        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
645    }
646
647    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
648        bytes::BytesMut::reserve(self, len);
649        Ok(())
650    }
651
652    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
653        if self.capacity() - self.len() >= len {
654            return Ok(());
655        }
656
657        bytes::BytesMut::reserve(self, len);
658
659        if self.capacity() - self.len() != len {
660            Err(ReserveExactError::ExactSizeMismatch {
661                reserved: self.capacity() - self.len(),
662                expected: len,
663            })
664        } else {
665            Ok(())
666        }
667    }
668}
669
670#[cfg(feature = "read_buf")]
671impl IoBufMut for std::io::BorrowedBuf<'static, u8> {
672    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
673        let total_cap = self.capacity();
674
675        // SAFETY: We reconstruct the full buffer from the filled portion pointer.
676        // BorrowedBuf guarantees that the underlying buffer has capacity bytes.
677        unsafe {
678            let filled_ptr = self.filled().as_ptr() as *mut MaybeUninit<u8>;
679            std::slice::from_raw_parts_mut(filled_ptr, total_cap)
680        }
681    }
682}
683
684#[cfg(feature = "arrayvec")]
685impl<const N: usize> IoBufMut for arrayvec::ArrayVec<u8, N> {
686    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
687        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
688        // SAFETY: ArrayVec guarantees that the pointer is valid for N bytes
689        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
690    }
691}
692
693#[cfg(feature = "smallvec")]
694impl<const N: usize> IoBufMut for smallvec::SmallVec<[u8; N]>
695where
696    [u8; N]: smallvec::Array<Item = u8>,
697{
698    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
699        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
700        let cap = self.capacity();
701        // SAFETY: SmallVec guarantees that the pointer is valid for `capacity` bytes
702        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
703    }
704
705    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
706        if let Err(e) = smallvec::SmallVec::try_reserve(self, len) {
707            return Err(ReserveError::ReserveFailed(Box::new(
708                smallvec_err::SmallVecErr(e),
709            )));
710        }
711        Ok(())
712    }
713
714    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
715        if self.capacity() - self.len() >= len {
716            return Ok(());
717        }
718
719        if let Err(e) = smallvec::SmallVec::try_reserve_exact(self, len) {
720            return Err(ReserveExactError::ReserveFailed(Box::new(
721                smallvec_err::SmallVecErr(e),
722            )));
723        }
724
725        if self.capacity() - self.len() != len {
726            return Err(ReserveExactError::ExactSizeMismatch {
727                reserved: self.capacity() - self.len(),
728                expected: len,
729            });
730        }
731        Ok(())
732    }
733}
734
735#[cfg(feature = "memmap2")]
736impl IoBufMut for memmap2::MmapMut {
737    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
738        // Safety: &mut [u8] is valid &mut [MaybeUninit<u8>]
739        unsafe { std::mem::transmute(self.as_mut_slice()) }
740    }
741}
742
743/// A helper trait for `set_len` like methods.
744pub trait SetLen {
745    /// Set the buffer length.
746    ///
747    /// # Safety
748    ///
749    /// * `len` must be less or equal than `as_uninit().len()`.
750    /// * The bytes in the range `[buf_len(), len)` must be initialized.
751    unsafe fn set_len(&mut self, len: usize);
752}
753
754/// A static assertion that [`SetLen`] is dyn-compatible (object-safe).
755const _: [&dyn SetLen; 0] = [];
756
757/// Extension trait for `set_len` like methods.
758pub trait SetLenExt: SetLen {
759    /// Advance the buffer length by `len`.
760    ///
761    /// # Safety
762    ///
763    /// * The bytes in the range `[buf_len(), buf_len() + len)` must be
764    ///   initialized.
765    unsafe fn advance(&mut self, len: usize)
766    where
767        Self: IoBuf,
768    {
769        let current_len = (*self).buf_len();
770        let new_len = current_len.checked_add(len).expect("length overflow");
771        unsafe { self.set_len(new_len) };
772    }
773
774    /// Set the buffer length to `len`. If `len` is less than the current
775    /// length, this operation is a no-op.
776    ///
777    /// # Safety
778    ///
779    /// * `len` must be less or equal than `as_uninit().len()`.
780    /// * The bytes in the range `[buf_len(), len)` must be initialized.
781    unsafe fn advance_to(&mut self, len: usize)
782    where
783        Self: IoBuf,
784    {
785        let current_len = (*self).buf_len();
786        if len > current_len {
787            unsafe { self.set_len(len) };
788        }
789    }
790
791    /// Set the vector buffer's total length to `len`. If `len` is less than the
792    /// current total length, this operation is a no-op.
793    ///
794    /// # Safety
795    ///
796    /// * `len` must be less or equal than `total_len()`.
797    /// * The bytes in the range `[total_len(), len)` must be initialized.
798    unsafe fn advance_vec_to(&mut self, len: usize)
799    where
800        Self: IoVectoredBuf,
801    {
802        let current_len = (*self).total_len();
803        if len > current_len {
804            unsafe { self.set_len(len) };
805        }
806    }
807
808    /// Clear the buffer, setting its length to 0 without touching its content
809    /// or capacity.
810    fn clear(&mut self)
811    where
812        Self: IoBuf,
813    {
814        // SAFETY: setting length to 0 is always valid
815        unsafe { self.set_len(0) };
816    }
817}
818
819impl<B: SetLen + ?Sized> SetLenExt for B {}
820
821impl<B: SetLen + ?Sized> SetLen for &'static mut B {
822    unsafe fn set_len(&mut self, len: usize) {
823        unsafe { (**self).set_len(len) }
824    }
825}
826
827impl<B: SetLen + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen
828    for t_alloc!(Box, B, A)
829{
830    unsafe fn set_len(&mut self, len: usize) {
831        unsafe { (**self).set_len(len) }
832    }
833}
834
835impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen for t_alloc!(Vec, u8, A) {
836    unsafe fn set_len(&mut self, len: usize) {
837        unsafe { self.set_len(len) };
838    }
839}
840
841impl SetLen for [u8] {
842    unsafe fn set_len(&mut self, len: usize) {
843        debug_assert!(len <= self.len());
844    }
845}
846
847impl<const N: usize> SetLen for [u8; N] {
848    unsafe fn set_len(&mut self, len: usize) {
849        debug_assert!(len <= N);
850    }
851}
852
853#[cfg(feature = "bytes")]
854impl SetLen for bytes::BytesMut {
855    unsafe fn set_len(&mut self, len: usize) {
856        unsafe { self.set_len(len) };
857    }
858}
859
860#[cfg(feature = "read_buf")]
861impl SetLen for std::io::BorrowedBuf<'static, u8> {
862    unsafe fn set_len(&mut self, len: usize) {
863        debug_assert!(self.capacity() >= len);
864
865        // SAFETY: `len` range is initialized guaranteed by invariant of `set_len`
866        #[allow(unused_unsafe)]
867        unsafe {
868            self.clear().unfilled().advance(len)
869        };
870    }
871}
872
873#[cfg(feature = "arrayvec")]
874impl<const N: usize> SetLen for arrayvec::ArrayVec<u8, N> {
875    unsafe fn set_len(&mut self, len: usize) {
876        if (**self).buf_len() < len {
877            unsafe { self.set_len(len) };
878        }
879    }
880}
881
882#[cfg(feature = "smallvec")]
883impl<const N: usize> SetLen for smallvec::SmallVec<[u8; N]>
884where
885    [u8; N]: smallvec::Array<Item = u8>,
886{
887    unsafe fn set_len(&mut self, len: usize) {
888        if (**self).buf_len() < len {
889            unsafe { self.set_len(len) };
890        }
891    }
892}
893
894#[cfg(feature = "memmap2")]
895impl SetLen for memmap2::MmapMut {
896    unsafe fn set_len(&mut self, len: usize) {
897        debug_assert!(len <= self.len())
898    }
899}
900
901impl<T: IoBufMut> SetLen for [T] {
902    unsafe fn set_len(&mut self, len: usize) {
903        unsafe { default_set_len(self.iter_mut(), len) }
904    }
905}
906
907impl<T: IoBufMut, const N: usize> SetLen for [T; N] {
908    unsafe fn set_len(&mut self, len: usize) {
909        unsafe { default_set_len(self.iter_mut(), len) }
910    }
911}
912
913impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen
914    for t_alloc!(Vec, T, A)
915{
916    unsafe fn set_len(&mut self, len: usize) {
917        unsafe { default_set_len(self.iter_mut(), len) }
918    }
919}
920
921#[cfg(feature = "arrayvec")]
922impl<T: IoBufMut, const N: usize> SetLen for arrayvec::ArrayVec<T, N> {
923    unsafe fn set_len(&mut self, len: usize) {
924        unsafe { default_set_len(self.iter_mut(), len) }
925    }
926}
927
928#[cfg(feature = "smallvec")]
929impl<T: IoBufMut, const N: usize> SetLen for smallvec::SmallVec<[T; N]>
930where
931    [T; N]: smallvec::Array<Item = T>,
932{
933    unsafe fn set_len(&mut self, len: usize) {
934        unsafe { default_set_len(self.iter_mut(), len) }
935    }
936}
937
938/// # Safety
939/// * `len` should be less or equal than the sum of `buf_capacity()` of all
940///   buffers.
941/// * The bytes in the range `[buf_len(), new_len)` of each buffer must be
942///   initialized
943unsafe fn default_set_len<'a, B: IoBufMut>(
944    iter: impl IntoIterator<Item = &'a mut B>,
945    mut len: usize,
946) {
947    let mut iter = iter.into_iter();
948    while len > 0 {
949        let Some(curr) = iter.next() else { return };
950        let sub = (*curr).buf_capacity().min(len);
951        unsafe { curr.set_len(sub) };
952        len -= sub;
953    }
954}
955
956#[cfg(test)]
957mod test {
958    use crate::{IoBufMut, IoBufMutExt};
959
960    #[test]
961    fn test_vec_reserve() {
962        let mut buf = Vec::new();
963        IoBufMut::reserve(&mut buf, 10).unwrap();
964        assert!(buf.capacity() >= 10);
965
966        let mut buf = Vec::new();
967        IoBufMut::reserve_exact(&mut buf, 10).unwrap();
968        assert!(buf.capacity() == 10);
969
970        let mut buf = Box::new(Vec::new());
971        IoBufMut::reserve_exact(&mut buf, 10).unwrap();
972        assert!(buf.capacity() == 10);
973    }
974
975    #[test]
976    #[cfg(feature = "bytes")]
977    fn test_bytes_reserve() {
978        let mut buf = bytes::BytesMut::new();
979        IoBufMut::reserve(&mut buf, 10).unwrap();
980        assert!(buf.capacity() >= 10);
981    }
982
983    #[test]
984    #[cfg(feature = "smallvec")]
985    fn test_smallvec_reserve() {
986        let mut buf = smallvec::SmallVec::<[u8; 8]>::new();
987        IoBufMut::reserve(&mut buf, 10).unwrap();
988        assert!(buf.capacity() >= 10);
989    }
990
991    #[test]
992    #[cfg(feature = "memmap2")]
993    fn tests_memmap2() {
994        use std::{
995            fs::{OpenOptions, remove_file},
996            io::{Seek, SeekFrom, Write},
997        };
998
999        use memmap2::MmapOptions;
1000
1001        use super::*;
1002
1003        let path = std::env::temp_dir().join("compio_buf_mmap_mut_test");
1004
1005        let mut file = OpenOptions::new()
1006            .read(true)
1007            .write(true)
1008            .create(true)
1009            .truncate(true)
1010            .open(&path)
1011            .unwrap();
1012        let data = b"hello memmap2";
1013        file.write_all(data).unwrap();
1014        file.flush().unwrap();
1015        file.seek(SeekFrom::Start(0)).unwrap();
1016        let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
1017
1018        let uninit_slice = mmap.as_init();
1019        assert_eq!(uninit_slice, data);
1020
1021        remove_file(path).unwrap();
1022    }
1023
1024    #[test]
1025    fn test_other_reserve() {
1026        let mut buf = [1, 1, 4, 5, 1, 4];
1027        let res = IoBufMut::reserve(&mut buf, 10);
1028        assert!(res.is_err_and(|x| x.is_not_supported()));
1029        assert!(buf.buf_capacity() == 6);
1030    }
1031
1032    #[test]
1033    fn test_extend() {
1034        let mut buf = Vec::from(b"hello");
1035        IoBufMutExt::extend_from_slice(&mut buf, b" world").unwrap();
1036        assert_eq!(buf.as_slice(), b"hello world");
1037
1038        let mut buf = [];
1039        let res = IoBufMutExt::extend_from_slice(&mut buf, b" ");
1040        assert!(res.is_err_and(|x| x.is_not_supported()));
1041    }
1042}
compio_buf/io_buf.rs

compio_buf/
io_buf.rs
