GNU Linux-libre 6.8.7-gnu
[releases.git] / rust / alloc / raw_vec.rs
1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2
3 #![unstable(feature = "raw_vec_internals", reason = "unstable const warnings", issue = "none")]
4
5 use core::alloc::LayoutError;
6 use core::cmp;
7 use core::intrinsics;
8 use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
9 use core::ptr::{self, NonNull, Unique};
10 use core::slice;
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 enum AllocInit {
23     /// The contents of the new memory are uninitialized.
24     Uninitialized,
25     /// The new memory is guaranteed to be zeroed.
26     #[allow(dead_code)]
27     Zeroed,
28 }
29
30 /// A low-level utility for more ergonomically allocating, reallocating, and deallocating
31 /// a buffer of memory on the heap without having to worry about all the corner cases
32 /// involved. This type is excellent for building your own data structures like Vec and VecDeque.
33 /// In particular:
34 ///
35 /// * Produces `Unique::dangling()` on zero-sized types.
36 /// * Produces `Unique::dangling()` on zero-length allocations.
37 /// * Avoids freeing `Unique::dangling()`.
38 /// * Catches all overflows in capacity computations (promotes them to "capacity overflow" panics).
39 /// * Guards against 32-bit systems allocating more than isize::MAX bytes.
40 /// * Guards against overflowing your length.
41 /// * Calls `handle_alloc_error` for fallible allocations.
42 /// * Contains a `ptr::Unique` and thus endows the user with all related benefits.
43 /// * Uses the excess returned from the allocator to use the largest available capacity.
44 ///
45 /// This type does not in anyway inspect the memory that it manages. When dropped it *will*
46 /// free its memory, but it *won't* try to drop its contents. It is up to the user of `RawVec`
47 /// to handle the actual things *stored* inside of a `RawVec`.
48 ///
49 /// Note that the excess of a zero-sized types is always infinite, so `capacity()` always returns
50 /// `usize::MAX`. This means that you need to be careful when round-tripping this type with a
51 /// `Box<[T]>`, since `capacity()` won't yield the length.
52 #[allow(missing_debug_implementations)]
53 pub(crate) struct RawVec<T, A: Allocator = Global> {
54     ptr: Unique<T>,
55     cap: usize,
56     alloc: A,
57 }
58
59 impl<T> RawVec<T, Global> {
60     /// HACK(Centril): This exists because stable `const fn` can only call stable `const fn`, so
61     /// they cannot call `Self::new()`.
62     ///
63     /// If you change `RawVec<T>::new` or dependencies, please take care to not introduce anything
64     /// that would truly const-call something unstable.
65     pub const NEW: Self = Self::new();
66
67     /// Creates the biggest possible `RawVec` (on the system heap)
68     /// without allocating. If `T` has positive size, then this makes a
69     /// `RawVec` with capacity `0`. If `T` is zero-sized, then it makes a
70     /// `RawVec` with capacity `usize::MAX`. Useful for implementing
71     /// delayed allocation.
72     #[must_use]
73     pub const fn new() -> Self {
74         Self::new_in(Global)
75     }
76
77     /// Creates a `RawVec` (on the system heap) with exactly the
78     /// capacity and alignment requirements for a `[T; capacity]`. This is
79     /// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is
80     /// zero-sized. Note that if `T` is zero-sized this means you will
81     /// *not* get a `RawVec` with the requested capacity.
82     ///
83     /// # Panics
84     ///
85     /// Panics if the requested capacity exceeds `isize::MAX` bytes.
86     ///
87     /// # Aborts
88     ///
89     /// Aborts on OOM.
90     #[cfg(not(any(no_global_oom_handling, test)))]
91     #[must_use]
92     #[inline]
93     pub fn with_capacity(capacity: usize) -> Self {
94         Self::with_capacity_in(capacity, Global)
95     }
96
97     /// Like `with_capacity`, but guarantees the buffer is zeroed.
98     #[cfg(not(any(no_global_oom_handling, test)))]
99     #[must_use]
100     #[inline]
101     pub fn with_capacity_zeroed(capacity: usize) -> Self {
102         Self::with_capacity_zeroed_in(capacity, Global)
103     }
104 }
105
106 impl<T, A: Allocator> RawVec<T, A> {
107     // Tiny Vecs are dumb. Skip to:
108     // - 8 if the element size is 1, because any heap allocators is likely
109     //   to round up a request of less than 8 bytes to at least 8 bytes.
110     // - 4 if elements are moderate-sized (<= 1 KiB).
111     // - 1 otherwise, to avoid wasting too much space for very short Vecs.
112     pub(crate) const MIN_NON_ZERO_CAP: usize = if mem::size_of::<T>() == 1 {
113         8
114     } else if mem::size_of::<T>() <= 1024 {
115         4
116     } else {
117         1
118     };
119
120     /// Like `new`, but parameterized over the choice of allocator for
121     /// the returned `RawVec`.
122     pub const fn new_in(alloc: A) -> Self {
123         // `cap: 0` means "unallocated". zero-sized types are ignored.
124         Self { ptr: Unique::dangling(), cap: 0, alloc }
125     }
126
127     /// Like `with_capacity`, but parameterized over the choice of
128     /// allocator for the returned `RawVec`.
129     #[cfg(not(no_global_oom_handling))]
130     #[inline]
131     pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
132         Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
133     }
134
135     /// Like `try_with_capacity`, but parameterized over the choice of
136     /// allocator for the returned `RawVec`.
137     #[inline]
138     pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {
139         Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc)
140     }
141
142     /// Like `with_capacity_zeroed`, but parameterized over the choice
143     /// of allocator for the returned `RawVec`.
144     #[cfg(not(no_global_oom_handling))]
145     #[inline]
146     pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
147         Self::allocate_in(capacity, AllocInit::Zeroed, alloc)
148     }
149
150     /// Converts the entire buffer into `Box<[MaybeUninit<T>]>` with the specified `len`.
151     ///
152     /// Note that this will correctly reconstitute any `cap` changes
153     /// that may have been performed. (See description of type for details.)
154     ///
155     /// # Safety
156     ///
157     /// * `len` must be greater than or equal to the most recently requested capacity, and
158     /// * `len` must be less than or equal to `self.capacity()`.
159     ///
160     /// Note, that the requested capacity and `self.capacity()` could differ, as
161     /// an allocator could overallocate and return a greater memory block than requested.
162     pub unsafe fn into_box(self, len: usize) -> Box<[MaybeUninit<T>], A> {
163         // Sanity-check one half of the safety requirement (we cannot check the other half).
164         debug_assert!(
165             len <= self.capacity(),
166             "`len` must be smaller than or equal to `self.capacity()`"
167         );
168
169         let me = ManuallyDrop::new(self);
170         unsafe {
171             let slice = slice::from_raw_parts_mut(me.ptr() as *mut MaybeUninit<T>, len);
172             Box::from_raw_in(slice, ptr::read(&me.alloc))
173         }
174     }
175
176     #[cfg(not(no_global_oom_handling))]
177     fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
178         // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.
179         if T::IS_ZST || capacity == 0 {
180             Self::new_in(alloc)
181         } else {
182             // We avoid `unwrap_or_else` here because it bloats the amount of
183             // LLVM IR generated.
184             let layout = match Layout::array::<T>(capacity) {
185                 Ok(layout) => layout,
186                 Err(_) => capacity_overflow(),
187             };
188             match alloc_guard(layout.size()) {
189                 Ok(_) => {}
190                 Err(_) => capacity_overflow(),
191             }
192             let result = match init {
193                 AllocInit::Uninitialized => alloc.allocate(layout),
194                 AllocInit::Zeroed => alloc.allocate_zeroed(layout),
195             };
196             let ptr = match result {
197                 Ok(ptr) => ptr,
198                 Err(_) => handle_alloc_error(layout),
199             };
200
201             // Allocators currently return a `NonNull<[u8]>` whose length
202             // matches the size requested. If that ever changes, the capacity
203             // here should change to `ptr.len() / mem::size_of::<T>()`.
204             Self {
205                 ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
206                 cap: capacity,
207                 alloc,
208             }
209         }
210     }
211
212     fn try_allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Result<Self, TryReserveError> {
213         // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.
214         if T::IS_ZST || capacity == 0 {
215             return Ok(Self::new_in(alloc));
216         }
217
218         let layout = Layout::array::<T>(capacity).map_err(|_| CapacityOverflow)?;
219         alloc_guard(layout.size())?;
220         let result = match init {
221             AllocInit::Uninitialized => alloc.allocate(layout),
222             AllocInit::Zeroed => alloc.allocate_zeroed(layout),
223         };
224         let ptr = result.map_err(|_| AllocError { layout, non_exhaustive: () })?;
225
226         // Allocators currently return a `NonNull<[u8]>` whose length
227         // matches the size requested. If that ever changes, the capacity
228         // here should change to `ptr.len() / mem::size_of::<T>()`.
229         Ok(Self {
230             ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
231             cap: capacity,
232             alloc,
233         })
234     }
235
236     /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator.
237     ///
238     /// # Safety
239     ///
240     /// The `ptr` must be allocated (via the given allocator `alloc`), and with the given
241     /// `capacity`.
242     /// The `capacity` cannot exceed `isize::MAX` for sized types. (only a concern on 32-bit
243     /// systems). ZST vectors may have a capacity up to `usize::MAX`.
244     /// If the `ptr` and `capacity` come from a `RawVec` created via `alloc`, then this is
245     /// guaranteed.
246     #[inline]
247     pub unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, alloc: A) -> Self {
248         Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap: capacity, alloc }
249     }
250
251     /// Gets a raw pointer to the start of the allocation. Note that this is
252     /// `Unique::dangling()` if `capacity == 0` or `T` is zero-sized. In the former case, you must
253     /// be careful.
254     #[inline]
255     pub fn ptr(&self) -> *mut T {
256         self.ptr.as_ptr()
257     }
258
259     /// Gets the capacity of the allocation.
260     ///
261     /// This will always be `usize::MAX` if `T` is zero-sized.
262     #[inline(always)]
263     pub fn capacity(&self) -> usize {
264         if T::IS_ZST { usize::MAX } else { self.cap }
265     }
266
267     /// Returns a shared reference to the allocator backing this `RawVec`.
268     pub fn allocator(&self) -> &A {
269         &self.alloc
270     }
271
272     fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
273         if T::IS_ZST || self.cap == 0 {
274             None
275         } else {
276             // We could use Layout::array here which ensures the absence of isize and usize overflows
277             // and could hypothetically handle differences between stride and size, but this memory
278             // has already been allocated so we know it can't overflow and currently rust does not
279             // support such types. So we can do better by skipping some checks and avoid an unwrap.
280             let _: () = const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
281             unsafe {
282                 let align = mem::align_of::<T>();
283                 let size = mem::size_of::<T>().unchecked_mul(self.cap);
284                 let layout = Layout::from_size_align_unchecked(size, align);
285                 Some((self.ptr.cast().into(), layout))
286             }
287         }
288     }
289
290     /// Ensures that the buffer contains at least enough space to hold `len +
291     /// additional` elements. If it doesn't already have enough capacity, will
292     /// reallocate enough space plus comfortable slack space to get amortized
293     /// *O*(1) behavior. Will limit this behavior if it would needlessly cause
294     /// itself to panic.
295     ///
296     /// If `len` exceeds `self.capacity()`, this may fail to actually allocate
297     /// the requested space. This is not really unsafe, but the unsafe
298     /// code *you* write that relies on the behavior of this function may break.
299     ///
300     /// This is ideal for implementing a bulk-push operation like `extend`.
301     ///
302     /// # Panics
303     ///
304     /// Panics if the new capacity exceeds `isize::MAX` bytes.
305     ///
306     /// # Aborts
307     ///
308     /// Aborts on OOM.
309     #[cfg(not(no_global_oom_handling))]
310     #[inline]
311     pub fn reserve(&mut self, len: usize, additional: usize) {
312         // Callers expect this function to be very cheap when there is already sufficient capacity.
313         // Therefore, we move all the resizing and error-handling logic from grow_amortized and
314         // handle_reserve behind a call, while making sure that this function is likely to be
315         // inlined as just a comparison and a call if the comparison fails.
316         #[cold]
317         fn do_reserve_and_handle<T, A: Allocator>(
318             slf: &mut RawVec<T, A>,
319             len: usize,
320             additional: usize,
321         ) {
322             handle_reserve(slf.grow_amortized(len, additional));
323         }
324
325         if self.needs_to_grow(len, additional) {
326             do_reserve_and_handle(self, len, additional);
327         }
328     }
329
330     /// A specialized version of `reserve()` used only by the hot and
331     /// oft-instantiated `Vec::push()`, which does its own capacity check.
332     #[cfg(not(no_global_oom_handling))]
333     #[inline(never)]
334     pub fn reserve_for_push(&mut self, len: usize) {
335         handle_reserve(self.grow_amortized(len, 1));
336     }
337
338     /// The same as `reserve`, but returns on errors instead of panicking or aborting.
339     pub fn try_reserve(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
340         if self.needs_to_grow(len, additional) {
341             self.grow_amortized(len, additional)
342         } else {
343             Ok(())
344         }
345     }
346
347     /// The same as `reserve_for_push`, but returns on errors instead of panicking or aborting.
348     #[inline(never)]
349     pub fn try_reserve_for_push(&mut self, len: usize) -> Result<(), TryReserveError> {
350         self.grow_amortized(len, 1)
351     }
352
353     /// Ensures that the buffer contains at least enough space to hold `len +
354     /// additional` elements. If it doesn't already, will reallocate the
355     /// minimum possible amount of memory necessary. Generally this will be
356     /// exactly the amount of memory necessary, but in principle the allocator
357     /// is free to give back more than we asked for.
358     ///
359     /// If `len` exceeds `self.capacity()`, this may fail to actually allocate
360     /// the requested space. This is not really unsafe, but the unsafe code
361     /// *you* write that relies on the behavior of this function may break.
362     ///
363     /// # Panics
364     ///
365     /// Panics if the new capacity exceeds `isize::MAX` bytes.
366     ///
367     /// # Aborts
368     ///
369     /// Aborts on OOM.
370     #[cfg(not(no_global_oom_handling))]
371     pub fn reserve_exact(&mut self, len: usize, additional: usize) {
372         handle_reserve(self.try_reserve_exact(len, additional));
373     }
374
375     /// The same as `reserve_exact`, but returns on errors instead of panicking or aborting.
376     pub fn try_reserve_exact(
377         &mut self,
378         len: usize,
379         additional: usize,
380     ) -> Result<(), TryReserveError> {
381         if self.needs_to_grow(len, additional) { self.grow_exact(len, additional) } else { Ok(()) }
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     pub fn shrink_to_fit(&mut self, cap: usize) {
396         handle_reserve(self.shrink(cap));
397     }
398 }
399
400 impl<T, A: Allocator> RawVec<T, A> {
401     /// Returns if the buffer needs to grow to fulfill the needed extra capacity.
402     /// Mainly used to make inlining reserve-calls possible without inlining `grow`.
403     fn needs_to_grow(&self, len: usize, additional: usize) -> bool {
404         additional > self.capacity().wrapping_sub(len)
405     }
406
407     fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) {
408         // Allocators currently return a `NonNull<[u8]>` whose length matches
409         // the size requested. If that ever changes, the capacity here should
410         // change to `ptr.len() / mem::size_of::<T>()`.
411         self.ptr = unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) };
412         self.cap = cap;
413     }
414
415     // This method is usually instantiated many times. So we want it to be as
416     // small as possible, to improve compile times. But we also want as much of
417     // its contents to be statically computable as possible, to make the
418     // generated code run faster. Therefore, this method is carefully written
419     // so that all of the code that depends on `T` is within it, while as much
420     // of the code that doesn't depend on `T` as possible is in functions that
421     // are non-generic over `T`.
422     fn grow_amortized(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
423         // This is ensured by the calling contexts.
424         debug_assert!(additional > 0);
425
426         if T::IS_ZST {
427             // Since we return a capacity of `usize::MAX` when `elem_size` is
428             // 0, getting to here necessarily means the `RawVec` is overfull.
429             return Err(CapacityOverflow.into());
430         }
431
432         // Nothing we can really do about these checks, sadly.
433         let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
434
435         // This guarantees exponential growth. The doubling cannot overflow
436         // because `cap <= isize::MAX` and the type of `cap` is `usize`.
437         let cap = cmp::max(self.cap * 2, required_cap);
438         let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap);
439
440         let new_layout = Layout::array::<T>(cap);
441
442         // `finish_grow` is non-generic over `T`.
443         let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
444         self.set_ptr_and_cap(ptr, cap);
445         Ok(())
446     }
447
448     // The constraints on this method are much the same as those on
449     // `grow_amortized`, but this method is usually instantiated less often so
450     // it's less critical.
451     fn grow_exact(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
452         if T::IS_ZST {
453             // Since we return a capacity of `usize::MAX` when the type size is
454             // 0, getting to here necessarily means the `RawVec` is overfull.
455             return Err(CapacityOverflow.into());
456         }
457
458         let cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
459         let new_layout = Layout::array::<T>(cap);
460
461         // `finish_grow` is non-generic over `T`.
462         let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
463         self.set_ptr_and_cap(ptr, cap);
464         Ok(())
465     }
466
467     #[cfg(not(no_global_oom_handling))]
468     fn shrink(&mut self, cap: usize) -> Result<(), TryReserveError> {
469         assert!(cap <= self.capacity(), "Tried to shrink to a larger capacity");
470
471         let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) };
472         // See current_memory() why this assert is here
473         let _: () = const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
474
475         // If shrinking to 0, deallocate the buffer. We don't reach this point
476         // for the T::IS_ZST case since current_memory() will have returned
477         // None.
478         if cap == 0 {
479             unsafe { self.alloc.deallocate(ptr, layout) };
480             self.ptr = Unique::dangling();
481             self.cap = 0;
482         } else {
483             let ptr = unsafe {
484                 // `Layout::array` cannot overflow here because it would have
485                 // overflowed earlier when capacity was larger.
486                 let new_size = mem::size_of::<T>().unchecked_mul(cap);
487                 let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
488                 self.alloc
489                     .shrink(ptr, layout, new_layout)
490                     .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?
491             };
492             self.set_ptr_and_cap(ptr, cap);
493         }
494         Ok(())
495     }
496 }
497
498 // This function is outside `RawVec` to minimize compile times. See the comment
499 // above `RawVec::grow_amortized` for details. (The `A` parameter isn't
500 // significant, because the number of different `A` types seen in practice is
501 // much smaller than the number of `T` types.)
502 #[inline(never)]
503 fn finish_grow<A>(
504     new_layout: Result<Layout, LayoutError>,
505     current_memory: Option<(NonNull<u8>, Layout)>,
506     alloc: &mut A,
507 ) -> Result<NonNull<[u8]>, TryReserveError>
508 where
509     A: Allocator,
510 {
511     // Check for the error here to minimize the size of `RawVec::grow_*`.
512     let new_layout = new_layout.map_err(|_| CapacityOverflow)?;
513
514     alloc_guard(new_layout.size())?;
515
516     let memory = if let Some((ptr, old_layout)) = current_memory {
517         debug_assert_eq!(old_layout.align(), new_layout.align());
518         unsafe {
519             // The allocator checks for alignment equality
520             intrinsics::assume(old_layout.align() == new_layout.align());
521             alloc.grow(ptr, old_layout, new_layout)
522         }
523     } else {
524         alloc.allocate(new_layout)
525     };
526
527     memory.map_err(|_| AllocError { layout: new_layout, non_exhaustive: () }.into())
528 }
529
530 unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
531     /// Frees the memory owned by the `RawVec` *without* trying to drop its contents.
532     fn drop(&mut self) {
533         if let Some((ptr, layout)) = self.current_memory() {
534             unsafe { self.alloc.deallocate(ptr, layout) }
535         }
536     }
537 }
538
539 // Central function for reserve error handling.
540 #[cfg(not(no_global_oom_handling))]
541 #[inline]
542 fn handle_reserve(result: Result<(), TryReserveError>) {
543     match result.map_err(|e| e.kind()) {
544         Err(CapacityOverflow) => capacity_overflow(),
545         Err(AllocError { layout, .. }) => handle_alloc_error(layout),
546         Ok(()) => { /* yay */ }
547     }
548 }
549
550 // We need to guarantee the following:
551 // * We don't ever allocate `> isize::MAX` byte-size objects.
552 // * We don't overflow `usize::MAX` and actually allocate too little.
553 //
554 // On 64-bit we just need to check for overflow since trying to allocate
555 // `> isize::MAX` bytes will surely fail. On 32-bit and 16-bit we need to add
556 // an extra guard for this in case we're running on a platform which can use
557 // all 4GB in user-space, e.g., PAE or x32.
558
559 #[inline]
560 fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> {
561     if usize::BITS < 64 && alloc_size > isize::MAX as usize {
562         Err(CapacityOverflow.into())
563     } else {
564         Ok(())
565     }
566 }
567
568 // One central function responsible for reporting capacity overflows. This'll
569 // ensure that the code generation related to these panics is minimal as there's
570 // only one location which panics rather than a bunch throughout the module.
571 #[cfg(not(no_global_oom_handling))]
572 fn capacity_overflow() -> ! {
573     panic!("capacity overflow");
574 }