Line data    Source code

       1              : use std::ops::{Deref, Range, RangeBounds};
       2              : use std::sync::Arc;
       3              : 
       4              : use super::alignment::Alignment;
       5              : use super::raw::RawAlignedBuffer;
       6              : use super::{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        22312 :     pub(super) fn from_raw(raw: RawAlignedBuffer<A>, range: Range<usize>) -> Self {
      20        22312 :         AlignedBuffer {
      21        22312 :             raw: Arc::new(raw),
      22        22312 :             range,
      23        22312 :         }
      24        22312 :     }
      25              : 
      26              :     /// Returns the number of bytes in the buffer, also referred to as its 'length'.
      27              :     #[inline]
      28     22518837 :     pub fn len(&self) -> usize {
      29     22518837 :         self.range.len()
      30     22518837 :     }
      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     11366398 :     fn as_ptr(&self) -> *const u8 {
      40     11366398 :         // SAFETY: `self.range.start` is guaranteed to be within [0, self.len()).
      41     11366398 :         unsafe { self.raw.as_ptr().add(self.range.start) }
      42     11366398 :     }
      43              : 
      44              :     /// Extracts a slice containing the entire buffer.
      45              :     ///
      46              :     /// Equivalent to `&s[..]`.
      47              :     #[inline]
      48       508484 :     fn as_slice(&self) -> &[u8] {
      49       508484 :         &self.raw.as_slice()[self.range.start..self.range.end]
      50       508484 :     }
      51              : 
      52              :     /// Returns a slice of self for the index range `[begin..end)`.
      53      2192758 :     pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
      54              :         use core::ops::Bound;
      55      2192758 :         let len = self.len();
      56              : 
      57      2192758 :         let begin = match range.start_bound() {
      58      2192758 :             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      2192758 :         let end = match range.end_bound() {
      64            0 :             Bound::Included(&n) => n.checked_add(1).expect("out of range"),
      65      2192758 :             Bound::Excluded(&n) => n,
      66            0 :             Bound::Unbounded => len,
      67              :         };
      68              : 
      69      2192758 :         assert!(
      70      2192758 :             begin <= end,
      71            0 :             "range start must not be greater than end: {:?} <= {:?}",
      72              :             begin,
      73              :             end,
      74              :         );
      75      2192758 :         assert!(
      76      2192758 :             end <= len,
      77            0 :             "range end out of bounds: {:?} <= {:?}",
      78              :             end,
      79              :             len,
      80              :         );
      81              : 
      82      2192758 :         let begin = self.range.start + begin;
      83      2192758 :         let end = self.range.start + end;
      84      2192758 : 
      85      2192758 :         AlignedBuffer {
      86      2192758 :             raw: Arc::clone(&self.raw),
      87      2192758 :             range: begin..end,
      88      2192758 :         }
      89      2192758 :     }
      90              : 
      91              :     /// Returns the mutable aligned buffer, if the immutable aligned buffer
      92              :     /// has exactly one strong reference. Otherwise returns `None`.
      93        11805 :     pub fn into_mut(self) -> Option<AlignedBufferMut<A>> {
      94        11805 :         let raw = Arc::into_inner(self.raw)?;
      95        11805 :         Some(AlignedBufferMut::from_raw(raw))
      96        11805 :     }
      97              : }
      98              : 
      99              : impl<A: Alignment> Deref for AlignedBuffer<A> {
     100              :     type Target = [u8];
     101              : 
     102       508484 :     fn deref(&self) -> &Self::Target {
     103       508484 :         self.as_slice()
     104       508484 :     }
     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            0 :     fn from(value: &[u8; N]) -> Self {
     121            0 :         let mut buf = AlignedBufferMut::with_capacity(N);
     122            0 :         buf.extend_from_slice(value);
     123            0 :         buf.freeze()
     124            0 :     }
     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     11366398 :     fn stable_ptr(&self) -> *const u8 {
     130     11366398 :         self.as_ptr()
     131     11366398 :     }
     132              : 
     133     13632213 :     fn bytes_init(&self) -> usize {
     134     13632213 :         self.len()
     135     13632213 :     }
     136              : 
     137      4472332 :     fn bytes_total(&self) -> usize {
     138      4472332 :         self.len()
     139      4472332 :     }
     140              : }