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 8872 : pub(super) fn from_raw(raw: RawAlignedBuffer<A>, range: Range<usize>) -> Self {
20 8872 : AlignedBuffer {
21 8872 : raw: Arc::new(raw),
22 8872 : range,
23 8872 : }
24 8872 : }
25 :
26 : /// Returns the number of bytes in the buffer, also referred to as its 'length'.
27 : #[inline]
28 70627614 : pub fn len(&self) -> usize {
29 70627614 : self.range.len()
30 70627614 : }
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 31086674 : fn as_ptr(&self) -> *const u8 {
40 31086674 : // SAFETY: `self.range.start` is guaranteed to be within [0, self.len()).
41 31086674 : unsafe { self.raw.as_ptr().add(self.range.start) }
42 31086674 : }
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 4385516 : pub fn slice(&self, range: impl RangeBounds<usize>) -> Self {
54 : use core::ops::Bound;
55 4385516 : let len = self.len();
56 :
57 4385516 : let begin = match range.start_bound() {
58 4385516 : 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 4385516 : let end = match range.end_bound() {
64 0 : Bound::Included(&n) => n.checked_add(1).expect("out of range"),
65 4385516 : Bound::Excluded(&n) => n,
66 0 : Bound::Unbounded => len,
67 : };
68 :
69 4385516 : assert!(
70 4385516 : begin <= end,
71 0 : "range start must not be greater than end: {:?} <= {:?}",
72 : begin,
73 : end,
74 : );
75 4385516 : assert!(
76 4385516 : end <= len,
77 0 : "range end out of bounds: {:?} <= {:?}",
78 : end,
79 : len,
80 : );
81 :
82 4385516 : let begin = self.range.start + begin;
83 4385516 : let end = self.range.start + end;
84 4385516 :
85 4385516 : AlignedBuffer {
86 4385516 : raw: Arc::clone(&self.raw),
87 4385516 : range: begin..end,
88 4385516 : }
89 4385516 : }
90 :
91 : /// Returns the mutable aligned buffer, if the immutable aligned buffer
92 : /// has exactly one strong reference. Otherwise returns `None`.
93 6619 : pub fn into_mut(self) -> Option<AlignedBufferMut<A>> {
94 6619 : let raw = Arc::into_inner(self.raw)?;
95 6619 : Some(AlignedBufferMut::from_raw(raw))
96 6619 : }
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 8 : fn from(value: &[u8; N]) -> Self {
121 8 : let mut buf = AlignedBufferMut::with_capacity(N);
122 8 : buf.extend_from_slice(value);
123 8 : buf.freeze()
124 8 : }
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 31086674 : fn stable_ptr(&self) -> *const u8 {
130 31086674 : self.as_ptr()
131 31086674 : }
132 :
133 48656477 : fn bytes_init(&self) -> usize {
134 48656477 : self.len()
135 48656477 : }
136 :
137 13192201 : fn bytes_total(&self) -> usize {
138 13192201 : self.len()
139 13192201 : }
140 : }
|