Commit 51d3a25a authored by Miguel Ojeda's avatar Miguel Ojeda
Browse files

rust: alloc: add `RawVec::try_with_capacity_in()` constructor 



Add the `RawVec::try_with_capacity_in()` constructor as the fallible
version of `RawVec::with_capacity_in()`.

The implementation follows the original.

The infallible constructor is implemented in terms of the private
`RawVec::allocate_in()` constructor, thus also add the private
`RawVec::try_allocate_in()` constructor following the other.

It will be used to implement `Vec::try_with_capacity{,_in}()` in
the next patch.

Reviewed-by: default avatarGary Guo <gary@garyguo.net>
Signed-off-by: default avatarMiguel Ojeda <ojeda@kernel.org>
parent 25d176a4

rust/alloc/raw_vec.rs

+33 −1
Original line number Diff line number Diff line
@@ -20,11 +20,11 @@ 
#[cfg(test)]

mod tests;



#[cfg(not(no_global_oom_handling))]

enum AllocInit {

    /// The contents of the new memory are uninitialized.

    Uninitialized,

    /// The new memory is guaranteed to be zeroed.

    #[allow(dead_code)]

    Zeroed,

}
@@ -133,6 +133,14 @@ pub fn with_capacity_in(capacity: usize, alloc: A) -> Self { 
        Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)

    }



    /// Like `try_with_capacity`, but parameterized over the choice of

    /// allocator for the returned `RawVec`.

    #[allow(dead_code)]

    #[inline]

    pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> {

        Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc)

    }



    /// Like `with_capacity_zeroed`, but parameterized over the choice

    /// of allocator for the returned `RawVec`.

    #[cfg(not(no_global_oom_handling))]
@@ -203,6 +211,30 @@ fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self { 
        }

    }



    fn try_allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Result<Self, TryReserveError> {

        // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.

        if mem::size_of::<T>() == 0 || capacity == 0 {

            return Ok(Self::new_in(alloc));

        }



        let layout = Layout::array::<T>(capacity).map_err(|_| CapacityOverflow)?;

        alloc_guard(layout.size())?;

        let result = match init {

            AllocInit::Uninitialized => alloc.allocate(layout),

            AllocInit::Zeroed => alloc.allocate_zeroed(layout),

        };

        let ptr = result.map_err(|_| AllocError { layout, non_exhaustive: () })?;



        // Allocators currently return a `NonNull<[u8]>` whose length

        // matches the size requested. If that ever changes, the capacity

        // here should change to `ptr.len() / mem::size_of::<T>()`.

        Ok(Self {

            ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) },

            cap: capacity,

            alloc,

        })

    }



    /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator.

    ///

    /// # Safety