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