Line data Source code
1 : use core::slice;
2 : use std::{
3 : alloc::{self, Layout},
4 : cmp,
5 : mem::ManuallyDrop,
6 : };
7 :
8 : use super::alignment::{Alignment, ConstAlign};
9 :
10 : #[derive(Debug)]
11 : struct AlignedBufferPtr(*mut u8);
12 :
13 : // SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
14 : unsafe impl Send for AlignedBufferPtr {}
15 :
16 : // SAFETY: We gurantees no one besides `IoBufferPtr` itself has the raw pointer.
17 : unsafe impl Sync for AlignedBufferPtr {}
18 :
19 : /// An aligned buffer type.
20 : #[derive(Debug)]
21 : pub struct RawAlignedBuffer<A: Alignment> {
22 : ptr: AlignedBufferPtr,
23 : capacity: usize,
24 : len: usize,
25 : align: A,
26 : }
27 :
28 : impl<const A: usize> RawAlignedBuffer<ConstAlign<A>> {
29 : /// Constructs a new, empty `IoBufferMut` with at least the specified capacity and alignment.
30 : ///
31 : /// The buffer will be able to hold at most `capacity` elements and will never resize.
32 : ///
33 : ///
34 : /// # Panics
35 : ///
36 : /// Panics if the new capacity exceeds `isize::MAX` _bytes_, or if the following alignment requirement is not met:
37 : /// * `align` must not be zero,
38 : ///
39 : /// * `align` must be a power of two,
40 : ///
41 : /// * `capacity`, when rounded up to the nearest multiple of `align`,
42 : /// must not overflow isize (i.e., the rounded value must be
43 : /// less than or equal to `isize::MAX`).
44 1440396 : pub fn with_capacity(capacity: usize) -> Self {
45 1440396 : let align = ConstAlign::<A>;
46 1440396 : let layout = Layout::from_size_align(capacity, align.align()).expect("Invalid layout");
47 :
48 : // SAFETY: Making an allocation with a sized and aligned layout. The memory is manually freed with the same layout.
49 1440396 : let ptr = unsafe {
50 1440396 : let ptr = alloc::alloc(layout);
51 1440396 : if ptr.is_null() {
52 0 : alloc::handle_alloc_error(layout);
53 1440396 : }
54 1440396 : AlignedBufferPtr(ptr)
55 1440396 : };
56 1440396 :
57 1440396 : RawAlignedBuffer {
58 1440396 : ptr,
59 1440396 : capacity,
60 1440396 : len: 0,
61 1440396 : align,
62 1440396 : }
63 1440396 : }
64 : }
65 :
66 : impl<A: Alignment> RawAlignedBuffer<A> {
67 : /// Returns the total number of bytes the buffer can hold.
68 : #[inline]
69 74389305 : pub fn capacity(&self) -> usize {
70 74389305 : self.capacity
71 74389305 : }
72 :
73 : /// Returns the alignment of the buffer.
74 : #[inline]
75 68 : pub fn align(&self) -> usize {
76 68 : self.align.align()
77 68 : }
78 :
79 : /// Returns the number of bytes in the buffer, also referred to as its 'length'.
80 : #[inline]
81 95045920 : pub fn len(&self) -> usize {
82 95045920 : self.len
83 95045920 : }
84 :
85 : /// Force the length of the buffer to `new_len`.
86 : #[inline]
87 11062530 : pub unsafe fn set_len(&mut self, new_len: usize) {
88 11062530 : debug_assert!(new_len <= self.capacity());
89 11062530 : self.len = new_len;
90 11062530 : }
91 :
92 : #[inline]
93 77110571 : pub fn as_ptr(&self) -> *const u8 {
94 77110571 : self.ptr.0
95 77110571 : }
96 :
97 : #[inline]
98 14380596 : pub fn as_mut_ptr(&mut self) -> *mut u8 {
99 14380596 : self.ptr.0
100 14380596 : }
101 :
102 : /// Extracts a slice containing the entire buffer.
103 : ///
104 : /// Equivalent to `&s[..]`.
105 : #[inline]
106 8367853 : pub fn as_slice(&self) -> &[u8] {
107 8367853 : // SAFETY: The pointer is valid and `len` bytes are initialized.
108 8367853 : unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
109 8367853 : }
110 :
111 : /// Extracts a mutable slice of the entire buffer.
112 : ///
113 : /// Equivalent to `&mut s[..]`.
114 0 : pub fn as_mut_slice(&mut self) -> &mut [u8] {
115 0 : // SAFETY: The pointer is valid and `len` bytes are initialized.
116 0 : unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
117 0 : }
118 :
119 : /// Drops the all the contents of the buffer, setting its length to `0`.
120 : #[inline]
121 13298 : pub fn clear(&mut self) {
122 13298 : self.len = 0;
123 13298 : }
124 :
125 : /// Reserves capacity for at least `additional` more bytes to be inserted
126 : /// in the given `IoBufferMut`. The collection may reserve more space to
127 : /// speculatively avoid frequent reallocations. After calling `reserve`,
128 : /// capacity will be greater than or equal to `self.len() + additional`.
129 : /// Does nothing if capacity is already sufficient.
130 : ///
131 : /// # Panics
132 : ///
133 : /// Panics if the new capacity exceeds `isize::MAX` _bytes_.
134 9623128 : pub fn reserve(&mut self, additional: usize) {
135 9623128 : if additional > self.capacity() - self.len() {
136 12 : self.reserve_inner(additional);
137 9623116 : }
138 9623128 : }
139 :
140 12 : fn reserve_inner(&mut self, additional: usize) {
141 12 : let Some(required_cap) = self.len().checked_add(additional) else {
142 0 : capacity_overflow()
143 : };
144 :
145 12 : let old_capacity = self.capacity();
146 12 : let align = self.align();
147 12 : // This guarantees exponential growth. The doubling cannot overflow
148 12 : // because `cap <= isize::MAX` and the type of `cap` is `usize`.
149 12 : let cap = cmp::max(old_capacity * 2, required_cap);
150 12 :
151 12 : if !is_valid_alloc(cap) {
152 0 : capacity_overflow()
153 12 : }
154 12 : let new_layout = Layout::from_size_align(cap, self.align()).expect("Invalid layout");
155 12 :
156 12 : let old_ptr = self.as_mut_ptr();
157 :
158 : // SAFETY: old allocation was allocated with std::alloc::alloc with the same layout,
159 : // and we panics on null pointer.
160 12 : let (ptr, cap) = unsafe {
161 12 : let old_layout = Layout::from_size_align_unchecked(old_capacity, align);
162 12 : let ptr = alloc::realloc(old_ptr, old_layout, new_layout.size());
163 12 : if ptr.is_null() {
164 0 : alloc::handle_alloc_error(new_layout);
165 12 : }
166 12 : (AlignedBufferPtr(ptr), cap)
167 12 : };
168 12 :
169 12 : self.ptr = ptr;
170 12 : self.capacity = cap;
171 12 : }
172 :
173 : /// Shortens the buffer, keeping the first len bytes.
174 20 : pub fn truncate(&mut self, len: usize) {
175 20 : if len > self.len {
176 0 : return;
177 20 : }
178 20 : self.len = len;
179 20 : }
180 :
181 : /// Consumes and leaks the `IoBufferMut`, returning a mutable reference to the contents, &'a mut [u8].
182 184 : pub fn leak<'a>(self) -> &'a mut [u8] {
183 184 : let mut buf = ManuallyDrop::new(self);
184 184 : // SAFETY: leaking the buffer as intended.
185 184 : unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len) }
186 184 : }
187 : }
188 :
189 0 : fn capacity_overflow() -> ! {
190 0 : panic!("capacity overflow")
191 : }
192 :
193 : // We need to guarantee the following:
194 : // * We don't ever allocate `> isize::MAX` byte-size objects.
195 : // * We don't overflow `usize::MAX` and actually allocate too little.
196 : //
197 : // On 64-bit we just need to check for overflow since trying to allocate
198 : // `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add
199 : // an extra guard for this in case we're running on a platform which can use
200 : // all 4GB in user-space, e.g., PAE or x32.
201 : #[inline]
202 12 : fn is_valid_alloc(alloc_size: usize) -> bool {
203 12 : !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
204 12 : }
205 :
206 : impl<A: Alignment> Drop for RawAlignedBuffer<A> {
207 1439737 : fn drop(&mut self) {
208 1439737 : // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
209 1439737 : unsafe {
210 1439737 : alloc::dealloc(
211 1439737 : self.as_mut_ptr(),
212 1439737 : Layout::from_size_align_unchecked(self.capacity, self.align.align()),
213 1439737 : )
214 1439737 : }
215 1439737 : }
216 : }
|
