Line data    Source code

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