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