cl_generic_vec/raw/heap/
nightly.rs

1use crate::raw::{AllocError, AllocResult, Storage, StorageWithCapacity};
2
3use core::{alloc::Layout, ptr::NonNull};
4use std::{alloc::handle_alloc_error, mem::MaybeUninit};
5
6use std::alloc::Allocator;
7
8enum OnFailure {
9    Abort,
10    Error,
11}
12
13type Heap<T, A> = Box<[MaybeUninit<T>], A>;
14
15/// Create a new `Heap<T>`storage from the given pointer and capacity
16///
17/// # Safety
18///
19/// If the capacity is non-zero
20/// * You must have allocated the pointer from the given allocator
21/// * The pointer must be valid to read-write for the range `ptr..ptr.add(capacity)`
22pub(crate) unsafe fn box_from_raw_parts_in<T, A: Allocator>(
23    ptr: NonNull<T>,
24    capacity: usize,
25    allocator: A,
26) -> Heap<T, A> {
27    unsafe {
28        let ptr = std::ptr::slice_from_raw_parts_mut(ptr.as_ptr().cast(), capacity);
29        Box::from_raw_in(ptr, allocator)
30    }
31}
32
33/// Convert a `Heap` storage into a pointer and capacity, without
34/// deallocating the storage
35pub(crate) fn box_into_raw_parts_with_alloc<T, A: Allocator>(b: Heap<T, A>) -> (NonNull<T>, usize, A) {
36    let (ptr, alloc) = Box::into_raw_with_allocator(b);
37    unsafe {
38        let (ptr, capacity) = ptr.to_raw_parts();
39        (NonNull::new_unchecked(ptr.cast()), capacity, alloc)
40    }
41}
42
43unsafe impl<T, A: Allocator> Storage for Heap<T, A> {
44    type Item = T;
45
46    fn reserve(&mut self, new_capacity: usize) {
47        if self.len() < new_capacity {
48            let _ = reserve_slow(self, new_capacity, OnFailure::Abort);
49        }
50    }
51
52    fn try_reserve(&mut self, new_capacity: usize) -> AllocResult {
53        if self.len() < new_capacity {
54            reserve_slow(self, new_capacity, OnFailure::Error)
55        } else {
56            Ok(())
57        }
58    }
59}
60
61unsafe impl<T, A: Default + Allocator> StorageWithCapacity for Heap<T, A> {
62    fn with_capacity(cap: usize) -> Self { Box::new_uninit_slice_in(cap, A::default()) }
63}
64
65#[cold]
66#[inline(never)]
67fn reserve_slow<T, A: Allocator>(b: &mut Heap<T, A>, new_capacity: usize, on_failure: OnFailure) -> AllocResult {
68    assert!(new_capacity > b.len());
69
70    // taking a copy of the box so we can get it's contents and then update it later
71    // Safety:
72    // we forget the box just as soon we we copy it, so we have no risk of double-free
73    let (ptr, cap, alloc) = unsafe { box_into_raw_parts_with_alloc(std::ptr::read(b)) };
74
75    // grow by at least doubling
76    let new_capacity = new_capacity
77        .max(cap.checked_mul(2).expect("Could not grow further"))
78        .max(super::INIT_ALLOC_CAPACITY);
79    let layout = Layout::new::<T>().repeat(new_capacity).expect("Invalid layout").0;
80
81    let ptr = if cap == 0 {
82        unsafe { alloc.allocate(layout) }
83    } else {
84        let new_layout = layout;
85        let old_layout = Layout::new::<T>().repeat(cap).expect("Invalid layout").0;
86
87        unsafe { alloc.grow(ptr.cast(), old_layout, new_layout) }
88    };
89
90    let ptr = match (ptr, on_failure) {
91        (Ok(ptr), _) => ptr,
92        (Err(_), OnFailure::Abort) => handle_alloc_error(layout),
93        (Err(_), OnFailure::Error) => return Err(AllocError),
94    };
95
96    // Creating a new Heap using the re-alloced pointer.
97    // Replacing the existing heap and forgetting it so
98    // that no drop code happens, avoiding the
99    unsafe {
100        let new = box_from_raw_parts_in(ptr.cast(), new_capacity, alloc);
101        let old = std::mem::replace(b, new);
102        std::mem::forget(old);
103    }
104
105    Ok(())
106}
cl_generic_vec/raw/heap/nightly.rs

cl_generic_vec/raw/heap/
nightly.rs
