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