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        22424 :     pub(super) fn from_raw(raw: RawAlignedBuffer<A>, range: Range<usize>) -> Self {
      20        22424 :         AlignedBuffer {
      21        22424 :             raw: Arc::new(raw),
      22        22424 :             range,
      23        22424 :         }
      24        22424 :     }
      25              : 
      26              :     /// Returns the number of bytes in the buffer, also referred to as its 'length'.
      27              :     #[inline]
      28     22520609 :     pub fn len(&self) -> usize {
      29     22520609 :         self.range.len()
      30     22520609 :     }
      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     11367537 :     fn as_ptr(&self) -> *const u8 {
      40              :         // SAFETY: `self.range.start` is guaranteed to be within [0, self.len()).
      41     11367537 :         unsafe { self.raw.as_ptr().add(self.range.start) }
      42     11367537 :     }
      43              : 
      44              :     /// Extracts a slice containing the entire buffer.
      45              :     ///
      46              :     /// Equivalent to `&s[..]`.
      47              :     #[inline]
      48       508346 :     fn as_slice(&self) -> &[u8] {
      49       508346 :         &self.raw.as_slice()[self.range.start..self.range.end]
      50       508346 :     }
      51              : 
      52              :     /// Returns a slice of self for the index range `[begin..end)`.
      53      2192760 :     pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
      54              :         use core::ops::Bound;
      55      2192760 :         let len = self.len();
      56              : 
      57      2192760 :         let begin = match range.start_bound() {
      58      2192760 :             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      2192760 :         let end = match range.end_bound() {
      64            0 :             Bound::Included(&n) => n.checked_add(1).expect("out of range"),
      65      2192760 :             Bound::Excluded(&n) => n,
      66            0 :             Bound::Unbounded => len,
      67              :         };
      68              : 
      69      2192760 :         assert!(
      70      2192760 :             begin <= end,
      71            0 :             "range start must not be greater than end: {begin:?} <= {end:?}",
      72              :         );
      73      2192760 :         assert!(end <= len, "range end out of bounds: {end:?} <= {len:?}",);
      74              : 
      75      2192760 :         let begin = self.range.start + begin;
      76      2192760 :         let end = self.range.start + end;
      77              : 
      78      2192760 :         AlignedBuffer {
      79      2192760 :             raw: Arc::clone(&self.raw),
      80      2192760 :             range: begin..end,
      81      2192760 :         }
      82      2192760 :     }
      83              : 
      84              :     /// Returns the mutable aligned buffer, if the immutable aligned buffer
      85              :     /// has exactly one strong reference. Otherwise returns `None`.
      86        11817 :     pub fn into_mut(self) -> Option<AlignedBufferMut<A>> {
      87        11817 :         let raw = Arc::into_inner(self.raw)?;
      88        11817 :         Some(AlignedBufferMut::from_raw(raw))
      89        11817 :     }
      90              : }
      91              : 
      92              : impl<A: Alignment> Deref for AlignedBuffer<A> {
      93              :     type Target = [u8];
      94              : 
      95       508346 :     fn deref(&self) -> &Self::Target {
      96       508346 :         self.as_slice()
      97       508346 :     }
      98              : }
      99              : 
     100              : impl<A: Alignment> AsRef<[u8]> for AlignedBuffer<A> {
     101            0 :     fn as_ref(&self) -> &[u8] {
     102            0 :         self.as_slice()
     103            0 :     }
     104              : }
     105              : 
     106              : impl<A: Alignment> PartialEq<[u8]> for AlignedBuffer<A> {
     107            0 :     fn eq(&self, other: &[u8]) -> bool {
     108            0 :         self.as_slice().eq(other)
     109            0 :     }
     110              : }
     111              : 
     112              : impl<const A: usize, const N: usize> From<&[u8; N]> for AlignedBuffer<ConstAlign<A>> {
     113            0 :     fn from(value: &[u8; N]) -> Self {
     114            0 :         let mut buf = AlignedBufferMut::with_capacity(N);
     115            0 :         buf.extend_from_slice(value);
     116            0 :         buf.freeze()
     117            0 :     }
     118              : }
     119              : 
     120              : /// SAFETY: the underlying buffer references a stable memory region.
     121              : unsafe impl<A: Alignment> tokio_epoll_uring::IoBuf for AlignedBuffer<A> {
     122     11367537 :     fn stable_ptr(&self) -> *const u8 {
     123     11367537 :         self.as_ptr()
     124     11367537 :     }
     125              : 
     126     13633561 :     fn bytes_init(&self) -> usize {
     127     13633561 :         self.len()
     128     13633561 :     }
     129              : 
     130      4472624 :     fn bytes_total(&self) -> usize {
     131      4472624 :         self.len()
     132      4472624 :     }
     133              : }