@@ -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