1use crate::Array;
8use core::mem::ManuallyDrop;
9use core::ops::{Deref, DerefMut};
10use kalloc::{AllocError, Allocator, Box, DefaultAllocator};
11
12#[macro_export]
36macro_rules! try_vec {
37 ($($x:expr),* $(,)?) => {
38 {
39 let mut v = $crate::Vector::new();
40 let count = 0 $(+ { let _ = stringify!($x); 1 })*;
41 let f = || {
42 if count > 0 {
43 v.reserve(count)?;
44 }
45 $(
46 v.push_back($x)?;
47 )*
48 Ok(v)
49 };
50 f()
51 }
52 };
53
54 ($elem:expr; $n:expr) => {
55 {
56 let mut v = $crate::Vector::new();
57 let n = $n;
58 let f = || {
59 if n > 0 {
60 v.reserve(n)?;
61 for _ in 0..n {
62 v.push_back($elem.clone())?;
63 }
64 }
65 Ok(v)
66 };
67 f()
68 }
69 };
70}
71
72pub struct Vector<T, A: Allocator = DefaultAllocator> {
81 buf: Box<[core::mem::MaybeUninit<T>], A>,
82
83 size: usize,
88}
89
90const CAPACITY_MINIMUM: usize = 16;
91const CAPACITY_GROWTH_FACTOR: usize = 2;
92const CAPACITY_SHRINK_FACTOR: usize = 4;
93
94zr::static_assert!(core::mem::size_of::<Vector<u32>>() == 24);
96zr::static_assert!(core::mem::align_of::<Vector<u32>>() == 8);
97
98impl<T, A: Allocator> Vector<T, A> {
99 pub const fn new_in(allocator: A) -> Self {
101 Vector { buf: Box::empty_slice_in(allocator), size: 0 }
102 }
103
104 pub fn len(&self) -> usize {
106 self.size
107 }
108
109 pub fn capacity(&self) -> usize {
111 self.buf.len()
112 }
113
114 pub fn is_empty(&self) -> bool {
116 self.size == 0
117 }
118
119 pub fn reserve(&mut self, new_capacity: usize) -> Result<(), AllocError> {
121 if new_capacity <= self.buf.len() {
122 return Ok(());
123 }
124 self.reallocate(new_capacity)
125 }
126
127 pub fn clear(&mut self) {
129 self.truncate(0);
130 }
131
132 pub fn swap(&mut self, other: &mut Self) {
134 core::mem::swap(self, other);
135 }
136
137 pub fn push_back(&mut self, value: T) -> Result<(), AllocError> {
139 self.grow_for_new_element()?;
140 self.buf[self.size].write(value);
141 self.size += 1;
142 Ok(())
143 }
144
145 pub fn pop_back(&mut self) -> Option<T> {
147 if self.is_empty() {
148 return None;
149 }
150 self.size -= 1;
151 let val = unsafe { self.buf[self.size].assume_init_read() };
154 self.consider_shrinking();
155 Some(val)
156 }
157
158 pub fn insert(&mut self, index: usize, value: T) -> Result<(), AllocError> {
160 assert!(index <= self.size);
161 self.push_back(value)?;
162 let size = self.size;
163 self[index..size].rotate_right(1);
164 Ok(())
165 }
166
167 pub fn erase(&mut self, index: usize) -> T {
169 assert!(index < self.size);
170 let size = self.size;
171 self[index..size].rotate_left(1);
172 self.pop_back().unwrap()
173 }
174
175 pub fn truncate(&mut self, new_len: usize) {
178 if new_len >= self.size {
179 return;
180 }
181 let old_size = self.size;
182 self.size = new_len;
183 unsafe {
185 core::ptr::drop_in_place(self.buf[new_len..old_size].assume_init_mut());
186 }
187 self.consider_shrinking();
188 }
189
190 pub fn resize_with_default(&mut self, new_size: usize) -> Result<(), AllocError>
195 where
196 T: Default,
197 {
198 self.resize_with(new_size, T::default)
199 }
200
201 pub fn resize(&mut self, new_size: usize, value: T) -> Result<(), AllocError>
206 where
207 T: Clone,
208 {
209 self.resize_with(new_size, || value.clone())
210 }
211
212 pub fn resize_with<F>(&mut self, new_size: usize, mut f: F) -> Result<(), AllocError>
217 where
218 F: FnMut() -> T,
219 {
220 if new_size <= self.size {
221 self.truncate(new_size);
222 } else {
223 self.reserve(new_size)?;
224 while self.size < new_size {
225 self.push_back(f())?;
226 }
227 }
228 Ok(())
229 }
230
231 fn reallocate(&mut self, new_capacity: usize) -> Result<(), AllocError> {
233 assert!(new_capacity > 0);
234 assert!(new_capacity >= self.size);
235
236 if new_capacity > self.buf.len() {
237 Box::try_grow(&mut self.buf, new_capacity)?;
238 } else if new_capacity < self.buf.len() {
239 unsafe {
242 Box::try_shrink(&mut self.buf, new_capacity)?;
243 }
244 }
245 Ok(())
246 }
247
248 fn grow_for_new_element(&mut self) -> Result<(), AllocError> {
251 if self.size == self.buf.len() {
252 let new_capacity = if self.buf.is_empty() {
253 CAPACITY_MINIMUM
254 } else {
255 self.buf.len() * CAPACITY_GROWTH_FACTOR
256 };
257 self.reallocate(new_capacity)?;
258 }
259 Ok(())
260 }
261
262 fn consider_shrinking(&mut self) {
264 if self.size * CAPACITY_SHRINK_FACTOR < self.buf.len() && self.buf.len() > CAPACITY_MINIMUM
265 {
266 let new_capacity =
267 core::cmp::max(CAPACITY_MINIMUM, self.buf.len() / CAPACITY_SHRINK_FACTOR);
268 let _ = self.reallocate(new_capacity);
270 }
271 }
272
273 pub fn try_from_iter_in<I: IntoIterator<Item = T>>(
278 iter: I,
279 allocator: A,
280 ) -> Result<Self, AllocError> {
281 let mut v = Vector::new_in(allocator);
282 let iter = iter.into_iter();
283
284 let (lower, _) = iter.size_hint();
285 if lower > 0 {
286 v.reserve(lower)?;
287 }
288
289 for item in iter {
290 v.push_back(item)?;
291 }
292 Ok(v)
293 }
294
295 pub fn try_into_array(mut self) -> Result<Array<T, A>, AllocError> {
301 if self.size < self.buf.len() {
302 unsafe { Box::try_shrink(&mut self.buf, self.size)? };
305 }
306
307 let this = ManuallyDrop::new(self);
308 let buf = unsafe { core::ptr::read(&this.buf) };
310 Ok(Array::from_box(unsafe { buf.assume_init() }))
312 }
313}
314
315impl<T, A: Allocator> Deref for Vector<T, A> {
316 type Target = [T];
317
318 fn deref(&self) -> &Self::Target {
319 unsafe { self.buf[0..self.size].assume_init_ref() }
321 }
322}
323
324impl<T, A: Allocator> DerefMut for Vector<T, A> {
325 fn deref_mut(&mut self) -> &mut Self::Target {
326 unsafe { self.buf[0..self.size].assume_init_mut() }
328 }
329}
330
331impl<T, A: Allocator> Drop for Vector<T, A> {
332 fn drop(&mut self) {
333 unsafe { core::ptr::drop_in_place(self.buf[..self.size].assume_init_mut()) }
335 }
336}
337
338impl<T> Vector<T, DefaultAllocator> {
339 pub const fn new() -> Self {
341 Vector { buf: Box::empty_slice(), size: 0 }
342 }
343
344 pub fn try_from_iter<I: IntoIterator<Item = T>>(iter: I) -> Result<Self, AllocError> {
346 let mut v = Vector::new();
347 let iter = iter.into_iter();
348
349 let (lower, _) = iter.size_hint();
350 if lower > 0 {
351 v.reserve(lower)?;
352 }
353
354 for item in iter {
355 v.push_back(item)?;
356 }
357 Ok(v)
358 }
359}
360
361impl<T> Default for Vector<T, DefaultAllocator> {
362 fn default() -> Self {
363 Self::new()
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 extern crate std;
370
371 use super::*;
372
373 use core::cell::Cell;
374 use core::ptr::NonNull;
375
376 #[derive(Debug, PartialEq, Eq)]
377 struct TestState {
378 live_obj_count: Cell<usize>,
379 ctor_count: Cell<usize>,
380 dtor_count: Cell<usize>,
381 alloc_count: Cell<usize>,
382 fail_threshold: Cell<usize>,
383 }
384
385 impl Default for TestState {
386 fn default() -> Self {
387 Self {
388 live_obj_count: Cell::new(0),
389 ctor_count: Cell::new(0),
390 dtor_count: Cell::new(0),
391 alloc_count: Cell::new(0),
392 fail_threshold: Cell::new(usize::MAX),
393 }
394 }
395 }
396
397 #[derive(Clone)]
398 struct TestAllocator<'a> {
399 state: &'a TestState,
400 }
401
402 impl<'a> Allocator for TestAllocator<'a> {
403 fn allocate(&self, layout: core::alloc::Layout) -> Result<NonNull<[u8]>, AllocError> {
404 let current = self.state.alloc_count.get();
405 self.state.alloc_count.set(current + 1);
406 if current >= self.state.fail_threshold.get() {
407 return Err(AllocError);
408 }
409 DefaultAllocator::default().allocate(layout)
410 }
411
412 unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: core::alloc::Layout) {
413 unsafe { DefaultAllocator::default().deallocate(ptr, layout) }
414 }
415
416 unsafe fn grow(
417 &self,
418 ptr: NonNull<u8>,
419 old_layout: core::alloc::Layout,
420 new_layout: core::alloc::Layout,
421 ) -> Result<NonNull<[u8]>, AllocError> {
422 let current = self.state.alloc_count.get();
423 self.state.alloc_count.set(current + 1);
424 if current >= self.state.fail_threshold.get() {
425 return Err(AllocError);
426 }
427 unsafe { DefaultAllocator::default().grow(ptr, old_layout, new_layout) }
428 }
429
430 unsafe fn shrink(
431 &self,
432 ptr: NonNull<u8>,
433 old_layout: core::alloc::Layout,
434 new_layout: core::alloc::Layout,
435 ) -> Result<NonNull<[u8]>, AllocError> {
436 let current = self.state.alloc_count.get();
437 self.state.alloc_count.set(current + 1);
438 if current >= self.state.fail_threshold.get() {
439 return Err(AllocError);
440 }
441 unsafe { DefaultAllocator::default().shrink(ptr, old_layout, new_layout) }
442 }
443
444 fn allocate_zeroed(
445 &self,
446 layout: core::alloc::Layout,
447 ) -> Result<NonNull<[u8]>, AllocError> {
448 let current = self.state.alloc_count.get();
449 self.state.alloc_count.set(current + 1);
450 if current >= self.state.fail_threshold.get() {
451 return Err(AllocError);
452 }
453 DefaultAllocator::default().allocate_zeroed(layout)
454 }
455 }
456
457 #[derive(Debug, Eq, PartialEq)]
458 struct TestObject<'a> {
459 val: usize,
460 alive: bool,
461 state: &'a TestState,
462 }
463
464 impl<'a> TestObject<'a> {
465 fn new(val: usize, state: &'a TestState) -> Self {
466 state.live_obj_count.set(state.live_obj_count.get() + 1);
467 state.ctor_count.set(state.ctor_count.get() + 1);
468 TestObject { val, alive: true, state }
469 }
470 }
471
472 impl<'a> Drop for TestObject<'a> {
473 fn drop(&mut self) {
474 if self.alive {
475 self.state.live_obj_count.set(self.state.live_obj_count.get() - 1);
476 self.state.dtor_count.set(self.state.dtor_count.get() + 1);
477 }
478 }
479 }
480
481 #[test]
482 fn test_empty() {
483 let v: Vector<u32> = Vector::new();
484 assert_eq!(v.len(), 0);
485 assert_eq!(v.capacity(), 0);
486 assert!(v.is_empty());
487 }
488
489 #[test]
490 fn test_push_pop() {
491 let mut v: Vector<u32> = Vector::new();
492 v.push_back(1).unwrap();
493 v.push_back(2).unwrap();
494 v.push_back(3).unwrap();
495
496 assert_eq!(v.len(), 3);
497 assert_eq!(v[0], 1);
498 assert_eq!(v[1], 2);
499 assert_eq!(v[2], 3);
500
501 assert_eq!(v.pop_back(), Some(3));
502 assert_eq!(v.pop_back(), Some(2));
503 assert_eq!(v.pop_back(), Some(1));
504 assert_eq!(v.pop_back(), None);
505 }
506
507 #[test]
508 fn test_insert_erase() {
509 let mut v: Vector<u32> = Vector::new();
510 v.push_back(1).unwrap();
511 v.push_back(3).unwrap();
512
513 v.insert(1, 2).unwrap();
514 assert_eq!(v.len(), 3);
515 assert_eq!(v[0], 1);
516 assert_eq!(v[1], 2);
517 assert_eq!(v[2], 3);
518
519 assert_eq!(v.erase(1), 2);
520 assert_eq!(v.len(), 2);
521 assert_eq!(v[0], 1);
522 assert_eq!(v[1], 3);
523 }
524
525 #[test]
526 fn test_resize() {
527 let mut v: Vector<u32> = Vector::new();
528 v.resize_with_default(5).unwrap();
529 assert_eq!(v.len(), 5);
530 for i in 0..5 {
531 assert_eq!(v[i], 0);
532 }
533
534 v.resize_with_default(2).unwrap();
535 assert_eq!(v.len(), 2);
536 }
537
538 #[test]
539 fn test_drop_behavior() {
540 let state = TestState::default();
541 {
542 let mut v: Vector<TestObject<'_>, TestAllocator<'_>> =
543 Vector::new_in(TestAllocator { state: &state });
544 v.push_back(TestObject::new(1, &state)).unwrap();
545 v.push_back(TestObject::new(2, &state)).unwrap();
546 assert_eq!(state.live_obj_count.get(), 2);
547 }
548 assert_eq!(state.live_obj_count.get(), 0);
549 assert_eq!(state.ctor_count.get(), 2);
550 assert_eq!(state.dtor_count.get(), 2);
551 }
552
553 #[test]
554 fn test_counting_allocator() {
555 let state = TestState::default();
556 {
557 let mut v: Vector<u32, TestAllocator<'_>> =
558 Vector::new_in(TestAllocator { state: &state });
559 v.push_back(1).unwrap(); assert_eq!(state.alloc_count.get(), 1);
561 }
562 }
563
564 #[test]
565 fn test_failing_allocator() {
566 let state = TestState::default();
567 state.fail_threshold.set(1); let mut v: Vector<u32, TestAllocator<'_>> = Vector::new_in(TestAllocator { state: &state });
570 v.push_back(1).unwrap(); for i in 2..=16 {
574 v.push_back(i).unwrap();
575 }
576 assert_eq!(v.len(), 16);
577 assert_eq!(v.capacity(), 16);
578
579 assert_eq!(v.push_back(17), Err(AllocError));
583 assert_eq!(v.len(), 16); for i in 0..16 {
587 assert_eq!(v[i], (i + 1) as u32);
588 }
589 }
590
591 #[test]
592 fn test_truncate() {
593 let mut v: Vector<u32> = Vector::new();
594 v.push_back(1).unwrap();
595 v.push_back(2).unwrap();
596 v.push_back(3).unwrap();
597
598 v.truncate(2);
599 assert_eq!(v.len(), 2);
600 assert_eq!(v[0], 1);
601 assert_eq!(v[1], 2);
602
603 v.truncate(5); assert_eq!(v.len(), 2);
605 }
606
607 #[test]
608 fn test_truncate_drops_elements() {
609 let state = TestState::default();
610 {
611 let mut v: Vector<TestObject<'_>, TestAllocator<'_>> =
612 Vector::new_in(TestAllocator { state: &state });
613 v.push_back(TestObject::new(1, &state)).unwrap();
614 v.push_back(TestObject::new(2, &state)).unwrap();
615 v.push_back(TestObject::new(3, &state)).unwrap();
616
617 assert_eq!(state.live_obj_count.get(), 3);
618
619 v.truncate(1);
620 assert_eq!(v.len(), 1);
621 assert_eq!(state.live_obj_count.get(), 1);
622 assert_eq!(state.dtor_count.get(), 2);
623 }
624 assert_eq!(state.live_obj_count.get(), 0);
625 }
626
627 #[test]
628 fn test_iterator() {
629 let mut v: Vector<u32> = Vector::new();
630 v.push_back(1).unwrap();
631 v.push_back(2).unwrap();
632
633 let mut it = v.iter();
634 assert_eq!(it.next(), Some(&1));
635 assert_eq!(it.next(), Some(&2));
636 assert_eq!(it.next(), None);
637
638 for x in v.iter_mut() {
639 *x += 10;
640 }
641
642 assert_eq!(v[0], 11);
643 assert_eq!(v[1], 12);
644 }
645
646 #[test]
647 fn test_box() {
648 let state = TestState::default();
649 {
650 let mut v: Vector<Box<TestObject<'_>, TestAllocator<'_>>, TestAllocator<'_>> =
651 Vector::new_in(TestAllocator { state: &state });
652 v.push_back(
653 Box::try_new_in(TestObject::new(1, &state), TestAllocator { state: &state })
654 .unwrap(),
655 )
656 .unwrap();
657 assert_eq!(v.len(), 1);
658 assert_eq!(v[0].val, 1);
659 }
660 }
661
662 #[test]
663 fn test_try_from_iter() {
664 let items = [1, 2, 3, 4, 5];
665 let v: Vector<u32> = Vector::try_from_iter(items.iter().copied()).unwrap();
666 assert_eq!(v.len(), 5);
667 assert_eq!(v[0], 1);
668 assert_eq!(v[4], 5);
669 }
670
671 #[test]
672 fn test_try_from_iter_failing() {
673 let state = TestState::default();
674 state.fail_threshold.set(0); let items = [1, 2, 3, 4, 5];
677 let v: Result<Vector<u32, TestAllocator<'_>>, AllocError> =
678 Vector::try_from_iter_in(items.iter().copied(), TestAllocator { state: &state });
679 assert!(v.is_err());
680 }
681
682 #[test]
683 fn test_try_vec_macro() {
684 let v: Result<Vector<u32>, AllocError> = try_vec![1, 2, 3];
685 let v = v.unwrap();
686 assert_eq!(v.len(), 3);
687 assert_eq!(v[0], 1);
688 assert_eq!(v[2], 3);
689
690 let v2: Result<Vector<u32>, AllocError> = try_vec![0; 5];
691 let v2 = v2.unwrap();
692 assert_eq!(v2.len(), 5);
693 for i in 0..5 {
694 assert_eq!(v2[i], 0);
695 }
696 }
697
698 #[test]
699 fn test_try_vec_macro_failing() {
700 let state = TestState::default();
701 state.fail_threshold.set(0); let mut v: Vector<u32, TestAllocator<'_>> = Vector::new_in(TestAllocator { state: &state });
704 assert!(v.push_back(1).is_err());
705 }
706
707 #[test]
708 fn test_swap() {
709 let v1: Result<Vector<u32>, AllocError> = try_vec![1, 2, 3];
710 let mut v1 = v1.unwrap();
711 let v2: Result<Vector<u32>, AllocError> = try_vec![4, 5];
712 let mut v2 = v2.unwrap();
713
714 v1.swap(&mut v2);
715
716 assert_eq!(v1.len(), 2);
717 assert_eq!(v1[0], 4);
718 assert_eq!(v1[1], 5);
719
720 assert_eq!(v2.len(), 3);
721 assert_eq!(v2[0], 1);
722 assert_eq!(v2[1], 2);
723 assert_eq!(v2[2], 3);
724 }
725
726 #[test]
727 fn test_resize_with_value() {
728 let mut v: Vector<u32> = Vector::new();
729 v.resize(3, 42).unwrap();
730 assert_eq!(v.len(), 3);
731 assert_eq!(v[0], 42);
732 assert_eq!(v[1], 42);
733 assert_eq!(v[2], 42);
734
735 v.resize(1, 10).unwrap();
736 assert_eq!(v.len(), 1);
737 assert_eq!(v[0], 42); }
739
740 #[test]
741 fn test_resize_with() {
742 let mut v: Vector<u32> = Vector::new();
743 let mut c = 0;
744 v.resize_with(3, || {
745 c += 1;
746 c
747 })
748 .unwrap();
749 assert_eq!(v.len(), 3);
750 assert_eq!(v[0], 1);
751 assert_eq!(v[1], 2);
752 assert_eq!(v[2], 3);
753 }
754
755 #[test]
756 fn test_try_into_array() {
757 let mut v: Vector<u32> = Vector::new();
758 v.reserve(3).unwrap();
759 v.push_back(1).unwrap();
760 v.push_back(2).unwrap();
761 v.push_back(3).unwrap();
762 assert_eq!(v.capacity(), 3);
763
764 let ptr = v.as_ptr();
765 let a = v.try_into_array().unwrap();
766 assert_eq!(&a[..], &[1, 2, 3]);
767 assert_eq!(a.as_ptr(), ptr);
768 }
769
770 #[test]
771 fn test_try_into_array_shrinks() {
772 let mut v: Vector<u32> = Vector::new();
773 v.push_back(1).unwrap();
774 v.push_back(2).unwrap();
775 assert!(v.capacity() > v.len());
776
777 let a = v.try_into_array().unwrap();
778 assert_eq!(&a[..], &[1, 2]);
779 }
780
781 #[test]
782 fn test_try_into_array_empty() {
783 let v: Vector<u32> = Vector::new();
784 let a = v.try_into_array().unwrap();
785 assert!(a.is_empty());
786
787 let mut v: Vector<u32> = Vector::new();
788 v.reserve(4).unwrap();
789 let a = v.try_into_array().unwrap();
790 assert!(a.is_empty());
791 }
792
793 #[test]
794 fn test_try_into_array_drop_behavior() {
795 let state = TestState::default();
796 {
797 let mut v: Vector<TestObject<'_>, TestAllocator<'_>> =
798 Vector::new_in(TestAllocator { state: &state });
799 v.push_back(TestObject::new(1, &state)).unwrap();
800 v.push_back(TestObject::new(2, &state)).unwrap();
801
802 let a = v.try_into_array().unwrap();
803 assert_eq!(state.live_obj_count.get(), 2);
804 assert_eq!(state.dtor_count.get(), 0);
805 assert_eq!(a[0].val, 1);
806 assert_eq!(a[1].val, 2);
807 }
808 assert_eq!(state.live_obj_count.get(), 0);
809 assert_eq!(state.dtor_count.get(), 2);
810 }
811
812 #[test]
813 fn test_try_into_array_shrink_failure() {
814 let state = TestState::default();
815 {
816 let mut v: Vector<TestObject<'_>, TestAllocator<'_>> =
817 Vector::new_in(TestAllocator { state: &state });
818 v.push_back(TestObject::new(1, &state)).unwrap();
819
820 state.fail_threshold.set(state.alloc_count.get());
822 assert!(v.try_into_array().is_err());
823 }
824 assert_eq!(state.live_obj_count.get(), 0);
825 assert_eq!(state.dtor_count.get(), 1);
826 }
827}