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