Line data    Source code

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