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}