CapacityHolder WIP

Author: scottmcm

library/alloc/src/raw_vec.rs

@@ -33,11 +33,33 @@ enum AllocInit {
     Zeroed,
 }
 
+pub trait CapacityHolder: Copy {
+    fn as_usize<T>(self) -> usize {
+        if T::IS_ZST {
+            usize::MAX
+        } else {
+            Self::as_usize_raw(self)
+        }
+    }
+    fn as_usize_raw(c: Self) -> usize;
+}
+
+#[derive(Copy, Clone)]
+pub struct FixedCapacity<const N: usize>;
+impl<const N: usize> CapacityHolder for FixedCapacity<N> {
+    fn as_usize_raw(_: Self) -> usize { N }
+}
+
 #[repr(transparent)]
+#[derive(Copy, Clone)]
 #[cfg_attr(target_pointer_width = "16", rustc_layout_scalar_valid_range_end(0x7fff))]
 #[cfg_attr(target_pointer_width = "32", rustc_layout_scalar_valid_range_end(0x7fff_ffff))]
 #[cfg_attr(target_pointer_width = "64", rustc_layout_scalar_valid_range_end(0x7fff_ffff_ffff_ffff))]
-struct Cap(usize);
+pub(crate) struct Cap(usize);
+
+impl CapacityHolder for Cap {
+    fn as_usize_raw(c: Self) -> usize { c.0 }
+}
 
 impl Cap {
     const ZERO: Cap = unsafe { Cap(0) };
@@ -66,14 +88,9 @@ impl Cap {
 /// `usize::MAX`. This means that you need to be careful when round-tripping this type with a
 /// `Box<[T]>`, since `capacity()` won't yield the length.
 #[allow(missing_debug_implementations)]
-pub(crate) struct RawVec<T, A: Allocator = Global> {
+pub(crate) struct RawVec<T, A: Allocator = Global, C: CapacityHolder = Cap> {
     ptr: Unique<T>,
-    /// Never used for ZSTs; it's `capacity()`'s responsibility to return usize::MAX in that case.
-    ///
-    /// # Safety
-    ///
-    /// `cap` must be in the `0..=isize::MAX` range.
-    cap: Cap,
+    cap: C,
     alloc: A,
 }
 
@@ -129,6 +146,27 @@ impl<T> RawVec<T, Global> {
     }
 }
 
+impl<T, A: Allocator, C: CapacityHolder> RawVec<T, A, C> {
+    fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
+        let cap = self.cap.as_usize::<T>();
+        if T::IS_ZST || cap == 0 {
+            None
+        } else {
+            // We could use Layout::array here which ensures the absence of isize and usize overflows
+            // and could hypothetically handle differences between stride and size, but this memory
+            // has already been allocated so we know it can't overflow and currently Rust does not
+            // support such types. So we can do better by skipping some checks and avoid an unwrap.
+            const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
+            unsafe {
+                let align = mem::align_of::<T>();
+                let size = mem::size_of::<T>().unchecked_mul(cap);
+                let layout = Layout::from_size_align_unchecked(size, align);
+                Some((self.ptr.cast().into(), layout))
+            }
+        }
+    }
+}
+
 impl<T, A: Allocator> RawVec<T, A> {
     // Tiny Vecs are dumb. Skip to:
     // - 8 if the element size is 1, because any heap allocators is likely
@@ -296,24 +334,6 @@ impl<T, A: Allocator> RawVec<T, A> {
         &self.alloc
     }
 
-    fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
-        if T::IS_ZST || self.cap.0 == 0 {
-            None
-        } else {
-            // We could use Layout::array here which ensures the absence of isize and usize overflows
-            // and could hypothetically handle differences between stride and size, but this memory
-            // has already been allocated so we know it can't overflow and currently Rust does not
-            // support such types. So we can do better by skipping some checks and avoid an unwrap.
-            const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
-            unsafe {
-                let align = mem::align_of::<T>();
-                let size = mem::size_of::<T>().unchecked_mul(self.cap.0);
-                let layout = Layout::from_size_align_unchecked(size, align);
-                Some((self.ptr.cast().into(), layout))
-            }
-        }
-    }
-
     /// Ensures that the buffer contains at least enough space to hold `len +
     /// additional` elements. If it doesn't already have enough capacity, will
     /// reallocate enough space plus comfortable slack space to get amortized
@@ -576,7 +596,7 @@ where
     memory.map_err(|_| AllocError { layout: new_layout, non_exhaustive: () }.into())
 }
 
-unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
+unsafe impl<#[may_dangle] T, A: Allocator, C: CapacityHolder> Drop for RawVec<T, A, C> {
     /// Frees the memory owned by the `RawVec` *without* trying to drop its contents.
     fn drop(&mut self) {
         if let Some((ptr, layout)) = self.current_memory() {

library/alloc/src/vec/into_iter.rs

@@ -46,12 +46,7 @@ pub struct IntoIter<
     T,
     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
 > {
-    pub(super) buf: NonNull<T>,
-    pub(super) phantom: PhantomData<T>,
-    pub(super) cap: usize,
-    // the drop impl reconstructs a RawVec from buf, cap and alloc
-    // to avoid dropping the allocator twice we need to wrap it into ManuallyDrop
-    pub(super) alloc: ManuallyDrop<A>,
+    pub(super) buf: RawVec<T, A>,
     pub(super) ptr: NonNull<T>,
     /// If T is a ZST, this is actually ptr+len. This encoding is picked so that
     /// ptr == end is a quick test for the Iterator being empty, that works
@@ -106,7 +101,7 @@ impl<T, A: Allocator> IntoIter<T, A> {
     #[unstable(feature = "allocator_api", issue = "32838")]
     #[inline]
     pub fn allocator(&self) -> &A {
-        &self.alloc
+        self.buf.allocator()
     }
 
     fn as_raw_mut_slice(&mut self) -> *mut [T] {
@@ -136,9 +131,7 @@ impl<T, A: Allocator> IntoIter<T, A> {
         // struct and then overwriting &mut self.
         // this creates less assembly
         self.cap = 0;
-        self.buf = RawVec::NEW.non_null();
-        self.ptr = self.buf;
-        self.end = self.buf.as_ptr();
+        unsafe { self.buf.set_ptr_and_cap(NonNull::dangling(), 0) };
 
         // Dropping the remaining elements can panic, so this needs to be
         // done only after updating the other fields.