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    /// Returns a pointer to the first element of the buffer.
61    pub const fn as_ptr(&self) -> *const T {
62        self.array.as_ptr().cast()
63    }
64
65    /// Returns a mutable pointer to the first element of the buffer.
66    pub const fn as_mut_ptr(&mut self) -> *mut T {
67        self.array.as_mut_ptr().cast()
68    }
69
70    /// Appends an element to the buffer.
71    ///
72    /// # Safety
73    ///
74    /// The caller must ensure that the buffer is not full.
75    pub unsafe fn push_unchecked(&mut self, item: T) {
76        assert_hint(self.len() < N, "Tried to push to a full array buffer");
77
78        self.array[self.len].write(item);
79        self.len += 1;
80    }
81
82    /// Appends an element to the buffer or returns `Err(value)` if the buffer is full.
83    pub fn push(&mut self, item: T) -> Result<(), T> {
84        if unlikely(self.len == self.capacity()) {
85            return Err(item);
86        }
87
88        unsafe { self.push_unchecked(item) };
89
90        Ok(())
91    }
92
93    /// Pops an element from the buffer or returns `None` if the buffer is empty.
94    pub fn pop(&mut self) -> Option<T> {
95        if unlikely(self.len == 0) {
96            return None;
97        }
98
99        self.len -= 1;
100
101        Some(unsafe { self.array[self.len].as_ptr().read() })
102    }
103
104    /// Clears with calling the provided function on each element.
105    pub fn clear_with<F>(&mut self, mut f: F)
106    where
107        F: FnMut(T),
108    {
109        for i in 0..self.len {
110            f(unsafe { self.array[i].as_ptr().read() });
111        }
112
113        self.len = 0;
114    }
115
116    /// Drops all elements in the buffer and set the length to 0.
117    pub fn clear(&mut self) {
118        if mem::needs_drop::<T>() {
119            for i in 0..self.len {
120                drop(unsafe { self.array[i].as_ptr().read() });
121            }
122        }
123
124        self.len = 0;
125    }
126
127    /// Returns a reference iterator over the buffer.
128    pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
129        struct Iter<'array_buffer, T, const N: usize> {
130            buffer: &'array_buffer ArrayBuffer<T, N>,
131            current: *const T,
132            end: *const T,
133        }
134
135        impl<'array_buffer, T, const N: usize> Iterator for Iter<'array_buffer, T, N> {
136            type Item = &'array_buffer T;
137
138            fn next(&mut self) -> Option<Self::Item> {
139                if likely(self.current < self.end) {
140                    let item = unsafe { &*self.current };
141
142                    unsafe {
143                        self.current = self.current.add(1);
144                    }
145
146                    Some(item)
147                } else {
148                    None
149                }
150            }
151
152            fn size_hint(&self) -> (usize, Option<usize>) {
153                let size = (self.end as usize - self.current as usize) / size_of::<T>();
154
155                (size, Some(size))
156            }
157        }
158
159        impl<T, const N: usize> ExactSizeIterator for Iter<'_, T, N> {
160            fn len(&self) -> usize {
161                self.buffer.len
162            }
163        }
164
165        let current = (&raw const self.array[0]).cast();
166
167        Iter {
168            buffer: self,
169            current,
170            end: unsafe { current.add(self.len) },
171        }
172    }
173
174    /// Returns a mutable reference iterator over the buffer.
175    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
176        struct IterMut<'array_buffer, T, const N: usize> {
177            buffer: &'array_buffer mut ArrayBuffer<T, N>,
178            current: *mut T,
179            end: *mut T,
180        }
181
182        impl<'array_buffer, T, const N: usize> Iterator for IterMut<'array_buffer, T, N> {
183            type Item = &'array_buffer mut T;
184
185            fn next(&mut self) -> Option<Self::Item> {
186                if likely(self.current < self.end) {
187                    let item = unsafe { &mut *self.current };
188
189                    unsafe {
190                        self.current = self.current.add(1);
191                    }
192
193                    Some(item)
194                } else {
195                    None
196                }
197            }
198
199            fn size_hint(&self) -> (usize, Option<usize>) {
200                let size = (self.end as usize - self.current as usize) / size_of::<T>();
201
202                (size, Some(size))
203            }
204        }
205
206        impl<T, const N: usize> ExactSizeIterator for IterMut<'_, T, N> {
207            fn len(&self) -> usize {
208                self.buffer.len
209            }
210        }
211
212        let current: *mut T = (&raw mut self.array[0]).cast();
213        let end = unsafe { current.add(self.len) };
214
215        IterMut {
216            buffer: self,
217            current,
218            end,
219        }
220    }
221
222    /// Refills the buffer with elements provided by the function.
223    ///
224    /// # Safety
225    ///
226    /// The caller must ensure that the buffer is empty before refilling.
227    pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
228        debug_assert!(
229            self.is_empty(),
230            "ArrayBuffer should be empty before refilling"
231        );
232
233        let filled = f(&mut self.array);
234
235        debug_assert!(filled <= N, "Filled more than the capacity");
236
237        self.len = filled;
238    }
239    /// Returns a pointer to the underlying array.
240    fn as_slice_ptr(&self) -> *const [T] {
241        slice_from_raw_parts(self.as_ptr(), self.len)
242    }
243
244    /// Returns a mutable pointer to the underlying array.
245    fn as_mut_slice_ptr(&mut self) -> *mut [T] {
246        slice_from_raw_parts_mut(self.as_mut_ptr(), self.len)
247    }
248}
249
250impl<T, const N: usize> Deref for ArrayBuffer<T, N> {
251    type Target = [T];
252
253    fn deref(&self) -> &Self::Target {
254        unsafe { &*self.as_slice_ptr() }
255    }
256}
257
258impl<T, const N: usize> AsRef<[T]> for ArrayBuffer<T, N> {
259    fn as_ref(&self) -> &[T] {
260        unsafe { &*self.as_slice_ptr() }
261    }
262}
263
264impl<T, const N: usize> DerefMut for ArrayBuffer<T, N> {
265    fn deref_mut(&mut self) -> &mut Self::Target {
266        unsafe { &mut *self.as_mut_slice_ptr() }
267    }
268}
269
270impl<T, const N: usize> AsMut<[T]> for ArrayBuffer<T, N> {
271    fn as_mut(&mut self) -> &mut [T] {
272        unsafe { &mut *self.as_mut_slice_ptr() }
273    }
274}
275
276impl<T, const N: usize> Default for ArrayBuffer<T, N> {
277    fn default() -> Self {
278        Self::new()
279    }
280}
281
282impl<T, const N: usize> From<[T; N]> for ArrayBuffer<T, N> {
283    fn from(array: [T; N]) -> Self {
284        Self {
285            array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
286            len: N,
287        }
288    }
289}
290
291impl<T, const N: usize> Drop for ArrayBuffer<T, N> {
292    fn drop(&mut self) {
293        self.clear();
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use alloc::vec;
301    use alloc::vec::Vec;
302
303    #[allow(
304        clippy::explicit_auto_deref,
305        reason = "We test deref and deref_mut methods"
306    )]
307    #[test]
308    fn test_array_buffer_pop_push_len() {
309        let mut buffer = ArrayBuffer::<u32, 4>::new();
310
311        unsafe {
312            buffer.push_unchecked(1);
313            assert_eq!(buffer.len(), 1);
314            assert_eq!((*buffer).len(), 1);
315
316            buffer.push_unchecked(2);
317            assert_eq!(buffer.len(), 2);
318            assert_eq!((*buffer).len(), 2);
319
320            buffer.push(3).unwrap();
321            assert_eq!(buffer.len(), 3);
322            assert_eq!(buffer.as_ref().len(), 3);
323
324            assert_eq!(buffer.pop(), Some(3));
325            assert_eq!(buffer.len(), 2);
326            assert_eq!(buffer.as_mut().len(), 2);
327
328            buffer.push_unchecked(4);
329            assert_eq!(buffer.len(), 3);
330            assert_eq!(buffer.deref_mut().len(), 3);
331
332            buffer.push_unchecked(5);
333            assert_eq!(buffer.len(), 4);
334            assert_eq!(buffer.deref_mut().len(), 4);
335
336            assert_eq!(buffer.push(6), Err(6));
337
338            assert_eq!(buffer.pop(), Some(5));
339            assert_eq!(buffer.pop(), Some(4));
340            assert_eq!(buffer.pop(), Some(2));
341            assert_eq!(buffer.pop(), Some(1));
342            assert_eq!(buffer.pop(), None);
343        }
344    }
345
346    #[test]
347    fn test_array_buffer_iterators() {
348        let mut buffer = ArrayBuffer::<u32, 4>::new();
349
350        unsafe {
351            buffer.push_unchecked(1);
352            buffer.push_unchecked(2);
353            buffer.push_unchecked(3);
354            buffer.push_unchecked(4);
355        }
356
357        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
358        assert_eq!(
359            buffer.iter_mut().collect::<Vec<_>>(),
360            vec![&mut 1, &mut 2, &mut 3, &mut 4]
361        );
362    }
363
364    #[test]
365    fn test_array_buffer_refill_with() {
366        let mut buffer = ArrayBuffer::<u32, 4>::new();
367
368        unsafe {
369            buffer.refill_with(|array| {
370                array.copy_from_slice(&[
371                    MaybeUninit::new(1),
372                    MaybeUninit::new(2),
373                    MaybeUninit::new(3),
374                    MaybeUninit::new(4),
375                ]);
376
377                4
378            });
379        };
380
381        assert_eq!(buffer.len(), 4);
382        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
383    }
384}
orengine_utils/array_buffer.rs

orengine_utils/
array_buffer.rs
