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