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 405605 : pub fn with_capacity(capacity: usize) -> Self {
43 405605 : let align = ConstAlign::<A>;
44 405605 : 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 405605 : let ptr = unsafe {
48 405605 : let ptr = alloc::alloc(layout);
49 405605 : if ptr.is_null() {
50 0 : alloc::handle_alloc_error(layout);
51 405605 : }
52 405605 : AlignedBufferPtr(ptr)
53 : };
54 :
55 405605 : RawAlignedBuffer {
56 405605 : ptr,
57 405605 : capacity,
58 405605 : len: 0,
59 405605 : align,
60 405605 : }
61 405605 : }
62 : }
63 :
64 : impl<A: Alignment> RawAlignedBuffer<A> {
65 : /// Returns the total number of bytes the buffer can hold.
66 : #[inline]
67 65129902 : pub fn capacity(&self) -> usize {
68 65129902 : self.capacity
69 65129902 : }
70 :
71 : /// Returns the alignment of the buffer.
72 : #[inline]
73 17 : pub fn align(&self) -> usize {
74 17 : self.align.align()
75 17 : }
76 :
77 : /// Returns the number of bytes in the buffer, also referred to as its 'length'.
78 : #[inline]
79 78312468 : pub fn len(&self) -> usize {
80 78312468 : self.len
81 78312468 : }
82 :
83 : /// Force the length of the buffer to `new_len`.
84 : #[inline]
85 9459961 : pub unsafe fn set_len(&mut self, new_len: usize) {
86 9459961 : debug_assert!(new_len <= self.capacity());
87 9459961 : self.len = new_len;
88 9459961 : }
89 :
90 : #[inline]
91 16912194 : pub fn as_ptr(&self) -> *const u8 {
92 16912194 : self.ptr.0
93 16912194 : }
94 :
95 : #[inline]
96 10747679 : pub fn as_mut_ptr(&mut self) -> *mut u8 {
97 10747679 : self.ptr.0
98 10747679 : }
99 :
100 : /// Extracts a slice containing the entire buffer.
101 : ///
102 : /// Equivalent to `&s[..]`.
103 : #[inline]
104 3222812 : pub fn as_slice(&self) -> &[u8] {
105 : // SAFETY: The pointer is valid and `len` bytes are initialized.
106 3222812 : unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
107 3222812 : }
108 :
109 : /// Extracts a mutable slice of the entire buffer.
110 : ///
111 : /// Equivalent to `&mut s[..]`.
112 203 : pub fn as_mut_slice(&mut self) -> &mut [u8] {
113 : // SAFETY: The pointer is valid and `len` bytes are initialized.
114 203 : unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
115 203 : }
116 :
117 : /// Drops the all the contents of the buffer, setting its length to `0`.
118 : #[inline]
119 11830 : pub fn clear(&mut self) {
120 11830 : self.len = 0;
121 11830 : }
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 9007222 : pub fn reserve(&mut self, additional: usize) {
133 9007222 : if additional > self.capacity() - self.len() {
134 3 : self.reserve_inner(additional);
135 9007219 : }
136 9007222 : }
137 :
138 3 : fn reserve_inner(&mut self, additional: usize) {
139 3 : let Some(required_cap) = self.len().checked_add(additional) else {
140 0 : capacity_overflow()
141 : };
142 :
143 3 : let old_capacity = self.capacity();
144 3 : let align = self.align();
145 : // This guarantees exponential growth. The doubling cannot overflow
146 : // because `cap <= isize::MAX` and the type of `cap` is `usize`.
147 3 : let cap = cmp::max(old_capacity * 2, required_cap);
148 :
149 3 : if !is_valid_alloc(cap) {
150 0 : capacity_overflow()
151 3 : }
152 3 : let new_layout = Layout::from_size_align(cap, self.align()).expect("Invalid layout");
153 :
154 3 : 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 3 : let (ptr, cap) = unsafe {
159 3 : let old_layout = Layout::from_size_align_unchecked(old_capacity, align);
160 3 : let ptr = alloc::realloc(old_ptr, old_layout, new_layout.size());
161 3 : if ptr.is_null() {
162 0 : alloc::handle_alloc_error(new_layout);
163 3 : }
164 3 : (AlignedBufferPtr(ptr), cap)
165 : };
166 :
167 3 : self.ptr = ptr;
168 3 : self.capacity = cap;
169 3 : }
170 :
171 : /// Shortens the buffer, keeping the first len bytes.
172 5 : pub fn truncate(&mut self, len: usize) {
173 5 : if len > self.len {
174 0 : return;
175 5 : }
176 5 : self.len = len;
177 5 : }
178 :
179 : /// Consumes and leaks the `IoBufferMut`, returning a mutable reference to the contents, &'a mut [u8].
180 50 : pub fn leak<'a>(self) -> &'a mut [u8] {
181 50 : let mut buf = ManuallyDrop::new(self);
182 : // SAFETY: leaking the buffer as intended.
183 50 : unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len) }
184 50 : }
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 3 : fn is_valid_alloc(alloc_size: usize) -> bool {
201 3 : !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
202 3 : }
203 :
204 : impl<A: Alignment> Drop for RawAlignedBuffer<A> {
205 405421 : fn drop(&mut self) {
206 : // SAFETY: memory was allocated with std::alloc::alloc with the same layout.
207 : unsafe {
208 405421 : alloc::dealloc(
209 405421 : self.as_mut_ptr(),
210 405421 : Layout::from_size_align_unchecked(self.capacity, self.align.align()),
211 : )
212 : }
213 405421 : }
214 : }
|
