Skip to main content

orengine_utils/
array_buffer.rs

1//! This module contains the [`ArrayBuffer`].
2use crate::hints::{assert_hint, likely, unlikely};
3use core::mem;
4use core::mem::MaybeUninit;
5use core::ops::{Deref, DerefMut};
6use core::ptr::{slice_from_raw_parts, slice_from_raw_parts_mut};
7
8/// `ArrayBuffer` is a fixed-sized array-based buffer.
9///
10/// # Example
11///
12/// ```rust
13/// use core::mem::MaybeUninit;
14/// use orengine_utils::ArrayBuffer;
15///
16/// let mut buffer = ArrayBuffer::<u16, 4>::new();
17///
18/// unsafe {
19///     buffer.refill_with(|buf| {
20///         buf[0..2].copy_from_slice(&[MaybeUninit::new(22), MaybeUninit::new(23)]);
21///
22///         2
23///     });
24/// }
25///
26/// buffer[1] = 21;
27///
28/// assert_eq!(buffer.pop(), Some(21));
29/// assert_eq!(buffer.pop(), Some(22));
30/// ```
31pub struct ArrayBuffer<T, const N: usize> {
32    array: [MaybeUninit<T>; N],
33    len: usize,
34}
35
36impl<T, const N: usize> ArrayBuffer<T, N> {
37    /// Creates a new ` ArrayBuffer `.
38    pub const fn new() -> Self {
39        Self {
40            array: [const { MaybeUninit::uninit() }; N],
41            len: 0,
42        }
43    }
44
45    /// Returns the capacity of the buffer.
46    pub const fn capacity(&self) -> usize {
47        N
48    }
49
50    /// Returns the number of elements in the buffer.
51    pub fn len(&self) -> usize {
52        self.len
53    }
54
55    /// Returns `true` if the buffer is empty.
56    pub fn is_empty(&self) -> bool {
57        self.len == 0
58    }
59
60    /// Forces the length of the buffer to `new_len`.
61    ///
62    /// # Safety
63    ///
64    /// - `new_len` must be less than or equal to `N`.
65    /// - The elements at `old_len..new_len` must be initialized.
66    pub unsafe fn set_len(&mut self, new_len: usize) {
67        debug_assert!(
68            new_len <= self.capacity(),
69            "provided len is more than the capacity: {new_len} > {N}"
70        );
71
72        self.len = new_len;
73    }
74
75    /// Returns a pointer to the first element of the buffer.
76    pub const fn as_ptr(&self) -> *const T {
77        self.array.as_ptr().cast()
78    }
79
80    /// Returns a mutable pointer to the first element of the buffer.
81    pub const fn as_mut_ptr(&mut self) -> *mut T {
82        self.array.as_mut_ptr().cast()
83    }
84
85    /// Appends an element to the buffer.
86    ///
87    /// # Safety
88    ///
89    /// The caller must ensure that the buffer is not full.
90    pub unsafe fn push_unchecked(&mut self, item: T) {
91        assert_hint(self.len() < N, "Tried to push to a full array buffer");
92
93        self.array[self.len].write(item);
94        self.len += 1;
95    }
96
97    /// Appends an element to the buffer or returns `Err(value)` if the buffer is full.
98    pub fn push(&mut self, item: T) -> Result<(), T> {
99        if unlikely(self.len == self.capacity()) {
100            return Err(item);
101        }
102
103        unsafe { self.push_unchecked(item) };
104
105        Ok(())
106    }
107
108    /// Pops an element from the buffer or returns `None` if the buffer is empty.
109    pub fn pop(&mut self) -> Option<T> {
110        if unlikely(self.len == 0) {
111            return None;
112        }
113
114        self.len -= 1;
115
116        Some(unsafe { self.array[self.len].as_ptr().read() })
117    }
118
119    /// Clears with calling the provided function on each element.
120    pub fn clear_with<F>(&mut self, mut f: F)
121    where
122        F: FnMut(T),
123    {
124        for i in 0..self.len {
125            f(unsafe { self.array[i].as_ptr().read() });
126        }
127
128        self.len = 0;
129    }
130
131    /// Drops all elements in the buffer and set the length to 0.
132    pub fn clear(&mut self) {
133        if mem::needs_drop::<T>() {
134            for i in 0..self.len {
135                drop(unsafe { self.array[i].as_ptr().read() });
136            }
137        }
138
139        self.len = 0;
140    }
141
142    /// Returns a reference iterator over the buffer.
143    pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
144        struct Iter<'array_buffer, T> {
145            current: *const T,
146            end: *const T,
147            _marker: core::marker::PhantomData<&'array_buffer T>,
148        }
149
150        impl<'array_buffer, T> Iterator for Iter<'array_buffer, T> {
151            type Item = &'array_buffer T;
152
153            fn next(&mut self) -> Option<Self::Item> {
154                if likely(self.current < self.end) {
155                    let item = unsafe { &*self.current };
156
157                    self.current = unsafe { self.current.add(1) };
158
159                    Some(item)
160                } else {
161                    None
162                }
163            }
164
165            fn size_hint(&self) -> (usize, Option<usize>) {
166                #[allow(clippy::cast_sign_loss, reason = "It is impossible")]
167                let size = unsafe { self.end.offset_from(self.current) as usize };
168
169                (size, Some(size))
170            }
171        }
172
173        impl<T> ExactSizeIterator for Iter<'_, T> {
174            fn len(&self) -> usize {
175                #[allow(clippy::cast_sign_loss, reason = "It is impossible")]
176                unsafe {
177                    self.end.offset_from(self.current) as usize
178                }
179            }
180        }
181
182        let current = self.as_ptr();
183        let end = unsafe { current.add(self.len) };
184
185        Iter {
186            current,
187            end,
188            _marker: core::marker::PhantomData,
189        }
190    }
191
192    /// Returns a mutable reference iterator over the buffer.
193    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
194        struct IterMut<'array_buffer, T> {
195            current: *mut T,
196            end: *mut T,
197            _marker: core::marker::PhantomData<&'array_buffer T>,
198        }
199
200        impl<'array_buffer, T> Iterator for IterMut<'array_buffer, T> {
201            type Item = &'array_buffer mut T;
202
203            fn next(&mut self) -> Option<Self::Item> {
204                if likely(self.current < self.end) {
205                    let item = unsafe { &mut *self.current };
206
207                    self.current = unsafe { self.current.add(1) };
208
209                    Some(item)
210                } else {
211                    None
212                }
213            }
214
215            fn size_hint(&self) -> (usize, Option<usize>) {
216                #[allow(clippy::cast_sign_loss, reason = "It is impossible")]
217                let size = unsafe { self.end.offset_from(self.current) as usize };
218
219                (size, Some(size))
220            }
221        }
222
223        impl<T> ExactSizeIterator for IterMut<'_, T> {
224            fn len(&self) -> usize {
225                #[allow(clippy::cast_sign_loss, reason = "It is impossible")]
226                unsafe {
227                    self.end.offset_from(self.current) as usize
228                }
229            }
230        }
231
232        let current = self.as_mut_ptr();
233        let end = unsafe { current.add(self.len) };
234
235        IterMut {
236            current,
237            end,
238            _marker: core::marker::PhantomData,
239        }
240    }
241
242    /// Refills the buffer with elements provided by the function.
243    ///
244    /// # Safety
245    ///
246    /// The caller must ensure that the buffer is empty before refilling.
247    pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
248        debug_assert!(
249            self.is_empty(),
250            "ArrayBuffer should be empty before refilling"
251        );
252
253        let filled = f(&mut self.array);
254
255        debug_assert!(filled <= N, "Filled more than the capacity");
256
257        self.len = filled;
258    }
259    /// Returns a pointer to the underlying array.
260    fn as_slice_ptr(&self) -> *const [T] {
261        slice_from_raw_parts(self.as_ptr(), self.len)
262    }
263
264    /// Returns a mutable pointer to the underlying array.
265    fn as_mut_slice_ptr(&mut self) -> *mut [T] {
266        slice_from_raw_parts_mut(self.as_mut_ptr(), self.len)
267    }
268}
269
270impl<T, const N: usize> Deref for ArrayBuffer<T, N> {
271    type Target = [T];
272
273    fn deref(&self) -> &Self::Target {
274        unsafe { &*self.as_slice_ptr() }
275    }
276}
277
278impl<T, const N: usize> AsRef<[T]> for ArrayBuffer<T, N> {
279    fn as_ref(&self) -> &[T] {
280        unsafe { &*self.as_slice_ptr() }
281    }
282}
283
284impl<T, const N: usize> DerefMut for ArrayBuffer<T, N> {
285    fn deref_mut(&mut self) -> &mut Self::Target {
286        unsafe { &mut *self.as_mut_slice_ptr() }
287    }
288}
289
290impl<T, const N: usize> AsMut<[T]> for ArrayBuffer<T, N> {
291    fn as_mut(&mut self) -> &mut [T] {
292        unsafe { &mut *self.as_mut_slice_ptr() }
293    }
294}
295
296impl<T, const N: usize> Default for ArrayBuffer<T, N> {
297    fn default() -> Self {
298        Self::new()
299    }
300}
301
302impl<T, const N: usize> From<[T; N]> for ArrayBuffer<T, N> {
303    fn from(array: [T; N]) -> Self {
304        Self {
305            array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
306            len: N,
307        }
308    }
309}
310
311impl<T: Clone, const N: usize> Clone for ArrayBuffer<T, N> {
312    fn clone(&self) -> Self {
313        let mut res = Self {
314            array: [const { MaybeUninit::uninit() }; N],
315            len: self.len,
316        };
317
318        for (i, item) in self.as_ref().iter().enumerate() {
319            res.array[i] = MaybeUninit::new(item.clone());
320        }
321
322        res
323    }
324}
325
326impl<T, const N: usize> Drop for ArrayBuffer<T, N> {
327    fn drop(&mut self) {
328        self.clear();
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335    use alloc::vec;
336    use alloc::vec::Vec;
337
338    #[allow(
339        clippy::explicit_auto_deref,
340        reason = "We test deref and deref_mut methods"
341    )]
342    #[test]
343    fn test_array_buffer_pop_push_len() {
344        let mut buffer = ArrayBuffer::<u32, 4>::new();
345
346        unsafe {
347            buffer.push_unchecked(1);
348            assert_eq!(buffer.len(), 1);
349            assert_eq!((*buffer).len(), 1);
350
351            buffer.push_unchecked(2);
352            assert_eq!(buffer.len(), 2);
353            assert_eq!((*buffer).len(), 2);
354
355            buffer.push(3).unwrap();
356            assert_eq!(buffer.len(), 3);
357            assert_eq!(buffer.as_ref().len(), 3);
358
359            assert_eq!(buffer.pop(), Some(3));
360            assert_eq!(buffer.len(), 2);
361            assert_eq!(buffer.as_mut().len(), 2);
362
363            buffer.push_unchecked(4);
364            assert_eq!(buffer.len(), 3);
365            assert_eq!(buffer.deref_mut().len(), 3);
366
367            buffer.push_unchecked(5);
368            assert_eq!(buffer.len(), 4);
369            assert_eq!(buffer.deref_mut().len(), 4);
370
371            assert_eq!(buffer.push(6), Err(6));
372
373            assert_eq!(buffer.pop(), Some(5));
374            assert_eq!(buffer.pop(), Some(4));
375            assert_eq!(buffer.pop(), Some(2));
376            assert_eq!(buffer.pop(), Some(1));
377            assert_eq!(buffer.pop(), None);
378        }
379    }
380
381    #[test]
382    fn test_array_buffer_iterators() {
383        let mut buffer = ArrayBuffer::<u32, 4>::new();
384
385        unsafe {
386            buffer.push_unchecked(1);
387            buffer.push_unchecked(2);
388            buffer.push_unchecked(3);
389            buffer.push_unchecked(4);
390        }
391
392        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
393        assert_eq!(
394            buffer.iter_mut().collect::<Vec<_>>(),
395            vec![&mut 1, &mut 2, &mut 3, &mut 4]
396        );
397    }
398
399    #[test]
400    fn test_array_buffer_refill_with() {
401        let mut buffer = ArrayBuffer::<u32, 4>::new();
402
403        unsafe {
404            buffer.refill_with(|array| {
405                array.copy_from_slice(&[
406                    MaybeUninit::new(1),
407                    MaybeUninit::new(2),
408                    MaybeUninit::new(3),
409                    MaybeUninit::new(4),
410                ]);
411
412                4
413            });
414        };
415
416        assert_eq!(buffer.len(), 4);
417        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
418    }
419
420    #[test]
421    fn test_array_buffer_set_len() {
422        let mut buffer = ArrayBuffer::<u32, 4>::new();
423
424        unsafe { buffer.set_len(1) };
425
426        assert_eq!(buffer.len(), 1);
427    }
428}