1use crate::Alignment;
2use core::{
3 marker::PhantomData,
4 mem::{align_of, size_of},
5 ptr::{null_mut, NonNull},
6};
7use alloc::alloc::{alloc, dealloc, handle_alloc_error, realloc, Layout};
8
9pub struct ARawVec<T, A: Alignment> {
10 pub ptr: NonNull<T>,
11 pub capacity: usize,
12 pub align: A,
13 _marker: PhantomData<T>,
14}
15
16impl<T, A: Alignment> Drop for ARawVec<T, A> {
17 #[inline]
18 fn drop(&mut self) {
19 // this can't overflow since we already have this much stored in a slice
20 let size_bytes: usize = self.capacity * size_of::<T>();
21 if size_bytes > 0 {
22 // SAFETY: memory was allocated with alloc::alloc::alloc
23 unsafe {
24 dealloc(
25 self.ptr.as_ptr() as *mut u8,
26 Layout::from_size_align_unchecked(
27 size_bytes,
28 self.align.alignment(minimum_align:align_of::<T>()),
29 ),
30 )
31 }
32 }
33 }
34}
35
36pub fn capacity_overflow() -> ! {
37 panic!("capacity overflow")
38}
39
40impl<T, A: Alignment> ARawVec<T, A> {
41 /// # Safety
42 ///
43 /// `align` must be a power of two.
44 /// `align` must be greater than or equal to `core::mem::align_of::<T>()`.
45 #[inline]
46 pub unsafe fn new_unchecked(align: usize) -> Self {
47 let cap = if size_of::<T>() == 0 { usize::MAX } else { 0 };
48 Self::from_raw_parts(null_mut::<u8>().wrapping_add(align) as *mut T, cap, align)
49 }
50
51 /// # Safety
52 ///
53 /// `align` must be a power of two.
54 /// `align` must be greater than or equal to `core::mem::align_of::<T>()`.
55 #[inline]
56 pub unsafe fn with_capacity_unchecked(capacity: usize, align: usize) -> Self {
57 if capacity == 0 || size_of::<T>() == 0 {
58 Self::new_unchecked(align)
59 } else {
60 Self {
61 ptr: NonNull::<T>::new_unchecked(with_capacity_unchecked(
62 capacity,
63 align,
64 size_of::<T>(),
65 ) as *mut T),
66 capacity,
67 align: A::new(align, align_of::<T>()),
68 _marker: PhantomData,
69 }
70 }
71 }
72
73 const MIN_NON_ZERO_CAP: usize = if size_of::<T>() == 1 {
74 8
75 } else if size_of::<T>() <= 1024 {
76 4
77 } else {
78 1
79 };
80
81 pub unsafe fn grow_amortized(&mut self, len: usize, additional: usize) {
82 debug_assert!(additional > 0);
83 if self.capacity == 0 {
84 *self = Self::with_capacity_unchecked(
85 additional.max(Self::MIN_NON_ZERO_CAP),
86 self.align.alignment(align_of::<T>()),
87 );
88 return;
89 }
90
91 if size_of::<T>() == 0 {
92 debug_assert_eq!(self.capacity, usize::MAX);
93 capacity_overflow();
94 }
95
96 let new_cap = match len.checked_add(additional) {
97 Some(cap) => cap,
98 None => capacity_overflow(),
99 };
100
101 // self.cap * 2 can't overflow because it's less than isize::MAX
102 let new_cap = new_cap.max(self.capacity * 2);
103 let new_cap = new_cap.max(Self::MIN_NON_ZERO_CAP);
104
105 let ptr = {
106 grow_unchecked(
107 self.as_mut_ptr() as *mut u8,
108 self.capacity,
109 new_cap,
110 self.align.alignment(align_of::<T>()),
111 size_of::<T>(),
112 ) as *mut T
113 };
114
115 self.capacity = new_cap;
116 self.ptr = NonNull::<T>::new_unchecked(ptr);
117 }
118
119 pub unsafe fn grow_exact(&mut self, len: usize, additional: usize) {
120 debug_assert!(additional > 0);
121 if size_of::<T>() == 0 {
122 debug_assert_eq!(self.capacity, usize::MAX);
123 capacity_overflow();
124 }
125
126 if self.capacity == 0 {
127 *self =
128 Self::with_capacity_unchecked(additional, self.align.alignment(align_of::<T>()));
129 return;
130 }
131
132 let new_cap = match len.checked_add(additional) {
133 Some(cap) => cap,
134 None => capacity_overflow(),
135 };
136
137 let ptr = grow_unchecked(
138 self.as_mut_ptr() as *mut u8,
139 self.capacity,
140 new_cap,
141 self.align.alignment(align_of::<T>()),
142 size_of::<T>(),
143 ) as *mut T;
144
145 self.capacity = new_cap;
146 self.ptr = NonNull::<T>::new_unchecked(ptr);
147 }
148
149 pub unsafe fn shrink_to(&mut self, len: usize) {
150 if size_of::<T>() == 0 {
151 return;
152 }
153
154 debug_assert!(len < self.capacity());
155 let size_of = size_of::<T>();
156 let old_capacity = self.capacity;
157 let align = self.align;
158 let old_ptr = self.ptr.as_ptr() as *mut u8;
159
160 // this cannot overflow or exceed isize::MAX bytes since len < cap and the same was true
161 // for cap
162 let new_size_bytes = len * size_of;
163 let old_size_bytes = old_capacity * size_of;
164 let old_layout =
165 Layout::from_size_align_unchecked(old_size_bytes, align.alignment(align_of::<T>()));
166
167 let ptr = realloc(old_ptr, old_layout, new_size_bytes);
168 let ptr = ptr as *mut T;
169 self.capacity = len;
170 self.ptr = NonNull::<T>::new_unchecked(ptr);
171 }
172
173 #[inline]
174 pub unsafe fn from_raw_parts(ptr: *mut T, capacity: usize, align: usize) -> Self {
175 Self {
176 ptr: NonNull::<T>::new_unchecked(ptr),
177 capacity,
178 align: A::new(align, align_of::<T>()),
179 _marker: PhantomData,
180 }
181 }
182
183 /// Returns the capacity of the vector.
184 #[inline]
185 pub fn capacity(&self) -> usize {
186 self.capacity
187 }
188
189 #[inline]
190 pub fn align(&self) -> usize {
191 self.align.alignment(align_of::<T>())
192 }
193
194 #[inline]
195 pub fn as_ptr(&self) -> *const T {
196 self.ptr.as_ptr()
197 }
198
199 #[inline]
200 pub fn as_mut_ptr(&mut self) -> *mut T {
201 self.ptr.as_ptr()
202 }
203}
204
205pub unsafe fn with_capacity_unchecked(capacity: usize, align: usize, size_of: usize) -> *mut u8 {
206 let size_bytes: usize = match capacity.checked_mul(size_of) {
207 Some(size_bytes: usize) => size_bytes,
208 None => capacity_overflow(),
209 };
210 debug_assert!(size_bytes > 0);
211 let will_overflow: bool = size_bytes > usize::MAX - (align - 1);
212 if will_overflow || !is_valid_alloc(alloc_size:size_bytes) {
213 capacity_overflow();
214 }
215
216 let layout: Layout = Layout::from_size_align_unchecked(size_bytes, align);
217 let ptr: *mut u8 = alloc(layout);
218 if ptr.is_null() {
219 handle_alloc_error(layout);
220 }
221 ptr
222}
223
224unsafe fn grow_unchecked(
225 old_ptr: *mut u8,
226 old_capacity: usize,
227 new_capacity: usize,
228 align: usize,
229 size_of: usize,
230) -> *mut u8 {
231 let new_size_bytes: usize = match new_capacity.checked_mul(size_of) {
232 Some(size_bytes: usize) => size_bytes,
233 None => capacity_overflow(),
234 };
235 let will_overflow: bool = new_size_bytes > usize::MAX - (align - 1);
236 if will_overflow || !is_valid_alloc(alloc_size:new_size_bytes) {
237 capacity_overflow();
238 }
239
240 // can't overflow because we already allocated this much
241 let old_size_bytes: usize = old_capacity * size_of;
242 let old_layout: Layout = Layout::from_size_align_unchecked(size:old_size_bytes, align);
243
244 let ptr: *mut u8 = realloc(old_ptr, old_layout, new_size_bytes);
245
246 if ptr.is_null() {
247 let new_layout: Layout = Layout::from_size_align_unchecked(size:old_size_bytes, align);
248 handle_alloc_error(new_layout);
249 }
250
251 ptr
252}
253
254#[inline]
255fn is_valid_alloc(alloc_size: usize) -> bool {
256 !(usize::BITS < 64 && alloc_size > isize::MAX as usize)
257}
258