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

compio_buf/
io_buf.rs
