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

compio_buf/
io_buf.rs
