1use crate::hints::{assert_hint, likely, unlikely};
3use std::mem::MaybeUninit;
4use std::ops::{Deref, DerefMut};
5use std::mem;
6
7pub struct ArrayBuffer<T, const N: usize> {
31    array: [MaybeUninit<T>; N],
32    len: usize,
33}
34
35impl<T, const N: usize> ArrayBuffer<T, N> {
36    pub fn new() -> Self {
38        #[allow(
39            clippy::uninit_assumed_init,
40            reason = "We guarantee that the array is initialized, when reading from it"
41        )]
42        {
43            Self {
44                array: [const { MaybeUninit::uninit() }; N],
45                len: 0,
46            }
47        }
48    }
49
50    pub const fn capacity(&self) -> usize {
52        N
53    }
54
55    pub fn len(&self) -> usize {
57        self.len
58    }
59
60    pub fn is_empty(&self) -> bool {
62        self.len == 0
63    }
64
65
66    pub unsafe fn push_unchecked(&mut self, item: T) {
72        assert_hint(self.len() < N, "Tried to push to a full array buffer");
73
74        self.array[self.len].write(item);
75        self.len += 1;
76    }
77
78    pub fn push(&mut self, item: T) -> Result<(), T> {
80        if unlikely(self.len == self.capacity()) {
81            return Err(item);
82        }
83
84        unsafe { self.push_unchecked(item) };
85
86        Ok(())
87    }
88
89    pub fn pop(&mut self) -> Option<T> {
91        if unlikely(self.len == 0) {
92            return None;
93        }
94
95        self.len -= 1;
96
97        Some(unsafe { self.array[self.len].as_ptr().read() })
98    }
99
100    pub fn clear_with<F>(&mut self, mut f: F)
102    where
103        F: FnMut(T),
104    {
105        for i in 0..self.len {
106            f(unsafe { self.array[i].as_ptr().read() });
107        }
108
109        self.len = 0;
110    }
111
112    pub fn clear(&mut self) {
114        if mem::needs_drop::<T>() {
115            for i in 0..self.len {
116                drop(unsafe { self.array[i].as_ptr().read() });
117            }
118        }
119
120        self.len = 0;
121    }
122
123    pub fn iter(&self) -> impl ExactSizeIterator<Item = &T> {
125        struct Iter<'array_buffer, T, const N: usize> {
126            buffer: &'array_buffer ArrayBuffer<T, N>,
127            current: *const T,
128            end: *const T,
129        }
130
131        impl<'array_buffer, T, const N: usize> Iterator for Iter<'array_buffer, T, N> {
132            type Item = &'array_buffer T;
133
134            fn next(&mut self) -> Option<Self::Item> {
135                if likely(self.current < self.end) {
136                    let item = unsafe { &*self.current };
137
138                    unsafe { self.current = self.current.add(1); }
139
140                    Some(item)
141                } else {
142                    None
143                }
144            }
145
146            fn size_hint(&self) -> (usize, Option<usize>) {
147                let size = (self.end as usize - self.current as usize) / size_of::<T>();
148
149                (size, Some(size))
150            }
151        }
152
153        impl<T, const N: usize> ExactSizeIterator for Iter<'_, T, N> {
154            fn len(&self) -> usize {
155                self.buffer.len
156            }
157        }
158
159        let current = (&raw const self.array[0]).cast();
160
161        Iter {
162            buffer: self,
163            current,
164            end: unsafe { current.add(self.len) },
165        }
166    }
167
168    pub fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = &mut T> {
170        struct IterMut<'array_buffer, T, const N: usize> {
171            buffer: &'array_buffer mut ArrayBuffer<T, N>,
172            current: *mut T,
173            end: *mut T,
174        }
175
176        impl<'array_buffer, T, const N: usize> Iterator for IterMut<'array_buffer, T, N> {
177            type Item = &'array_buffer mut T;
178
179            fn next(&mut self) -> Option<Self::Item> {
180                if likely(self.current < self.end) {
181                    let item = unsafe { &mut *self.current };
182
183                    unsafe { self.current = self.current.add(1); }
184
185                    Some(item)
186                } else {
187                    None
188                }
189            }
190
191            fn size_hint(&self) -> (usize, Option<usize>) {
192                let size = (self.end as usize - self.current as usize) / size_of::<T>();
193
194                (size, Some(size))
195            }
196        }
197
198        impl<T, const N: usize> ExactSizeIterator for IterMut<'_, T, N> {
199            fn len(&self) -> usize {
200                self.buffer.len
201            }
202        }
203
204        let current: *mut T = (&raw mut self.array[0]).cast();
205        let end = unsafe { current.add(self.len) };
206
207        IterMut {
208            buffer: self,
209            current,
210            end,
211        }
212    }
213
214    pub unsafe fn refill_with(&mut self, f: impl FnOnce(&mut [MaybeUninit<T>; N]) -> usize) {
220        debug_assert!(self.is_empty(), "ArrayBuffer should be empty before refilling");
221
222        let filled = f(&mut self.array);
223
224        debug_assert!(filled <= N, "Filled more than the capacity");
225
226        self.len = filled;
227    }
228    fn as_slice_ptr(&self) -> *const [T; N] {
230        (&raw const self.array).cast()
231    }
232
233    fn as_mut_slice_ptr(&mut self) -> *mut [T; N] {
235        (&raw mut self.array).cast()
236    }
237
238}
239
240impl<T, const N: usize> Deref for ArrayBuffer<T, N> {
241    type Target = [T];
242
243    fn deref(&self) -> &Self::Target {
244        unsafe { &*self.as_slice_ptr() }
245    }
246}
247
248impl<T, const N: usize> AsRef<[T]> for ArrayBuffer<T, N> {
249    fn as_ref(&self) -> &[T] {
250        unsafe { &*self.as_slice_ptr() }
251    }
252}
253
254impl<T, const N: usize> DerefMut for ArrayBuffer<T, N> {
255    fn deref_mut(&mut self) -> &mut Self::Target {
256        unsafe { &mut *self.as_mut_slice_ptr() }
257    }
258}
259
260impl<T, const N: usize> AsMut<[T]> for ArrayBuffer<T, N> {
261    fn as_mut(&mut self) -> &mut [T] {
262        unsafe { &mut *self.as_mut_slice_ptr() }
263    }
264}
265
266impl<T, const N: usize> Default for ArrayBuffer<T, N> {
267    fn default() -> Self {
268        Self::new()
269    }
270}
271
272impl<T, const N: usize> From<[T; N]> for ArrayBuffer<T, N> {
273    fn from(array: [T; N]) -> Self {
274        Self {
275            array: unsafe { (&raw const array).cast::<[MaybeUninit<T>; N]>().read() },
276            len: N,
277        }
278    }
279}
280
281impl<T, const N: usize> Drop for ArrayBuffer<T, N> {
282    fn drop(&mut self) {
283        self.clear();
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn test_array_buffer_pop_push_len() {
293        let mut buffer = ArrayBuffer::<u32, 4>::new();
294
295        unsafe {
296            buffer.push_unchecked(1);
297            assert_eq!(buffer.len(), 1);
298
299            buffer.push_unchecked(2);
300            assert_eq!(buffer.len(), 2);
301
302            buffer.push(3).unwrap();
303            assert_eq!(buffer.len(), 3);
304
305            assert_eq!(buffer.pop(), Some(3));
306            assert_eq!(buffer.len(), 2);
307
308            buffer.push_unchecked(4);
309            assert_eq!(buffer.len(), 3);
310
311            buffer.push_unchecked(5);
312            assert_eq!(buffer.len(), 4);
313
314            assert_eq!(buffer.push(6), Err(6));
315
316            assert_eq!(buffer.pop(), Some(5));
317            assert_eq!(buffer.pop(), Some(4));
318            assert_eq!(buffer.pop(), Some(2));
319            assert_eq!(buffer.pop(), Some(1));
320            assert_eq!(buffer.pop(), None);
321        }
322    }
323
324    #[test]
325    fn test_array_buffer_iterators() {
326        let mut buffer = ArrayBuffer::<u32, 4>::new();
327
328        unsafe {
329            buffer.push_unchecked(1);
330            buffer.push_unchecked(2);
331            buffer.push_unchecked(3);
332            buffer.push_unchecked(4);
333        }
334
335        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
336        assert_eq!(buffer.iter_mut().collect::<Vec<_>>(), vec![&mut 1, &mut 2, &mut 3, &mut 4]);
337    }
338
339    #[test]
340    fn test_array_buffer_refill_with() {
341        let mut buffer = ArrayBuffer::<u32, 4>::new();
342
343        unsafe {
344        buffer.refill_with(|array| {
345                array.copy_from_slice(&[
346                    MaybeUninit::new(1),
347                    MaybeUninit::new(2),
348                    MaybeUninit::new(3),
349                    MaybeUninit::new(4)
350                ]);
351
352                4
353            });
354        };
355
356        assert_eq!(buffer.len(), 4);
357        assert_eq!(buffer.iter().collect::<Vec<_>>(), vec![&1, &2, &3, &4]);
358    }
359}