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
475            //   one allocation
476            // - src is valid for len bytes
477            // - ptr is valid for len bytes
478            // - &mut self guarantees that src cannot overlap with dst
479            std::ptr::copy_nonoverlapping(src.as_ptr() as _, ptr, len);
480
481            // SAFETY: the bytes in range [init, init + len) are initialized now
482            self.advance_to(init + len);
483        }
484
485        Ok(())
486    }
487
488    /// Like [`slice::copy_within`], copy a range of bytes within the buffer to
489    /// another location in the same buffer. This will count in both initialized
490    /// and uninitialized bytes.
491    ///
492    /// # Panics
493    ///
494    /// This method will panic if the source or destination range is out of
495    /// bounds.
496    ///
497    /// [`slice::copy_within`]: https://doc.rust-lang.org/std/primitive.slice.html#method.copy_within
498    fn copy_within<R>(&mut self, src: R, dest: usize)
499    where
500        R: RangeBounds<usize>,
501    {
502        self.as_uninit().copy_within(src, dest);
503    }
504
505    /// Returns an [`Uninit`], which is a [`Slice`] that only exposes
506    /// uninitialized bytes.
507    ///
508    /// It will always point to the uninitialized area of a [`IoBufMut`] even
509    /// after reading in some bytes, which is done by [`SetLen`]. This
510    /// is useful for writing data into buffer without overwriting any
511    /// existing bytes.
512    ///
513    /// # Examples
514    ///
515    /// ```
516    /// use compio_buf::{IoBuf, IoBufMut, IoBufMutExt};
517    ///
518    /// let mut buf = Vec::from(b"hello world");
519    /// buf.reserve_exact(10);
520    /// let mut slice = buf.uninit();
521    ///
522    /// assert_eq!(slice.as_init(), b"");
523    /// assert_eq!(slice.buf_capacity(), 10);
524    /// ```
525    fn uninit(self) -> Uninit<Self>
526    where
527        Self: Sized,
528    {
529        Uninit::new(self)
530    }
531
532    /// Create a [`Writer`] from this buffer, which implements
533    /// [`std::io::Write`].
534    fn into_writer(self) -> Writer<Self>
535    where
536        Self: Sized,
537    {
538        Writer::new(self)
539    }
540
541    /// Create a [`Writer`] from a mutable reference of the buffer, which
542    /// implements [`std::io::Write`].
543    fn as_writer(&mut self) -> WriterRef<'_, Self> {
544        WriterRef::new(self)
545    }
546
547    /// Indicate whether the buffer has been filled (uninit portion is empty)
548    fn is_filled(&mut self) -> bool {
549        let len = (*self).as_init().len();
550        let cap = (*self).buf_capacity();
551        len == cap
552    }
553}
554
555impl<B: IoBufMut + ?Sized> IoBufMutExt for B {}
556
557impl<B: IoBufMut + ?Sized> IoBufMut for &'static mut B {
558    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
559        (**self).as_uninit()
560    }
561
562    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
563        (**self).reserve(len)
564    }
565
566    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
567        (**self).reserve_exact(len)
568    }
569}
570
571impl<B: IoBufMut + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
572    for t_alloc!(Box, B, A)
573{
574    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
575        (**self).as_uninit()
576    }
577
578    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
579        (**self).reserve(len)
580    }
581
582    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
583        (**self).reserve_exact(len)
584    }
585}
586
587impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut for t_alloc!(Vec, u8, A) {
588    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
589        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
590        let cap = self.capacity();
591        // SAFETY: Vec guarantees that the pointer is valid for `capacity` bytes
592        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
593    }
594
595    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
596        if let Err(e) = Vec::try_reserve(self, len) {
597            return Err(ReserveError::ReserveFailed(Box::new(e)));
598        }
599
600        Ok(())
601    }
602
603    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
604        if self.capacity() - self.len() >= len {
605            return Ok(());
606        }
607
608        if let Err(e) = Vec::try_reserve_exact(self, len) {
609            return Err(ReserveExactError::ReserveFailed(Box::new(e)));
610        }
611
612        if self.capacity() - self.len() != len {
613            return Err(ReserveExactError::ExactSizeMismatch {
614                reserved: self.capacity() - self.len(),
615                expected: len,
616            });
617        }
618        Ok(())
619    }
620}
621
622impl IoBufMut for [u8] {
623    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
624        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
625        let len = self.len();
626        // SAFETY: slice is fully initialized, so treating it as MaybeUninit is
627        // safe
628        unsafe { std::slice::from_raw_parts_mut(ptr, len) }
629    }
630}
631
632impl<const N: usize> IoBufMut for [u8; N] {
633    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
634        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
635        // SAFETY: array is fully initialized, so treating it as MaybeUninit is
636        // safe
637        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
638    }
639}
640
641#[cfg(feature = "bytes")]
642impl IoBufMut for bytes::BytesMut {
643    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
644        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
645        let cap = self.capacity();
646        // SAFETY: BytesMut guarantees that the pointer is valid for `capacity`
647        // bytes
648        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
649    }
650
651    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
652        bytes::BytesMut::reserve(self, len);
653        Ok(())
654    }
655
656    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
657        if self.capacity() - self.len() >= len {
658            return Ok(());
659        }
660
661        bytes::BytesMut::reserve(self, len);
662
663        if self.capacity() - self.len() != len {
664            Err(ReserveExactError::ExactSizeMismatch {
665                reserved: self.capacity() - self.len(),
666                expected: len,
667            })
668        } else {
669            Ok(())
670        }
671    }
672}
673
674#[cfg(feature = "read_buf")]
675impl IoBufMut for std::io::BorrowedBuf<'static, u8> {
676    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
677        let total_cap = self.capacity();
678
679        // SAFETY: We reconstruct the full buffer from the filled portion
680        // pointer. BorrowedBuf guarantees that the underlying buffer
681        // has capacity bytes.
682        unsafe {
683            let filled_ptr = self.filled().as_ptr() as *mut MaybeUninit<u8>;
684            std::slice::from_raw_parts_mut(filled_ptr, total_cap)
685        }
686    }
687}
688
689#[cfg(feature = "arrayvec")]
690impl<const N: usize> IoBufMut for arrayvec::ArrayVec<u8, N> {
691    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
692        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
693        // SAFETY: ArrayVec guarantees that the pointer is valid for N bytes
694        unsafe { std::slice::from_raw_parts_mut(ptr, N) }
695    }
696}
697
698#[cfg(feature = "smallvec")]
699impl<const N: usize> IoBufMut for smallvec::SmallVec<[u8; N]>
700where
701    [u8; N]: smallvec::Array<Item = u8>,
702{
703    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
704        let ptr = self.as_mut_ptr() as *mut MaybeUninit<u8>;
705        let cap = self.capacity();
706        // SAFETY: SmallVec guarantees that the pointer is valid for `capacity`
707        // bytes
708        unsafe { std::slice::from_raw_parts_mut(ptr, cap) }
709    }
710
711    fn reserve(&mut self, len: usize) -> Result<(), ReserveError> {
712        if let Err(e) = smallvec::SmallVec::try_reserve(self, len) {
713            return Err(ReserveError::ReserveFailed(Box::new(
714                smallvec_err::SmallVecErr(e),
715            )));
716        }
717        Ok(())
718    }
719
720    fn reserve_exact(&mut self, len: usize) -> Result<(), ReserveExactError> {
721        if self.capacity() - self.len() >= len {
722            return Ok(());
723        }
724
725        if let Err(e) = smallvec::SmallVec::try_reserve_exact(self, len) {
726            return Err(ReserveExactError::ReserveFailed(Box::new(
727                smallvec_err::SmallVecErr(e),
728            )));
729        }
730
731        if self.capacity() - self.len() != len {
732            return Err(ReserveExactError::ExactSizeMismatch {
733                reserved: self.capacity() - self.len(),
734                expected: len,
735            });
736        }
737        Ok(())
738    }
739}
740
741#[cfg(feature = "memmap2")]
742impl IoBufMut for memmap2::MmapMut {
743    fn as_uninit(&mut self) -> &mut [MaybeUninit<u8>] {
744        // Safety: &mut [u8] is valid &mut [MaybeUninit<u8>]
745        unsafe { std::mem::transmute(self.as_mut()) }
746    }
747}
748
749/// A helper trait for `set_len` like methods.
750pub trait SetLen {
751    /// Set the buffer length.
752    ///
753    /// # Safety
754    ///
755    /// * `len` must be less or equal than `as_uninit().len()`.
756    /// * The bytes in the range `[buf_len(), len)` must be initialized.
757    unsafe fn set_len(&mut self, len: usize);
758}
759
760/// A static assertion that [`SetLen`] is dyn-compatible (object-safe).
761const _: [&dyn SetLen; 0] = [];
762
763/// Extension trait for `set_len` like methods.
764pub trait SetLenExt: SetLen {
765    /// Advance the buffer length by `len`.
766    ///
767    /// # Safety
768    ///
769    /// * The bytes in the range `[buf_len(), buf_len() + len)` must be
770    ///   initialized.
771    unsafe fn advance(&mut self, len: usize)
772    where
773        Self: IoBuf,
774    {
775        let current_len = (*self).buf_len();
776        let new_len = current_len.checked_add(len).expect("length overflow");
777        unsafe { self.set_len(new_len) };
778    }
779
780    /// Set the buffer length to `len`. If `len` is less than the current
781    /// length, this operation is a no-op.
782    ///
783    /// # Safety
784    ///
785    /// * `len` must be less or equal than `as_uninit().len()`.
786    /// * The bytes in the range `[buf_len(), len)` must be initialized.
787    unsafe fn advance_to(&mut self, len: usize)
788    where
789        Self: IoBuf,
790    {
791        let current_len = (*self).buf_len();
792        if len > current_len {
793            unsafe { self.set_len(len) };
794        }
795    }
796
797    /// Set the vector buffer's total length to `len`. If `len` is less than the
798    /// current total length, this operation is a no-op.
799    ///
800    /// # Safety
801    ///
802    /// * `len` must be less or equal than `total_len()`.
803    /// * The bytes in the range `[total_len(), len)` must be initialized.
804    unsafe fn advance_vec_to(&mut self, len: usize)
805    where
806        Self: IoVectoredBuf,
807    {
808        let current_len = (*self).total_len();
809        if len > current_len {
810            unsafe { self.set_len(len) };
811        }
812    }
813
814    /// Clear the buffer, setting its length to 0 without touching its content
815    /// or capacity.
816    fn clear(&mut self)
817    where
818        Self: IoBuf,
819    {
820        // SAFETY: setting length to 0 is always valid
821        unsafe { self.set_len(0) };
822    }
823}
824
825impl<B: SetLen + ?Sized> SetLenExt for B {}
826
827impl<B: SetLen + ?Sized> SetLen for &'static mut B {
828    unsafe fn set_len(&mut self, len: usize) {
829        unsafe { (**self).set_len(len) }
830    }
831}
832
833impl<B: SetLen + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen
834    for t_alloc!(Box, B, A)
835{
836    unsafe fn set_len(&mut self, len: usize) {
837        unsafe { (**self).set_len(len) }
838    }
839}
840
841impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen for t_alloc!(Vec, u8, A) {
842    unsafe fn set_len(&mut self, len: usize) {
843        unsafe { self.set_len(len) };
844    }
845}
846
847impl SetLen for [u8] {
848    unsafe fn set_len(&mut self, len: usize) {
849        debug_assert!(len <= self.len());
850    }
851}
852
853impl<const N: usize> SetLen for [u8; N] {
854    unsafe fn set_len(&mut self, len: usize) {
855        debug_assert!(len <= N);
856    }
857}
858
859#[cfg(feature = "bytes")]
860impl SetLen for bytes::BytesMut {
861    unsafe fn set_len(&mut self, len: usize) {
862        unsafe { self.set_len(len) };
863    }
864}
865
866#[cfg(feature = "read_buf")]
867impl SetLen for std::io::BorrowedBuf<'static, u8> {
868    unsafe fn set_len(&mut self, len: usize) {
869        debug_assert!(self.capacity() >= len);
870
871        // SAFETY: `len` range is initialized guaranteed by invariant of
872        // `set_len`
873        #[allow(unused_unsafe)]
874        unsafe {
875            self.clear().unfilled().advance(len)
876        };
877    }
878}
879
880#[cfg(feature = "arrayvec")]
881impl<const N: usize> SetLen for arrayvec::ArrayVec<u8, N> {
882    unsafe fn set_len(&mut self, len: usize) {
883        if (**self).buf_len() < len {
884            unsafe { self.set_len(len) };
885        }
886    }
887}
888
889#[cfg(feature = "smallvec")]
890impl<const N: usize> SetLen for smallvec::SmallVec<[u8; N]>
891where
892    [u8; N]: smallvec::Array<Item = u8>,
893{
894    unsafe fn set_len(&mut self, len: usize) {
895        if (**self).buf_len() < len {
896            unsafe { self.set_len(len) };
897        }
898    }
899}
900
901#[cfg(feature = "memmap2")]
902impl SetLen for memmap2::MmapMut {
903    unsafe fn set_len(&mut self, len: usize) {
904        debug_assert!(len <= self.len())
905    }
906}
907
908impl<T: IoBufMut> SetLen for [T] {
909    unsafe fn set_len(&mut self, len: usize) {
910        unsafe { default_set_len(self.iter_mut(), len) }
911    }
912}
913
914impl<T: IoBufMut, const N: usize> SetLen for [T; N] {
915    unsafe fn set_len(&mut self, len: usize) {
916        unsafe { default_set_len(self.iter_mut(), len) }
917    }
918}
919
920impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetLen
921    for t_alloc!(Vec, T, A)
922{
923    unsafe fn set_len(&mut self, len: usize) {
924        unsafe { default_set_len(self.iter_mut(), len) }
925    }
926}
927
928#[cfg(feature = "arrayvec")]
929impl<T: IoBufMut, const N: usize> SetLen for arrayvec::ArrayVec<T, N> {
930    unsafe fn set_len(&mut self, len: usize) {
931        unsafe { default_set_len(self.iter_mut(), len) }
932    }
933}
934
935#[cfg(feature = "smallvec")]
936impl<T: IoBufMut, const N: usize> SetLen for smallvec::SmallVec<[T; N]>
937where
938    [T; N]: smallvec::Array<Item = T>,
939{
940    unsafe fn set_len(&mut self, len: usize) {
941        unsafe { default_set_len(self.iter_mut(), len) }
942    }
943}
944
945/// # Safety
946/// * `len` should be less or equal than the sum of `buf_capacity()` of all
947///   buffers.
948/// * The bytes in the range `[buf_len(), new_len)` of each buffer must be
949///   initialized
950unsafe fn default_set_len<'a, B: IoBufMut>(
951    iter: impl IntoIterator<Item = &'a mut B>,
952    mut len: usize,
953) {
954    let mut iter = iter.into_iter();
955    while len > 0 {
956        let Some(curr) = iter.next() else { return };
957        let sub = (*curr).buf_capacity().min(len);
958        unsafe { curr.set_len(sub) };
959        len -= sub;
960    }
961}
962
963#[cfg(test)]
964mod test {
965    use crate::{IoBufMut, IoBufMutExt};
966
967    #[test]
968    fn test_vec_reserve() {
969        let mut buf = Vec::new();
970        IoBufMut::reserve(&mut buf, 10).unwrap();
971        assert!(buf.capacity() >= 10);
972
973        let mut buf = Vec::new();
974        IoBufMut::reserve_exact(&mut buf, 10).unwrap();
975        assert!(buf.capacity() == 10);
976
977        let mut buf = Box::new(Vec::new());
978        IoBufMut::reserve_exact(&mut buf, 10).unwrap();
979        assert!(buf.capacity() == 10);
980    }
981
982    #[test]
983    #[cfg(feature = "bytes")]
984    fn test_bytes_reserve() {
985        let mut buf = bytes::BytesMut::new();
986        IoBufMut::reserve(&mut buf, 10).unwrap();
987        assert!(buf.capacity() >= 10);
988    }
989
990    #[test]
991    #[cfg(feature = "smallvec")]
992    fn test_smallvec_reserve() {
993        let mut buf = smallvec::SmallVec::<[u8; 8]>::new();
994        IoBufMut::reserve(&mut buf, 10).unwrap();
995        assert!(buf.capacity() >= 10);
996    }
997
998    #[test]
999    #[cfg(feature = "memmap2")]
1000    fn tests_memmap2() {
1001        use std::{
1002            fs::{OpenOptions, remove_file},
1003            io::{Seek, SeekFrom, Write},
1004        };
1005
1006        use memmap2::MmapOptions;
1007
1008        use super::*;
1009
1010        let path = std::env::temp_dir().join("compio_buf_mmap_mut_test");
1011
1012        let mut file = OpenOptions::new()
1013            .read(true)
1014            .write(true)
1015            .create(true)
1016            .truncate(true)
1017            .open(&path)
1018            .unwrap();
1019        let data = b"hello memmap2";
1020        file.write_all(data).unwrap();
1021        file.flush().unwrap();
1022        file.seek(SeekFrom::Start(0)).unwrap();
1023        {
1024            let mmap = unsafe { MmapOptions::new().map(&file).unwrap() };
1025
1026            let init_slice = mmap.as_init();
1027            assert_eq!(init_slice, data);
1028        }
1029
1030        {
1031            let mut mmap = unsafe { MmapOptions::new().map_mut(&file).unwrap() };
1032
1033            let init_slice = mmap.as_mut_slice();
1034            assert_eq!(init_slice, data);
1035        }
1036
1037        remove_file(path).unwrap();
1038    }
1039
1040    #[test]
1041    fn test_other_reserve() {
1042        let mut buf = [1, 1, 4, 5, 1, 4];
1043        let res = IoBufMut::reserve(&mut buf, 10);
1044        assert!(res.is_err_and(|x| x.is_not_supported()));
1045        assert!(buf.buf_capacity() == 6);
1046    }
1047
1048    #[test]
1049    fn test_extend() {
1050        let mut buf = Vec::from(b"hello");
1051        IoBufMutExt::extend_from_slice(&mut buf, b" world").unwrap();
1052        assert_eq!(buf.as_slice(), b"hello world");
1053
1054        let mut buf = [];
1055        let res = IoBufMutExt::extend_from_slice(&mut buf, b" ");
1056        assert!(res.is_err_and(|x| x.is_not_supported()));
1057    }
1058}
compio_buf/io_buf.rs

compio_buf/
io_buf.rs
