Skip to main content

osom_lib_arrays/inline_array/
inline_array.rs

1use core::{
2    alloc::Layout,
3    mem::{ManuallyDrop, MaybeUninit},
4    ops::DerefMut,
5    ptr::NonNull,
6};
7
8use osom_lib_alloc::traits::Allocator;
9use osom_lib_primitives::length::Length;
10use osom_lib_reprc::traits::ReprC;
11
12use crate::errors::ArrayError;
13
14/// Represents a dynamic array, where first `TCAPACITY` items are inlined in the struct itself.
15///
16/// In other words, this array allocates memory only when its length exceeds `TCAPACITY`. In which
17/// case it allocates data on heap, and becomes pretty much a [`DynamicArray`][crate::dynamic_array::DynamicArray].
18#[repr(C)]
19#[must_use]
20pub struct InlineArray<const TCAPACITY: usize, T, TAllocator>
21where
22    T: Sized,
23    TAllocator: Allocator,
24{
25    pub(super) internal: InlineArrayUnion<TCAPACITY, T>,
26    pub(super) size: Length,
27    pub(super) capacity: Length,
28    pub(super) allocator: TAllocator,
29}
30
31unsafe impl<const TCAPACITY: usize, T, TAllocator> ReprC for InlineArray<TCAPACITY, T, TAllocator>
32where
33    T: Sized + ReprC,
34    TAllocator: Allocator,
35{
36    const CHECK: () = const {
37        let () = <T as ReprC>::CHECK;
38        let () = <TAllocator as ReprC>::CHECK;
39        let () = <Length as ReprC>::CHECK;
40        let () = <InlineArrayUnion<TCAPACITY, T> as ReprC>::CHECK;
41    };
42}
43
44#[repr(C)]
45pub(super) union InlineArrayUnion<const TCAPACITY: usize, T>
46where
47    T: Sized,
48{
49    pub inlined: ManuallyDrop<MaybeUninit<[T; TCAPACITY]>>,
50    pub ptr: *mut T,
51}
52
53unsafe impl<const TCAPACITY: usize, T> ReprC for InlineArrayUnion<TCAPACITY, T>
54where
55    T: ReprC,
56{
57    const CHECK: () = const {
58        let () = <T as ReprC>::CHECK;
59        let () = <*mut T as ReprC>::CHECK;
60        let () = <ManuallyDrop<MaybeUninit<[T; TCAPACITY]>> as ReprC>::CHECK;
61    };
62}
63
64impl<const TCAPACITY: usize, T, TAllocator> InlineArray<TCAPACITY, T, TAllocator>
65where
66    T: Sized,
67    TAllocator: Allocator,
68{
69    /// Creates a new empty [`InlineArray`] with the default `TAllocator`.
70    #[inline(always)]
71    pub fn new() -> Self {
72        Self::with_allocator(TAllocator::default())
73    }
74
75    /// Creates a new empty [`InlineArray`] with the given `TAllocator`.
76    #[inline]
77    pub const fn with_allocator(allocator: TAllocator) -> Self {
78        let inlined = ManuallyDrop::new(MaybeUninit::uninit());
79        Self {
80            internal: InlineArrayUnion { inlined },
81            size: Length::ZERO,
82            capacity: static_capacity::<TCAPACITY>(),
83            allocator,
84        }
85    }
86
87    /// Creates a new [`InlineArray`] with the default `TAllocator`. This method allocates memory
88    /// if `capacity` exceeds `TCAPACITY`.
89    ///
90    /// # Errors
91    ///
92    /// For details see [`ArrayError`].
93    #[inline(always)]
94    pub fn with_capacity(capacity: Length) -> Result<Self, ArrayError> {
95        Self::with_capacity_and_allocator(capacity, TAllocator::default())
96    }
97
98    /// Creates a new [`InlineArray`] with the given `TAllocator`. This method allocates memory
99    /// if `capacity` exceeds `TCAPACITY`.
100    ///
101    /// # Errors
102    ///
103    /// For details see [`ArrayError`].
104    pub fn with_capacity_and_allocator(capacity: Length, allocator: TAllocator) -> Result<Self, ArrayError> {
105        if capacity.as_usize() <= TCAPACITY {
106            return Ok(Self::with_allocator(allocator));
107        }
108
109        if capacity.as_usize() > Length::MAX_LENGTH.as_usize() {
110            return Err(ArrayError::LengthLimitExceeded);
111        }
112
113        let new_ptr = allocator
114            .allocate(Self::layout_for_size(capacity))
115            .map_err(Into::into)?
116            .cast::<T>();
117
118        Ok(Self {
119            internal: InlineArrayUnion { ptr: new_ptr.as_ptr() },
120            size: Length::ZERO,
121            capacity: capacity,
122            allocator: allocator,
123        })
124    }
125
126    pub(super) const fn layout_for_size(size: Length) -> Layout {
127        let tsize = size_of::<T>();
128        let Some(real_size) = tsize.checked_mul(size.as_usize()) else {
129            panic!("Tried to allocate array of size outside of usize range");
130        };
131        unsafe { Layout::from_size_align_unchecked(real_size, align_of::<T>()) }
132    }
133
134    #[inline(always)]
135    pub(super) const fn current_layout(&self) -> Layout {
136        Self::layout_for_size(self.capacity)
137    }
138
139    #[inline(always)]
140    pub(super) fn is_inlined(&self) -> bool {
141        self.capacity.as_usize() <= TCAPACITY
142    }
143
144    pub(super) fn current_ptr_mut(&mut self) -> *mut T {
145        unsafe {
146            if self.is_inlined() {
147                self.internal.inlined.deref_mut().as_mut_ptr().cast::<T>()
148            } else {
149                self.internal.ptr
150            }
151        }
152    }
153
154    pub(super) fn as_slice_internal(&self) -> &[T] {
155        unsafe {
156            let ptr = if self.is_inlined() {
157                self.internal.inlined.as_ptr().cast::<T>()
158            } else {
159                self.internal.ptr
160            };
161
162            core::slice::from_raw_parts(ptr, self.size.as_usize())
163        }
164    }
165
166    pub(super) fn as_slice_mut_internal(&mut self) -> &mut [T] {
167        unsafe { core::slice::from_raw_parts_mut(self.current_ptr_mut(), self.size.as_usize()) }
168    }
169
170    #[allow(clippy::cast_possible_truncation)]
171    pub(super) fn reserve_if_needed(&mut self, new_length: u32) -> Result<(), ArrayError> {
172        let capacity = self.capacity.as_u32();
173        if new_length <= capacity {
174            return Ok(());
175        }
176
177        let new_capacity = {
178            let upper_bound = ((u64::from(new_length) * 3) / 2) + 1;
179            let capped = core::cmp::min(upper_bound, Length::MAX_LENGTH.as_usize() as u64) as u32;
180            unsafe { Length::new_unchecked(capped) }
181        };
182
183        let new_ptr = unsafe {
184            let new_layout = Self::layout_for_size(new_capacity);
185            if self.is_inlined() {
186                let ptr = self
187                    .allocator
188                    .allocate(new_layout)
189                    .map_err(Into::into)?
190                    .cast::<T>()
191                    .as_ptr();
192                let current_length = self.size.as_usize();
193                if current_length > 0 {
194                    let inlined_ptr = self.internal.inlined.deref_mut().as_mut_ptr().cast::<T>();
195                    ptr.copy_from_nonoverlapping(inlined_ptr, current_length);
196                }
197                ptr
198            } else {
199                let old_layout = self.current_layout();
200                let raw_ptr = NonNull::new_unchecked(self.internal.ptr);
201                self.allocator
202                    .resize(raw_ptr.cast(), old_layout, new_layout)
203                    .map_err(Into::into)?
204                    .cast()
205                    .as_ptr()
206            }
207        };
208
209        self.internal = InlineArrayUnion { ptr: new_ptr };
210        self.capacity = new_capacity;
211        Ok(())
212    }
213}
214
215impl<const TCAPACITY: usize, T, TAllocator> Drop for InlineArray<TCAPACITY, T, TAllocator>
216where
217    T: Sized,
218    TAllocator: Allocator,
219{
220    fn drop(&mut self) {
221        unsafe {
222            if core::mem::needs_drop::<T>() {
223                for item in self.as_slice_mut_internal() {
224                    core::ptr::drop_in_place(item);
225                }
226            }
227
228            if self.is_inlined() {
229                return;
230            }
231
232            let ptr = NonNull::new_unchecked(self.internal.ptr);
233            self.allocator.deallocate(ptr.cast(), self.current_layout());
234        }
235    }
236}
237
238#[inline(always)]
239#[allow(clippy::cast_possible_truncation)]
240pub(super) const fn static_capacity<const TCAPACITY: usize>() -> Length {
241    const {
242        assert!(
243            TCAPACITY > 0,
244            "TCAPACITY cannot be zero. It reduces the InlineArray into a less efficient DynamicArray. Use latter instead."
245        );
246        assert!(
247            TCAPACITY <= Length::MAX_LENGTH.as_usize(),
248            "TCAPACITY cannot exceed Length::MAX_LENGTH"
249        );
250    }
251
252    unsafe { Length::new_unchecked(TCAPACITY as u32) }
253}