Line data    Source code

       1              : use std::{
       2              :     ops::{Deref, Range, RangeBounds},
       3              :     sync::Arc,
       4              : };
       5              : 
       6              : use super::{alignment::Alignment, raw::RawAlignedBuffer, AlignedBufferMut, ConstAlign};
       7              : 
       8              : /// An shared, immutable aligned buffer type.
       9              : #[derive(Clone, Debug)]
      10              : pub struct AlignedBuffer<A: Alignment> {
      11              :     /// Shared raw buffer.
      12              :     raw: Arc<RawAlignedBuffer<A>>,
      13              :     /// Range that specifies the current slice.
      14              :     range: Range<usize>,
      15              : }
      16              : 
      17              : impl<A: Alignment> AlignedBuffer<A> {
      18              :     /// Creates an immutable `IoBuffer` from the raw buffer
      19        17813 :     pub(super) fn from_raw(raw: RawAlignedBuffer<A>, range: Range<usize>) -> Self {
      20        17813 :         AlignedBuffer {
      21        17813 :             raw: Arc::new(raw),
      22        17813 :             range,
      23        17813 :         }
      24        17813 :     }
      25              : 
      26              :     /// Returns the number of bytes in the buffer, also referred to as its 'length'.
      27              :     #[inline]
      28    141255171 :     pub fn len(&self) -> usize {
      29    141255171 :         self.range.len()
      30    141255171 :     }
      31              : 
      32              :     /// Returns the alignment of the buffer.
      33              :     #[inline]
      34            0 :     pub fn align(&self) -> usize {
      35            0 :         self.raw.align()
      36            0 :     }
      37              : 
      38              :     #[inline]
      39     62172967 :     fn as_ptr(&self) -> *const u8 {
      40     62172967 :         // SAFETY: `self.range.start` is guaranteed to be within [0, self.len()).
      41     62172967 :         unsafe { self.raw.as_ptr().add(self.range.start) }
      42     62172967 :     }
      43              : 
      44              :     /// Extracts a slice containing the entire buffer.
      45              :     ///
      46              :     /// Equivalent to `&s[..]`.
      47              :     #[inline]
      48            0 :     fn as_slice(&self) -> &[u8] {
      49            0 :         &self.raw.as_slice()[self.range.start..self.range.end]
      50            0 :     }
      51              : 
      52              :     /// Returns a slice of self for the index range `[begin..end)`.
      53      8771032 :     pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
      54              :         use core::ops::Bound;
      55      8771032 :         let len = self.len();
      56              : 
      57      8771032 :         let begin = match range.start_bound() {
      58      8771032 :             Bound::Included(&n) => n,
      59            0 :             Bound::Excluded(&n) => n.checked_add(1).expect("out of range"),
      60            0 :             Bound::Unbounded => 0,
      61              :         };
      62              : 
      63      8771032 :         let end = match range.end_bound() {
      64            0 :             Bound::Included(&n) => n.checked_add(1).expect("out of range"),
      65      8771032 :             Bound::Excluded(&n) => n,
      66            0 :             Bound::Unbounded => len,
      67              :         };
      68              : 
      69      8771032 :         assert!(
      70      8771032 :             begin <= end,
      71            0 :             "range start must not be greater than end: {:?} <= {:?}",
      72              :             begin,
      73              :             end,
      74              :         );
      75      8771032 :         assert!(
      76      8771032 :             end <= len,
      77            0 :             "range end out of bounds: {:?} <= {:?}",
      78              :             end,
      79              :             len,
      80              :         );
      81              : 
      82      8771032 :         let begin = self.range.start + begin;
      83      8771032 :         let end = self.range.start + end;
      84      8771032 : 
      85      8771032 :         AlignedBuffer {
      86      8771032 :             raw: Arc::clone(&self.raw),
      87      8771032 :             range: begin..end,
      88      8771032 :         }
      89      8771032 :     }
      90              : 
      91              :     /// Returns the mutable aligned buffer, if the immutable aligned buffer
      92              :     /// has exactly one strong reference. Otherwise returns `None`.
      93        13246 :     pub fn into_mut(self) -> Option<AlignedBufferMut<A>> {
      94        13246 :         let raw = Arc::into_inner(self.raw)?;
      95        13246 :         Some(AlignedBufferMut::from_raw(raw))
      96        13246 :     }
      97              : }
      98              : 
      99              : impl<A: Alignment> Deref for AlignedBuffer<A> {
     100              :     type Target = [u8];
     101              : 
     102            0 :     fn deref(&self) -> &Self::Target {
     103            0 :         self.as_slice()
     104            0 :     }
     105              : }
     106              : 
     107              : impl<A: Alignment> AsRef<[u8]> for AlignedBuffer<A> {
     108            0 :     fn as_ref(&self) -> &[u8] {
     109            0 :         self.as_slice()
     110            0 :     }
     111              : }
     112              : 
     113              : impl<A: Alignment> PartialEq<[u8]> for AlignedBuffer<A> {
     114            0 :     fn eq(&self, other: &[u8]) -> bool {
     115            0 :         self.as_slice().eq(other)
     116            0 :     }
     117              : }
     118              : 
     119              : impl<const A: usize, const N: usize> From<&[u8; N]> for AlignedBuffer<ConstAlign<A>> {
     120           16 :     fn from(value: &[u8; N]) -> Self {
     121           16 :         let mut buf = AlignedBufferMut::with_capacity(N);
     122           16 :         buf.extend_from_slice(value);
     123           16 :         buf.freeze()
     124           16 :     }
     125              : }
     126              : 
     127              : /// SAFETY: the underlying buffer references a stable memory region.
     128              : unsafe impl<A: Alignment> tokio_epoll_uring::IoBuf for AlignedBuffer<A> {
     129     62172967 :     fn stable_ptr(&self) -> *const u8 {
     130     62172967 :         self.as_ptr()
     131     62172967 :     }
     132              : 
     133     97312666 :     fn bytes_init(&self) -> usize {
     134     97312666 :         self.len()
     135     97312666 :     }
     136              : 
     137     26384564 :     fn bytes_total(&self) -> usize {
     138     26384564 :         self.len()
     139     26384564 :     }
     140              : }