File Index
Files and directories
Sign in
FILE

kvec.rs

File preview with syntax highlighting.

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