1 | #![unstable (feature = "raw_vec_internals" , reason = "unstable const warnings" , issue = "none" )] |
2 | #![cfg_attr (test, allow(dead_code))] |
3 | |
4 | // Note: This module is also included in the alloctests crate using #[path] to |
5 | // run the tests. See the comment there for an explanation why this is the case. |
6 | |
7 | use core::marker::PhantomData; |
8 | use core::mem::{ManuallyDrop, MaybeUninit, SizedTypeProperties}; |
9 | use core::ptr::{self, NonNull, Unique}; |
10 | use core::{cmp, hint}; |
11 | |
12 | #[cfg (not(no_global_oom_handling))] |
13 | use crate::alloc::handle_alloc_error; |
14 | use crate::alloc::{Allocator, Global, Layout}; |
15 | use crate::boxed::Box; |
16 | use crate::collections::TryReserveError; |
17 | use crate::collections::TryReserveErrorKind::*; |
18 | |
19 | #[cfg (test)] |
20 | mod tests; |
21 | |
22 | // One central function responsible for reporting capacity overflows. This'll |
23 | // ensure that the code generation related to these panics is minimal as there's |
24 | // only one location which panics rather than a bunch throughout the module. |
25 | #[cfg (not(no_global_oom_handling))] |
26 | #[cfg_attr (not(feature = "panic_immediate_abort" ), inline(never))] |
27 | #[track_caller ] |
28 | fn capacity_overflow() -> ! { |
29 | panic!("capacity overflow" ); |
30 | } |
31 | |
32 | enum AllocInit { |
33 | /// The contents of the new memory are uninitialized. |
34 | Uninitialized, |
35 | #[cfg (not(no_global_oom_handling))] |
36 | /// The new memory is guaranteed to be zeroed. |
37 | Zeroed, |
38 | } |
39 | |
40 | type Cap = core::num::niche_types::UsizeNoHighBit; |
41 | |
42 | const ZERO_CAP: Cap = unsafe { Cap::new_unchecked(val:0) }; |
43 | |
44 | /// `Cap(cap)`, except if `T` is a ZST then `Cap::ZERO`. |
45 | /// |
46 | /// # Safety: cap must be <= `isize::MAX`. |
47 | unsafe fn new_cap<T>(cap: usize) -> Cap { |
48 | if T::IS_ZST { ZERO_CAP } else { unsafe { Cap::new_unchecked(val:cap) } } |
49 | } |
50 | |
51 | /// A low-level utility for more ergonomically allocating, reallocating, and deallocating |
52 | /// a buffer of memory on the heap without having to worry about all the corner cases |
53 | /// involved. This type is excellent for building your own data structures like Vec and VecDeque. |
54 | /// In particular: |
55 | /// |
56 | /// * Produces `Unique::dangling()` on zero-sized types. |
57 | /// * Produces `Unique::dangling()` on zero-length allocations. |
58 | /// * Avoids freeing `Unique::dangling()`. |
59 | /// * Catches all overflows in capacity computations (promotes them to "capacity overflow" panics). |
60 | /// * Guards against 32-bit systems allocating more than `isize::MAX` bytes. |
61 | /// * Guards against overflowing your length. |
62 | /// * Calls `handle_alloc_error` for fallible allocations. |
63 | /// * Contains a `ptr::Unique` and thus endows the user with all related benefits. |
64 | /// * Uses the excess returned from the allocator to use the largest available capacity. |
65 | /// |
66 | /// This type does not in anyway inspect the memory that it manages. When dropped it *will* |
67 | /// free its memory, but it *won't* try to drop its contents. It is up to the user of `RawVec` |
68 | /// to handle the actual things *stored* inside of a `RawVec`. |
69 | /// |
70 | /// Note that the excess of a zero-sized types is always infinite, so `capacity()` always returns |
71 | /// `usize::MAX`. This means that you need to be careful when round-tripping this type with a |
72 | /// `Box<[T]>`, since `capacity()` won't yield the length. |
73 | #[allow (missing_debug_implementations)] |
74 | pub(crate) struct RawVec<T, A: Allocator = Global> { |
75 | inner: RawVecInner<A>, |
76 | _marker: PhantomData<T>, |
77 | } |
78 | |
79 | /// Like a `RawVec`, but only generic over the allocator, not the type. |
80 | /// |
81 | /// As such, all the methods need the layout passed-in as a parameter. |
82 | /// |
83 | /// Having this separation reduces the amount of code we need to monomorphize, |
84 | /// as most operations don't need the actual type, just its layout. |
85 | #[allow (missing_debug_implementations)] |
86 | struct RawVecInner<A: Allocator = Global> { |
87 | ptr: Unique<u8>, |
88 | /// Never used for ZSTs; it's `capacity()`'s responsibility to return usize::MAX in that case. |
89 | /// |
90 | /// # Safety |
91 | /// |
92 | /// `cap` must be in the `0..=isize::MAX` range. |
93 | cap: Cap, |
94 | alloc: A, |
95 | } |
96 | |
97 | impl<T> RawVec<T, Global> { |
98 | /// Creates the biggest possible `RawVec` (on the system heap) |
99 | /// without allocating. If `T` has positive size, then this makes a |
100 | /// `RawVec` with capacity `0`. If `T` is zero-sized, then it makes a |
101 | /// `RawVec` with capacity `usize::MAX`. Useful for implementing |
102 | /// delayed allocation. |
103 | #[must_use ] |
104 | pub(crate) const fn new() -> Self { |
105 | Self::new_in(Global) |
106 | } |
107 | |
108 | /// Creates a `RawVec` (on the system heap) with exactly the |
109 | /// capacity and alignment requirements for a `[T; capacity]`. This is |
110 | /// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is |
111 | /// zero-sized. Note that if `T` is zero-sized this means you will |
112 | /// *not* get a `RawVec` with the requested capacity. |
113 | /// |
114 | /// Non-fallible version of `try_with_capacity` |
115 | /// |
116 | /// # Panics |
117 | /// |
118 | /// Panics if the requested capacity exceeds `isize::MAX` bytes. |
119 | /// |
120 | /// # Aborts |
121 | /// |
122 | /// Aborts on OOM. |
123 | #[cfg (not(any(no_global_oom_handling, test)))] |
124 | #[must_use ] |
125 | #[inline ] |
126 | #[track_caller ] |
127 | pub(crate) fn with_capacity(capacity: usize) -> Self { |
128 | Self { inner: RawVecInner::with_capacity(capacity, T::LAYOUT), _marker: PhantomData } |
129 | } |
130 | |
131 | /// Like `with_capacity`, but guarantees the buffer is zeroed. |
132 | #[cfg (not(any(no_global_oom_handling, test)))] |
133 | #[must_use ] |
134 | #[inline ] |
135 | #[track_caller ] |
136 | pub(crate) fn with_capacity_zeroed(capacity: usize) -> Self { |
137 | Self { |
138 | inner: RawVecInner::with_capacity_zeroed_in(capacity, Global, T::LAYOUT), |
139 | _marker: PhantomData, |
140 | } |
141 | } |
142 | } |
143 | |
144 | impl RawVecInner<Global> { |
145 | #[cfg (not(any(no_global_oom_handling, test)))] |
146 | #[must_use ] |
147 | #[inline ] |
148 | #[track_caller ] |
149 | fn with_capacity(capacity: usize, elem_layout: Layout) -> Self { |
150 | match Self::try_allocate_in(capacity, init:AllocInit::Uninitialized, alloc:Global, elem_layout) { |
151 | Ok(res: RawVecInner) => res, |
152 | Err(err: TryReserveError) => handle_error(err), |
153 | } |
154 | } |
155 | } |
156 | |
157 | // Tiny Vecs are dumb. Skip to: |
158 | // - 8 if the element size is 1, because any heap allocators is likely |
159 | // to round up a request of less than 8 bytes to at least 8 bytes. |
160 | // - 4 if elements are moderate-sized (<= 1 KiB). |
161 | // - 1 otherwise, to avoid wasting too much space for very short Vecs. |
162 | const fn min_non_zero_cap(size: usize) -> usize { |
163 | if size == 1 { |
164 | 8 |
165 | } else if size <= 1024 { |
166 | 4 |
167 | } else { |
168 | 1 |
169 | } |
170 | } |
171 | |
172 | impl<T, A: Allocator> RawVec<T, A> { |
173 | #[cfg (not(no_global_oom_handling))] |
174 | pub(crate) const MIN_NON_ZERO_CAP: usize = min_non_zero_cap(size_of::<T>()); |
175 | |
176 | /// Like `new`, but parameterized over the choice of allocator for |
177 | /// the returned `RawVec`. |
178 | #[inline ] |
179 | pub(crate) const fn new_in(alloc: A) -> Self { |
180 | Self { inner: RawVecInner::new_in(alloc, align_of::<T>()), _marker: PhantomData } |
181 | } |
182 | |
183 | /// Like `with_capacity`, but parameterized over the choice of |
184 | /// allocator for the returned `RawVec`. |
185 | #[cfg (not(no_global_oom_handling))] |
186 | #[inline ] |
187 | #[track_caller ] |
188 | pub(crate) fn with_capacity_in(capacity: usize, alloc: A) -> Self { |
189 | Self { |
190 | inner: RawVecInner::with_capacity_in(capacity, alloc, T::LAYOUT), |
191 | _marker: PhantomData, |
192 | } |
193 | } |
194 | |
195 | /// Like `try_with_capacity`, but parameterized over the choice of |
196 | /// allocator for the returned `RawVec`. |
197 | #[inline ] |
198 | pub(crate) fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> { |
199 | match RawVecInner::try_with_capacity_in(capacity, alloc, T::LAYOUT) { |
200 | Ok(inner) => Ok(Self { inner, _marker: PhantomData }), |
201 | Err(e) => Err(e), |
202 | } |
203 | } |
204 | |
205 | /// Like `with_capacity_zeroed`, but parameterized over the choice |
206 | /// of allocator for the returned `RawVec`. |
207 | #[cfg (not(no_global_oom_handling))] |
208 | #[inline ] |
209 | #[track_caller ] |
210 | pub(crate) fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self { |
211 | Self { |
212 | inner: RawVecInner::with_capacity_zeroed_in(capacity, alloc, T::LAYOUT), |
213 | _marker: PhantomData, |
214 | } |
215 | } |
216 | |
217 | /// Converts the entire buffer into `Box<[MaybeUninit<T>]>` with the specified `len`. |
218 | /// |
219 | /// Note that this will correctly reconstitute any `cap` changes |
220 | /// that may have been performed. (See description of type for details.) |
221 | /// |
222 | /// # Safety |
223 | /// |
224 | /// * `len` must be greater than or equal to the most recently requested capacity, and |
225 | /// * `len` must be less than or equal to `self.capacity()`. |
226 | /// |
227 | /// Note, that the requested capacity and `self.capacity()` could differ, as |
228 | /// an allocator could overallocate and return a greater memory block than requested. |
229 | pub(crate) unsafe fn into_box(self, len: usize) -> Box<[MaybeUninit<T>], A> { |
230 | // Sanity-check one half of the safety requirement (we cannot check the other half). |
231 | debug_assert!( |
232 | len <= self.capacity(), |
233 | "`len` must be smaller than or equal to `self.capacity()`" |
234 | ); |
235 | |
236 | let me = ManuallyDrop::new(self); |
237 | unsafe { |
238 | let slice = ptr::slice_from_raw_parts_mut(me.ptr() as *mut MaybeUninit<T>, len); |
239 | Box::from_raw_in(slice, ptr::read(&me.inner.alloc)) |
240 | } |
241 | } |
242 | |
243 | /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator. |
244 | /// |
245 | /// # Safety |
246 | /// |
247 | /// The `ptr` must be allocated (via the given allocator `alloc`), and with the given |
248 | /// `capacity`. |
249 | /// The `capacity` cannot exceed `isize::MAX` for sized types. (only a concern on 32-bit |
250 | /// systems). For ZSTs capacity is ignored. |
251 | /// If the `ptr` and `capacity` come from a `RawVec` created via `alloc`, then this is |
252 | /// guaranteed. |
253 | #[inline ] |
254 | pub(crate) unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, alloc: A) -> Self { |
255 | // SAFETY: Precondition passed to the caller |
256 | unsafe { |
257 | let ptr = ptr.cast(); |
258 | let capacity = new_cap::<T>(capacity); |
259 | Self { |
260 | inner: RawVecInner::from_raw_parts_in(ptr, capacity, alloc), |
261 | _marker: PhantomData, |
262 | } |
263 | } |
264 | } |
265 | |
266 | /// A convenience method for hoisting the non-null precondition out of [`RawVec::from_raw_parts_in`]. |
267 | /// |
268 | /// # Safety |
269 | /// |
270 | /// See [`RawVec::from_raw_parts_in`]. |
271 | #[inline ] |
272 | pub(crate) unsafe fn from_nonnull_in(ptr: NonNull<T>, capacity: usize, alloc: A) -> Self { |
273 | // SAFETY: Precondition passed to the caller |
274 | unsafe { |
275 | let ptr = ptr.cast(); |
276 | let capacity = new_cap::<T>(capacity); |
277 | Self { inner: RawVecInner::from_nonnull_in(ptr, capacity, alloc), _marker: PhantomData } |
278 | } |
279 | } |
280 | |
281 | /// Gets a raw pointer to the start of the allocation. Note that this is |
282 | /// `Unique::dangling()` if `capacity == 0` or `T` is zero-sized. In the former case, you must |
283 | /// be careful. |
284 | #[inline ] |
285 | pub(crate) const fn ptr(&self) -> *mut T { |
286 | self.inner.ptr() |
287 | } |
288 | |
289 | #[inline ] |
290 | pub(crate) fn non_null(&self) -> NonNull<T> { |
291 | self.inner.non_null() |
292 | } |
293 | |
294 | /// Gets the capacity of the allocation. |
295 | /// |
296 | /// This will always be `usize::MAX` if `T` is zero-sized. |
297 | #[inline ] |
298 | pub(crate) const fn capacity(&self) -> usize { |
299 | self.inner.capacity(size_of::<T>()) |
300 | } |
301 | |
302 | /// Returns a shared reference to the allocator backing this `RawVec`. |
303 | #[inline ] |
304 | pub(crate) fn allocator(&self) -> &A { |
305 | self.inner.allocator() |
306 | } |
307 | |
308 | /// Ensures that the buffer contains at least enough space to hold `len + |
309 | /// additional` elements. If it doesn't already have enough capacity, will |
310 | /// reallocate enough space plus comfortable slack space to get amortized |
311 | /// *O*(1) behavior. Will limit this behavior if it would needlessly cause |
312 | /// itself to panic. |
313 | /// |
314 | /// If `len` exceeds `self.capacity()`, this may fail to actually allocate |
315 | /// the requested space. This is not really unsafe, but the unsafe |
316 | /// code *you* write that relies on the behavior of this function may break. |
317 | /// |
318 | /// This is ideal for implementing a bulk-push operation like `extend`. |
319 | /// |
320 | /// # Panics |
321 | /// |
322 | /// Panics if the new capacity exceeds `isize::MAX` _bytes_. |
323 | /// |
324 | /// # Aborts |
325 | /// |
326 | /// Aborts on OOM. |
327 | #[cfg (not(no_global_oom_handling))] |
328 | #[inline ] |
329 | #[track_caller ] |
330 | pub(crate) fn reserve(&mut self, len: usize, additional: usize) { |
331 | self.inner.reserve(len, additional, T::LAYOUT) |
332 | } |
333 | |
334 | /// A specialized version of `self.reserve(len, 1)` which requires the |
335 | /// caller to ensure `len == self.capacity()`. |
336 | #[cfg (not(no_global_oom_handling))] |
337 | #[inline (never)] |
338 | #[track_caller ] |
339 | pub(crate) fn grow_one(&mut self) { |
340 | self.inner.grow_one(T::LAYOUT) |
341 | } |
342 | |
343 | /// The same as `reserve`, but returns on errors instead of panicking or aborting. |
344 | pub(crate) fn try_reserve( |
345 | &mut self, |
346 | len: usize, |
347 | additional: usize, |
348 | ) -> Result<(), TryReserveError> { |
349 | self.inner.try_reserve(len, additional, T::LAYOUT) |
350 | } |
351 | |
352 | /// Ensures that the buffer contains at least enough space to hold `len + |
353 | /// additional` elements. If it doesn't already, will reallocate the |
354 | /// minimum possible amount of memory necessary. Generally this will be |
355 | /// exactly the amount of memory necessary, but in principle the allocator |
356 | /// is free to give back more than we asked for. |
357 | /// |
358 | /// If `len` exceeds `self.capacity()`, this may fail to actually allocate |
359 | /// the requested space. This is not really unsafe, but the unsafe code |
360 | /// *you* write that relies on the behavior of this function may break. |
361 | /// |
362 | /// # Panics |
363 | /// |
364 | /// Panics if the new capacity exceeds `isize::MAX` _bytes_. |
365 | /// |
366 | /// # Aborts |
367 | /// |
368 | /// Aborts on OOM. |
369 | #[cfg (not(no_global_oom_handling))] |
370 | #[track_caller ] |
371 | pub(crate) fn reserve_exact(&mut self, len: usize, additional: usize) { |
372 | self.inner.reserve_exact(len, additional, T::LAYOUT) |
373 | } |
374 | |
375 | /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting. |
376 | pub(crate) fn try_reserve_exact( |
377 | &mut self, |
378 | len: usize, |
379 | additional: usize, |
380 | ) -> Result<(), TryReserveError> { |
381 | self.inner.try_reserve_exact(len, additional, T::LAYOUT) |
382 | } |
383 | |
384 | /// Shrinks the buffer down to the specified capacity. If the given amount |
385 | /// is 0, actually completely deallocates. |
386 | /// |
387 | /// # Panics |
388 | /// |
389 | /// Panics if the given amount is *larger* than the current capacity. |
390 | /// |
391 | /// # Aborts |
392 | /// |
393 | /// Aborts on OOM. |
394 | #[cfg (not(no_global_oom_handling))] |
395 | #[track_caller ] |
396 | #[inline ] |
397 | pub(crate) fn shrink_to_fit(&mut self, cap: usize) { |
398 | self.inner.shrink_to_fit(cap, T::LAYOUT) |
399 | } |
400 | } |
401 | |
402 | unsafe impl<#[may_dangle ] T, A: Allocator> Drop for RawVec<T, A> { |
403 | /// Frees the memory owned by the `RawVec` *without* trying to drop its contents. |
404 | fn drop(&mut self) { |
405 | // SAFETY: We are in a Drop impl, self.inner will not be used again. |
406 | unsafe { self.inner.deallocate(T::LAYOUT) } |
407 | } |
408 | } |
409 | |
410 | impl<A: Allocator> RawVecInner<A> { |
411 | #[inline ] |
412 | const fn new_in(alloc: A, align: usize) -> Self { |
413 | let ptr = unsafe { core::mem::transmute(align) }; |
414 | // `cap: 0` means "unallocated". zero-sized types are ignored. |
415 | Self { ptr, cap: ZERO_CAP, alloc } |
416 | } |
417 | |
418 | #[cfg (not(no_global_oom_handling))] |
419 | #[inline ] |
420 | #[track_caller ] |
421 | fn with_capacity_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self { |
422 | match Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout) { |
423 | Ok(this) => { |
424 | unsafe { |
425 | // Make it more obvious that a subsequent Vec::reserve(capacity) will not allocate. |
426 | hint::assert_unchecked(!this.needs_to_grow(0, capacity, elem_layout)); |
427 | } |
428 | this |
429 | } |
430 | Err(err) => handle_error(err), |
431 | } |
432 | } |
433 | |
434 | #[inline ] |
435 | fn try_with_capacity_in( |
436 | capacity: usize, |
437 | alloc: A, |
438 | elem_layout: Layout, |
439 | ) -> Result<Self, TryReserveError> { |
440 | Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc, elem_layout) |
441 | } |
442 | |
443 | #[cfg (not(no_global_oom_handling))] |
444 | #[inline ] |
445 | #[track_caller ] |
446 | fn with_capacity_zeroed_in(capacity: usize, alloc: A, elem_layout: Layout) -> Self { |
447 | match Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc, elem_layout) { |
448 | Ok(res) => res, |
449 | Err(err) => handle_error(err), |
450 | } |
451 | } |
452 | |
453 | fn try_allocate_in( |
454 | capacity: usize, |
455 | init: AllocInit, |
456 | alloc: A, |
457 | elem_layout: Layout, |
458 | ) -> Result<Self, TryReserveError> { |
459 | // We avoid `unwrap_or_else` here because it bloats the amount of |
460 | // LLVM IR generated. |
461 | let layout = match layout_array(capacity, elem_layout) { |
462 | Ok(layout) => layout, |
463 | Err(_) => return Err(CapacityOverflow.into()), |
464 | }; |
465 | |
466 | // Don't allocate here because `Drop` will not deallocate when `capacity` is 0. |
467 | if layout.size() == 0 { |
468 | return Ok(Self::new_in(alloc, elem_layout.align())); |
469 | } |
470 | |
471 | if let Err(err) = alloc_guard(layout.size()) { |
472 | return Err(err); |
473 | } |
474 | |
475 | let result = match init { |
476 | AllocInit::Uninitialized => alloc.allocate(layout), |
477 | #[cfg (not(no_global_oom_handling))] |
478 | AllocInit::Zeroed => alloc.allocate_zeroed(layout), |
479 | }; |
480 | let ptr = match result { |
481 | Ok(ptr) => ptr, |
482 | Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()), |
483 | }; |
484 | |
485 | // Allocators currently return a `NonNull<[u8]>` whose length |
486 | // matches the size requested. If that ever changes, the capacity |
487 | // here should change to `ptr.len() / size_of::<T>()`. |
488 | Ok(Self { |
489 | ptr: Unique::from(ptr.cast()), |
490 | cap: unsafe { Cap::new_unchecked(capacity) }, |
491 | alloc, |
492 | }) |
493 | } |
494 | |
495 | #[inline ] |
496 | unsafe fn from_raw_parts_in(ptr: *mut u8, cap: Cap, alloc: A) -> Self { |
497 | Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc } |
498 | } |
499 | |
500 | #[inline ] |
501 | unsafe fn from_nonnull_in(ptr: NonNull<u8>, cap: Cap, alloc: A) -> Self { |
502 | Self { ptr: Unique::from(ptr), cap, alloc } |
503 | } |
504 | |
505 | #[inline ] |
506 | const fn ptr<T>(&self) -> *mut T { |
507 | self.non_null::<T>().as_ptr() |
508 | } |
509 | |
510 | #[inline ] |
511 | const fn non_null<T>(&self) -> NonNull<T> { |
512 | self.ptr.cast().as_non_null_ptr() |
513 | } |
514 | |
515 | #[inline ] |
516 | const fn capacity(&self, elem_size: usize) -> usize { |
517 | if elem_size == 0 { usize::MAX } else { self.cap.as_inner() } |
518 | } |
519 | |
520 | #[inline ] |
521 | fn allocator(&self) -> &A { |
522 | &self.alloc |
523 | } |
524 | |
525 | #[inline ] |
526 | fn current_memory(&self, elem_layout: Layout) -> Option<(NonNull<u8>, Layout)> { |
527 | if elem_layout.size() == 0 || self.cap.as_inner() == 0 { |
528 | None |
529 | } else { |
530 | // We could use Layout::array here which ensures the absence of isize and usize overflows |
531 | // and could hypothetically handle differences between stride and size, but this memory |
532 | // has already been allocated so we know it can't overflow and currently Rust does not |
533 | // support such types. So we can do better by skipping some checks and avoid an unwrap. |
534 | unsafe { |
535 | let alloc_size = elem_layout.size().unchecked_mul(self.cap.as_inner()); |
536 | let layout = Layout::from_size_align_unchecked(alloc_size, elem_layout.align()); |
537 | Some((self.ptr.into(), layout)) |
538 | } |
539 | } |
540 | } |
541 | |
542 | #[cfg (not(no_global_oom_handling))] |
543 | #[inline ] |
544 | #[track_caller ] |
545 | fn reserve(&mut self, len: usize, additional: usize, elem_layout: Layout) { |
546 | // Callers expect this function to be very cheap when there is already sufficient capacity. |
547 | // Therefore, we move all the resizing and error-handling logic from grow_amortized and |
548 | // handle_reserve behind a call, while making sure that this function is likely to be |
549 | // inlined as just a comparison and a call if the comparison fails. |
550 | #[cold ] |
551 | fn do_reserve_and_handle<A: Allocator>( |
552 | slf: &mut RawVecInner<A>, |
553 | len: usize, |
554 | additional: usize, |
555 | elem_layout: Layout, |
556 | ) { |
557 | if let Err(err) = slf.grow_amortized(len, additional, elem_layout) { |
558 | handle_error(err); |
559 | } |
560 | } |
561 | |
562 | if self.needs_to_grow(len, additional, elem_layout) { |
563 | do_reserve_and_handle(self, len, additional, elem_layout); |
564 | } |
565 | } |
566 | |
567 | #[cfg (not(no_global_oom_handling))] |
568 | #[inline ] |
569 | #[track_caller ] |
570 | fn grow_one(&mut self, elem_layout: Layout) { |
571 | if let Err(err) = self.grow_amortized(self.cap.as_inner(), 1, elem_layout) { |
572 | handle_error(err); |
573 | } |
574 | } |
575 | |
576 | fn try_reserve( |
577 | &mut self, |
578 | len: usize, |
579 | additional: usize, |
580 | elem_layout: Layout, |
581 | ) -> Result<(), TryReserveError> { |
582 | if self.needs_to_grow(len, additional, elem_layout) { |
583 | self.grow_amortized(len, additional, elem_layout)?; |
584 | } |
585 | unsafe { |
586 | // Inform the optimizer that the reservation has succeeded or wasn't needed |
587 | hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout)); |
588 | } |
589 | Ok(()) |
590 | } |
591 | |
592 | #[cfg (not(no_global_oom_handling))] |
593 | #[track_caller ] |
594 | fn reserve_exact(&mut self, len: usize, additional: usize, elem_layout: Layout) { |
595 | if let Err(err) = self.try_reserve_exact(len, additional, elem_layout) { |
596 | handle_error(err); |
597 | } |
598 | } |
599 | |
600 | fn try_reserve_exact( |
601 | &mut self, |
602 | len: usize, |
603 | additional: usize, |
604 | elem_layout: Layout, |
605 | ) -> Result<(), TryReserveError> { |
606 | if self.needs_to_grow(len, additional, elem_layout) { |
607 | self.grow_exact(len, additional, elem_layout)?; |
608 | } |
609 | unsafe { |
610 | // Inform the optimizer that the reservation has succeeded or wasn't needed |
611 | hint::assert_unchecked(!self.needs_to_grow(len, additional, elem_layout)); |
612 | } |
613 | Ok(()) |
614 | } |
615 | |
616 | #[cfg (not(no_global_oom_handling))] |
617 | #[inline ] |
618 | #[track_caller ] |
619 | fn shrink_to_fit(&mut self, cap: usize, elem_layout: Layout) { |
620 | if let Err(err) = self.shrink(cap, elem_layout) { |
621 | handle_error(err); |
622 | } |
623 | } |
624 | |
625 | #[inline ] |
626 | fn needs_to_grow(&self, len: usize, additional: usize, elem_layout: Layout) -> bool { |
627 | additional > self.capacity(elem_layout.size()).wrapping_sub(len) |
628 | } |
629 | |
630 | #[inline ] |
631 | unsafe fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) { |
632 | // Allocators currently return a `NonNull<[u8]>` whose length matches |
633 | // the size requested. If that ever changes, the capacity here should |
634 | // change to `ptr.len() / size_of::<T>()`. |
635 | self.ptr = Unique::from(ptr.cast()); |
636 | self.cap = unsafe { Cap::new_unchecked(cap) }; |
637 | } |
638 | |
639 | fn grow_amortized( |
640 | &mut self, |
641 | len: usize, |
642 | additional: usize, |
643 | elem_layout: Layout, |
644 | ) -> Result<(), TryReserveError> { |
645 | // This is ensured by the calling contexts. |
646 | debug_assert!(additional > 0); |
647 | |
648 | if elem_layout.size() == 0 { |
649 | // Since we return a capacity of `usize::MAX` when `elem_size` is |
650 | // 0, getting to here necessarily means the `RawVec` is overfull. |
651 | return Err(CapacityOverflow.into()); |
652 | } |
653 | |
654 | // Nothing we can really do about these checks, sadly. |
655 | let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?; |
656 | |
657 | // This guarantees exponential growth. The doubling cannot overflow |
658 | // because `cap <= isize::MAX` and the type of `cap` is `usize`. |
659 | let cap = cmp::max(self.cap.as_inner() * 2, required_cap); |
660 | let cap = cmp::max(min_non_zero_cap(elem_layout.size()), cap); |
661 | |
662 | let new_layout = layout_array(cap, elem_layout)?; |
663 | |
664 | let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?; |
665 | // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than `isize::MAX` items |
666 | |
667 | unsafe { self.set_ptr_and_cap(ptr, cap) }; |
668 | Ok(()) |
669 | } |
670 | |
671 | fn grow_exact( |
672 | &mut self, |
673 | len: usize, |
674 | additional: usize, |
675 | elem_layout: Layout, |
676 | ) -> Result<(), TryReserveError> { |
677 | if elem_layout.size() == 0 { |
678 | // Since we return a capacity of `usize::MAX` when the type size is |
679 | // 0, getting to here necessarily means the `RawVec` is overfull. |
680 | return Err(CapacityOverflow.into()); |
681 | } |
682 | |
683 | let cap = len.checked_add(additional).ok_or(CapacityOverflow)?; |
684 | let new_layout = layout_array(cap, elem_layout)?; |
685 | |
686 | let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?; |
687 | // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than `isize::MAX` items |
688 | unsafe { |
689 | self.set_ptr_and_cap(ptr, cap); |
690 | } |
691 | Ok(()) |
692 | } |
693 | |
694 | #[cfg (not(no_global_oom_handling))] |
695 | #[inline ] |
696 | fn shrink(&mut self, cap: usize, elem_layout: Layout) -> Result<(), TryReserveError> { |
697 | assert!(cap <= self.capacity(elem_layout.size()), "Tried to shrink to a larger capacity" ); |
698 | // SAFETY: Just checked this isn't trying to grow |
699 | unsafe { self.shrink_unchecked(cap, elem_layout) } |
700 | } |
701 | |
702 | /// `shrink`, but without the capacity check. |
703 | /// |
704 | /// This is split out so that `shrink` can inline the check, since it |
705 | /// optimizes out in things like `shrink_to_fit`, without needing to |
706 | /// also inline all this code, as doing that ends up failing the |
707 | /// `vec-shrink-panic` codegen test when `shrink_to_fit` ends up being too |
708 | /// big for LLVM to be willing to inline. |
709 | /// |
710 | /// # Safety |
711 | /// `cap <= self.capacity()` |
712 | #[cfg (not(no_global_oom_handling))] |
713 | unsafe fn shrink_unchecked( |
714 | &mut self, |
715 | cap: usize, |
716 | elem_layout: Layout, |
717 | ) -> Result<(), TryReserveError> { |
718 | let (ptr, layout) = |
719 | if let Some(mem) = self.current_memory(elem_layout) { mem } else { return Ok(()) }; |
720 | |
721 | // If shrinking to 0, deallocate the buffer. We don't reach this point |
722 | // for the T::IS_ZST case since current_memory() will have returned |
723 | // None. |
724 | if cap == 0 { |
725 | unsafe { self.alloc.deallocate(ptr, layout) }; |
726 | self.ptr = |
727 | unsafe { Unique::new_unchecked(ptr::without_provenance_mut(elem_layout.align())) }; |
728 | self.cap = ZERO_CAP; |
729 | } else { |
730 | let ptr = unsafe { |
731 | // Layout cannot overflow here because it would have |
732 | // overflowed earlier when capacity was larger. |
733 | let new_size = elem_layout.size().unchecked_mul(cap); |
734 | let new_layout = Layout::from_size_align_unchecked(new_size, layout.align()); |
735 | self.alloc |
736 | .shrink(ptr, layout, new_layout) |
737 | .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })? |
738 | }; |
739 | // SAFETY: if the allocation is valid, then the capacity is too |
740 | unsafe { |
741 | self.set_ptr_and_cap(ptr, cap); |
742 | } |
743 | } |
744 | Ok(()) |
745 | } |
746 | |
747 | /// # Safety |
748 | /// |
749 | /// This function deallocates the owned allocation, but does not update `ptr` or `cap` to |
750 | /// prevent double-free or use-after-free. Essentially, do not do anything with the caller |
751 | /// after this function returns. |
752 | /// Ideally this function would take `self` by move, but it cannot because it exists to be |
753 | /// called from a `Drop` impl. |
754 | unsafe fn deallocate(&mut self, elem_layout: Layout) { |
755 | if let Some((ptr, layout)) = self.current_memory(elem_layout) { |
756 | unsafe { |
757 | self.alloc.deallocate(ptr, layout); |
758 | } |
759 | } |
760 | } |
761 | } |
762 | |
763 | // not marked inline(never) since we want optimizers to be able to observe the specifics of this |
764 | // function, see tests/codegen/vec-reserve-extend.rs. |
765 | #[cold ] |
766 | fn finish_grow<A>( |
767 | new_layout: Layout, |
768 | current_memory: Option<(NonNull<u8>, Layout)>, |
769 | alloc: &mut A, |
770 | ) -> Result<NonNull<[u8]>, TryReserveError> |
771 | where |
772 | A: Allocator, |
773 | { |
774 | alloc_guard(alloc_size:new_layout.size())?; |
775 | |
776 | let memory: Result, AllocError> = if let Some((ptr: NonNull, old_layout: Layout)) = current_memory { |
777 | debug_assert_eq!(old_layout.align(), new_layout.align()); |
778 | unsafe { |
779 | // The allocator checks for alignment equality |
780 | hint::assert_unchecked(cond:old_layout.align() == new_layout.align()); |
781 | alloc.grow(ptr, old_layout, new_layout) |
782 | } |
783 | } else { |
784 | alloc.allocate(new_layout) |
785 | }; |
786 | |
787 | memory.map_err(|_| AllocError { layout: new_layout, non_exhaustive: () }.into()) |
788 | } |
789 | |
790 | // Central function for reserve error handling. |
791 | #[cfg (not(no_global_oom_handling))] |
792 | #[cold ] |
793 | #[optimize (size)] |
794 | #[track_caller ] |
795 | fn handle_error(e: TryReserveError) -> ! { |
796 | match e.kind() { |
797 | CapacityOverflow => capacity_overflow(), |
798 | AllocError { layout: Layout, .. } => handle_alloc_error(layout), |
799 | } |
800 | } |
801 | |
802 | // We need to guarantee the following: |
803 | // * We don't ever allocate `> isize::MAX` byte-size objects. |
804 | // * We don't overflow `usize::MAX` and actually allocate too little. |
805 | // |
806 | // On 64-bit we just need to check for overflow since trying to allocate |
807 | // `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add |
808 | // an extra guard for this in case we're running on a platform which can use |
809 | // all 4GB in user-space, e.g., PAE or x32. |
810 | #[inline ] |
811 | fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> { |
812 | if usize::BITS < 64 && alloc_size > isize::MAX as usize { |
813 | Err(CapacityOverflow.into()) |
814 | } else { |
815 | Ok(()) |
816 | } |
817 | } |
818 | |
819 | #[inline ] |
820 | fn layout_array(cap: usize, elem_layout: Layout) -> Result<Layout, TryReserveError> { |
821 | elem_layout.repeat(cap).map(|(layout, _pad)| layout).map_err(|_| CapacityOverflow.into()) |
822 | } |
823 | |