1use 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
8pub 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 pub const fn new() -> Self {
39 Self {
40 array: [const { MaybeUninit::uninit() }; N],
41 len: 0,
42 }
43 }
44
45 pub const fn capacity(&self) -> usize {
47 N
48 }
49
50 pub fn len(&self) -> usize {
52 self.len
53 }
54
55 pub fn is_empty(&self) -> bool {
57 self.len == 0
58 }
59
60 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 pub const fn as_ptr(&self) -> *const T {
77 self.array.as_ptr().cast()
78 }
79
80 pub const fn as_mut_ptr(&mut self) -> *mut T {
82 self.array.as_mut_ptr().cast()
83 }
84
85 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 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 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 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 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 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 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 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 fn as_slice_ptr(&self) -> *const [T] {
261 slice_from_raw_parts(self.as_ptr(), self.len)
262 }
263
264 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}