@@ -25,6 +25,16 @@ enum AllocInit {
2525 Zeroed,
2626}
2727
28+ #[repr(transparent)]
29+ #[cfg_attr(target_pointer_width = "16", rustc_layout_scalar_valid_range_end(0x7fff))]
30+ #[cfg_attr(target_pointer_width = "32", rustc_layout_scalar_valid_range_end(0x7fff_ffff))]
31+ #[cfg_attr(target_pointer_width = "64", rustc_layout_scalar_valid_range_end(0x7fff_ffff_ffff_ffff))]
32+ struct Cap(usize);
33+
34+ impl Cap {
35+ const ZERO: Cap = unsafe { Cap(0) };
36+ }
37+
2838/// A low-level utility for more ergonomically allocating, reallocating, and deallocating
2939/// a buffer of memory on the heap without having to worry about all the corner cases
3040/// involved. This type is excellent for building your own data structures like Vec and VecDeque.
@@ -50,7 +60,12 @@ enum AllocInit {
5060#[allow(missing_debug_implementations)]
5161pub(crate) struct RawVec<T, A: Allocator = Global> {
5262 ptr: Unique<T>,
53- cap: usize,
63+ /// Never used for ZSTs; it's `capacity()`'s responsibility to return usize::MAX in that case.
64+ ///
65+ /// # Safety
66+ ///
67+ /// `cap` must be in the `0..=isize::MAX` range.
68+ cap: Cap,
5469 alloc: A,
5570}
5671
@@ -119,7 +134,7 @@ impl<T, A: Allocator> RawVec<T, A> {
119134 /// the returned `RawVec`.
120135 pub const fn new_in(alloc: A) -> Self {
121136 // `cap: 0` means "unallocated". zero-sized types are ignored.
122- Self { ptr: Unique::dangling(), cap: 0 , alloc }
137+ Self { ptr: Unique::dangling(), cap: Cap::ZERO , alloc }
123138 }
124139
125140 /// Like `with_capacity`, but parameterized over the choice of
@@ -194,7 +209,7 @@ impl<T, A: Allocator> RawVec<T, A> {
194209 // here should change to `ptr.len() / mem::size_of::<T>()`.
195210 Self {
196211 ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },
197- cap: capacity,
212+ cap: unsafe { Cap( capacity) } ,
198213 alloc,
199214 }
200215 }
@@ -207,12 +222,13 @@ impl<T, A: Allocator> RawVec<T, A> {
207222 /// The `ptr` must be allocated (via the given allocator `alloc`), and with the given
208223 /// `capacity`.
209224 /// The `capacity` cannot exceed `isize::MAX` for sized types. (only a concern on 32-bit
210- /// systems). ZST vectors may have a capacity up to `usize::MAX` .
225+ /// systems). For ZSTs capacity is ignored .
211226 /// If the `ptr` and `capacity` come from a `RawVec` created via `alloc`, then this is
212227 /// guaranteed.
213228 #[inline]
214229 pub unsafe fn from_raw_parts_in(ptr: *mut T, capacity: usize, alloc: A) -> Self {
215- Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap: capacity, alloc }
230+ let cap = if T::IS_ZST { Cap::ZERO } else { unsafe { Cap(capacity) } };
231+ Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc }
216232 }
217233
218234 /// Gets a raw pointer to the start of the allocation. Note that this is
@@ -228,7 +244,7 @@ impl<T, A: Allocator> RawVec<T, A> {
228244 /// This will always be `usize::MAX` if `T` is zero-sized.
229245 #[inline(always)]
230246 pub fn capacity(&self) -> usize {
231- if T::IS_ZST { usize::MAX } else { self.cap }
247+ if T::IS_ZST { usize::MAX } else { self.cap.0 }
232248 }
233249
234250 /// Returns a shared reference to the allocator backing this `RawVec`.
@@ -237,7 +253,7 @@ impl<T, A: Allocator> RawVec<T, A> {
237253 }
238254
239255 fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
240- if T::IS_ZST || self.cap == 0 {
256+ if T::IS_ZST || self.cap.0 == 0 {
241257 None
242258 } else {
243259 // We could use Layout::array here which ensures the absence of isize and usize overflows
@@ -247,7 +263,7 @@ impl<T, A: Allocator> RawVec<T, A> {
247263 let _: () = const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
248264 unsafe {
249265 let align = mem::align_of::<T>();
250- let size = mem::size_of::<T>().unchecked_mul(self.cap);
266+ let size = mem::size_of::<T>().unchecked_mul(self.cap.0 );
251267 let layout = Layout::from_size_align_unchecked(size, align);
252268 Some((self.ptr.cast().into(), layout))
253269 }
@@ -375,12 +391,15 @@ impl<T, A: Allocator> RawVec<T, A> {
375391 additional > self.capacity().wrapping_sub(len)
376392 }
377393
378- fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) {
394+ /// # Safety:
395+ ///
396+ /// `cap` must not exceed `isize::MAX`.
397+ unsafe fn set_ptr_and_cap(&mut self, ptr: NonNull<[u8]>, cap: usize) {
379398 // Allocators currently return a `NonNull<[u8]>` whose length matches
380399 // the size requested. If that ever changes, the capacity here should
381400 // change to `ptr.len() / mem::size_of::<T>()`.
382401 self.ptr = unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) };
383- self.cap = cap;
402+ self.cap = unsafe { Cap( cap) } ;
384403 }
385404
386405 // This method is usually instantiated many times. So we want it to be as
@@ -405,14 +424,15 @@ impl<T, A: Allocator> RawVec<T, A> {
405424
406425 // This guarantees exponential growth. The doubling cannot overflow
407426 // because `cap <= isize::MAX` and the type of `cap` is `usize`.
408- let cap = cmp::max(self.cap * 2, required_cap);
427+ let cap = cmp::max(self.cap.0 * 2, required_cap);
409428 let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap);
410429
411430 let new_layout = Layout::array::<T>(cap);
412431
413432 // `finish_grow` is non-generic over `T`.
414433 let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
415- self.set_ptr_and_cap(ptr, cap);
434+ // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than isize::MAX items
435+ unsafe { self.set_ptr_and_cap(ptr, cap) };
416436 Ok(())
417437 }
418438
@@ -431,7 +451,10 @@ impl<T, A: Allocator> RawVec<T, A> {
431451
432452 // `finish_grow` is non-generic over `T`.
433453 let ptr = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
434- self.set_ptr_and_cap(ptr, cap);
454+ // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than isize::MAX items
455+ unsafe {
456+ self.set_ptr_and_cap(ptr, cap);
457+ }
435458 Ok(())
436459 }
437460
@@ -449,7 +472,7 @@ impl<T, A: Allocator> RawVec<T, A> {
449472 if cap == 0 {
450473 unsafe { self.alloc.deallocate(ptr, layout) };
451474 self.ptr = Unique::dangling();
452- self.cap = 0 ;
475+ self.cap = Cap::ZERO ;
453476 } else {
454477 let ptr = unsafe {
455478 // `Layout::array` cannot overflow here because it would have
@@ -460,7 +483,10 @@ impl<T, A: Allocator> RawVec<T, A> {
460483 .shrink(ptr, layout, new_layout)
461484 .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?
462485 };
463- self.set_ptr_and_cap(ptr, cap);
486+ // SAFETY: if the allocation is valid, then the capacity is too
487+ unsafe {
488+ self.set_ptr_and_cap(ptr, cap);
489+ }
464490 }
465491 Ok(())
466492 }
0 commit comments