fastvec/small.rs
1use alloc::alloc as malloc;
2use alloc::boxed::Box;
3use alloc::vec::Vec;
4use core::alloc::Layout;
5use core::fmt::Debug;
6use core::iter::FusedIterator;
7use core::mem::{ManuallyDrop, MaybeUninit};
8use core::num::NonZeroUsize;
9use core::panic::{RefUnwindSafe, UnwindSafe};
10use core::ptr::{self, NonNull};
11use core::slice;
12
13use super::utils::{IsZST, min_cap, split_range_bound};
14
15// -----------------------------------------------------------------------------
16// Utils
17
18const MAX_LEN: usize = usize::MAX >> 1;
19const MARKER: usize = usize::MAX ^ MAX_LEN;
20
21union Data<T, const N: usize> {
22 heap: (NonNull<T>, usize),
23 cache: [ManuallyDrop<T>; N],
24}
25
26// -----------------------------------------------------------------------------
27// SmallVec
28
29/// A vector with inline storage that spills to heap when capacity is exceeded.
30///
31/// This type implements small-buffer optimization (SBO) while prioritizing the
32/// space efficiency of the container itself:
33/// - Storage mode (inline/heap) is encoded in a single bit of the length field.
34/// - Heap metadata is stored in a compact union with the inline buffer.
35///
36/// Unlike `FastVec`, this type does not cache a data pointer, meaning all data
37/// accesses require an additional branch to determine the storage location.
38/// However, this design makes the container more compact. For example, `SmallVec<u64, 2>`
39/// can have the same size as `Vec<u64>`.
40///
41/// Hot paths are optimized with cold path annotations to minimize the overhead
42/// of branch checks, especially when data is inlined.
43///
44/// # Panics
45/// Any operation that would cause `capacity > isize::MAX` will panic.
46///
47/// # Examples
48///
49/// ```
50/// use fastvec::SmallVec;
51///
52/// // Allocate inline space for 2 elements (uninitialized)
53/// let mut vec: SmallVec<String, 2> = SmallVec::new();
54///
55/// assert_eq!(vec.len(), 0);
56/// assert_eq!(vec.capacity(), 2);
57///
58/// // Use it like a normal `Vec`
59/// vec.push("Hello".to_string());
60/// vec.push(", world!".to_string());
61///
62/// assert_eq!(vec, ["Hello", ", world!"]);
63///
64/// // Convert into a standard `Vec`
65/// let vec: Vec<String> = vec.into_vec();
66/// ```
67pub struct SmallVec<T, const N: usize> {
68 // The highest bit stores the location flag: 0 means cache, 1 means heap.
69 // The remaining bits store length.
70 len_and_flag: usize,
71 data: Data<T, N>,
72}
73
74// -----------------------------------------------------------------------------
75// Marker
76
77unsafe impl<T: Sync, const N: usize> Sync for SmallVec<T, N> {}
78unsafe impl<T: Send, const N: usize> Send for SmallVec<T, N> {}
79impl<T, const N: usize> UnwindSafe for SmallVec<T, N> where T: UnwindSafe {}
80impl<T, const N: usize> RefUnwindSafe for SmallVec<T, N> where T: RefUnwindSafe {}
81
82// -----------------------------------------------------------------------------
83// Basic
84
85impl<T, const N: usize> Default for SmallVec<T, N> {
86 /// Constructs a new, empty `SmallVec<T>`.
87 ///
88 /// # Panic
89 /// Panic if the generic param `N` > `isize::MAX`.
90 #[inline(always)]
91 fn default() -> Self {
92 Self::new()
93 }
94}
95
96impl<T, const N: usize> Drop for SmallVec<T, N> {
97 fn drop(&mut self) {
98 self.clear();
99 self.dealloc();
100 }
101}
102
103impl<T, const N: usize> SmallVec<T, N> {
104 #[inline(always)]
105 const fn in_heap(&self) -> bool {
106 self.len_and_flag & MARKER != 0
107 }
108
109 /// Returns current capacity.
110 ///
111 /// Before spilling to heap this equals `N`.
112 ///
113 /// # Examples
114 ///
115 /// ```
116 /// # use fastvec::SmallVec;
117 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
118 /// assert_eq!(vec.capacity(), 2);
119 /// vec.extend([1, 2, 3]);
120 /// assert!(vec.capacity() >= 3);
121 /// ```
122 #[inline(always)]
123 pub const fn capacity(&self) -> usize {
124 if self.in_heap() {
125 crate::utils::cold_path();
126 unsafe { self.data.heap.1 }
127 } else {
128 N
129 }
130 }
131
132 /// Returns the number of initialized elements in the vector.
133 ///
134 /// # Examples
135 ///
136 /// ```
137 /// # use fastvec::SmallVec;
138 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
139 /// assert_eq!(vec.len(), 0);
140 /// vec.push(1);
141 /// assert_eq!(vec.len(), 1);
142 /// ```
143 #[inline(always)]
144 pub const fn len(&self) -> usize {
145 self.len_and_flag & MAX_LEN
146 }
147
148 /// Returns `true` if the vector contains no elements.
149 ///
150 /// # Examples
151 ///
152 /// ```
153 /// # use fastvec::SmallVec;
154 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
155 /// assert!(vec.is_empty());
156 /// vec.push(1);
157 /// assert!(!vec.is_empty());
158 /// ```
159 #[inline(always)]
160 pub const fn is_empty(&self) -> bool {
161 self.len_and_flag & MAX_LEN == 0
162 }
163
164 /// Returns a raw pointer to the vector's buffer, or a dangling raw
165 /// pointer valid for zero sized reads if the vector didn't allocate.
166 ///
167 /// # Examples
168 ///
169 /// ```
170 /// # use fastvec::SmallVec;
171 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
172 /// vec.push(7);
173 /// let ptr = vec.as_ptr();
174 /// assert_eq!(unsafe { *ptr }, 7);
175 /// ```
176 #[inline(always)]
177 pub const fn as_ptr(&self) -> *const T {
178 if self.in_heap() {
179 crate::utils::cold_path();
180 unsafe { self.data.heap.0.as_ptr() }
181 } else {
182 unsafe { self.data.cache.as_ptr() as *const T }
183 }
184 }
185
186 /// Returns a raw mutable pointer to the vector's buffer, or a dangling
187 /// raw pointer valid for zero sized reads if the vector didn't allocate.
188 ///
189 /// # Examples
190 ///
191 /// ```
192 /// # use fastvec::SmallVec;
193 /// let mut vec: SmallVec<i32, 2> = SmallVec::from([1, 2]);
194 /// let ptr = vec.as_mut_ptr();
195 /// unsafe { *ptr.add(1) = 9; }
196 /// assert_eq!(vec.as_slice(), &[1, 9]);
197 /// ```
198 #[inline(always)]
199 pub const fn as_mut_ptr(&mut self) -> *mut T {
200 if self.in_heap() {
201 crate::utils::cold_path();
202 unsafe { self.data.heap.0.as_ptr() }
203 } else {
204 unsafe { self.data.cache.as_mut_ptr() as *mut T }
205 }
206 }
207
208 /// Extracts a slice containing the entire vector.
209 ///
210 /// # Examples
211 ///
212 /// ```
213 /// # use fastvec::SmallVec;
214 /// let vec: SmallVec<_, 4> = [1, 2, 3].into();
215 /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
216 /// ```
217 #[inline]
218 pub const fn as_slice(&self) -> &[T] {
219 let data = self.as_ptr();
220 let len = self.len();
221 unsafe { slice::from_raw_parts(data, len) }
222 }
223
224 /// Extracts a mutable slice of the entire vector.
225 ///
226 /// # Examples
227 ///
228 /// ```
229 /// # use fastvec::SmallVec;
230 /// let mut vec: SmallVec<_, 4> = [1, 2, 3].into();
231 /// vec.as_mut_slice()[1] = 9;
232 /// assert_eq!(vec.as_slice(), &[1, 9, 3]);
233 /// ```
234 #[inline]
235 pub const fn as_mut_slice(&mut self) -> &mut [T] {
236 let data = self.as_mut_ptr();
237 let len = self.len();
238 unsafe { slice::from_raw_parts_mut(data, len) }
239 }
240
241 /// Forces the length of the vector to `new_len`.
242 ///
243 /// This is a low-level operation that does not initialize or drop elements.
244 ///
245 /// # Safety
246 /// - `new_len <= capacity()`.
247 /// - Elements in the range `old_len..new_len` must already be initialized.
248 /// - Elements in the range `new_len..old_len` are considered logically removed
249 /// and will not be dropped by this call.
250 ///
251 /// # Examples
252 ///
253 /// ```
254 /// # use fastvec::SmallVec;
255 /// let mut vec: SmallVec<i32, 4> = SmallVec::new();
256 /// unsafe {
257 /// let ptr = vec.as_mut_ptr();
258 /// ptr.write(10);
259 /// ptr.add(1).write(20);
260 /// vec.set_len(2);
261 /// }
262 /// assert_eq!(vec.as_slice(), &[10, 20]);
263 /// ```
264 #[inline(always)]
265 pub const unsafe fn set_len(&mut self, new_len: usize) {
266 debug_assert!(new_len <= self.capacity());
267 self.len_and_flag = (self.len_and_flag & MARKER) | new_len;
268 }
269
270 /// Clears the vector, removing all values.
271 ///
272 /// Note that this method has no effect on allocated capacity.
273 ///
274 /// # Examples
275 ///
276 /// ```
277 /// # use fastvec::SmallVec;
278 /// let mut vec: SmallVec<_, 2> = [1, 2, 3].into();
279 /// let cap = vec.capacity();
280 /// vec.clear();
281 /// assert!(vec.is_empty());
282 /// assert_eq!(vec.capacity(), cap);
283 /// ```
284 pub fn clear(&mut self) {
285 if core::mem::needs_drop::<T>() {
286 unsafe {
287 let slice: &mut [T] = self.as_mut_slice();
288 ptr::drop_in_place::<[T]>(slice);
289 }
290 }
291 self.len_and_flag &= MARKER;
292 }
293
294 /// Shortens the vector, keeping the first `len` elements
295 /// and dropping the rest.
296 ///
297 /// If `len >= self.len()`, this has no effect.
298 ///
299 /// # Examples
300 ///
301 /// ```
302 /// # use fastvec::SmallVec;
303 /// let mut vec: SmallVec<_, 4> = [1, 2, 3, 4].into();
304 /// vec.truncate(2);
305 /// assert_eq!(vec.as_slice(), &[1, 2]);
306 /// ```
307 pub fn truncate(&mut self, len: usize) {
308 let old = self.len_and_flag;
309 let old_len = old & MAX_LEN;
310
311 if old_len > len {
312 if core::mem::needs_drop::<T>() {
313 unsafe {
314 let data = self.as_mut_ptr().add(len);
315 let len = old_len - len;
316 let to_drop = ptr::slice_from_raw_parts_mut(data, len);
317 ptr::drop_in_place::<[T]>(to_drop)
318 }
319 }
320
321 self.len_and_flag = (old & MARKER) | len;
322 }
323 }
324}
325
326// -----------------------------------------------------------------------------
327// Memory
328
329impl<T, const N: usize> SmallVec<T, N> {
330 #[inline(never)]
331 fn realloc(&mut self, cap: usize) {
332 let len = self.len();
333
334 debug_assert!(cap >= len);
335 assert!(
336 cap <= MAX_LEN,
337 "the capacity of SmallVec cannot exceed isize::MAX"
338 );
339
340 if cap <= N {
341 debug_assert!(self.in_heap());
342 if !T::IS_ZST {
343 let (ptr, old_cap) = unsafe { self.data.heap };
344 let old_layout = Layout::array::<T>(old_cap).unwrap();
345 self.len_and_flag = 0; // Ensure the safety during panic
346 unsafe {
347 let dst = self.data.cache.as_mut_ptr() as *mut T;
348 ptr::copy_nonoverlapping::<T>(ptr.as_ptr(), dst, len);
349 malloc::dealloc(ptr.as_ptr() as *mut u8, old_layout);
350 }
351 }
352 self.len_and_flag = len & MAX_LEN;
353 } else if self.in_heap() {
354 let (mut ptr, old_cap) = unsafe { self.data.heap };
355 if !T::IS_ZST {
356 let old_layout = Layout::array::<T>(old_cap).unwrap();
357 let new_layout = Layout::array::<T>(cap).unwrap();
358 let new_size = new_layout.size();
359 let raw_ptr = ptr.as_ptr() as *mut u8;
360 unsafe {
361 ptr = NonNull::new(malloc::realloc(raw_ptr, old_layout, new_size) as *mut T)
362 .unwrap_or_else(|| malloc::handle_alloc_error(new_layout));
363 }
364 }
365 self.data.heap = (ptr, cap);
366 } else {
367 let ptr: NonNull<T> = if !T::IS_ZST {
368 let layout = Layout::array::<T>(cap).unwrap();
369 NonNull::new(unsafe { malloc::alloc(layout) as *mut T })
370 .unwrap_or_else(|| malloc::handle_alloc_error(layout))
371 } else {
372 let align = ::core::mem::align_of::<T>();
373 debug_assert!(align != 0);
374 NonNull::<T>::without_provenance(unsafe { NonZeroUsize::new_unchecked(align) })
375 };
376 if !T::IS_ZST {
377 unsafe {
378 let src = self.data.cache.as_ptr() as *const T;
379 ptr::copy_nonoverlapping::<T>(src, ptr.as_ptr(), len);
380 }
381 }
382 self.len_and_flag = len | MARKER;
383 self.data.heap = (ptr, cap);
384 }
385 }
386
387 fn dealloc(&mut self) {
388 if !T::IS_ZST && self.in_heap() {
389 let (ptr, cap) = unsafe { self.data.heap };
390 let layout = Layout::array::<T>(cap).unwrap();
391 unsafe {
392 malloc::dealloc(ptr.as_ptr() as *mut u8, layout);
393 }
394 }
395 }
396}
397
398impl<T, const N: usize> SmallVec<T, N> {
399 /// Constructs a new, empty `SmallVec<T>`.
400 ///
401 /// The vector will not allocate until the number of elements exceed `N`.
402 ///
403 /// # Panic
404 /// Panic if the generic param `N` > `isize::MAX`.
405 ///
406 /// # Examples
407 ///
408 /// ```
409 /// # use fastvec::SmallVec;
410 /// let mut vec: SmallVec<i32, 4> = SmallVec::new();
411 /// vec.push(1);
412 /// assert_eq!(vec.as_slice(), &[1]);
413 /// ```
414 #[must_use]
415 #[inline(always)]
416 pub const fn new() -> Self {
417 // This expressions can be eliminated at compile time.
418 assert!(
419 N <= MAX_LEN,
420 "The capacity of SmallVec can not exceed isize::MAX"
421 );
422
423 Self {
424 data: Data {
425 // SAFETY: Full buffer uninitialized to internal uninitialized is safe.
426 #[expect(clippy::uninit_assumed_init)]
427 cache: unsafe { MaybeUninit::<[ManuallyDrop<T>; N]>::uninit().assume_init() },
428 },
429 len_and_flag: 0,
430 }
431 }
432
433 /// Constructs a new, empty `SmallVec` with at least the specified capacity.
434 ///
435 /// If the specified capacity is less than or equal to `N`, this is equivalent
436 /// to [`new`](SmallVec::new), and no heap memory will be allocated.
437 ///
438 /// # Panics
439 /// Panics if `capacity > isize::MAX` elements.
440 ///
441 /// # Examples
442 ///
443 /// ```
444 /// # use fastvec::SmallVec;
445 /// let vec: SmallVec<i32, 4> = SmallVec::with_capacity(16);
446 /// assert!(vec.capacity() >= 16);
447 /// ```
448 #[inline]
449 #[must_use]
450 pub fn with_capacity(capacity: usize) -> Self {
451 let mut this = Self::new();
452
453 if capacity > N {
454 this.realloc(capacity);
455 }
456 this
457 }
458
459 /// Creates a `SmallVec<T>` directly from a pointer, a length, and a capacity.
460 ///
461 /// The data is still stored using pointers and will not be moved.
462 ///
463 /// # Safety
464 /// - `ptr` must be allocated with the global allocator for `capacity` elements of `T`.
465 /// - `length <= capacity`.
466 /// - The first `length` elements at `ptr` must be initialized.
467 /// - The allocation must be uniquely owned and valid for deallocation by `SmallVec`.
468 /// - `capacity` and `length` must both be `<= isize::MAX`.
469 ///
470 /// # Examples
471 ///
472 /// ```
473 /// # use fastvec::SmallVec;
474 /// let mut source = vec![1, 2, 3];
475 /// let ptr = source.as_mut_ptr();
476 /// let (len, cap) = (source.len(), source.capacity());
477 /// core::mem::forget(source);
478 ///
479 /// let small = unsafe { SmallVec::<i32, 2>::from_raw_parts(ptr, len, cap) };
480 /// assert_eq!(small.as_slice(), &[1, 2, 3]);
481 /// ```
482 #[inline]
483 pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
484 debug_assert!(length <= capacity && capacity & MARKER == 0);
485 let ptr = NonNull::new(ptr).unwrap();
486
487 Self {
488 data: Data {
489 heap: (ptr, capacity),
490 },
491 len_and_flag: length | MARKER,
492 }
493 }
494
495 /// Creates a `SmallVec` by copying all elements from a slice.
496 ///
497 /// If `slice.len() <= N`, data is stored inline, otherwise it is allocated on heap.
498 ///
499 /// # Examples
500 ///
501 /// ```
502 /// # use fastvec::SmallVec;
503 /// let vec = SmallVec::<i32, 2>::from_slice(&[10, 20, 30]);
504 /// assert_eq!(vec.as_slice(), &[10, 20, 30]);
505 /// ```
506 #[inline]
507 pub fn from_slice(slice: &[T]) -> Self
508 where
509 T: Copy,
510 {
511 let mut this = Self::with_capacity(slice.len());
512 unsafe {
513 if !T::IS_ZST {
514 let ptr = this.as_mut_ptr();
515 ptr::copy_nonoverlapping(slice.as_ptr(), ptr, slice.len());
516 }
517 this.set_len(slice.len());
518 }
519 this
520 }
521
522 /// Reserves capacity for at least `additional`
523 /// more elements to be inserted in the given `SmallVec<T>`.
524 ///
525 /// This may reserve more than requested to reduce future reallocations.
526 ///
527 /// # Panics
528 /// Panics if the new capacity exceeds `isize::MAX`.
529 ///
530 /// # Examples
531 ///
532 /// ```
533 /// # use fastvec::SmallVec;
534 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
535 /// vec.reserve(10);
536 /// assert!(vec.capacity() >= 10);
537 /// ```
538 pub fn reserve(&mut self, additional: usize) {
539 let cap = self.capacity();
540 let len = self.len();
541 let target = len.saturating_add(additional);
542 if target > cap {
543 self.realloc(target.min(MARKER).next_power_of_two());
544 }
545 }
546
547 /// Reserves the minimum capacity for exactly `additional` more elements.
548 ///
549 /// Unlike [`reserve`](Self::reserve), this does not intentionally over-allocate.
550 ///
551 /// # Panics
552 /// Panics if the new capacity exceeds `isize::MAX`.
553 ///
554 /// # Examples
555 ///
556 /// ```
557 /// # use fastvec::SmallVec;
558 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
559 /// vec.reserve_exact(5);
560 /// assert!(vec.capacity() >= 5);
561 /// ```
562 pub fn reserve_exact(&mut self, additional: usize) {
563 let cap = self.capacity();
564 let len = self.len();
565 let target = len.saturating_add(additional);
566 if target > cap {
567 self.realloc(target);
568 }
569 }
570
571 /// Shrinks the capacity of the vector as much as possible.
572 ///
573 /// If `len <= N`, data may move back to inline storage.
574 ///
575 /// # Examples
576 ///
577 /// ```
578 /// # use fastvec::SmallVec;
579 /// let mut vec: SmallVec<_, 2> = SmallVec::with_capacity(16);
580 /// vec.extend([1, 2, 3]);
581 /// vec.shrink_to_fit();
582 /// assert!(vec.capacity() >= vec.len());
583 /// ```
584 pub fn shrink_to_fit(&mut self) {
585 if self.in_heap() {
586 let len = self.len();
587 let cap = self.capacity();
588 if cap > len {
589 self.realloc(len);
590 }
591 }
592 }
593
594 /// Shrinks the capacity of the vector as much as possible.
595 ///
596 /// The resulting capacity will be at least `max(self.len(), min_capacity)`.
597 ///
598 /// # Examples
599 ///
600 /// ```
601 /// # use fastvec::SmallVec;
602 /// let mut vec: SmallVec<_, 2> = SmallVec::with_capacity(32);
603 /// vec.extend([1, 2, 3]);
604 /// vec.shrink_to(4);
605 /// assert!(vec.capacity() >= 4);
606 /// ```
607 pub fn shrink_to(&mut self, min_capacity: usize) {
608 if self.in_heap() {
609 let len = self.len();
610 let cap = self.capacity();
611 if min_capacity >= len && min_capacity < cap {
612 self.realloc(min_capacity);
613 }
614 }
615 }
616
617 /// Converts the `SmallVec` into `Vec<T>`.
618 ///
619 /// If the data is already in the heap, transfer pointer directly.
620 ///
621 /// # Examples
622 ///
623 /// ```
624 /// # use fastvec::SmallVec;
625 /// let small: SmallVec<_, 2> = [1, 2, 3].into();
626 /// let vec = small.into_vec();
627 /// assert_eq!(vec, vec![1, 2, 3]);
628 /// ```
629 #[inline]
630 pub fn into_vec(self) -> Vec<T> {
631 if self.in_heap() {
632 let (ptr, cap) = unsafe { self.data.heap };
633 let len = self.len();
634 ::core::mem::forget(self);
635 unsafe { Vec::from_raw_parts(ptr.as_ptr(), len, cap) }
636 } else {
637 // !in_heap, self.len_and_flag == self.len()
638 let len = self.len_and_flag;
639 let mut ret = Vec::<T>::with_capacity(len);
640 unsafe {
641 if !T::IS_ZST {
642 let src = self.data.cache.as_ptr() as *const T;
643 let dst = ret.as_mut_ptr();
644 ptr::copy_nonoverlapping(src, dst, len);
645 }
646
647 ::core::mem::forget(self);
648 ret.set_len(len);
649 }
650 ret
651 }
652 }
653
654 /// Converts the `SmallVec` into `Box<[T]>`.
655 ///
656 /// If the data is already in the heap, transfer pointer directly.
657 ///
658 /// # Examples
659 ///
660 /// ```
661 /// # use fastvec::SmallVec;
662 /// let small: SmallVec<_, 2> = [1, 2, 3].into();
663 /// let boxed = small.into_boxed_slice();
664 /// assert_eq!(&*boxed, &[1, 2, 3]);
665 /// ```
666 #[inline]
667 pub fn into_boxed_slice(self) -> Box<[T]> {
668 self.into_vec().into_boxed_slice()
669 }
670
671 /// Appends an element to the back of a collection.
672 ///
673 /// # Panics
674 ///
675 /// Panics if the new capacity exceeds `isize::MAX`.
676 ///
677 /// # Examples
678 ///
679 /// ```
680 /// # use fastvec::SmallVec;
681 /// let mut vec: SmallVec<_, 2> = SmallVec::new();
682 /// vec.push(1);
683 /// vec.push(2);
684 /// vec.push(3);
685 /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
686 /// ```
687 pub fn push(&mut self, value: T) {
688 if self.in_heap() {
689 crate::utils::cold_path();
690 let len = self.len();
691 let cap = unsafe { self.data.heap.1 };
692
693 if len == cap {
694 crate::utils::cold_path();
695 let new_cap = (cap << 1).max(min_cap::<T>());
696 self.realloc(new_cap);
697 }
698 unsafe {
699 let ptr = self.data.heap.0.as_ptr();
700 ptr::write(ptr.add(len), value);
701 self.len_and_flag += 1;
702 }
703 } else {
704 // !in_heap, self.len_and_flag == self.len()
705 let len = self.len_and_flag;
706 let ptr: *mut T = if len == N {
707 crate::utils::cold_path();
708 let new_cap = (N << 1).max(min_cap::<T>());
709 self.realloc(new_cap);
710 unsafe { self.data.heap.0.as_ptr() }
711 } else {
712 unsafe { self.data.cache.as_mut_ptr() as *mut T }
713 };
714 unsafe {
715 ptr::write(ptr.add(len), value);
716 }
717 self.len_and_flag += 1;
718 }
719 }
720
721 /// Appends an element to the back of the vector without checking capacity.
722 ///
723 /// # Safety
724 /// `self.len() < self.capacity()` must hold before calling this method.
725 ///
726 /// # Examples
727 ///
728 /// ```
729 /// # use fastvec::SmallVec;
730 /// let mut vec: SmallVec<i32, 2> = SmallVec::new();
731 /// unsafe {
732 /// vec.push_unchecked(1);
733 /// vec.push_unchecked(2);
734 /// }
735 /// assert_eq!(vec.as_slice(), &[1, 2]);
736 /// ```
737 #[inline(always)]
738 pub unsafe fn push_unchecked(&mut self, value: T) {
739 let ptr = self.as_mut_ptr();
740 let len = self.len();
741 unsafe {
742 ptr::write(ptr.add(len), value);
743 }
744 self.len_and_flag += 1;
745 }
746
747 /// Removes the last element and returns it, or `None` if empty.
748 ///
749 /// # Examples
750 ///
751 /// ```
752 /// # use fastvec::SmallVec;
753 /// let mut vec: SmallVec<_, 2> = [1, 2].into();
754 /// assert_eq!(vec.pop(), Some(2));
755 /// assert_eq!(vec.pop(), Some(1));
756 /// assert_eq!(vec.pop(), None);
757 /// ```
758 #[inline]
759 pub fn pop(&mut self) -> Option<T> {
760 let len = self.len();
761 if len != 0 {
762 unsafe {
763 self.len_and_flag -= 1;
764 Some(ptr::read(self.as_ptr().add(len - 1)))
765 }
766 } else {
767 crate::utils::cold_path();
768 None
769 }
770 }
771
772 /// Removes and returns the last element if `predicate` returns `true`.
773 ///
774 /// Returns `None` when the vector is empty or predicate returns `false`.
775 ///
776 /// # Examples
777 ///
778 /// ```
779 /// # use fastvec::SmallVec;
780 /// let mut vec: SmallVec<_, 4> = [1, 2, 3, 4].into();
781 /// assert_eq!(vec.pop_if(|v| *v % 2 == 0), Some(4));
782 /// assert_eq!(vec.pop_if(|v| *v % 2 == 0), None);
783 /// ```
784 #[inline]
785 pub fn pop_if(&mut self, predicate: impl FnOnce(&mut T) -> bool) -> Option<T> {
786 let last = self.last_mut()?;
787 if predicate(last) { self.pop() } else { None }
788 }
789
790 /// Inserts `element` at `index`, shifting all elements after it to the right.
791 ///
792 /// # Panics
793 /// Panics if `index > len`.
794 ///
795 /// # Examples
796 ///
797 /// ```
798 /// # use fastvec::SmallVec;
799 /// let mut vec: SmallVec<_, 2> = [1, 3].into();
800 /// vec.insert(1, 2);
801 /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
802 /// ```
803 pub fn insert(&mut self, index: usize, element: T) {
804 #[cold]
805 #[inline(never)]
806 fn assert_failed(index: usize, len: usize) -> ! {
807 panic!("insertion index (is {index}) should be <= len (is {len})");
808 }
809
810 let len = self.len();
811 if index > len {
812 assert_failed(index, len);
813 }
814
815 // space for the new element
816 if len == self.capacity() {
817 crate::utils::cold_path();
818 self.reserve(1);
819 }
820
821 unsafe {
822 let p = self.as_mut_ptr().add(index);
823 if index < len {
824 ptr::copy(p, p.add(1), len - index);
825 }
826 ptr::write(p, element);
827 self.len_and_flag += 1;
828 }
829 }
830
831 /// Removes and returns the element at `index`, shifting later elements left.
832 ///
833 /// # Panics
834 /// Panics if `index >= len`.
835 ///
836 /// # Examples
837 ///
838 /// ```
839 /// # use fastvec::SmallVec;
840 /// let mut vec: SmallVec<_, 4> = [10, 20, 30].into();
841 /// assert_eq!(vec.remove(1), 20);
842 /// assert_eq!(vec.as_slice(), &[10, 30]);
843 /// ```
844 pub fn remove(&mut self, index: usize) -> T {
845 #[cold]
846 #[inline(never)]
847 fn assert_failed(index: usize, len: usize) -> ! {
848 panic!("removal index (is {index}) should be < len (is {len})");
849 }
850
851 let len = self.len();
852 if index >= len {
853 assert_failed(index, len);
854 }
855
856 unsafe {
857 let ptr = self.as_mut_ptr().add(index);
858 let ret = ptr::read(ptr);
859 ptr::copy(ptr.add(1), ptr, len - index - 1);
860 self.len_and_flag -= 1;
861 ret
862 }
863 }
864
865 /// Removes and returns the element at `index`.
866 ///
867 /// The last element is moved into `index`, so ordering is not preserved.
868 ///
869 /// # Panics
870 /// Panics if `index >= len`.
871 ///
872 /// # Examples
873 ///
874 /// ```
875 /// # use fastvec::SmallVec;
876 /// let mut vec: SmallVec<_, 4> = [10, 20, 30, 40].into();
877 /// let removed = vec.swap_remove(1);
878 /// assert_eq!(removed, 20);
879 /// assert_eq!(vec.len(), 3);
880 /// ```
881 pub fn swap_remove(&mut self, index: usize) -> T {
882 #[cold]
883 #[inline(never)]
884 fn assert_failed(index: usize, len: usize) -> ! {
885 panic!("swap_remove index (is {index}) should be < len (is {len})");
886 }
887
888 let len = self.len();
889 if index >= len {
890 assert_failed(index, len);
891 }
892
893 unsafe {
894 let ptr = self.as_mut_ptr();
895 let value = ptr::read(ptr.add(index));
896 ptr::copy(ptr.add(len - 1), ptr.add(index), 1);
897 self.len_and_flag -= 1;
898 value
899 }
900 }
901
902 /// Moves all elements from `other` to the end of `self`.
903 ///
904 /// `other` is emptied after the call.
905 ///
906 /// # Examples
907 ///
908 /// ```
909 /// # use fastvec::SmallVec;
910 /// let mut a: SmallVec<_, 2> = [1, 2].into();
911 /// let mut b: SmallVec<_, 4> = [3, 4].into();
912 /// a.append(&mut b);
913 /// assert_eq!(a.as_slice(), &[1, 2, 3, 4]);
914 /// assert!(b.is_empty());
915 /// ```
916 pub fn append<const P: usize>(&mut self, other: &mut SmallVec<T, P>) {
917 unsafe {
918 let slice = other.as_slice();
919 let count = slice.len();
920 self.reserve(count);
921
922 if !T::IS_ZST {
923 let len = self.len();
924 let dst = self.as_mut_ptr().add(len);
925 ptr::copy_nonoverlapping::<T>(slice.as_ptr(), dst, count);
926 }
927
928 self.len_and_flag += count;
929 other.set_len(0);
930 }
931 }
932
933 /// Resizes the vector to `new_len` using `f` to generate new values.
934 ///
935 /// If `new_len < len`, this truncates the vector.
936 ///
937 /// # Examples
938 ///
939 /// ```
940 /// # use fastvec::SmallVec;
941 /// let mut vec: SmallVec<_, 2> = [1].into();
942 /// vec.resize_with(4, || 7);
943 /// assert_eq!(vec.as_slice(), &[1, 7, 7, 7]);
944 /// ```
945 pub fn resize_with<F>(&mut self, new_len: usize, mut f: F)
946 where
947 F: FnMut() -> T,
948 {
949 let len = self.len();
950 if new_len > len {
951 self.reserve(new_len - len);
952 let ptr = self.as_mut_ptr();
953 (len..new_len).for_each(|idx| unsafe {
954 ptr::write(ptr.add(idx), f());
955 });
956 unsafe {
957 self.set_len(new_len);
958 }
959 } else {
960 self.truncate(new_len);
961 }
962 }
963
964 /// Retains only elements for which `f` returns `true`, passing each item mutably.
965 ///
966 /// # Examples
967 ///
968 /// ```
969 /// # use fastvec::SmallVec;
970 /// let mut vec: SmallVec<_, 4> = [1, 2, 3, 4].into();
971 /// vec.retain_mut(|v| {
972 /// *v *= 2;
973 /// *v > 4
974 /// });
975 /// assert_eq!(vec.as_slice(), &[6, 8]);
976 /// ```
977 pub fn retain_mut<F: FnMut(&mut T) -> bool>(&mut self, mut f: F) {
978 let base_ptr = self.as_mut_ptr();
979 let len = self.len_and_flag & MAX_LEN;
980 let flag = self.len_and_flag & MARKER;
981 self.len_and_flag = 0;
982 let mut count = 0usize;
983
984 for index in 0..len {
985 unsafe {
986 let dst = base_ptr.add(index);
987 if f(&mut *dst) {
988 ptr::copy(dst, base_ptr.add(count), 1);
989 count += 1;
990 } else {
991 ptr::drop_in_place(dst);
992 }
993 }
994 }
995 self.len_and_flag = count | flag;
996 }
997
998 /// Retains only elements for which `f` returns `true`.
999 ///
1000 /// # Examples
1001 ///
1002 /// ```
1003 /// # use fastvec::SmallVec;
1004 /// let mut vec: SmallVec<_, 4> = [1, 2, 3, 4].into();
1005 /// vec.retain(|v| *v % 2 == 0);
1006 /// assert_eq!(vec.as_slice(), &[2, 4]);
1007 /// ```
1008 #[inline]
1009 pub fn retain<F: FnMut(&T) -> bool>(&mut self, mut f: F) {
1010 self.retain_mut(|v| f(v));
1011 }
1012
1013 /// Removes consecutive repeated elements according to `same_bucket`.
1014 ///
1015 /// # Examples
1016 ///
1017 /// ```
1018 /// # use fastvec::SmallVec;
1019 /// let mut vec: SmallVec<_, 8> = [10, 20, 21, 30, 20].into();
1020 /// vec.dedup_by(|a, b| *a / 10 == *b / 10);
1021 /// assert_eq!(vec.as_slice(), &[10, 20, 30, 20]);
1022 /// ```
1023 pub fn dedup_by<F: FnMut(&mut T, &mut T) -> bool>(&mut self, mut same_bucket: F) {
1024 let len = self.len();
1025 if len <= 1 {
1026 return;
1027 }
1028
1029 let ptr = self.as_mut_ptr();
1030 let mut left = 0usize;
1031
1032 unsafe {
1033 let mut p_l = ptr.add(left);
1034 for right in 1..len {
1035 let p_r = ptr.add(right);
1036 if !same_bucket(&mut *p_r, &mut *p_l) {
1037 left += 1;
1038 p_l = ptr.add(left);
1039 if right != left {
1040 ptr::swap(p_r, p_l);
1041 }
1042 }
1043 }
1044 }
1045 self.truncate(left + 1);
1046 }
1047
1048 /// Removes consecutive repeated elements that map to the same key.
1049 ///
1050 /// # Examples
1051 ///
1052 /// ```
1053 /// # use fastvec::SmallVec;
1054 /// let mut vec: SmallVec<_, 8> = [10, 20, 21, 30, 20].into();
1055 /// vec.dedup_by_key(|v| *v / 10);
1056 /// assert_eq!(vec.as_slice(), &[10, 20, 30, 20]);
1057 /// ```
1058 #[inline]
1059 pub fn dedup_by_key<F, K>(&mut self, mut key: F)
1060 where
1061 F: FnMut(&mut T) -> K,
1062 K: PartialEq,
1063 {
1064 self.dedup_by(|a, b| key(a) == key(b));
1065 }
1066
1067 /// Returns the remaining spare capacity as a slice of `MaybeUninit<T>`.
1068 ///
1069 /// # Examples
1070 ///
1071 /// ```
1072 /// # use core::mem::MaybeUninit;
1073 /// # use fastvec::SmallVec;
1074 /// let mut vec: SmallVec<i32, 4> = SmallVec::new();
1075 /// let spare: &mut [MaybeUninit<i32>] = vec.spare_capacity_mut();
1076 /// spare[0].write(11);
1077 /// unsafe { vec.set_len(1); }
1078 /// assert_eq!(vec.as_slice(), &[11]);
1079 /// ```
1080 #[inline]
1081 pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
1082 let len = self.len();
1083 let cap = self.capacity();
1084 unsafe {
1085 let data = self.as_mut_ptr().add(len);
1086 slice::from_raw_parts_mut(data as *mut MaybeUninit<T>, cap - len)
1087 }
1088 }
1089}
1090
1091// -----------------------------------------------------------------------------
1092// Common
1093
1094super::utils::impl_common_traits!(SmallVec<T, N>);
1095
1096impl<T, U, const N: usize> PartialEq<SmallVec<U, N>> for SmallVec<T, N>
1097where
1098 T: PartialEq<U>,
1099{
1100 #[inline]
1101 fn eq(&self, other: &SmallVec<U, N>) -> bool {
1102 PartialEq::eq(self.as_slice(), other.as_slice())
1103 }
1104}
1105
1106impl<T: PartialEq, const N: usize> SmallVec<T, N> {
1107 /// Removes consecutive repeated elements using `PartialEq`.
1108 ///
1109 /// # Examples
1110 ///
1111 /// ```
1112 /// # use fastvec::SmallVec;
1113 /// let mut vec: SmallVec<_, 8> = [1, 1, 2, 2, 3].into();
1114 /// vec.dedup();
1115 /// assert_eq!(vec.as_slice(), &[1, 2, 3]);
1116 /// ```
1117 #[inline]
1118 pub fn dedup(&mut self) {
1119 self.dedup_by(|x, y| PartialEq::eq(x, y));
1120 }
1121}
1122
1123impl<T: Clone, const N: usize> Clone for SmallVec<T, N> {
1124 fn clone(&self) -> Self {
1125 let slice = self.as_slice();
1126
1127 let mut this = Self::with_capacity(slice.len());
1128 slice.iter().for_each(|item| unsafe {
1129 this.push_unchecked(item.clone());
1130 });
1131 this
1132 }
1133
1134 fn clone_from(&mut self, source: &Self) {
1135 let slice = source.as_slice();
1136 self.clear();
1137
1138 self.reserve(slice.len());
1139 slice.iter().for_each(|item| unsafe {
1140 self.push_unchecked(item.clone());
1141 });
1142 }
1143}
1144
1145impl<T: Clone, const N: usize> SmallVec<T, N> {
1146 /// Resizes the vector to `new_len` by cloning `value` when growing.
1147 ///
1148 /// # Examples
1149 ///
1150 /// ```
1151 /// # use fastvec::SmallVec;
1152 /// let mut vec: SmallVec<_, 2> = [1, 2].into();
1153 /// vec.resize(4, 9);
1154 /// assert_eq!(vec.as_slice(), &[1, 2, 9, 9]);
1155 /// vec.resize(1, 0);
1156 /// assert_eq!(vec.as_slice(), &[1]);
1157 /// ```
1158 pub fn resize(&mut self, new_len: usize, value: T) {
1159 let len = self.len();
1160
1161 if new_len > len {
1162 self.reserve(new_len - len);
1163 (len..new_len - 1).for_each(|_| unsafe {
1164 self.push_unchecked(value.clone());
1165 });
1166 unsafe {
1167 self.push_unchecked(value);
1168 }
1169 } else {
1170 self.truncate(new_len);
1171 }
1172 }
1173
1174 /// Extends the vector by cloning all elements from `other`.
1175 ///
1176 /// # Examples
1177 ///
1178 /// ```
1179 /// # use fastvec::SmallVec;
1180 /// let mut vec: SmallVec<_, 2> = [1].into();
1181 /// vec.extend_from_slice(&[2, 3, 4]);
1182 /// assert_eq!(vec.as_slice(), &[1, 2, 3, 4]);
1183 /// ```
1184 pub fn extend_from_slice(&mut self, other: &[T]) {
1185 self.reserve(other.len());
1186 other.iter().for_each(|item| unsafe {
1187 self.push_unchecked(item.clone());
1188 });
1189 }
1190}
1191
1192impl<'a, T: 'a + Clone + 'a, const N: usize> Extend<&'a T> for SmallVec<T, N> {
1193 fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
1194 let iter = iter.into_iter();
1195 self.reserve(iter.size_hint().0);
1196
1197 iter.for_each(|item| {
1198 self.push(item.clone());
1199 });
1200 }
1201}
1202
1203impl<T, const N: usize> Extend<T> for SmallVec<T, N> {
1204 fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
1205 let iter = iter.into_iter();
1206 self.reserve(iter.size_hint().0);
1207
1208 iter.for_each(|item| {
1209 self.push(item);
1210 });
1211 }
1212}
1213
1214impl<T, const N: usize> FromIterator<T> for SmallVec<T, N> {
1215 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
1216 let iter = iter.into_iter();
1217 let mut vec = Self::with_capacity(iter.size_hint().0);
1218 iter.for_each(|item| {
1219 vec.push(item);
1220 });
1221 vec
1222 }
1223}
1224
1225// -----------------------------------------------------------------------------
1226// From/Into
1227
1228impl<T, const N: usize> From<SmallVec<T, N>> for Vec<T> {
1229 fn from(v: SmallVec<T, N>) -> Self {
1230 v.into_vec()
1231 }
1232}
1233
1234impl<T, const N: usize> From<SmallVec<T, N>> for Box<[T]> {
1235 fn from(v: SmallVec<T, N>) -> Self {
1236 v.into_boxed_slice()
1237 }
1238}
1239
1240impl<T, const N: usize, const P: usize> TryFrom<SmallVec<T, N>> for [T; P] {
1241 type Error = SmallVec<T, N>;
1242
1243 fn try_from(mut vec: SmallVec<T, N>) -> Result<[T; P], SmallVec<T, N>> {
1244 if vec.len() != P {
1245 return Err(vec);
1246 }
1247
1248 let src = vec.as_ptr();
1249 unsafe { vec.set_len(0) };
1250 let array = unsafe { ptr::read(src as *const [T; P]) };
1251 Ok(array)
1252 }
1253}
1254
1255impl<T: Clone, const N: usize> From<&[T]> for SmallVec<T, N> {
1256 fn from(s: &[T]) -> SmallVec<T, N> {
1257 let mut vec = SmallVec::<T, N>::with_capacity(s.len());
1258 s.iter().for_each(|item| unsafe {
1259 vec.push_unchecked(item.clone());
1260 });
1261 vec
1262 }
1263}
1264
1265impl<T: Clone, const N: usize> From<&mut [T]> for SmallVec<T, N> {
1266 fn from(s: &mut [T]) -> SmallVec<T, N> {
1267 let mut vec = SmallVec::<T, N>::with_capacity(s.len());
1268 s.iter().for_each(|item| unsafe {
1269 vec.push_unchecked(item.clone());
1270 });
1271 vec
1272 }
1273}
1274
1275impl<T: Clone, const N: usize, const P: usize> From<&[T; N]> for SmallVec<T, P> {
1276 fn from(s: &[T; N]) -> SmallVec<T, P> {
1277 Self::from(s.as_slice())
1278 }
1279}
1280
1281impl<T: Clone, const N: usize, const P: usize> From<&mut [T; N]> for SmallVec<T, P> {
1282 fn from(s: &mut [T; N]) -> Self {
1283 Self::from(s.as_mut_slice())
1284 }
1285}
1286
1287impl<T, const N: usize, const P: usize> From<[T; N]> for SmallVec<T, P> {
1288 fn from(s: [T; N]) -> Self {
1289 if N <= P {
1290 let mut this = Self::new();
1291 let ptr = unsafe { this.data.cache.as_mut_ptr() as *mut T };
1292 let s = ManuallyDrop::new(s);
1293 let len = s.len();
1294 unsafe {
1295 ptr::copy_nonoverlapping(s.as_ptr(), ptr, len);
1296 this.set_len(len);
1297 }
1298 this
1299 } else {
1300 let vec = Vec::<T>::from(s);
1301 SmallVec::from(vec)
1302 }
1303 }
1304}
1305
1306impl<T, const N: usize> From<Vec<T>> for SmallVec<T, N> {
1307 fn from(s: Vec<T>) -> Self {
1308 let (p, l, c) = s.into_raw_parts();
1309 unsafe { SmallVec::from_raw_parts(p, l, c) }
1310 }
1311}
1312
1313impl<T, const N: usize> From<Box<[T]>> for SmallVec<T, N> {
1314 fn from(s: Box<[T]>) -> Self {
1315 Self::from(s.into_vec())
1316 }
1317}
1318
1319// -----------------------------------------------------------------------------
1320// IntoIterator
1321
1322/// An iterator that consumes a [`SmallVec`] and yields its items by value.
1323///
1324/// # Examples
1325///
1326/// ```
1327/// # use fastvec::SmallVec;
1328///
1329/// let vec = SmallVec::<_, 5>::from(["1", "2", "3"]);
1330/// let mut iter = vec.into_iter();
1331///
1332/// assert_eq!(iter.next(), Some("1"));
1333///
1334/// let vec: Vec<&'static str> = iter.collect();
1335/// assert_eq!(vec, ["2", "3"]);
1336/// ```
1337pub struct IntoIter<T, const N: usize> {
1338 vec: ManuallyDrop<SmallVec<T, N>>,
1339 index: usize,
1340}
1341
1342unsafe impl<T: Send, const N: usize> Send for IntoIter<T, N> {}
1343unsafe impl<T: Sync, const N: usize> Sync for IntoIter<T, N> {}
1344
1345impl<T, const N: usize> IntoIterator for SmallVec<T, N> {
1346 type Item = T;
1347 type IntoIter = IntoIter<T, N>;
1348
1349 #[inline]
1350 fn into_iter(self) -> Self::IntoIter {
1351 IntoIter {
1352 vec: ManuallyDrop::new(self),
1353 index: 0,
1354 }
1355 }
1356}
1357
1358impl<T, const N: usize> Iterator for IntoIter<T, N> {
1359 type Item = T;
1360 #[inline]
1361 fn next(&mut self) -> Option<Self::Item> {
1362 if self.index < self.vec.len() {
1363 self.index += 1;
1364 unsafe { Some(ptr::read(self.vec.as_ptr().add(self.index - 1))) }
1365 } else {
1366 None
1367 }
1368 }
1369
1370 #[inline]
1371 fn size_hint(&self) -> (usize, Option<usize>) {
1372 let v = self.vec.len() - self.index;
1373 (v, Some(v))
1374 }
1375}
1376
1377impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
1378 #[inline]
1379 fn next_back(&mut self) -> Option<Self::Item> {
1380 let len = self.vec.len();
1381 if self.index < len {
1382 unsafe {
1383 self.vec.set_len(len - 1);
1384 }
1385 unsafe { Some(ptr::read(self.vec.as_ptr().add(len - 1))) }
1386 } else {
1387 None
1388 }
1389 }
1390}
1391
1392impl<T, const N: usize> ExactSizeIterator for IntoIter<T, N> {
1393 #[inline]
1394 fn len(&self) -> usize {
1395 self.vec.len() - self.index
1396 }
1397}
1398
1399impl<T, const N: usize> FusedIterator for IntoIter<T, N> {}
1400
1401impl<T: Debug, const N: usize> Debug for IntoIter<T, N> {
1402 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1403 f.debug_tuple("IntoIter")
1404 .field(&self.vec.as_slice())
1405 .finish()
1406 }
1407}
1408
1409impl<T, const N: usize> Drop for IntoIter<T, N> {
1410 fn drop(&mut self) {
1411 let len = self.vec.len();
1412 if self.index < len {
1413 unsafe {
1414 let ptr = self.vec.as_mut_ptr().add(self.index);
1415 let count = len - self.index;
1416 let to_drop = ptr::slice_from_raw_parts_mut(ptr, count);
1417 ptr::drop_in_place(to_drop);
1418 }
1419 }
1420 self.vec.dealloc();
1421 }
1422}
1423
1424// -----------------------------------------------------------------------------
1425// Drain
1426
1427/// An iterator that removes the items from a [`SmallVec`]
1428/// and yields them by value.
1429///
1430/// See [`SmallVec::drain`] .
1431pub struct Drain<'a, T: 'a, const N: usize> {
1432 tail_start: usize,
1433 tail_len: usize,
1434 iter: slice::Iter<'a, T>,
1435 vec: NonNull<SmallVec<T, N>>,
1436}
1437
1438impl<T, const N: usize> SmallVec<T, N> {
1439 /// Removes the subslice indicated by the given range from the vector,
1440 /// returning a double-ended iterator over the removed subslice.
1441 ///
1442 /// If the iterator is dropped before being fully consumed, it drops the remaining removed elements.
1443 ///
1444 /// The returned iterator keeps a mutable borrow on the vector to optimize its implementation.
1445 ///
1446 /// See more information in [`Vec::drain`].
1447 ///
1448 /// # Panics
1449 ///
1450 /// Panics if the range is out of bounds.
1451 ///
1452 /// # Examples
1453 ///
1454 /// ```
1455 /// # use fastvec::SmallVec;
1456 /// let mut v = SmallVec::<_, 3>::from([1, 2, 3]);
1457 /// let u: Vec<_> = v.drain(1..).collect();
1458 /// assert_eq!(v, [1]);
1459 /// assert_eq!(u, [2, 3]);
1460 ///
1461 /// // A full range clears the vector, like `clear()` does
1462 /// v.drain(..);
1463 /// assert_eq!(v, []);
1464 /// ```
1465 pub fn drain<R: core::ops::RangeBounds<usize>>(&mut self, range: R) -> Drain<'_, T, N> {
1466 let len = self.len();
1467 let (start, end) = split_range_bound(&range, len);
1468
1469 unsafe {
1470 self.set_len(start);
1471 let data = self.as_ptr().add(start);
1472 let range_slice = slice::from_raw_parts(data, end - start);
1473
1474 Drain {
1475 tail_start: end,
1476 tail_len: len - end,
1477 iter: range_slice.iter(),
1478 vec: NonNull::new_unchecked(self as *mut _),
1479 }
1480 }
1481 }
1482}
1483
1484impl<T: Debug, const N: usize> Debug for Drain<'_, T, N> {
1485 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1486 f.debug_tuple("Drain").field(&self.iter.as_slice()).finish()
1487 }
1488}
1489
1490impl<T, const N: usize> Iterator for Drain<'_, T, N> {
1491 type Item = T;
1492
1493 #[inline]
1494 fn next(&mut self) -> Option<T> {
1495 self.iter
1496 .next()
1497 .map(|reference| unsafe { ptr::read(reference) })
1498 }
1499
1500 #[inline]
1501 fn size_hint(&self) -> (usize, Option<usize>) {
1502 self.iter.size_hint()
1503 }
1504}
1505
1506impl<T, const N: usize> DoubleEndedIterator for Drain<'_, T, N> {
1507 #[inline]
1508 fn next_back(&mut self) -> Option<Self::Item> {
1509 self.iter
1510 .next_back()
1511 .map(|reference| unsafe { ptr::read(reference) })
1512 }
1513}
1514
1515impl<T, const N: usize> ExactSizeIterator for Drain<'_, T, N> {
1516 #[inline]
1517 fn len(&self) -> usize {
1518 self.iter.len()
1519 }
1520}
1521
1522impl<T, const N: usize> FusedIterator for Drain<'_, T, N> {}
1523
1524impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
1525 fn drop(&mut self) {
1526 /// Moves back the un-`Drain`ed elements to restore the original `Vec`.
1527 struct DropGuard<'r, 'a, T, const N: usize>(&'r mut Drain<'a, T, N>);
1528
1529 impl<'r, 'a, T, const N: usize> Drop for DropGuard<'r, 'a, T, N> {
1530 fn drop(&mut self) {
1531 if self.0.tail_len > 0 {
1532 unsafe {
1533 let source_vec = self.0.vec.as_mut();
1534 // memmove back untouched tail, update to new length
1535 let start = source_vec.len();
1536 let tail = self.0.tail_start;
1537 if tail != start {
1538 let base = source_vec.as_mut_ptr();
1539 let src = base.add(tail);
1540 let dst = base.add(start);
1541 ptr::copy(src, dst, self.0.tail_len);
1542 }
1543 source_vec.set_len(start + self.0.tail_len);
1544 }
1545 }
1546 }
1547 }
1548
1549 let iter = core::mem::take(&mut self.iter);
1550 let drop_len = iter.len();
1551
1552 let mut vec = self.vec;
1553
1554 if T::IS_ZST {
1555 // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
1556 // this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
1557 unsafe {
1558 let vec = vec.as_mut();
1559 let old_len = vec.len();
1560 vec.set_len(old_len + drop_len + self.tail_len);
1561 vec.truncate(old_len + self.tail_len);
1562 }
1563
1564 return;
1565 }
1566
1567 // ensure elements are moved back into their appropriate places, even when drop_in_place panics
1568 let _guard = DropGuard(self);
1569
1570 if drop_len == 0 {
1571 return;
1572 }
1573
1574 // as_slice() must only be called when iter.len() is > 0 because
1575 // it also gets touched by vec::Splice which may turn it into a dangling pointer
1576 // which would make it and the vec pointer point to different allocations which would
1577 // lead to invalid pointer arithmetic below.
1578 let drop_ptr = iter.as_slice().as_ptr();
1579
1580 unsafe {
1581 // drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place
1582 // a pointer with mutable provenance is necessary. Therefore we must reconstruct
1583 // it from the original vec but also avoid creating a &mut to the front since that could
1584 // invalidate raw pointers to it which some unsafe code might rely on.
1585 let vec_ptr = vec.as_mut().as_mut_ptr();
1586 let drop_offset = drop_ptr.offset_from_unsigned(vec_ptr);
1587 let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
1588 ptr::drop_in_place(to_drop);
1589 }
1590 }
1591}
1592
1593// -----------------------------------------------------------------------------
1594// ExtractIf
1595
1596/// An iterator that removes elements matching a predicate from a range.
1597///
1598/// This yields removed items by value while compacting retained elements in place.
1599///
1600/// See [`SmallVec::extract_if`] .
1601pub struct ExtractIf<'a, T, F: FnMut(&mut T) -> bool, const N: usize> {
1602 vec: &'a mut SmallVec<T, N>,
1603 idx: usize,
1604 end: usize,
1605 del: usize,
1606 old_len: usize,
1607 pred: F,
1608}
1609
1610impl<T, const N: usize> SmallVec<T, N> {
1611 /// Creates an iterator which uses a closure to determine
1612 /// if an element in the range should be removed.
1613 ///
1614 /// See more information in [`Vec::extract_if`].
1615 ///
1616 /// # Panics
1617 /// Panics if the range is out of bounds.
1618 ///
1619 ///
1620 /// # Examples
1621 ///
1622 /// Splitting a vector into even and odd values, reusing the original vector:
1623 ///
1624 /// ```
1625 /// # use fastvec::SmallVec;
1626 /// let mut numbers = SmallVec::<_, 5>::from([1, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
1627 ///
1628 /// let evens = numbers.extract_if(.., |x| *x % 2 == 0).collect::<SmallVec<_, 10>>();
1629 /// let odds = numbers;
1630 ///
1631 /// assert_eq!(evens, [2, 4, 6, 8, 14]);
1632 /// assert_eq!(odds, [1, 3, 5, 9, 11, 13, 15]);
1633 /// ```
1634 ///
1635 /// Using the range argument to only process a part of the vector:
1636 ///
1637 /// ```
1638 /// # use fastvec::SmallVec;
1639 /// let mut items = SmallVec::<_, 5>::from([0, 0, 0, 0, 0, 0, 0, 1, 2, 1, 2, 1, 2]);
1640 /// let ones = items.extract_if(7.., |x| *x == 1).collect::<Vec<_>>();
1641 /// assert_eq!(items, [0, 0, 0, 0, 0, 0, 0, 2, 2, 2]);
1642 /// assert_eq!(ones.len(), 3);
1643 /// ```
1644 pub fn extract_if<F, R>(&mut self, range: R, filter: F) -> ExtractIf<'_, T, F, N>
1645 where
1646 F: FnMut(&mut T) -> bool,
1647 R: core::ops::RangeBounds<usize>,
1648 {
1649 let old_len = self.len();
1650 let (start, end) = split_range_bound(&range, old_len);
1651
1652 // Guard against the vec getting leaked (leak amplification)
1653 unsafe {
1654 self.set_len(0);
1655 }
1656
1657 ExtractIf {
1658 vec: self,
1659 idx: start,
1660 del: 0,
1661 end,
1662 old_len,
1663 pred: filter,
1664 }
1665 }
1666}
1667
1668impl<T, F: FnMut(&mut T) -> bool, const N: usize> Iterator for ExtractIf<'_, T, F, N> {
1669 type Item = T;
1670
1671 fn next(&mut self) -> Option<T> {
1672 while self.idx < self.end {
1673 let i = self.idx;
1674 // SAFETY:
1675 // We know that `i < self.end` from the if guard and that `self.end <= self.old_len` from
1676 // the validity of `Self`. Therefore `i` points to an element within `vec`.
1677 //
1678 // Additionally, the i-th element is valid because each element is visited at most once
1679 // and it is the first time we access vec[i].
1680 //
1681 // Note: we can't use `vec.get_unchecked_mut(i)` here since the precondition for that
1682 // function is that i < vec.len(), but we've set vec's length to zero.
1683 let cur = unsafe { &mut *self.vec.as_mut_ptr().add(i) };
1684 let drained = (self.pred)(cur);
1685 // Update the index *after* the predicate is called. If the index
1686 // is updated prior and the predicate panics, the element at this
1687 // index would be leaked.
1688 self.idx += 1;
1689 if drained {
1690 self.del += 1;
1691 // SAFETY: We never touch this element again after returning it.
1692 return Some(unsafe { ptr::read(cur) });
1693 } else if self.del > 0 {
1694 // SAFETY: `self.del` > 0, so the hole slot must not overlap with current element.
1695 // We use copy for move, and never touch this element again.
1696 unsafe {
1697 let hole_slot = self.vec.as_mut_ptr().add(i - self.del);
1698 ptr::copy_nonoverlapping(cur, hole_slot, 1);
1699 }
1700 }
1701 }
1702 None
1703 }
1704
1705 fn size_hint(&self) -> (usize, Option<usize>) {
1706 (0, Some(self.end - self.idx))
1707 }
1708}
1709
1710impl<T, F: FnMut(&mut T) -> bool, const N: usize> Drop for ExtractIf<'_, T, F, N> {
1711 fn drop(&mut self) {
1712 if !T::IS_ZST && self.del > 0 {
1713 // SAFETY: Trailing unchecked items must be valid since we never touch them.
1714 unsafe {
1715 let base = self.vec.as_mut_ptr();
1716 ptr::copy(
1717 base.add(self.idx),
1718 base.add(self.idx - self.del),
1719 self.old_len - self.idx,
1720 );
1721 }
1722 }
1723 // SAFETY: After filling holes, all items are in contiguous memory.
1724 unsafe {
1725 self.vec.set_len(self.old_len - self.del);
1726 }
1727 }
1728}
1729
1730impl<T: Debug, F: FnMut(&mut T) -> bool, const N: usize> Debug for ExtractIf<'_, T, F, N> {
1731 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1732 let peek = if self.idx < self.end {
1733 self.vec.get(self.idx)
1734 } else {
1735 None
1736 };
1737 f.debug_struct("ExtractIf")
1738 .field("peek", &peek)
1739 .finish_non_exhaustive()
1740 }
1741}
1742
1743// -----------------------------------------------------------------------------
1744// Tests
1745
1746#[cfg(test)]
1747mod tests {
1748 use super::SmallVec;
1749 use core::sync::atomic::{AtomicUsize, Ordering};
1750
1751 macro_rules! define_tracker {
1752 () => {
1753 static DROPS: AtomicUsize = AtomicUsize::new(0);
1754
1755 struct Tracker;
1756 impl Drop for Tracker {
1757 fn drop(&mut self) {
1758 DROPS.fetch_add(1, Ordering::SeqCst);
1759 }
1760 }
1761 };
1762 }
1763
1764 #[test]
1765 fn drop_zst() {
1766 define_tracker!();
1767
1768 DROPS.store(0, Ordering::SeqCst);
1769
1770 let mut vec = SmallVec::<Tracker, 0>::new();
1771 vec.push(Tracker);
1772 vec.push(Tracker);
1773 vec.push(Tracker);
1774 vec.push(Tracker);
1775 vec.push(Tracker);
1776
1777 {
1778 let mut drain = vec.drain(1..4);
1779 let one = drain.next_back().unwrap();
1780 drop(one);
1781 assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1782 }
1783
1784 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1785
1786 drop(vec);
1787 assert_eq!(DROPS.load(Ordering::SeqCst), 5);
1788 }
1789
1790 #[test]
1791 fn drop_inline_and_heap() {
1792 define_tracker!();
1793
1794 DROPS.store(0, Ordering::SeqCst);
1795 {
1796 let mut vec = SmallVec::<Tracker, 4>::new();
1797 vec.push(Tracker);
1798 vec.push(Tracker);
1799 vec.push(Tracker);
1800 assert_eq!(DROPS.load(Ordering::SeqCst), 0);
1801 }
1802 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1803
1804 DROPS.store(0, Ordering::SeqCst);
1805 {
1806 let mut vec = SmallVec::<Tracker, 2>::new();
1807 vec.push(Tracker);
1808 vec.push(Tracker);
1809 vec.push(Tracker);
1810 assert_eq!(DROPS.load(Ordering::SeqCst), 0);
1811 }
1812 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1813 }
1814
1815 #[test]
1816 fn drop_pop_remove() {
1817 define_tracker!();
1818
1819 DROPS.store(0, Ordering::SeqCst);
1820
1821 let mut vec = SmallVec::<Tracker, 2>::new();
1822 vec.push(Tracker);
1823 vec.push(Tracker);
1824 vec.push(Tracker);
1825
1826 let popped = vec.pop().unwrap();
1827 assert_eq!(DROPS.load(Ordering::SeqCst), 0);
1828 drop(popped);
1829 assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1830
1831 let removed = vec.remove(0);
1832 assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1833 drop(removed);
1834 assert_eq!(DROPS.load(Ordering::SeqCst), 2);
1835
1836 drop(vec);
1837 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1838 }
1839
1840 #[test]
1841 fn drop_into_iter() {
1842 define_tracker!();
1843
1844 DROPS.store(0, Ordering::SeqCst);
1845
1846 let mut vec = SmallVec::<Tracker, 2>::new();
1847 vec.push(Tracker);
1848 vec.push(Tracker);
1849 vec.push(Tracker);
1850
1851 let mut iter = vec.into_iter();
1852 let first = iter.next().unwrap();
1853 drop(first);
1854 assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1855
1856 drop(iter);
1857 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1858 }
1859
1860 #[test]
1861 fn drop_drain() {
1862 define_tracker!();
1863
1864 DROPS.store(0, Ordering::SeqCst);
1865
1866 let mut vec = SmallVec::<Tracker, 2>::new();
1867 vec.push(Tracker);
1868 vec.push(Tracker);
1869 vec.push(Tracker);
1870 vec.push(Tracker);
1871 vec.push(Tracker);
1872
1873 {
1874 let mut drain = vec.drain(1..4);
1875 let first = drain.next().unwrap();
1876 drop(first);
1877 assert_eq!(DROPS.load(Ordering::SeqCst), 1);
1878 }
1879
1880 // 1 consumed + 2 still in drained range
1881 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1882
1883 drop(vec);
1884 assert_eq!(DROPS.load(Ordering::SeqCst), 5);
1885 }
1886
1887 #[test]
1888 fn drop_extract_if() {
1889 static DROPS: AtomicUsize = AtomicUsize::new(0);
1890
1891 struct Tracker {
1892 id: usize,
1893 }
1894 impl Drop for Tracker {
1895 fn drop(&mut self) {
1896 DROPS.fetch_add(1, Ordering::SeqCst);
1897 }
1898 }
1899
1900 DROPS.store(0, Ordering::SeqCst);
1901
1902 let mut vec = SmallVec::<Tracker, 2>::new();
1903 for id in 0..6 {
1904 vec.push(Tracker { id });
1905 }
1906
1907 let removed: SmallVec<Tracker, 8> = vec.extract_if(.., |t| t.id % 2 == 0).collect();
1908 assert_eq!(DROPS.load(Ordering::SeqCst), 0);
1909
1910 drop(removed);
1911 assert_eq!(DROPS.load(Ordering::SeqCst), 3);
1912
1913 drop(vec);
1914 assert_eq!(DROPS.load(Ordering::SeqCst), 6);
1915 }
1916}