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 std::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, const N: usize> {
145            buffer: &'array_buffer ArrayBuffer<T, N>,
146            current: *const T,
147            end: *const T,
148        }
149
150        impl<'array_buffer, T, const N: usize> Iterator for Iter<'array_buffer, T, N> {
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                    unsafe {
158                        self.current = self.current.add(1);
159                    }
160
161                    Some(item)
162                } else {
163                    None
164                }
165            }
166
167            fn size_hint(&self) -> (usize, Option<usize>) {
168                let size = (self.end as usize - self.current as usize) / size_of::<T>();
169
170                (size, Some(size))
171            }
172        }
173
174        impl<T, const N: usize> ExactSizeIterator for Iter<'_, T, N> {
175            fn len(&self) -> usize {
176                self.buffer.len
177            }
178        }
179
180        let current = (&raw const self.array[0]).cast();
181
182        Iter {
183            buffer: self,
184            current,
185            end: unsafe { current.add(self.len) },
186        }
187    }
188
189    /// Returns a mutable reference iterator over the buffer.
190    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
191        struct IterMut<'array_buffer, T, const N: usize> {
192            buffer: &'array_buffer mut ArrayBuffer<T, N>,
193            current: *mut T,
194            end: *mut T,
195        }
196
197        impl<'array_buffer, T, const N: usize> Iterator for IterMut<'array_buffer, T, N> {
198            type Item = &'array_buffer mut T;
199
200            fn next(&mut self) -> Option<Self::Item> {
201                if likely(self.current < self.end) {
202                    let item = unsafe { &mut *self.current };
203
204                    unsafe {
205                        self.current = self.current.add(1);
206                    }
207
208                    Some(item)
209                } else {
210                    None
211                }
212            }
213
214            fn size_hint(&self) -> (usize, Option<usize>) {
215                let size = (self.end as usize - self.current as usize) / size_of::<T>();
216
217                (size, Some(size))
218            }
219        }
220
221        impl<T, const N: usize> ExactSizeIterator for IterMut<'_, T, N> {
222            fn len(&self) -> usize {
223                self.buffer.len
224            }
225        }
226
227        let current: *mut T = (&raw mut self.array[0]).cast();
228        let end = unsafe { current.add(self.len) };
229
230        IterMut {
231            buffer: self,
232            current,
233            end,
234        }
235    }
236
237    /// Refills the buffer with elements provided by the function.
238    ///
239    /// # Safety
240    ///
241    /// The caller must ensure that the buffer is empty before refilling.
242    pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
243        debug_assert!(
244            self.is_empty(),
245            "ArrayBuffer should be empty before refilling"
246        );
247
248        let filled = f(&mut self.array);
249
250        debug_assert!(filled <= N, "Filled more than the capacity");
251
252        self.len = filled;
253    }
254    /// Returns a pointer to the underlying array.
255    fn as_slice_ptr(&self) -> *const [T] {
256        slice_from_raw_parts(self.as_ptr(), self.len)
257    }
258
259    /// Returns a mutable pointer to the underlying array.
260    fn as_mut_slice_ptr(&mut self) -> *mut [T] {
261        slice_from_raw_parts_mut(self.as_mut_ptr(), self.len)
262    }
263}
264
265impl<T, const N: usize> Deref for ArrayBuffer<T, N> {
266    type Target = [T];
267
268    fn deref(&self) -> &Self::Target {
269        unsafe { &*self.as_slice_ptr() }
270    }
271}
272
273impl<T, const N: usize> AsRef<[T]> for ArrayBuffer<T, N> {
274    fn as_ref(&self) -> &[T] {
275        unsafe { &*self.as_slice_ptr() }
276    }
277}
278
279impl<T, const N: usize> DerefMut for ArrayBuffer<T, N> {
280    fn deref_mut(&mut self) -> &mut Self::Target {
281        unsafe { &mut *self.as_mut_slice_ptr() }
282    }
283}
284
285impl<T, const N: usize> AsMut<[T]> for ArrayBuffer<T, N> {
286    fn as_mut(&mut self) -> &mut [T] {
287        unsafe { &mut *self.as_mut_slice_ptr() }
288    }
289}
290
291impl<T, const N: usize> Default for ArrayBuffer<T, N> {
292    fn default() -> Self {
293        Self::new()
294    }
295}
296
297impl<T, const N: usize> From<[T; N]> for ArrayBuffer<T, N> {
298    fn from(array: [T; N]) -> Self {
299        Self {
300            array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
301            len: N,
302        }
303    }
304}
305
306impl<T, const N: usize> Drop for ArrayBuffer<T, N> {
307    fn drop(&mut self) {
308        self.clear();
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use alloc::vec;
316    use alloc::vec::Vec;
317
318    #[allow(
319        clippy::explicit_auto_deref,
320        reason = "We test deref and deref_mut methods"
321    )]
322    #[test]
323    fn test_array_buffer_pop_push_len() {
324        let mut buffer = ArrayBuffer::<u32, 4>::new();
325
326        unsafe {
327            buffer.push_unchecked(1);
328            assert_eq!(buffer.len(), 1);
329            assert_eq!((*buffer).len(), 1);
330
331            buffer.push_unchecked(2);
332            assert_eq!(buffer.len(), 2);
333            assert_eq!((*buffer).len(), 2);
334
335            buffer.push(3).unwrap();
336            assert_eq!(buffer.len(), 3);
337            assert_eq!(buffer.as_ref().len(), 3);
338
339            assert_eq!(buffer.pop(), Some(3));
340            assert_eq!(buffer.len(), 2);
341            assert_eq!(buffer.as_mut().len(), 2);
342
343            buffer.push_unchecked(4);
344            assert_eq!(buffer.len(), 3);
345            assert_eq!(buffer.deref_mut().len(), 3);
346
347            buffer.push_unchecked(5);
348            assert_eq!(buffer.len(), 4);
349            assert_eq!(buffer.deref_mut().len(), 4);
350
351            assert_eq!(buffer.push(6), Err(6));
352
353            assert_eq!(buffer.pop(), Some(5));
354            assert_eq!(buffer.pop(), Some(4));
355            assert_eq!(buffer.pop(), Some(2));
356            assert_eq!(buffer.pop(), Some(1));
357            assert_eq!(buffer.pop(), None);
358        }
359    }
360
361    #[test]
362    fn test_array_buffer_iterators() {
363        let mut buffer = ArrayBuffer::<u32, 4>::new();
364
365        unsafe {
366            buffer.push_unchecked(1);
367            buffer.push_unchecked(2);
368            buffer.push_unchecked(3);
369            buffer.push_unchecked(4);
370        }
371
372        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
373        assert_eq!(
374            buffer.iter_mut().collect::<Vec<_>>(),
375            vec![&mut 1, &mut 2, &mut 3, &mut 4]
376        );
377    }
378
379    #[test]
380    fn test_array_buffer_refill_with() {
381        let mut buffer = ArrayBuffer::<u32, 4>::new();
382
383        unsafe {
384            buffer.refill_with(|array| {
385                array.copy_from_slice(&[
386                    MaybeUninit::new(1),
387                    MaybeUninit::new(2),
388                    MaybeUninit::new(3),
389                    MaybeUninit::new(4),
390                ]);
391
392                4
393            });
394        };
395
396        assert_eq!(buffer.len(), 4);
397        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
398    }
399
400    #[test]
401    fn test_array_buffer_set_len() {
402        let mut buffer = ArrayBuffer::<u32, 4>::new();
403
404        unsafe { buffer.set_len(1) };
405
406        assert_eq!(buffer.len(), 1);
407    }
408}
orengine_utils/array_buffer.rs

orengine_utils/
array_buffer.rs
