compio_buf/
io_buf.rs

1#[cfg(feature = "allocator_api")]
2use std::alloc::Allocator;
3use std::{mem::MaybeUninit, rc::Rc, sync::Arc};
4
5use crate::*;
6
7/// A trait for buffers.
8///
9/// The `IoBuf` trait is implemented by buffer types that can be passed to
10/// compio operations. Users will not need to use this trait directly.
11///
12/// # Safety
13///
14/// The implementer should ensure the pointer, len and capacity are valid, so
15/// that the returned slice of [`IoBuf::as_slice`] is valid.
16pub unsafe trait IoBuf: 'static {
17    /// Returns a raw pointer to the vector’s buffer.
18    ///
19    /// This method is to be used by the `compio` runtime and it is not
20    /// expected for users to call it directly.
21    fn as_buf_ptr(&self) -> *const u8;
22
23    /// Number of initialized bytes.
24    ///
25    /// This method is to be used by the `compio` runtime and it is not
26    /// expected for users to call it directly.
27    ///
28    /// For [`Vec`], this is identical to `len()`.
29    fn buf_len(&self) -> usize;
30
31    /// Total size of the buffer, including uninitialized memory, if any.
32    ///
33    /// This method is to be used by the `compio` runtime and it is not
34    /// expected for users to call it directly.
35    ///
36    /// For [`Vec`], this is identical to `capacity()`.
37    fn buf_capacity(&self) -> usize;
38
39    /// Get the initialized part of the buffer.
40    fn as_slice(&self) -> &[u8] {
41        unsafe { std::slice::from_raw_parts(self.as_buf_ptr(), self.buf_len()) }
42    }
43
44    /// Create an [`IoSlice`] of this buffer.
45    ///
46    /// # Safety
47    ///
48    /// The return slice will not live longer than `Self`.
49    /// It is static to provide convenience from writing self-referenced
50    /// structure.
51    unsafe fn as_io_slice(&self) -> IoSlice {
52        unsafe { IoSlice::from_slice(self.as_slice()) }
53    }
54
55    /// Create an [`IoBuffer`] of this buffer.
56    ///
57    /// # Safety
58    ///
59    /// The returned [`IoBuffer`] will not live longer than `Self`.
60    /// It is static to provide convenience from writing self-referenced
61    /// structure.
62    unsafe fn as_io_buffer(&self) -> IoBuffer {
63        unsafe {
64            IoBuffer::new(
65                self.as_buf_ptr().cast_mut().cast(),
66                self.buf_len(),
67                self.buf_capacity(),
68            )
69        }
70    }
71
72    /// Returns a view of the buffer with the specified range.
73    ///
74    /// This method is similar to Rust's slicing (`&buf[..]`), but takes
75    /// ownership of the buffer.
76    ///
77    /// # Examples
78    ///
79    /// ```
80    /// use compio_buf::IoBuf;
81    ///
82    /// let buf = b"hello world";
83    /// assert_eq!(buf.slice(6..).as_slice(), b"world");
84    /// ```
85    fn slice(self, range: impl std::ops::RangeBounds<usize>) -> Slice<Self>
86    where
87        Self: Sized,
88    {
89        use std::ops::Bound;
90
91        let begin = match range.start_bound() {
92            Bound::Included(&n) => n,
93            Bound::Excluded(&n) => n + 1,
94            Bound::Unbounded => 0,
95        };
96
97        assert!(begin <= self.buf_capacity());
98
99        let end = match range.end_bound() {
100            Bound::Included(&n) => n.checked_add(1).expect("out of range"),
101            Bound::Excluded(&n) => n,
102            Bound::Unbounded => self.buf_capacity(),
103        };
104
105        assert!(end <= self.buf_capacity());
106        assert!(begin <= self.buf_len());
107
108        Slice::new(self, begin, end)
109    }
110
111    /// Returns an [`Uninit`], which is a [`Slice`] that only exposes
112    /// uninitialized bytes.
113    ///
114    /// It will always point to uninitialized area of a [`IoBuf`] even after
115    /// reading in some bytes, which is done by [`SetBufInit`]. This is useful
116    /// for writing data into buffer without overwriting any existing bytes.
117    ///
118    /// # Examples
119    ///
120    /// Creating an uninit slice
121    ///
122    /// ```
123    /// use compio_buf::IoBuf;
124    ///
125    /// let mut buf = Vec::from(b"hello world");
126    /// buf.reserve_exact(10);
127    /// let slice = buf.uninit();
128    ///
129    /// assert_eq!(slice.as_slice(), b"");
130    /// assert_eq!(slice.buf_capacity(), 10);
131    /// ```
132    fn uninit(self) -> Uninit<Self>
133    where
134        Self: Sized,
135    {
136        Uninit::new(self)
137    }
138
139    /// Indicate whether the buffer has been filled (uninit portion is empty)
140    fn filled(&self) -> bool {
141        self.buf_len() == self.buf_capacity()
142    }
143}
144
145unsafe impl<B: IoBuf + ?Sized> IoBuf for &'static B {
146    fn as_buf_ptr(&self) -> *const u8 {
147        (**self).as_buf_ptr()
148    }
149
150    fn buf_len(&self) -> usize {
151        (**self).buf_len()
152    }
153
154    fn buf_capacity(&self) -> usize {
155        (**self).buf_capacity()
156    }
157}
158
159unsafe impl<B: IoBuf + ?Sized> IoBuf for &'static mut B {
160    fn as_buf_ptr(&self) -> *const u8 {
161        (**self).as_buf_ptr()
162    }
163
164    fn buf_len(&self) -> usize {
165        (**self).buf_len()
166    }
167
168    fn buf_capacity(&self) -> usize {
169        (**self).buf_capacity()
170    }
171}
172
173unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
174    for t_alloc!(Box, B, A)
175{
176    fn as_buf_ptr(&self) -> *const u8 {
177        (**self).as_buf_ptr()
178    }
179
180    fn buf_len(&self) -> usize {
181        (**self).buf_len()
182    }
183
184    fn buf_capacity(&self) -> usize {
185        (**self).buf_capacity()
186    }
187}
188
189unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
190    for t_alloc!(Rc, B, A)
191{
192    fn as_buf_ptr(&self) -> *const u8 {
193        (**self).as_buf_ptr()
194    }
195
196    fn buf_len(&self) -> usize {
197        (**self).buf_len()
198    }
199
200    fn buf_capacity(&self) -> usize {
201        (**self).buf_capacity()
202    }
203}
204
205unsafe impl<B: IoBuf + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
206    for t_alloc!(Arc, B, A)
207{
208    fn as_buf_ptr(&self) -> *const u8 {
209        (**self).as_buf_ptr()
210    }
211
212    fn buf_len(&self) -> usize {
213        (**self).buf_len()
214    }
215
216    fn buf_capacity(&self) -> usize {
217        (**self).buf_capacity()
218    }
219}
220
221unsafe impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBuf
222    for t_alloc!(Vec, u8, A)
223{
224    fn as_buf_ptr(&self) -> *const u8 {
225        self.as_ptr()
226    }
227
228    fn buf_len(&self) -> usize {
229        self.len()
230    }
231
232    fn buf_capacity(&self) -> usize {
233        self.capacity()
234    }
235}
236
237unsafe impl IoBuf for String {
238    fn as_buf_ptr(&self) -> *const u8 {
239        self.as_ptr()
240    }
241
242    fn buf_len(&self) -> usize {
243        self.len()
244    }
245
246    fn buf_capacity(&self) -> usize {
247        self.capacity()
248    }
249}
250
251unsafe impl IoBuf for str {
252    fn as_buf_ptr(&self) -> *const u8 {
253        self.as_ptr()
254    }
255
256    fn buf_len(&self) -> usize {
257        self.len()
258    }
259
260    fn buf_capacity(&self) -> usize {
261        self.len()
262    }
263}
264
265unsafe impl IoBuf for [u8] {
266    fn as_buf_ptr(&self) -> *const u8 {
267        self.as_ptr()
268    }
269
270    fn buf_len(&self) -> usize {
271        self.len()
272    }
273
274    fn buf_capacity(&self) -> usize {
275        self.len()
276    }
277}
278
279unsafe impl<const N: usize> IoBuf for [u8; N] {
280    fn as_buf_ptr(&self) -> *const u8 {
281        self.as_ptr()
282    }
283
284    fn buf_len(&self) -> usize {
285        N
286    }
287
288    fn buf_capacity(&self) -> usize {
289        N
290    }
291}
292
293#[cfg(feature = "bytes")]
294unsafe impl IoBuf for bytes::Bytes {
295    fn as_buf_ptr(&self) -> *const u8 {
296        self.as_ptr()
297    }
298
299    fn buf_len(&self) -> usize {
300        self.len()
301    }
302
303    fn buf_capacity(&self) -> usize {
304        self.len()
305    }
306}
307
308#[cfg(feature = "bytes")]
309unsafe impl IoBuf for bytes::BytesMut {
310    fn as_buf_ptr(&self) -> *const u8 {
311        self.as_ptr()
312    }
313
314    fn buf_len(&self) -> usize {
315        self.len()
316    }
317
318    fn buf_capacity(&self) -> usize {
319        self.capacity()
320    }
321}
322
323#[cfg(feature = "read_buf")]
324unsafe impl IoBuf for std::io::BorrowedBuf<'static> {
325    fn as_buf_ptr(&self) -> *const u8 {
326        self.filled().as_ptr()
327    }
328
329    fn buf_len(&self) -> usize {
330        self.len()
331    }
332
333    fn buf_capacity(&self) -> usize {
334        self.capacity()
335    }
336}
337
338#[cfg(feature = "arrayvec")]
339unsafe impl<const N: usize> IoBuf for arrayvec::ArrayVec<u8, N> {
340    fn as_buf_ptr(&self) -> *const u8 {
341        self.as_ptr()
342    }
343
344    fn buf_len(&self) -> usize {
345        self.len()
346    }
347
348    fn buf_capacity(&self) -> usize {
349        self.capacity()
350    }
351}
352
353#[cfg(feature = "smallvec")]
354unsafe impl<const N: usize> IoBuf for smallvec::SmallVec<[u8; N]>
355where
356    [u8; N]: smallvec::Array<Item = u8>,
357{
358    fn as_buf_ptr(&self) -> *const u8 {
359        self.as_ptr()
360    }
361
362    fn buf_len(&self) -> usize {
363        self.len()
364    }
365
366    fn buf_capacity(&self) -> usize {
367        self.capacity()
368    }
369}
370
371/// A mutable compio compatible buffer.
372///
373/// The `IoBufMut` trait is implemented by buffer types that can be passed to
374/// compio operations. Users will not need to use this trait directly.
375///
376/// # Safety
377///
378/// Buffers passed to compio operations must reference a stable memory
379/// region. While the runtime holds ownership to a buffer, the pointer returned
380/// by `as_buf_mut_ptr` must remain valid even if the `IoBufMut` value is moved.
381pub unsafe trait IoBufMut: IoBuf + SetBufInit {
382    /// Returns a raw mutable pointer to the vector’s buffer.
383    ///
384    /// This method is to be used by the `compio` runtime and it is not
385    /// expected for users to call it directly.
386    fn as_buf_mut_ptr(&mut self) -> *mut u8;
387
388    /// Get the uninitialized part of the buffer.
389    fn as_mut_slice(&mut self) -> &mut [MaybeUninit<u8>] {
390        unsafe {
391            std::slice::from_raw_parts_mut(self.as_buf_mut_ptr().cast(), (*self).buf_capacity())
392        }
393    }
394
395    /// Create an [`IoSliceMut`] of the uninitialized part of the buffer.
396    ///
397    /// # Safety
398    ///
399    /// The return slice will not live longer than self.
400    /// It is static to provide convenience from writing self-referenced
401    /// structure.
402    unsafe fn as_io_slice_mut(&mut self) -> IoSliceMut {
403        unsafe { IoSliceMut::from_uninit(self.as_mut_slice()) }
404    }
405}
406
407unsafe impl<B: IoBufMut + ?Sized> IoBufMut for &'static mut B {
408    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
409        (**self).as_buf_mut_ptr()
410    }
411}
412
413unsafe impl<B: IoBufMut + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
414    for t_alloc!(Box, B, A)
415{
416    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
417        (**self).as_buf_mut_ptr()
418    }
419}
420
421unsafe impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> IoBufMut
422    for t_alloc!(Vec, u8, A)
423{
424    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
425        self.as_mut_ptr()
426    }
427}
428
429unsafe impl IoBufMut for [u8] {
430    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
431        self.as_mut_ptr()
432    }
433}
434
435unsafe impl<const N: usize> IoBufMut for [u8; N] {
436    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
437        self.as_mut_ptr()
438    }
439}
440
441#[cfg(feature = "bytes")]
442unsafe impl IoBufMut for bytes::BytesMut {
443    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
444        self.as_mut_ptr()
445    }
446}
447
448#[cfg(feature = "read_buf")]
449unsafe impl IoBufMut for std::io::BorrowedBuf<'static> {
450    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
451        (*self).filled().as_ptr() as _
452    }
453}
454
455#[cfg(feature = "arrayvec")]
456unsafe impl<const N: usize> IoBufMut for arrayvec::ArrayVec<u8, N> {
457    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
458        self.as_mut_ptr()
459    }
460}
461
462#[cfg(feature = "smallvec")]
463unsafe impl<const N: usize> IoBufMut for smallvec::SmallVec<[u8; N]>
464where
465    [u8; N]: smallvec::Array<Item = u8>,
466{
467    fn as_buf_mut_ptr(&mut self) -> *mut u8 {
468        self.as_mut_ptr()
469    }
470}
471
472/// A helper trait for `set_len` like methods.
473pub trait SetBufInit {
474    /// Set the buffer length. If `len` is less than the current length, nothing
475    /// should happen.
476    ///
477    /// # Safety
478    ///
479    /// * `len` should be less or equal than `buf_capacity()`.
480    /// * The bytes in the range `[buf_len(), len)` must be initialized.
481    unsafe fn set_buf_init(&mut self, len: usize);
482}
483
484impl<B: SetBufInit + ?Sized> SetBufInit for &'static mut B {
485    unsafe fn set_buf_init(&mut self, len: usize) {
486        unsafe { (**self).set_buf_init(len) }
487    }
488}
489
490impl<B: SetBufInit + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
491    for t_alloc!(Box, B, A)
492{
493    unsafe fn set_buf_init(&mut self, len: usize) {
494        unsafe { (**self).set_buf_init(len) }
495    }
496}
497
498impl<#[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit for t_alloc!(Vec, u8, A) {
499    unsafe fn set_buf_init(&mut self, len: usize) {
500        if (**self).buf_len() < len {
501            unsafe { self.set_len(len) };
502        }
503    }
504}
505
506impl SetBufInit for [u8] {
507    unsafe fn set_buf_init(&mut self, len: usize) {
508        debug_assert!(len <= self.len());
509    }
510}
511
512impl<const N: usize> SetBufInit for [u8; N] {
513    unsafe fn set_buf_init(&mut self, len: usize) {
514        debug_assert!(len <= N);
515    }
516}
517
518#[cfg(feature = "bytes")]
519impl SetBufInit for bytes::BytesMut {
520    unsafe fn set_buf_init(&mut self, len: usize) {
521        if (**self).buf_len() < len {
522            unsafe { self.set_len(len) };
523        }
524    }
525}
526
527#[cfg(feature = "read_buf")]
528impl SetBufInit for std::io::BorrowedBuf<'static> {
529    unsafe fn set_buf_init(&mut self, len: usize) {
530        let current_len = (*self).buf_len();
531        if current_len < len {
532            self.unfilled().advance(len - current_len);
533        }
534    }
535}
536
537#[cfg(feature = "arrayvec")]
538impl<const N: usize> SetBufInit for arrayvec::ArrayVec<u8, N> {
539    unsafe fn set_buf_init(&mut self, len: usize) {
540        if (**self).buf_len() < len {
541            unsafe { self.set_len(len) };
542        }
543    }
544}
545
546#[cfg(feature = "smallvec")]
547impl<const N: usize> SetBufInit for smallvec::SmallVec<[u8; N]>
548where
549    [u8; N]: smallvec::Array<Item = u8>,
550{
551    unsafe fn set_buf_init(&mut self, len: usize) {
552        if (**self).buf_len() < len {
553            unsafe { self.set_len(len) };
554        }
555    }
556}
557
558impl<T: IoBufMut> SetBufInit for [T] {
559    unsafe fn set_buf_init(&mut self, len: usize) {
560        unsafe { default_set_buf_init(self.iter_mut(), len) }
561    }
562}
563
564impl<T: IoBufMut, const N: usize> SetBufInit for [T; N] {
565    unsafe fn set_buf_init(&mut self, len: usize) {
566        unsafe { default_set_buf_init(self.iter_mut(), len) }
567    }
568}
569
570impl<T: IoBufMut, #[cfg(feature = "allocator_api")] A: Allocator + 'static> SetBufInit
571    for t_alloc!(Vec, T, A)
572{
573    unsafe fn set_buf_init(&mut self, len: usize) {
574        unsafe { default_set_buf_init(self.iter_mut(), len) }
575    }
576}
577
578#[cfg(feature = "arrayvec")]
579impl<T: IoBufMut, const N: usize> SetBufInit for arrayvec::ArrayVec<T, N> {
580    unsafe fn set_buf_init(&mut self, len: usize) {
581        unsafe { default_set_buf_init(self.iter_mut(), len) }
582    }
583}
584
585#[cfg(feature = "smallvec")]
586impl<T: IoBufMut, const N: usize> SetBufInit for smallvec::SmallVec<[T; N]>
587where
588    [T; N]: smallvec::Array<Item = T>,
589{
590    unsafe fn set_buf_init(&mut self, len: usize) {
591        unsafe { default_set_buf_init(self.iter_mut(), len) }
592    }
593}
594
595/// # Safety
596/// * `len` should be less or equal than the sum of `buf_capacity()` of all
597///   buffers.
598/// * The bytes in the range `[buf_len(), new_len)` of each buffer must be
599///   initialized
600unsafe fn default_set_buf_init<'a, B: IoBufMut>(
601    iter: impl IntoIterator<Item = &'a mut B>,
602    mut len: usize,
603) {
604    for buf in iter {
605        let capacity = (*buf).buf_capacity();
606        if len >= capacity {
607            unsafe { buf.set_buf_init(capacity) };
608            len -= capacity;
609        } else {
610            unsafe { buf.set_buf_init(len) };
611            len = 0;
612        }
613    }
614}
compio_buf/io_buf.rs

compio_buf/
io_buf.rs
