osom_tools_runtime/arrays/
inline_dynamic_array.rs

1#![allow(
2    clippy::needless_borrow,
3    clippy::uninit_assumed_init,
4    clippy::cast_possible_truncation,
5    clippy::cast_possible_wrap,
6    clippy::cast_sign_loss
7)]
8
9use core::alloc::Layout;
10use core::mem::ManuallyDrop;
11use core::ops::Deref;
12
13use crate::allocator::{AllocatedMemory, AllocationError, Allocator};
14
15use crate::Length;
16
17#[cfg(feature = "std_alloc")]
18use crate::allocator::StdAllocator;
19
20/// Represents an error that occurs when constructing new [`InlineDynamicArray`].
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22#[must_use]
23#[repr(u8)]
24pub enum InlineDynamicArrayConstructionError {
25    /// The allocator failed to allocate memory.
26    AllocationError,
27
28    /// The passed array is too long, it exceeds [`MAX_SIZE`][`InlineDynamicArray::MAX_SIZE`].
29    ArrayTooLong,
30}
31
32impl From<AllocationError> for InlineDynamicArrayConstructionError {
33    fn from(_: AllocationError) -> Self {
34        InlineDynamicArrayConstructionError::AllocationError
35    }
36}
37union InlineDynamicArrayUnion<T, const N: usize> {
38    stack_data: ManuallyDrop<[T; N]>,
39    heap_data: *mut T,
40}
41
42/// A struct similar to [`DynamicArray`][`super::DynamicArray`],
43/// but holds `N` items inlined. Meaning it won't allocate memory
44/// until it exceeds `N` elements.
45///
46/// Note that this struct can only grow, and will never shrink.
47#[must_use]
48pub struct InlineDynamicArray<
49    const N: usize,
50    T,
51    #[cfg(feature = "std_alloc")] TAllocator = StdAllocator,
52    #[cfg(not(feature = "std_alloc"))] TAllocator,
53> where
54    TAllocator: Allocator,
55{
56    data: InlineDynamicArrayUnion<T, N>,
57    length: Length,
58    capacity: Length,
59    allocator: TAllocator,
60}
61
62impl<const N: usize, T, TAllocator: Allocator> InlineDynamicArray<N, T, TAllocator> {
63    pub const MAX_SIZE: usize = Length::MAX;
64
65    const fn validate() {
66        assert!(
67            N > 0,
68            "N in InlineVec must be greater than 0. InlineVec with N == 0 is just Vec. Use Vec instead."
69        );
70
71        // Note: 2147482623 is (i32::MAX - 1024). This is definitely way too much,
72        // but we reserve some space, just in case.
73        assert!(
74            N < Self::MAX_SIZE,
75            "N in InlineVec must be at most 2147482623. Which likely already is waaaay too much."
76        );
77    }
78
79    #[inline(always)]
80    const fn layout(size: usize) -> Layout {
81        let real_size = size * size_of::<T>();
82        let alignment = align_of::<T>();
83        unsafe { Layout::from_size_align_unchecked(real_size, alignment) }
84    }
85
86    #[inline(always)]
87    fn allocate_memory(&self, size: usize) -> Result<*mut T, AllocationError> {
88        let new_memory = self.allocator.allocate(Self::layout(size))?;
89        let result = unsafe { new_memory.as_ptr().cast::<T>() };
90        assert!(result.is_aligned(), "Newly allocated memory is not aligned correctly.");
91        Ok(result)
92    }
93
94    #[inline(always)]
95    const fn is_inlined(&self) -> bool {
96        self.capacity.value() == N as i32
97    }
98
99    #[inline(always)]
100    fn data_ptr(&self) -> *mut T {
101        unsafe {
102            if self.is_inlined() {
103                (&self.data.stack_data).as_ptr().cast_mut()
104            } else {
105                self.data.heap_data
106            }
107        }
108    }
109
110    /// Pushes a value to the end of the [`InlineDynamicArray`].
111    ///
112    /// Note that the [`InlineDynamicArray`] data will be moved to the heap
113    /// only when length exceeds `N`. It won't come back from the
114    /// heap though.
115    ///
116    /// Also note that the [`InlineDynamicArray`] will never shrink, it will
117    /// only keep growing.
118    ///
119    /// # Errors
120    ///
121    /// For details see [`InlineDynamicArrayConstructionError`].
122    pub fn push(&mut self, value: T) -> Result<(), InlineDynamicArrayConstructionError> {
123        let capacity: usize = self.capacity.into();
124        unsafe {
125            if self.is_inlined() {
126                if self.length.value() < N as i32 {
127                    let data = &mut self.data.stack_data;
128                    data.as_mut_ptr().add(self.length.into()).write(value);
129                    self.length.inc(1);
130                    return Ok(());
131                }
132                let new_capacity = capacity * 2;
133                if new_capacity > Self::MAX_SIZE {
134                    return Err(InlineDynamicArrayConstructionError::ArrayTooLong);
135                }
136                let new_memory = self.allocate_memory(new_capacity)?;
137                let stack_data = (&self.data.stack_data).as_ptr();
138                new_memory.copy_from_nonoverlapping(stack_data, self.length.into());
139                new_memory.add(self.length.into()).write(value);
140                self.length.inc(1);
141                self.data = InlineDynamicArrayUnion { heap_data: new_memory };
142                self.capacity = Length::new_unchecked(new_capacity as i32);
143                return Ok(());
144            }
145
146            if self.length == self.capacity {
147                let new_capacity = capacity * 2;
148                if new_capacity > Self::MAX_SIZE {
149                    return Err(InlineDynamicArrayConstructionError::ArrayTooLong);
150                }
151                let old_layout = Self::layout(self.capacity.into());
152                let new_layout = Self::layout(new_capacity);
153                let old_memory = self.allocator.convert_raw_ptr(self.data.heap_data.cast());
154                let new_memory = old_memory.resize(old_layout, new_layout)?;
155                self.data.heap_data = new_memory.as_ptr().cast::<T>();
156                self.capacity = Length::new_unchecked(new_capacity as i32);
157            }
158
159            self.data.heap_data.add(self.length.into()).write(value);
160            self.length.inc(1);
161        }
162        Ok(())
163    }
164
165    /// Pops last element from the [`InlineDynamicArray`],
166    /// decreasing its size.
167    ///
168    /// # Returns
169    ///
170    /// * `Some(T)` if `self.len() > 0`
171    /// * `None` otherwise
172    pub fn pop(&mut self) -> Option<T> {
173        if self.is_empty() {
174            None
175        } else {
176            Some(unsafe { self.pop_unchecked() })
177        }
178    }
179
180    /// Unsafe variant of [`pop`][`Self::pop`].
181    ///
182    /// # Safety
183    ///
184    /// Returns `T` if `self.len() > 0` and decrease the
185    /// [`InlineDynamicArray`] size. The behaviour is undefined if
186    /// `self.len() == 0`.
187    #[inline]
188    pub unsafe fn pop_unchecked(&mut self) -> T {
189        debug_assert!(!self.is_empty(), "Tried pop_unchecked on length 0 InlineDynamicArray.");
190        unsafe {
191            let ptr = self.data_ptr();
192            self.length.dec(1);
193            ptr.add(self.length.into()).read()
194        }
195    }
196
197    /// Creates a new empty [`InlineDynamicArray`].
198    #[inline(always)]
199    pub fn new() -> Self {
200        Self::with_allocator(TAllocator::default())
201    }
202
203    /// Creates a new empty [`InlineDynamicArray`] with the given allocator.
204    #[inline]
205    pub fn with_allocator(allocator: TAllocator) -> Self {
206        const { Self::validate() };
207
208        Self {
209            data: InlineDynamicArrayUnion {
210                heap_data: unsafe { allocator.dangling::<T>().as_ptr().cast::<T>() },
211            },
212            length: Length::ZERO,
213            capacity: unsafe { Length::new_unchecked(N as i32) },
214            allocator: allocator,
215        }
216    }
217
218    /// Returns the number of elements in the [`InlineDynamicArray`].
219    #[inline(always)]
220    pub const fn len(&self) -> Length {
221        self.length
222    }
223
224    /// Returns `true` if the [`InlineDynamicArray`] is empty,
225    /// otherwise `false`.
226    #[inline(always)]
227    pub const fn is_empty(&self) -> bool {
228        self.length.value() == 0
229    }
230
231    /// Returns the capacity of the [`InlineDynamicArray`]. Note that
232    /// this is always at least `N`.
233    #[inline(always)]
234    pub const fn capacity(&self) -> Length {
235        self.capacity
236    }
237
238    /// Represents current [`InlineDynamicArray`] as a slice.
239    #[inline]
240    pub fn as_slice(&self) -> &[T] {
241        unsafe {
242            let ptr = self.data_ptr();
243            core::slice::from_raw_parts(ptr, self.len().into())
244        }
245    }
246}
247
248impl<const N: usize, T: Clone, TAllocator: Allocator> InlineDynamicArray<N, T, TAllocator> {
249    /// Tries to clone the [`InlineDynamicArray`].
250    ///
251    /// # Errors
252    ///
253    /// For details see [`InlineDynamicArrayConstructionError`].
254    pub fn try_clone(&self) -> Result<Self, InlineDynamicArrayConstructionError> {
255        let mut new = Self::with_allocator(self.allocator.clone());
256        let slice = self.as_slice();
257        for item in slice {
258            new.push(item.clone())?;
259        }
260        Ok(new)
261    }
262}
263
264impl<const N: usize, T, TAllocator: Allocator> Drop for InlineDynamicArray<N, T, TAllocator> {
265    fn drop(&mut self) {
266        unsafe {
267            if core::mem::needs_drop::<T>() {
268                let mut ptr = self.data_ptr();
269                let mut idx = 0;
270                while idx < self.len().into() {
271                    drop(ptr.read());
272                    ptr = ptr.add(1);
273                    idx += 1;
274                }
275            }
276
277            if !self.is_inlined() {
278                let layout = Self::layout(self.capacity.into());
279                let memory = self.allocator.convert_raw_ptr(self.data.heap_data.cast());
280                memory.deallocate(layout);
281            }
282        }
283    }
284}
285
286impl<const N: usize, T, TAllocator: Allocator> Default for InlineDynamicArray<N, T, TAllocator> {
287    fn default() -> Self {
288        Self::new()
289    }
290}
291
292impl<const N: usize, T: Clone, TAllocator: Allocator> Clone for InlineDynamicArray<N, T, TAllocator> {
293    fn clone(&self) -> Self {
294        self.try_clone().expect("Failed to clone the inline dynamic array")
295    }
296}
297
298impl<const N: usize, T, TAllocator: Allocator> core::fmt::Debug for InlineDynamicArray<N, T, TAllocator> {
299    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
300        f.debug_struct("InlineDynamicArray")
301            .field("N", &N)
302            .field("len", &self.len())
303            .field("capacity", &self.capacity())
304            .field("is_inlined", &self.is_inlined())
305            .field("raw_ptr", &self.data_ptr().addr())
306            .finish()
307    }
308}
309
310impl<const N1: usize, const N2: usize, T: PartialEq, TAllocator1: Allocator, TAllocator2: Allocator>
311    PartialEq<InlineDynamicArray<N1, T, TAllocator1>> for InlineDynamicArray<N2, T, TAllocator2>
312{
313    fn eq(&self, other: &InlineDynamicArray<N1, T, TAllocator1>) -> bool {
314        self.as_slice() == other.as_slice()
315    }
316}
317
318impl<const N: usize, T: Eq, TAllocator: Allocator> Eq for InlineDynamicArray<N, T, TAllocator> {}
319
320impl<const N: usize, T, TAllocator: Allocator> Deref for InlineDynamicArray<N, T, TAllocator> {
321    type Target = [T];
322
323    fn deref(&self) -> &Self::Target {
324        self.as_slice()
325    }
326}
327
328impl<const N: usize, T, TAllocator: Allocator> AsRef<[T]> for InlineDynamicArray<N, T, TAllocator> {
329    fn as_ref(&self) -> &[T] {
330        self.as_slice()
331    }
332}
333
334unsafe impl<const N: usize, T: Send, TAllocator: Allocator> Send for InlineDynamicArray<N, T, TAllocator> {}
335unsafe impl<const N: usize, T: Sync, TAllocator: Allocator> Sync for InlineDynamicArray<N, T, TAllocator> {}