projects/linux·blob master · refs

1688 lines · 51.3 KiB

1// SPDX-License-Identifier: GPL-2.0
2
3//! Implementation of [`Vec`].
4
5use super::{
6 allocator::{
7 KVmalloc,
8 Kmalloc,
9 Vmalloc,
10 VmallocPageIter, //
11 },
12 flags::__GFP_ZERO,
13 layout::ArrayLayout,
14 AllocError,
15 Allocator,
16 Box,
17 Flags,
18 NumaNode, //
19};
20
21use crate::{
22 fmt,
23 page::{
24 AsPageIter,
25 PAGE_SIZE, //
26 }, //
27};
28
29use core::{
30 borrow::{
31 Borrow,
32 BorrowMut, //
33 },
34 marker::PhantomData,
35 mem::{
36 ManuallyDrop,
37 MaybeUninit, //
38 },
39 ops::{
40 Deref,
41 DerefMut,
42 Index,
43 IndexMut, //
44 },
45 ptr::{
46 self,
47 NonNull, //
48 },
49 slice::{
50 self,
51 SliceIndex, //
52 }, //
53};
54
55use pin_init::Zeroable;
56
57mod errors;
58pub use self::errors::{InsertError, PushError, RemoveError};
59
60/// Create a [`KVec`] containing the arguments.
61///
62/// New memory is allocated with `GFP_KERNEL`.
63///
64/// # Examples
65///
66/// ```
67/// let mut v = kernel::kvec![];
68/// v.push(1, GFP_KERNEL)?;
69/// assert_eq!(v, [1]);
70///
71/// let mut v = kernel::kvec![1; 3]?;
72/// v.push(4, GFP_KERNEL)?;
73/// assert_eq!(v, [1, 1, 1, 4]);
74///
75/// let mut v = kernel::kvec![1, 2, 3]?;
76/// v.push(4, GFP_KERNEL)?;
77/// assert_eq!(v, [1, 2, 3, 4]);
78///
79/// # Ok::<(), Error>(())
80/// ```
81#[macro_export]
82macro_rules! kvec {
83 () => (
84 $crate::alloc::KVec::new()
85 );
86 ($elem:expr; $n:expr) => (
87 $crate::alloc::KVec::from_elem($elem, $n, GFP_KERNEL)
88 );
89 ($($x:expr),+ $(,)?) => (
90 match $crate::alloc::KBox::new_uninit(GFP_KERNEL) {
91 Ok(b) => Ok($crate::alloc::KVec::from($crate::alloc::KBox::write(b, [$($x),+]))),
92 Err(e) => Err(e),
93 }
94 );
95}
96
97/// The kernel's [`Vec`] type.
98///
99/// A contiguous growable array type with contents allocated with the kernel's allocators (e.g.
100/// [`Kmalloc`], [`Vmalloc`] or [`KVmalloc`]), written `Vec<T, A>`.
101///
102/// For non-zero-sized values, a [`Vec`] will use the given allocator `A` for its allocation. For
103/// the most common allocators the type aliases [`KVec`], [`VVec`] and [`KVVec`] exist.
104///
105/// For zero-sized types the [`Vec`]'s pointer must be `dangling_mut::<T>`; no memory is allocated.
106///
107/// Generally, [`Vec`] consists of a pointer that represents the vector's backing buffer, the
108/// capacity of the vector (the number of elements that currently fit into the vector), its length
109/// (the number of elements that are currently stored in the vector) and the `Allocator` type used
110/// to allocate (and free) the backing buffer.
111///
112/// A [`Vec`] can be deconstructed into and (re-)constructed from its previously named raw parts
113/// and manually modified.
114///
115/// [`Vec`]'s backing buffer gets, if required, automatically increased (re-allocated) when elements
116/// are added to the vector.
117///
118/// # Invariants
119///
120/// - `self.ptr` is always properly aligned and either points to memory allocated with `A` or, for
121/// zero-sized types, is a dangling, well aligned pointer.
122///
123/// - `self.len` always represents the exact number of elements stored in the vector.
124///
125/// - `self.layout` represents the absolute number of elements that can be stored within the vector
126/// without re-allocation. For ZSTs `self.layout`'s capacity is zero. However, it is legal for the
127/// backing buffer to be larger than `layout`.
128///
129/// - `self.len()` is always less than or equal to `self.capacity()`.
130///
131/// - The `Allocator` type `A` of the vector is the exact same `Allocator` type the backing buffer
132/// was allocated with (and must be freed with).
133pub struct Vec<T, A: Allocator> {
134 ptr: NonNull<T>,
135 /// Represents the actual buffer size as `cap` times `size_of::<T>` bytes.
136 ///
137 /// Note: This isn't quite the same as `Self::capacity`, which in contrast returns the number of
138 /// elements we can still store without reallocating.
139 layout: ArrayLayout<T>,
140 len: usize,
141 _p: PhantomData<A>,
142}
143
144/// Type alias for [`Vec`] with a [`Kmalloc`] allocator.
145///
146/// # Examples
147///
148/// ```
149/// let mut v = KVec::new();
150/// v.push(1, GFP_KERNEL)?;
151/// assert_eq!(&v, &[1]);
152///
153/// # Ok::<(), Error>(())
154/// ```
155pub type KVec<T> = Vec<T, Kmalloc>;
156
157/// Type alias for [`Vec`] with a [`Vmalloc`] allocator.
158///
159/// # Examples
160///
161/// ```
162/// let mut v = VVec::new();
163/// v.push(1, GFP_KERNEL)?;
164/// assert_eq!(&v, &[1]);
165///
166/// # Ok::<(), Error>(())
167/// ```
168pub type VVec<T> = Vec<T, Vmalloc>;
169
170/// Type alias for [`Vec`] with a [`KVmalloc`] allocator.
171///
172/// # Examples
173///
174/// ```
175/// let mut v = KVVec::new();
176/// v.push(1, GFP_KERNEL)?;
177/// assert_eq!(&v, &[1]);
178///
179/// # Ok::<(), Error>(())
180/// ```
181pub type KVVec<T> = Vec<T, KVmalloc>;
182
183// SAFETY: `Vec` is `Send` if `T` is `Send` because `Vec` owns its elements.
184unsafe impl<T, A> Send for Vec<T, A>
185where
186 T: Send,
187 A: Allocator,
188{
189}
190
191// SAFETY: `Vec` is `Sync` if `T` is `Sync` because `Vec` owns its elements.
192unsafe impl<T, A> Sync for Vec<T, A>
193where
194 T: Sync,
195 A: Allocator,
196{
197}
198
199impl<T, A> Vec<T, A>
200where
201 A: Allocator,
202{
203 #[inline]
204 const fn is_zst() -> bool {
205 core::mem::size_of::<T>() == 0
206 }
207
208 /// Returns the number of elements that can be stored within the vector without allocating
209 /// additional memory.
210 pub const fn capacity(&self) -> usize {
211 if const { Self::is_zst() } {
212 usize::MAX
213 } else {
214 self.layout.len()
215 }
216 }
217
218 /// Returns the number of elements stored within the vector.
219 #[inline]
220 pub const fn len(&self) -> usize {
221 self.len
222 }
223
224 /// Increments `self.len` by `additional`.
225 ///
226 /// # Safety
227 ///
228 /// - `additional` must be less than or equal to `self.capacity - self.len`.
229 /// - All elements within the interval [`self.len`,`self.len + additional`) must be initialized.
230 #[inline]
231 pub const unsafe fn inc_len(&mut self, additional: usize) {
232 // Guaranteed by the type invariant to never underflow.
233 debug_assert!(additional <= self.capacity() - self.len());
234 // INVARIANT: By the safety requirements of this method this represents the exact number of
235 // elements stored within `self`.
236 self.len += additional;
237 }
238
239 /// Decreases `self.len` by `count`.
240 ///
241 /// Returns a mutable slice to the elements forgotten by the vector. It is the caller's
242 /// responsibility to drop these elements if necessary.
243 ///
244 /// # Safety
245 ///
246 /// - `count` must be less than or equal to `self.len`.
247 unsafe fn dec_len(&mut self, count: usize) -> &mut [T] {
248 debug_assert!(count <= self.len());
249 // INVARIANT: We relinquish ownership of the elements within the range `[self.len - count,
250 // self.len)`, hence the updated value of `set.len` represents the exact number of elements
251 // stored within `self`.
252 self.len -= count;
253 // SAFETY: The memory after `self.len()` is guaranteed to contain `count` initialized
254 // elements of type `T`.
255 unsafe { slice::from_raw_parts_mut(self.as_mut_ptr().add(self.len), count) }
256 }
257
258 /// Returns a slice of the entire vector.
259 ///
260 /// # Examples
261 ///
262 /// ```
263 /// let mut v = KVec::new();
264 /// v.push(1, GFP_KERNEL)?;
265 /// v.push(2, GFP_KERNEL)?;
266 /// assert_eq!(v.as_slice(), &[1, 2]);
267 /// # Ok::<(), Error>(())
268 /// ```
269 #[inline]
270 pub fn as_slice(&self) -> &[T] {
271 self
272 }
273
274 /// Returns a mutable slice of the entire vector.
275 #[inline]
276 pub fn as_mut_slice(&mut self) -> &mut [T] {
277 self
278 }
279
280 /// Returns a mutable raw pointer to the vector's backing buffer, or, if `T` is a ZST, a
281 /// dangling raw pointer.
282 #[inline]
283 pub fn as_mut_ptr(&mut self) -> *mut T {
284 self.ptr.as_ptr()
285 }
286
287 /// Returns a raw pointer to the vector's backing buffer, or, if `T` is a ZST, a dangling raw
288 /// pointer.
289 #[inline]
290 pub const fn as_ptr(&self) -> *const T {
291 self.ptr.as_ptr()
292 }
293
294 /// Returns `true` if the vector contains no elements, `false` otherwise.
295 ///
296 /// # Examples
297 ///
298 /// ```
299 /// let mut v = KVec::new();
300 /// assert!(v.is_empty());
301 ///
302 /// v.push(1, GFP_KERNEL);
303 /// assert!(!v.is_empty());
304 /// ```
305 #[inline]
306 pub const fn is_empty(&self) -> bool {
307 self.len() == 0
308 }
309
310 /// Creates a new, empty `Vec<T, A>`.
311 ///
312 /// This method does not allocate by itself.
313 #[inline]
314 pub const fn new() -> Self {
315 // INVARIANT: Since this is a new, empty `Vec` with no backing memory yet,
316 // - `ptr` is a properly aligned dangling pointer for type `T`,
317 // - `layout` is an empty `ArrayLayout` (zero capacity)
318 // - `len` is zero, since no elements can be or have been stored,
319 // - `A` is always valid.
320 Self {
321 ptr: NonNull::dangling(),
322 layout: ArrayLayout::empty(),
323 len: 0,
324 _p: PhantomData::<A>,
325 }
326 }
327
328 /// Returns a slice of `MaybeUninit<T>` for the remaining spare capacity of the vector.
329 pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
330 // SAFETY:
331 // - `self.len` is smaller than `self.capacity` by the type invariant and hence, the
332 // resulting pointer is guaranteed to be part of the same allocated object.
333 // - `self.len` can not overflow `isize`.
334 let ptr = unsafe { self.as_mut_ptr().add(self.len) }.cast::<MaybeUninit<T>>();
335
336 // SAFETY: The memory between `self.len` and `self.capacity` is guaranteed to be allocated
337 // and valid, but uninitialized.
338 unsafe { slice::from_raw_parts_mut(ptr, self.capacity() - self.len) }
339 }
340
341 /// Appends an element to the back of the [`Vec`] instance.
342 ///
343 /// # Examples
344 ///
345 /// ```
346 /// let mut v = KVec::new();
347 /// v.push(1, GFP_KERNEL)?;
348 /// assert_eq!(&v, &[1]);
349 ///
350 /// v.push(2, GFP_KERNEL)?;
351 /// assert_eq!(&v, &[1, 2]);
352 /// # Ok::<(), Error>(())
353 /// ```
354 pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
355 self.reserve(1, flags)?;
356 // SAFETY: The call to `reserve` was successful, so the capacity is at least one greater
357 // than the length.
358 unsafe { self.push_within_capacity_unchecked(v) };
359 Ok(())
360 }
361
362 /// Appends an element to the back of the [`Vec`] instance without reallocating.
363 ///
364 /// Fails if the vector does not have capacity for the new element.
365 ///
366 /// # Examples
367 ///
368 /// ```
369 /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
370 /// for i in 0..10 {
371 /// v.push_within_capacity(i)?;
372 /// }
373 ///
374 /// assert!(v.push_within_capacity(10).is_err());
375 /// # Ok::<(), Error>(())
376 /// ```
377 pub fn push_within_capacity(&mut self, v: T) -> Result<(), PushError<T>> {
378 if self.len() < self.capacity() {
379 // SAFETY: The length is less than the capacity.
380 unsafe { self.push_within_capacity_unchecked(v) };
381 Ok(())
382 } else {
383 Err(PushError(v))
384 }
385 }
386
387 /// Appends an element to the back of the [`Vec`] instance without reallocating.
388 ///
389 /// # Safety
390 ///
391 /// The length must be less than the capacity.
392 unsafe fn push_within_capacity_unchecked(&mut self, v: T) {
393 let spare = self.spare_capacity_mut();
394
395 // SAFETY: By the safety requirements, `spare` is non-empty.
396 unsafe { spare.get_unchecked_mut(0) }.write(v);
397
398 // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
399 // by 1. We also know that the new length is <= capacity because the caller guarantees that
400 // the length is less than the capacity at the beginning of this function.
401 unsafe { self.inc_len(1) };
402 }
403
404 /// Inserts an element at the given index in the [`Vec`] instance.
405 ///
406 /// Fails if the vector does not have capacity for the new element. Panics if the index is out
407 /// of bounds.
408 ///
409 /// # Examples
410 ///
411 /// ```
412 /// use kernel::alloc::kvec::InsertError;
413 ///
414 /// let mut v = KVec::with_capacity(5, GFP_KERNEL)?;
415 /// for i in 0..5 {
416 /// v.insert_within_capacity(0, i)?;
417 /// }
418 ///
419 /// assert!(matches!(v.insert_within_capacity(0, 5), Err(InsertError::OutOfCapacity(_))));
420 /// assert!(matches!(v.insert_within_capacity(1000, 5), Err(InsertError::IndexOutOfBounds(_))));
421 /// assert_eq!(v, [4, 3, 2, 1, 0]);
422 /// # Ok::<(), Error>(())
423 /// ```
424 pub fn insert_within_capacity(
425 &mut self,
426 index: usize,
427 element: T,
428 ) -> Result<(), InsertError<T>> {
429 let len = self.len();
430 if index > len {
431 return Err(InsertError::IndexOutOfBounds(element));
432 }
433
434 if len >= self.capacity() {
435 return Err(InsertError::OutOfCapacity(element));
436 }
437
438 // SAFETY: This is in bounds since `index <= len < capacity`.
439 let p = unsafe { self.as_mut_ptr().add(index) };
440 // INVARIANT: This breaks the Vec invariants by making `index` contain an invalid element,
441 // but we restore the invariants below.
442 // SAFETY: Both the src and dst ranges end no later than one element after the length.
443 // Since the length is less than the capacity, both ranges are in bounds of the allocation.
444 unsafe { ptr::copy(p, p.add(1), len - index) };
445 // INVARIANT: This restores the Vec invariants.
446 // SAFETY: The pointer is in-bounds of the allocation.
447 unsafe { ptr::write(p, element) };
448 // SAFETY: Index `len` contains a valid element due to the above copy and write.
449 unsafe { self.inc_len(1) };
450 Ok(())
451 }
452
453 /// Removes the last element from a vector and returns it, or `None` if it is empty.
454 ///
455 /// # Examples
456 ///
457 /// ```
458 /// let mut v = KVec::new();
459 /// v.push(1, GFP_KERNEL)?;
460 /// v.push(2, GFP_KERNEL)?;
461 /// assert_eq!(&v, &[1, 2]);
462 ///
463 /// assert_eq!(v.pop(), Some(2));
464 /// assert_eq!(v.pop(), Some(1));
465 /// assert_eq!(v.pop(), None);
466 /// # Ok::<(), Error>(())
467 /// ```
468 pub fn pop(&mut self) -> Option<T> {
469 if self.is_empty() {
470 return None;
471 }
472
473 let removed: *mut T = {
474 // SAFETY: We just checked that the length is at least one.
475 let slice = unsafe { self.dec_len(1) };
476 // SAFETY: The argument to `dec_len` was 1 so this returns a slice of length 1.
477 unsafe { slice.get_unchecked_mut(0) }
478 };
479
480 // SAFETY: The guarantees of `dec_len` allow us to take ownership of this value.
481 Some(unsafe { removed.read() })
482 }
483
484 /// Removes the element at the given index.
485 ///
486 /// # Examples
487 ///
488 /// ```
489 /// let mut v = kernel::kvec![1, 2, 3]?;
490 /// assert_eq!(v.remove(1)?, 2);
491 /// assert_eq!(v, [1, 3]);
492 /// # Ok::<(), Error>(())
493 /// ```
494 pub fn remove(&mut self, i: usize) -> Result<T, RemoveError> {
495 let value = {
496 let value_ref = self.get(i).ok_or(RemoveError)?;
497 // INVARIANT: This breaks the invariants by invalidating the value at index `i`, but we
498 // restore the invariants below.
499 // SAFETY: The value at index `i` is valid, because otherwise we would have already
500 // failed with `RemoveError`.
501 unsafe { ptr::read(value_ref) }
502 };
503
504 // SAFETY: We checked that `i` is in-bounds.
505 let p = unsafe { self.as_mut_ptr().add(i) };
506
507 // INVARIANT: After this call, the invalid value is at the last slot, so the Vec invariants
508 // are restored after the below call to `dec_len(1)`.
509 // SAFETY: `p.add(1).add(self.len - i - 1)` is `i+1+len-i-1 == len` elements after the
510 // beginning of the vector, so this is in-bounds of the vector's allocation.
511 unsafe { ptr::copy(p.add(1), p, self.len - i - 1) };
512
513 // SAFETY: Since the check at the beginning of this call did not fail with `RemoveError`,
514 // the length is at least one.
515 unsafe { self.dec_len(1) };
516
517 Ok(value)
518 }
519
520 /// Creates a new [`Vec`] instance with at least the given capacity.
521 ///
522 /// # Examples
523 ///
524 /// ```
525 /// let v = KVec::<u32>::with_capacity(20, GFP_KERNEL)?;
526 ///
527 /// assert!(v.capacity() >= 20);
528 /// # Ok::<(), Error>(())
529 /// ```
530 pub fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
531 let mut v = Vec::new();
532
533 v.reserve(capacity, flags)?;
534
535 Ok(v)
536 }
537
538 /// Creates a new [`Vec`] with `n` zero-initialized elements.
539 ///
540 /// # Examples
541 ///
542 /// ```
543 /// let v = KVec::<u32>::zeroed(20, GFP_KERNEL)?;
544 ///
545 /// assert!(v.iter().all(|&x| x == 0));
546 /// # Ok::<(), Error>(())
547 /// ```
548 pub fn zeroed(n: usize, flags: Flags) -> Result<Self, AllocError>
549 where
550 T: Zeroable,
551 {
552 let mut v = Self::with_capacity(n, flags | __GFP_ZERO)?;
553
554 // SAFETY:
555 // - `n <= capacity - len`: `with_capacity(n)` guarantees capacity >= n, len is 0.
556 // - All elements in `[0, n)` are initialized: `__GFP_ZERO` zeroes the allocation,
557 // and `T: Zeroable` guarantees all-zeroes is a valid bit pattern.
558 unsafe { v.inc_len(n) };
559 Ok(v)
560 }
561
562 /// Creates a `Vec<T, A>` from a pointer, a length and a capacity using the allocator `A`.
563 ///
564 /// # Examples
565 ///
566 /// ```
567 /// let mut v = kernel::kvec![1, 2, 3]?;
568 /// v.reserve(1, GFP_KERNEL)?;
569 ///
570 /// let (mut ptr, mut len, cap) = v.into_raw_parts();
571 ///
572 /// // SAFETY: We've just reserved memory for another element.
573 /// unsafe { ptr.add(len).write(4) };
574 /// len += 1;
575 ///
576 /// // SAFETY: We only wrote an additional element at the end of the `KVec`'s buffer and
577 /// // correspondingly increased the length of the `KVec` by one. Otherwise, we construct it
578 /// // from the exact same raw parts.
579 /// let v = unsafe { KVec::from_raw_parts(ptr, len, cap) };
580 ///
581 /// assert_eq!(v, [1, 2, 3, 4]);
582 ///
583 /// # Ok::<(), Error>(())
584 /// ```
585 ///
586 /// # Safety
587 ///
588 /// If `T` is a ZST:
589 ///
590 /// - `ptr` must be a dangling, well aligned pointer.
591 ///
592 /// Otherwise:
593 ///
594 /// - `ptr` must have been allocated with the allocator `A`.
595 /// - `ptr` must satisfy or exceed the alignment requirements of `T`.
596 /// - `ptr` must point to memory with a size of at least `size_of::<T>() * capacity` bytes.
597 /// - The allocated size in bytes must not be larger than `isize::MAX`.
598 /// - `length` must be less than or equal to `capacity`.
599 /// - The first `length` elements must be initialized values of type `T`.
600 ///
601 /// It is also valid to create an empty `Vec` passing a dangling pointer for `ptr` and zero for
602 /// `cap` and `len`.
603 pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
604 let layout = if Self::is_zst() {
605 ArrayLayout::empty()
606 } else {
607 // SAFETY: By the safety requirements of this function, `capacity * size_of::<T>()` is
608 // smaller than `isize::MAX`.
609 unsafe { ArrayLayout::new_unchecked(capacity) }
610 };
611
612 // INVARIANT: For ZSTs, we store an empty `ArrayLayout`, all other type invariants are
613 // covered by the safety requirements of this function.
614 Self {
615 // SAFETY: By the safety requirements, `ptr` is either dangling or pointing to a valid
616 // memory allocation, allocated with `A`.
617 ptr: unsafe { NonNull::new_unchecked(ptr) },
618 layout,
619 len: length,
620 _p: PhantomData::<A>,
621 }
622 }
623
624 /// Consumes the `Vec<T, A>` and returns its raw components `pointer`, `length` and `capacity`.
625 ///
626 /// This will not run the destructor of the contained elements and for non-ZSTs the allocation
627 /// will stay alive indefinitely. Use [`Vec::from_raw_parts`] to recover the [`Vec`], drop the
628 /// elements and free the allocation, if any.
629 pub fn into_raw_parts(self) -> (*mut T, usize, usize) {
630 let mut me = ManuallyDrop::new(self);
631 let len = me.len();
632 let capacity = me.capacity();
633 let ptr = me.as_mut_ptr();
634 (ptr, len, capacity)
635 }
636
637 /// Clears the vector, removing all values.
638 ///
639 /// Note that this method has no effect on the allocated capacity
640 /// of the vector.
641 ///
642 /// # Examples
643 ///
644 /// ```
645 /// let mut v = kernel::kvec![1, 2, 3]?;
646 ///
647 /// v.clear();
648 ///
649 /// assert!(v.is_empty());
650 /// # Ok::<(), Error>(())
651 /// ```
652 #[inline]
653 pub fn clear(&mut self) {
654 self.truncate(0);
655 }
656
657 /// Ensures that the capacity exceeds the length by at least `additional` elements.
658 ///
659 /// # Examples
660 ///
661 /// ```
662 /// let mut v = KVec::new();
663 /// v.push(1, GFP_KERNEL)?;
664 ///
665 /// v.reserve(10, GFP_KERNEL)?;
666 /// let cap = v.capacity();
667 /// assert!(cap >= v.len() + 10);
668 ///
669 /// v.reserve(10, GFP_KERNEL)?;
670 /// let new_cap = v.capacity();
671 /// assert_eq!(new_cap, cap);
672 ///
673 /// # Ok::<(), Error>(())
674 /// ```
675 pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {
676 let len = self.len();
677 let cap = self.capacity();
678
679 if cap - len >= additional {
680 return Ok(());
681 }
682
683 if Self::is_zst() {
684 // The capacity is already `usize::MAX` for ZSTs, we can't go higher.
685 return Err(AllocError);
686 }
687
688 // We know that `cap <= isize::MAX` because of the type invariants of `Self`. So the
689 // multiplication by two won't overflow.
690 let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
691 let layout = ArrayLayout::new(new_cap).map_err(|_| AllocError)?;
692
693 // SAFETY:
694 // - `ptr` is valid because it's either `None` or comes from a previous call to
695 // `A::realloc`.
696 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
697 let ptr = unsafe {
698 A::realloc(
699 Some(self.ptr.cast()),
700 layout.into(),
701 self.layout.into(),
702 flags,
703 NumaNode::NO_NODE,
704 )?
705 };
706
707 // INVARIANT:
708 // - `layout` is some `ArrayLayout::<T>`,
709 // - `ptr` has been created by `A::realloc` from `layout`.
710 self.ptr = ptr.cast();
711 self.layout = layout;
712
713 Ok(())
714 }
715
716 /// Shortens the vector, setting the length to `len` and drops the removed values.
717 /// If `len` is greater than or equal to the current length, this does nothing.
718 ///
719 /// This has no effect on the capacity and will not allocate.
720 ///
721 /// # Examples
722 ///
723 /// ```
724 /// let mut v = kernel::kvec![1, 2, 3]?;
725 /// v.truncate(1);
726 /// assert_eq!(v.len(), 1);
727 /// assert_eq!(&v, &[1]);
728 ///
729 /// # Ok::<(), Error>(())
730 /// ```
731 pub fn truncate(&mut self, len: usize) {
732 if let Some(count) = self.len().checked_sub(len) {
733 // SAFETY: `count` is `self.len() - len` so it is guaranteed to be less than or
734 // equal to `self.len()`.
735 let ptr: *mut [T] = unsafe { self.dec_len(count) };
736
737 // SAFETY: the contract of `dec_len` guarantees that the elements in `ptr` are
738 // valid elements whose ownership has been transferred to the caller.
739 unsafe { ptr::drop_in_place(ptr) };
740 }
741 }
742
743 /// Takes ownership of all items in this vector without consuming the allocation.
744 ///
745 /// # Examples
746 ///
747 /// ```
748 /// let mut v = kernel::kvec![0, 1, 2, 3]?;
749 ///
750 /// for (i, j) in v.drain_all().enumerate() {
751 /// assert_eq!(i, j);
752 /// }
753 ///
754 /// assert!(v.capacity() >= 4);
755 /// # Ok::<(), Error>(())
756 /// ```
757 pub fn drain_all(&mut self) -> DrainAll<'_, T> {
758 // SAFETY: This does not underflow the length.
759 let elems = unsafe { self.dec_len(self.len()) };
760 // INVARIANT: The first `len` elements of the spare capacity are valid values, and as we
761 // just set the length to zero, we may transfer ownership to the `DrainAll` object.
762 DrainAll {
763 elements: elems.iter_mut(),
764 }
765 }
766
767 /// Removes all elements that don't match the provided closure.
768 ///
769 /// # Examples
770 ///
771 /// ```
772 /// let mut v = kernel::kvec![1, 2, 3, 4]?;
773 /// v.retain(|i| *i % 2 == 0);
774 /// assert_eq!(v, [2, 4]);
775 /// # Ok::<(), Error>(())
776 /// ```
777 pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
778 let mut num_kept = 0;
779 let mut next_to_check = 0;
780 while let Some(to_check) = self.get_mut(next_to_check) {
781 if f(to_check) {
782 self.swap(num_kept, next_to_check);
783 num_kept += 1;
784 }
785 next_to_check += 1;
786 }
787 self.truncate(num_kept);
788 }
789}
790// TODO: This is a temporary KVVec-specific implementation. It should be replaced with a generic
791// `shrink_to()` for `impl<T, A: Allocator> Vec<T, A>` that uses `A::realloc()` once the
792// underlying allocators properly support shrinking via realloc.
793impl<T> Vec<T, KVmalloc> {
794 /// Shrinks the capacity of the vector with a lower bound.
795 ///
796 /// The capacity will remain at least as large as both the length and the supplied value.
797 /// If the current capacity is less than the lower limit, this is a no-op.
798 ///
799 /// For `kmalloc` allocations, this delegates to `realloc()`, which decides whether
800 /// shrinking is worthwhile. For `vmalloc` allocations, shrinking only occurs if the
801 /// operation would free at least one page of memory, and performs a deep copy since
802 /// `vrealloc` does not yet support in-place shrinking.
803 ///
804 /// # Examples
805 ///
806 /// ```
807 /// // Allocate enough capacity to span multiple pages.
808 /// let elements_per_page = kernel::page::PAGE_SIZE / core::mem::size_of::<u32>();
809 /// let mut v = KVVec::with_capacity(elements_per_page * 4, GFP_KERNEL)?;
810 /// v.push(1, GFP_KERNEL)?;
811 /// v.push(2, GFP_KERNEL)?;
812 ///
813 /// v.shrink_to(0, GFP_KERNEL)?;
814 /// # Ok::<(), Error>(())
815 /// ```
816 pub fn shrink_to(&mut self, min_capacity: usize, flags: Flags) -> Result<(), AllocError> {
817 let target_cap = core::cmp::max(self.len(), min_capacity);
818
819 if self.capacity() <= target_cap {
820 return Ok(());
821 }
822
823 if Self::is_zst() {
824 return Ok(());
825 }
826
827 // For kmalloc allocations, delegate to realloc() and let the allocator decide
828 // whether shrinking is worthwhile.
829 //
830 // SAFETY: `self.ptr` points to a valid `KVmalloc` allocation.
831 if !unsafe { bindings::is_vmalloc_addr(self.ptr.as_ptr().cast()) } {
832 let new_layout = ArrayLayout::<T>::new(target_cap).map_err(|_| AllocError)?;
833
834 // SAFETY:
835 // - `self.ptr` is valid and was previously allocated with `KVmalloc`.
836 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
837 let ptr = unsafe {
838 KVmalloc::realloc(
839 Some(self.ptr.cast()),
840 new_layout.into(),
841 self.layout.into(),
842 flags,
843 NumaNode::NO_NODE,
844 )?
845 };
846
847 self.ptr = ptr.cast();
848 self.layout = new_layout;
849 return Ok(());
850 }
851
852 // Only shrink if we would free at least one page.
853 let current_size = self.capacity() * core::mem::size_of::<T>();
854 let target_size = target_cap * core::mem::size_of::<T>();
855 let current_pages = current_size.div_ceil(PAGE_SIZE);
856 let target_pages = target_size.div_ceil(PAGE_SIZE);
857
858 if current_pages <= target_pages {
859 return Ok(());
860 }
861
862 if target_cap == 0 {
863 if !self.layout.is_empty() {
864 // SAFETY:
865 // - `self.ptr` was previously allocated with `KVmalloc`.
866 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
867 unsafe { KVmalloc::free(self.ptr.cast(), self.layout.into()) };
868 }
869 self.ptr = NonNull::dangling();
870 self.layout = ArrayLayout::empty();
871 return Ok(());
872 }
873
874 // SAFETY: `target_cap <= self.capacity()` and original capacity was valid.
875 let new_layout = unsafe { ArrayLayout::<T>::new_unchecked(target_cap) };
876
877 let new_ptr = KVmalloc::alloc(new_layout.into(), flags, NumaNode::NO_NODE)?;
878
879 // SAFETY:
880 // - `self.as_ptr()` is valid for reads of `self.len()` elements of `T`.
881 // - `new_ptr` is valid for writes of at least `target_cap >= self.len()` elements.
882 // - The two allocations do not overlap since `new_ptr` is freshly allocated.
883 // - Both pointers are properly aligned for `T`.
884 unsafe {
885 ptr::copy_nonoverlapping(self.as_ptr(), new_ptr.as_ptr().cast::<T>(), self.len())
886 };
887
888 // SAFETY:
889 // - `self.ptr` was previously allocated with `KVmalloc`.
890 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
891 unsafe { KVmalloc::free(self.ptr.cast(), self.layout.into()) };
892
893 self.ptr = new_ptr.cast::<T>();
894 self.layout = new_layout;
895
896 Ok(())
897 }
898}
899
900impl<T: Clone, A: Allocator> Vec<T, A> {
901 /// Extend the vector by `n` clones of `value`.
902 ///
903 /// # Examples
904 ///
905 /// ```
906 /// let mut v = KVec::new();
907 /// v.push(1, GFP_KERNEL)?;
908 ///
909 /// v.extend_with(3, 5, GFP_KERNEL)?;
910 /// assert_eq!(&v, &[1, 5, 5, 5]);
911 ///
912 /// v.extend_with(2, 8, GFP_KERNEL)?;
913 /// assert_eq!(&v, &[1, 5, 5, 5, 8, 8]);
914 ///
915 /// v.extend_with(0, 3, GFP_KERNEL)?;
916 /// assert_eq!(&v, &[1, 5, 5, 5, 8, 8]);
917 ///
918 /// # Ok::<(), Error>(())
919 /// ```
920 pub fn extend_with(&mut self, n: usize, value: T, flags: Flags) -> Result<(), AllocError> {
921 if n == 0 {
922 return Ok(());
923 }
924
925 self.reserve(n, flags)?;
926
927 let spare = self.spare_capacity_mut();
928
929 for item in spare.iter_mut().take(n - 1) {
930 item.write(value.clone());
931 }
932
933 // We can write the last element directly without cloning needlessly.
934 spare[n - 1].write(value);
935
936 // SAFETY:
937 // - `self.len() + n <= self.capacity()` due to the call to reserve above,
938 // - the loop and the line above initialized the next `n` elements.
939 unsafe { self.inc_len(n) };
940
941 Ok(())
942 }
943
944 /// Pushes clones of the elements of slice into the [`Vec`] instance.
945 ///
946 /// # Examples
947 ///
948 /// ```
949 /// let mut v = KVec::new();
950 /// v.push(1, GFP_KERNEL)?;
951 ///
952 /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;
953 /// assert_eq!(&v, &[1, 20, 30, 40]);
954 ///
955 /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;
956 /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);
957 /// # Ok::<(), Error>(())
958 /// ```
959 pub fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError> {
960 self.reserve(other.len(), flags)?;
961 for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
962 slot.write(item.clone());
963 }
964
965 // SAFETY:
966 // - `other.len()` spare entries have just been initialized, so it is safe to increase
967 // the length by the same number.
968 // - `self.len() + other.len() <= self.capacity()` is guaranteed by the preceding `reserve`
969 // call.
970 unsafe { self.inc_len(other.len()) };
971 Ok(())
972 }
973
974 /// Create a new `Vec<T, A>` and extend it by `n` clones of `value`.
975 pub fn from_elem(value: T, n: usize, flags: Flags) -> Result<Self, AllocError> {
976 let mut v = Self::with_capacity(n, flags)?;
977
978 v.extend_with(n, value, flags)?;
979
980 Ok(v)
981 }
982
983 /// Resizes the [`Vec`] so that `len` is equal to `new_len`.
984 ///
985 /// If `new_len` is smaller than `len`, the `Vec` is [`Vec::truncate`]d.
986 /// If `new_len` is larger, each new slot is filled with clones of `value`.
987 ///
988 /// # Examples
989 ///
990 /// ```
991 /// let mut v = kernel::kvec![1, 2, 3]?;
992 /// v.resize(1, 42, GFP_KERNEL)?;
993 /// assert_eq!(&v, &[1]);
994 ///
995 /// v.resize(3, 42, GFP_KERNEL)?;
996 /// assert_eq!(&v, &[1, 42, 42]);
997 ///
998 /// # Ok::<(), Error>(())
999 /// ```
1000 pub fn resize(&mut self, new_len: usize, value: T, flags: Flags) -> Result<(), AllocError> {
1001 match new_len.checked_sub(self.len()) {
1002 Some(n) => self.extend_with(n, value, flags),
1003 None => {
1004 self.truncate(new_len);
1005 Ok(())
1006 }
1007 }
1008 }
1009}
1010
1011impl<T, A> Drop for Vec<T, A>
1012where
1013 A: Allocator,
1014{
1015 fn drop(&mut self) {
1016 // SAFETY: `self.as_mut_ptr` is guaranteed to be valid by the type invariant.
1017 unsafe {
1018 ptr::drop_in_place(core::ptr::slice_from_raw_parts_mut(
1019 self.as_mut_ptr(),
1020 self.len,
1021 ))
1022 };
1023
1024 // SAFETY:
1025 // - `self.ptr` was previously allocated with `A`.
1026 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
1027 unsafe { A::free(self.ptr.cast(), self.layout.into()) };
1028 }
1029}
1030
1031impl<T, A, const N: usize> From<Box<[T; N], A>> for Vec<T, A>
1032where
1033 A: Allocator,
1034{
1035 fn from(b: Box<[T; N], A>) -> Vec<T, A> {
1036 let len = b.len();
1037 let ptr = Box::into_raw(b);
1038
1039 // SAFETY:
1040 // - `b` has been allocated with `A`,
1041 // - `ptr` fulfills the alignment requirements for `T`,
1042 // - `ptr` points to memory with at least a size of `size_of::<T>() * len`,
1043 // - all elements within `b` are initialized values of `T`,
1044 // - `len` does not exceed `isize::MAX`.
1045 unsafe { Vec::from_raw_parts(ptr.cast(), len, len) }
1046 }
1047}
1048
1049impl<T, A: Allocator> Default for Vec<T, A> {
1050 #[inline]
1051 fn default() -> Self {
1052 Self::new()
1053 }
1054}
1055
1056impl<T: fmt::Debug, A: Allocator> fmt::Debug for Vec<T, A> {
1057 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1058 fmt::Debug::fmt(&**self, f)
1059 }
1060}
1061
1062impl<T, A> Deref for Vec<T, A>
1063where
1064 A: Allocator,
1065{
1066 type Target = [T];
1067
1068 #[inline]
1069 fn deref(&self) -> &[T] {
1070 // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len`
1071 // initialized elements of type `T`.
1072 unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
1073 }
1074}
1075
1076impl<T, A> DerefMut for Vec<T, A>
1077where
1078 A: Allocator,
1079{
1080 #[inline]
1081 fn deref_mut(&mut self) -> &mut [T] {
1082 // SAFETY: The memory behind `self.as_ptr()` is guaranteed to contain `self.len`
1083 // initialized elements of type `T`.
1084 unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
1085 }
1086}
1087
1088/// # Examples
1089///
1090/// ```
1091/// # use core::borrow::Borrow;
1092/// struct Foo<B: Borrow<[u32]>>(B);
1093///
1094/// // Owned array.
1095/// let owned_array = Foo([1, 2, 3]);
1096///
1097/// // Owned vector.
1098/// let owned_vec = Foo(KVec::from_elem(0, 3, GFP_KERNEL)?);
1099///
1100/// let arr = [1, 2, 3];
1101/// // Borrowed slice from `arr`.
1102/// let borrowed_slice = Foo(&arr[..]);
1103/// # Ok::<(), Error>(())
1104/// ```
1105impl<T, A> Borrow<[T]> for Vec<T, A>
1106where
1107 A: Allocator,
1108{
1109 fn borrow(&self) -> &[T] {
1110 self.as_slice()
1111 }
1112}
1113
1114/// # Examples
1115///
1116/// ```
1117/// # use core::borrow::BorrowMut;
1118/// struct Foo<B: BorrowMut<[u32]>>(B);
1119///
1120/// // Owned array.
1121/// let owned_array = Foo([1, 2, 3]);
1122///
1123/// // Owned vector.
1124/// let owned_vec = Foo(KVec::from_elem(0, 3, GFP_KERNEL)?);
1125///
1126/// let mut arr = [1, 2, 3];
1127/// // Borrowed slice from `arr`.
1128/// let borrowed_slice = Foo(&mut arr[..]);
1129/// # Ok::<(), Error>(())
1130/// ```
1131impl<T, A> BorrowMut<[T]> for Vec<T, A>
1132where
1133 A: Allocator,
1134{
1135 fn borrow_mut(&mut self) -> &mut [T] {
1136 self.as_mut_slice()
1137 }
1138}
1139
1140impl<T: Eq, A> Eq for Vec<T, A> where A: Allocator {}
1141
1142impl<T, I: SliceIndex<[T]>, A> Index<I> for Vec<T, A>
1143where
1144 A: Allocator,
1145{
1146 type Output = I::Output;
1147
1148 #[inline]
1149 fn index(&self, index: I) -> &Self::Output {
1150 Index::index(&**self, index)
1151 }
1152}
1153
1154impl<T, I: SliceIndex<[T]>, A> IndexMut<I> for Vec<T, A>
1155where
1156 A: Allocator,
1157{
1158 #[inline]
1159 fn index_mut(&mut self, index: I) -> &mut Self::Output {
1160 IndexMut::index_mut(&mut **self, index)
1161 }
1162}
1163
1164macro_rules! impl_slice_eq {
1165 ($([$($vars:tt)*] $lhs:ty, $rhs:ty,)*) => {
1166 $(
1167 impl<T, U, $($vars)*> PartialEq<$rhs> for $lhs
1168 where
1169 T: PartialEq<U>,
1170 {
1171 #[inline]
1172 fn eq(&self, other: &$rhs) -> bool { self[..] == other[..] }
1173 }
1174 )*
1175 }
1176}
1177
1178impl_slice_eq! {
1179 [A1: Allocator, A2: Allocator] Vec<T, A1>, Vec<U, A2>,
1180 [A: Allocator] Vec<T, A>, &[U],
1181 [A: Allocator] Vec<T, A>, &mut [U],
1182 [A: Allocator] &[T], Vec<U, A>,
1183 [A: Allocator] &mut [T], Vec<U, A>,
1184 [A: Allocator] Vec<T, A>, [U],
1185 [A: Allocator] [T], Vec<U, A>,
1186 [A: Allocator, const N: usize] Vec<T, A>, [U; N],
1187 [A: Allocator, const N: usize] Vec<T, A>, &[U; N],
1188}
1189
1190impl<'a, T, A> IntoIterator for &'a Vec<T, A>
1191where
1192 A: Allocator,
1193{
1194 type Item = &'a T;
1195 type IntoIter = slice::Iter<'a, T>;
1196
1197 fn into_iter(self) -> Self::IntoIter {
1198 self.iter()
1199 }
1200}
1201
1202impl<'a, T, A: Allocator> IntoIterator for &'a mut Vec<T, A>
1203where
1204 A: Allocator,
1205{
1206 type Item = &'a mut T;
1207 type IntoIter = slice::IterMut<'a, T>;
1208
1209 fn into_iter(self) -> Self::IntoIter {
1210 self.iter_mut()
1211 }
1212}
1213
1214/// # Examples
1215///
1216/// ```
1217/// use kernel::{
1218/// alloc::allocator::VmallocPageIter,
1219/// page::{
1220/// AsPageIter,
1221/// PAGE_SIZE, //
1222/// }, //
1223/// };
1224///
1225/// let mut vec = VVec::<u8>::new();
1226///
1227/// assert!(vec.page_iter().next().is_none());
1228///
1229/// vec.reserve(PAGE_SIZE, GFP_KERNEL)?;
1230///
1231/// let page = vec.page_iter().next().expect("At least one page should be available.\n");
1232///
1233/// // SAFETY: There is no concurrent read or write to the same page.
1234/// unsafe { page.fill_zero_raw(0, PAGE_SIZE)? };
1235/// # Ok::<(), Error>(())
1236/// ```
1237impl<T> AsPageIter for VVec<T> {
1238 type Iter<'a>
1239 = VmallocPageIter<'a>
1240 where
1241 T: 'a;
1242
1243 fn page_iter(&mut self) -> Self::Iter<'_> {
1244 let ptr = self.ptr.cast();
1245 let size = self.layout.size();
1246
1247 // SAFETY:
1248 // - `ptr` is a valid pointer to the beginning of a `Vmalloc` allocation.
1249 // - `ptr` is guaranteed to be valid for the lifetime of `'a`.
1250 // - `size` is the size of the `Vmalloc` allocation `ptr` points to.
1251 unsafe { VmallocPageIter::new(ptr, size) }
1252 }
1253}
1254
1255/// An [`Iterator`] implementation for [`Vec`] that moves elements out of a vector.
1256///
1257/// This structure is created by the [`Vec::into_iter`] method on [`Vec`] (provided by the
1258/// [`IntoIterator`] trait).
1259///
1260/// # Examples
1261///
1262/// ```
1263/// let v = kernel::kvec![0, 1, 2]?;
1264/// let iter = v.into_iter();
1265///
1266/// # Ok::<(), Error>(())
1267/// ```
1268pub struct IntoIter<T, A: Allocator> {
1269 ptr: *mut T,
1270 buf: NonNull<T>,
1271 len: usize,
1272 layout: ArrayLayout<T>,
1273 _p: PhantomData<A>,
1274}
1275
1276impl<T, A> IntoIter<T, A>
1277where
1278 A: Allocator,
1279{
1280 fn into_raw_parts(self) -> (*mut T, NonNull<T>, usize, usize) {
1281 let me = ManuallyDrop::new(self);
1282 let ptr = me.ptr;
1283 let buf = me.buf;
1284 let len = me.len;
1285 let cap = me.layout.len();
1286 (ptr, buf, len, cap)
1287 }
1288
1289 /// Same as `Iterator::collect` but specialized for `Vec`'s `IntoIter`.
1290 ///
1291 /// # Examples
1292 ///
1293 /// ```
1294 /// let v = kernel::kvec![1, 2, 3]?;
1295 /// let mut it = v.into_iter();
1296 ///
1297 /// assert_eq!(it.next(), Some(1));
1298 ///
1299 /// let v = it.collect(GFP_KERNEL);
1300 /// assert_eq!(v, [2, 3]);
1301 ///
1302 /// # Ok::<(), Error>(())
1303 /// ```
1304 ///
1305 /// # Implementation details
1306 ///
1307 /// Currently, we can't implement `FromIterator`. There are a couple of issues with this trait
1308 /// in the kernel, namely:
1309 ///
1310 /// - Rust's specialization feature is unstable. This prevents us to optimize for the special
1311 /// case where `I::IntoIter` equals `Vec`'s `IntoIter` type.
1312 /// - We also can't use `I::IntoIter`'s type ID either to work around this, since `FromIterator`
1313 /// doesn't require this type to be `'static`.
1314 /// - `FromIterator::from_iter` does return `Self` instead of `Result<Self, AllocError>`, hence
1315 /// we can't properly handle allocation failures.
1316 /// - Neither `Iterator::collect` nor `FromIterator::from_iter` can handle additional allocation
1317 /// flags.
1318 ///
1319 /// Instead, provide `IntoIter::collect`, such that we can at least convert a `IntoIter` into a
1320 /// `Vec` again.
1321 ///
1322 /// Note that `IntoIter::collect` doesn't require `Flags`, since it re-uses the existing backing
1323 /// buffer. However, this backing buffer may be shrunk to the actual count of elements.
1324 pub fn collect(self, flags: Flags) -> Vec<T, A> {
1325 let old_layout = self.layout;
1326 let (mut ptr, buf, len, mut cap) = self.into_raw_parts();
1327 let has_advanced = ptr != buf.as_ptr();
1328
1329 if has_advanced {
1330 // Copy the contents we have advanced to at the beginning of the buffer.
1331 //
1332 // SAFETY:
1333 // - `ptr` is valid for reads of `len * size_of::<T>()` bytes,
1334 // - `buf.as_ptr()` is valid for writes of `len * size_of::<T>()` bytes,
1335 // - `ptr` and `buf.as_ptr()` are not be subject to aliasing restrictions relative to
1336 // each other,
1337 // - both `ptr` and `buf.ptr()` are properly aligned.
1338 unsafe { ptr::copy(ptr, buf.as_ptr(), len) };
1339 ptr = buf.as_ptr();
1340
1341 // SAFETY: `len` is guaranteed to be smaller than `self.layout.len()` by the type
1342 // invariant.
1343 let layout = unsafe { ArrayLayout::<T>::new_unchecked(len) };
1344
1345 // SAFETY: `buf` points to the start of the backing buffer and `len` is guaranteed by
1346 // the type invariant to be smaller than `cap`. Depending on `realloc` this operation
1347 // may shrink the buffer or leave it as it is.
1348 ptr = match unsafe {
1349 A::realloc(
1350 Some(buf.cast()),
1351 layout.into(),
1352 old_layout.into(),
1353 flags,
1354 NumaNode::NO_NODE,
1355 )
1356 } {
1357 // If we fail to shrink, which likely can't even happen, continue with the existing
1358 // buffer.
1359 Err(_) => ptr,
1360 Ok(ptr) => {
1361 cap = len;
1362 ptr.as_ptr().cast()
1363 }
1364 };
1365 }
1366
1367 // SAFETY: If the iterator has been advanced, the advanced elements have been copied to
1368 // the beginning of the buffer and `len` has been adjusted accordingly.
1369 //
1370 // - `ptr` is guaranteed to point to the start of the backing buffer.
1371 // - `cap` is either the original capacity or, after shrinking the buffer, equal to `len`.
1372 // - `alloc` is guaranteed to be unchanged since `into_iter` has been called on the original
1373 // `Vec`.
1374 unsafe { Vec::from_raw_parts(ptr, len, cap) }
1375 }
1376}
1377
1378impl<T, A> Iterator for IntoIter<T, A>
1379where
1380 A: Allocator,
1381{
1382 type Item = T;
1383
1384 /// # Examples
1385 ///
1386 /// ```
1387 /// let v = kernel::kvec![1, 2, 3]?;
1388 /// let mut it = v.into_iter();
1389 ///
1390 /// assert_eq!(it.next(), Some(1));
1391 /// assert_eq!(it.next(), Some(2));
1392 /// assert_eq!(it.next(), Some(3));
1393 /// assert_eq!(it.next(), None);
1394 ///
1395 /// # Ok::<(), Error>(())
1396 /// ```
1397 fn next(&mut self) -> Option<T> {
1398 if self.len == 0 {
1399 return None;
1400 }
1401
1402 let current = self.ptr;
1403
1404 // SAFETY: We can't overflow; decreasing `self.len` by one every time we advance `self.ptr`
1405 // by one guarantees that.
1406 unsafe { self.ptr = self.ptr.add(1) };
1407
1408 self.len -= 1;
1409
1410 // SAFETY: `current` is guaranteed to point at a valid element within the buffer.
1411 Some(unsafe { current.read() })
1412 }
1413
1414 /// # Examples
1415 ///
1416 /// ```
1417 /// let v: KVec<u32> = kernel::kvec![1, 2, 3]?;
1418 /// let mut iter = v.into_iter();
1419 /// let size = iter.size_hint().0;
1420 ///
1421 /// iter.next();
1422 /// assert_eq!(iter.size_hint().0, size - 1);
1423 ///
1424 /// iter.next();
1425 /// assert_eq!(iter.size_hint().0, size - 2);
1426 ///
1427 /// iter.next();
1428 /// assert_eq!(iter.size_hint().0, size - 3);
1429 ///
1430 /// # Ok::<(), Error>(())
1431 /// ```
1432 fn size_hint(&self) -> (usize, Option<usize>) {
1433 (self.len, Some(self.len))
1434 }
1435}
1436
1437impl<T, A> Drop for IntoIter<T, A>
1438where
1439 A: Allocator,
1440{
1441 fn drop(&mut self) {
1442 // SAFETY: `self.ptr` is guaranteed to be valid by the type invariant.
1443 unsafe { ptr::drop_in_place(ptr::slice_from_raw_parts_mut(self.ptr, self.len)) };
1444
1445 // SAFETY:
1446 // - `self.buf` was previously allocated with `A`.
1447 // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
1448 unsafe { A::free(self.buf.cast(), self.layout.into()) };
1449 }
1450}
1451
1452impl<T, A> IntoIterator for Vec<T, A>
1453where
1454 A: Allocator,
1455{
1456 type Item = T;
1457 type IntoIter = IntoIter<T, A>;
1458
1459 /// Consumes the `Vec<T, A>` and creates an `Iterator`, which moves each value out of the
1460 /// vector (from start to end).
1461 ///
1462 /// # Examples
1463 ///
1464 /// ```
1465 /// let v = kernel::kvec![1, 2]?;
1466 /// let mut v_iter = v.into_iter();
1467 ///
1468 /// let first_element: Option<u32> = v_iter.next();
1469 ///
1470 /// assert_eq!(first_element, Some(1));
1471 /// assert_eq!(v_iter.next(), Some(2));
1472 /// assert_eq!(v_iter.next(), None);
1473 ///
1474 /// # Ok::<(), Error>(())
1475 /// ```
1476 ///
1477 /// ```
1478 /// let v = kernel::kvec![];
1479 /// let mut v_iter = v.into_iter();
1480 ///
1481 /// let first_element: Option<u32> = v_iter.next();
1482 ///
1483 /// assert_eq!(first_element, None);
1484 ///
1485 /// # Ok::<(), Error>(())
1486 /// ```
1487 #[inline]
1488 fn into_iter(self) -> Self::IntoIter {
1489 let buf = self.ptr;
1490 let layout = self.layout;
1491 let (ptr, len, _) = self.into_raw_parts();
1492
1493 IntoIter {
1494 ptr,
1495 buf,
1496 len,
1497 layout,
1498 _p: PhantomData::<A>,
1499 }
1500 }
1501}
1502
1503/// An iterator that owns all items in a vector, but does not own its allocation.
1504///
1505/// # Invariants
1506///
1507/// Every `&mut T` returned by the iterator references a `T` that the iterator may take ownership
1508/// of.
1509pub struct DrainAll<'vec, T> {
1510 elements: slice::IterMut<'vec, T>,
1511}
1512
1513impl<'vec, T> Iterator for DrainAll<'vec, T> {
1514 type Item = T;
1515
1516 fn next(&mut self) -> Option<T> {
1517 let elem: *mut T = self.elements.next()?;
1518 // SAFETY: By the type invariants, we may take ownership of this value.
1519 Some(unsafe { elem.read() })
1520 }
1521
1522 fn size_hint(&self) -> (usize, Option<usize>) {
1523 self.elements.size_hint()
1524 }
1525}
1526
1527impl<'vec, T> Drop for DrainAll<'vec, T> {
1528 fn drop(&mut self) {
1529 if core::mem::needs_drop::<T>() {
1530 let iter = core::mem::take(&mut self.elements);
1531 let ptr: *mut [T] = iter.into_slice();
1532 // SAFETY: By the type invariants, we own these values so we may destroy them.
1533 unsafe { ptr::drop_in_place(ptr) };
1534 }
1535 }
1536}
1537
1538#[cfg(CONFIG_RUST_KVEC_KUNIT_TEST)]
1539#[macros::kunit_tests(rust_kvec)]
1540mod tests {
1541 use super::*;
1542 use crate::prelude::*;
1543
1544 #[test]
1545 fn test_kvec_retain() {
1546 /// Verify correctness for one specific function.
1547 #[expect(clippy::needless_range_loop)]
1548 fn verify(c: &[bool]) {
1549 let mut vec1: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
1550 let mut vec2: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
1551
1552 for i in 0..c.len() {
1553 vec1.push_within_capacity(i).unwrap();
1554 if c[i] {
1555 vec2.push_within_capacity(i).unwrap();
1556 }
1557 }
1558
1559 vec1.retain(|i| c[*i]);
1560
1561 assert_eq!(vec1, vec2);
1562 }
1563
1564 /// Add one to a binary integer represented as a boolean array.
1565 fn add(value: &mut [bool]) {
1566 let mut carry = true;
1567 for v in value {
1568 let new_v = carry != *v;
1569 carry = carry && *v;
1570 *v = new_v;
1571 }
1572 }
1573
1574 // This boolean array represents a function from index to boolean. We check that `retain`
1575 // behaves correctly for all possible boolean arrays of every possible length less than
1576 // ten.
1577 let mut func = KVec::with_capacity(10, GFP_KERNEL).unwrap();
1578 for len in 0..10 {
1579 for _ in 0u32..1u32 << len {
1580 verify(&func);
1581 add(&mut func);
1582 }
1583 func.push_within_capacity(false).unwrap();
1584 }
1585 }
1586
1587 #[test]
1588 fn test_kvvec_shrink_to() {
1589 use crate::page::PAGE_SIZE;
1590
1591 // Create a vector with capacity spanning multiple pages.
1592 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap();
1593
1594 // Add a few elements.
1595 v.push(1, GFP_KERNEL).unwrap();
1596 v.push(2, GFP_KERNEL).unwrap();
1597 v.push(3, GFP_KERNEL).unwrap();
1598
1599 let initial_capacity = v.capacity();
1600 assert!(initial_capacity >= PAGE_SIZE * 4);
1601
1602 // Shrink to a capacity that would free at least one page.
1603 v.shrink_to(PAGE_SIZE, GFP_KERNEL).unwrap();
1604
1605 // Capacity should have been reduced.
1606 assert!(v.capacity() < initial_capacity);
1607 assert!(v.capacity() >= PAGE_SIZE);
1608
1609 // Elements should be preserved.
1610 assert_eq!(v.len(), 3);
1611 assert_eq!(v[0], 1);
1612 assert_eq!(v[1], 2);
1613 assert_eq!(v[2], 3);
1614
1615 // Shrink to zero (should shrink to len).
1616 v.shrink_to(0, GFP_KERNEL).unwrap();
1617
1618 // Capacity should be at least the length.
1619 assert!(v.capacity() >= v.len());
1620
1621 // Elements should still be preserved.
1622 assert_eq!(v.len(), 3);
1623 assert_eq!(v[0], 1);
1624 assert_eq!(v[1], 2);
1625 assert_eq!(v[2], 3);
1626 }
1627
1628 #[test]
1629 fn test_kvvec_shrink_to_empty() {
1630 use crate::page::PAGE_SIZE;
1631
1632 // Create a vector with large capacity but no elements.
1633 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap();
1634
1635 assert!(v.is_empty());
1636
1637 // Shrink empty vector to zero.
1638 v.shrink_to(0, GFP_KERNEL).unwrap();
1639
1640 // Should have freed the allocation.
1641 assert_eq!(v.capacity(), 0);
1642 assert!(v.is_empty());
1643 }
1644
1645 #[test]
1646 fn test_kvvec_shrink_to_no_op() {
1647 use crate::page::PAGE_SIZE;
1648
1649 // Create a small vector.
1650 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE, GFP_KERNEL).unwrap();
1651 v.push(1, GFP_KERNEL).unwrap();
1652
1653 let capacity_before = v.capacity();
1654
1655 // Try to shrink to a capacity larger than current - should be no-op.
1656 v.shrink_to(capacity_before + 100, GFP_KERNEL).unwrap();
1657
1658 assert_eq!(v.capacity(), capacity_before);
1659 assert_eq!(v.len(), 1);
1660 assert_eq!(v[0], 1);
1661 }
1662
1663 #[test]
1664 fn test_kvvec_shrink_to_respects_min_capacity() {
1665 use crate::page::PAGE_SIZE;
1666
1667 // Create a vector with large capacity.
1668 let mut v = KVVec::<u8>::with_capacity(PAGE_SIZE * 4, GFP_KERNEL).unwrap();
1669
1670 // Add some elements.
1671 for i in 0..10u8 {
1672 v.push(i, GFP_KERNEL).unwrap();
1673 }
1674
1675 // Shrink to a min_capacity larger than length.
1676 let min_cap = PAGE_SIZE * 2;
1677 v.shrink_to(min_cap, GFP_KERNEL).unwrap();
1678
1679 // Capacity should be at least min_capacity.
1680 assert!(v.capacity() >= min_cap);
1681
1682 // All elements preserved.
1683 assert_eq!(v.len(), 10);
1684 for i in 0..10u8 {
1685 assert_eq!(v[i as usize], i);
1686 }
1687 }
1688}