Line data    Source code

       1              : //! See [`FullSlice`].
       2              : 
       3              : use bytes::{Bytes, BytesMut};
       4              : use std::ops::{Deref, Range};
       5              : use tokio_epoll_uring::{BoundedBuf, IoBuf, Slice};
       6              : 
       7              : /// The true owned equivalent for Rust [`slice`]. Use this for the write path.
       8              : ///
       9              : /// Unlike [`tokio_epoll_uring::Slice`], which we unfortunately inherited from `tokio-uring`,
      10              : /// [`FullSlice`] is guaranteed to have all its bytes initialized. This means that
      11              : /// [`<FullSlice as Deref<Target = [u8]>>::len`] is equal to [`Slice::bytes_init`] and [`Slice::bytes_total`].
      12              : ///
      13              : pub struct FullSlice<B> {
      14              :     slice: Slice<B>,
      15              : }
      16              : 
      17              : impl<B> FullSlice<B>
      18              : where
      19              :     B: IoBuf,
      20              : {
      21     92982719 :     pub(crate) fn must_new(slice: Slice<B>) -> Self {
      22     92982719 :         assert_eq!(slice.bytes_init(), slice.bytes_total());
      23     92982719 :         FullSlice { slice }
      24     92982719 :     }
      25     71844746 :     pub(crate) fn into_raw_slice(self) -> Slice<B> {
      26     71844746 :         let FullSlice { slice: s } = self;
      27     71844746 :         s
      28     71844746 :     }
      29              : }
      30              : 
      31              : impl<B> Deref for FullSlice<B>
      32              : where
      33              :     B: IoBuf,
      34              : {
      35              :     type Target = [u8];
      36              : 
      37     22823767 :     fn deref(&self) -> &[u8] {
      38     22823767 :         let rust_slice = &self.slice[..];
      39     22823767 :         assert_eq!(rust_slice.len(), self.slice.bytes_init());
      40     22823767 :         assert_eq!(rust_slice.len(), self.slice.bytes_total());
      41     22823767 :         rust_slice
      42     22823767 :     }
      43              : }
      44              : 
      45              : pub(crate) trait IoBufExt {
      46              :     /// Get a [`FullSlice`] for the entire buffer, i.e., `self[..]` or `self[0..self.len()]`.
      47              :     fn slice_len(self) -> FullSlice<Self>
      48              :     where
      49              :         Self: Sized;
      50              : }
      51              : 
      52              : macro_rules! impl_io_buf_ext {
      53              :     ($T:ty) => {
      54              :         impl IoBufExt for $T {
      55              :             #[inline(always)]
      56     42284304 :             fn slice_len(self) -> FullSlice<Self> {
      57     42284304 :                 let len = self.len();
      58     42284304 :                 let s = if len == 0 {
      59              :                     // `BoundedBuf::slice(0..len)` or `BoundedBuf::slice(..)` has an incorrect assertion,
      60              :                     // causing a panic if len == 0.
      61              :                     // The Slice::from_buf_bounds has the correct assertion (<= instead of <).
      62              :                     // => https://github.com/neondatabase/tokio-epoll-uring/issues/46
      63          192 :                     let slice = self.slice_full();
      64          192 :                     let mut bounds: Range<_> = slice.bounds();
      65          192 :                     bounds.end = bounds.start;
      66          192 :                     Slice::from_buf_bounds(slice.into_inner(), bounds)
      67              :                 } else {
      68     42284112 :                     self.slice(0..len)
      69              :                 };
      70     42284304 :                 FullSlice::must_new(s)
      71     42284304 :             }
      72              :         }
      73              :     };
      74              : }
      75              : 
      76              : impl_io_buf_ext!(Bytes);
      77              : impl_io_buf_ext!(BytesMut);
      78              : impl_io_buf_ext!(Vec<u8>);