rust/kernel/alloc/vec_ext.rs

   1 // SPDX-License-Identifier: GPL-2.0
   2
   3 //! Extensions to [`Vec`] for fallible allocations.
   4
   5 use super::{AllocError, Flags};
   6 use alloc::vec::Vec;
   7 use core::ptr;
   8
   9 /// Extensions to [`Vec`].
  10 pub trait VecExt<T>: Sized {
  11     /// Creates a new [`Vec`] instance with at least the given capacity.
  12     ///
  13     /// # Examples
  14     ///
  15     /// ```
  16     /// let v = Vec::<u32>::with_capacity(20, GFP_KERNEL)?;
  17     ///
  18     /// assert!(v.capacity() >= 20);
  19     /// # Ok::<(), Error>(())
  20     /// ```
  21     fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError>;
  22
  23     /// Appends an element to the back of the [`Vec`] instance.
  24     ///
  25     /// # Examples
  26     ///
  27     /// ```
  28     /// let mut v = Vec::new();
  29     /// v.push(1, GFP_KERNEL)?;
  30     /// assert_eq!(&v, &[1]);
  31     ///
  32     /// v.push(2, GFP_KERNEL)?;
  33     /// assert_eq!(&v, &[1, 2]);
  34     /// # Ok::<(), Error>(())
  35     /// ```
  36     fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError>;
  37
  38     /// Pushes clones of the elements of slice into the [`Vec`] instance.
  39     ///
  40     /// # Examples
  41     ///
  42     /// ```
  43     /// let mut v = Vec::new();
  44     /// v.push(1, GFP_KERNEL)?;
  45     ///
  46     /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;
  47     /// assert_eq!(&v, &[1, 20, 30, 40]);
  48     ///
  49     /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;
  50     /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);
  51     /// # Ok::<(), Error>(())
  52     /// ```
  53     fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
  54     where
  55         T: Clone;
  56
  57     /// Ensures that the capacity exceeds the length by at least `additional` elements.
  58     ///
  59     /// # Examples
  60     ///
  61     /// ```
  62     /// let mut v = Vec::new();
  63     /// v.push(1, GFP_KERNEL)?;
  64     ///
  65     /// v.reserve(10, GFP_KERNEL)?;
  66     /// let cap = v.capacity();
  67     /// assert!(cap >= 10);
  68     ///
  69     /// v.reserve(10, GFP_KERNEL)?;
  70     /// let new_cap = v.capacity();
  71     /// assert_eq!(new_cap, cap);
  72     ///
  73     /// # Ok::<(), Error>(())
  74     /// ```
  75     fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError>;
  76 }
  77
  78 impl<T> VecExt<T> for Vec<T> {
  79     fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
  80         let mut v = Vec::new();
  81         <Self as VecExt<_>>::reserve(&mut v, capacity, flags)?;
  82         Ok(v)
  83     }
  84
  85     fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
  86         <Self as VecExt<_>>::reserve(self, 1, flags)?;
  87         let s = self.spare_capacity_mut();
  88         s[0].write(v);
  89
  90         // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
  91         // by 1. We also know that the new length is <= capacity because of the previous call to
  92         // `reserve` above.
  93         unsafe { self.set_len(self.len() + 1) };
  94         Ok(())
  95     }
  96
  97     fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
  98     where
  99         T: Clone,
 100     {
 101         <Self as VecExt<_>>::reserve(self, other.len(), flags)?;
 102         for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
 103             slot.write(item.clone());
 104         }
 105
 106         // SAFETY: We just initialised the `other.len()` spare entries, so it is safe to increase
 107         // the length by the same amount. We also know that the new length is <= capacity because
 108         // of the previous call to `reserve` above.
 109         unsafe { self.set_len(self.len() + other.len()) };
 110         Ok(())
 111     }
 112
 113     #[cfg(any(test, testlib))]
 114     fn reserve(&mut self, additional: usize, _flags: Flags) -> Result<(), AllocError> {
 115         Vec::reserve(self, additional);
 116         Ok(())
 117     }
 118
 119     #[cfg(not(any(test, testlib)))]
 120     fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {
 121         let len = self.len();
 122         let cap = self.capacity();
 123
 124         if cap - len >= additional {
 125             return Ok(());
 126         }
 127
 128         if core::mem::size_of::<T>() == 0 {
 129             // The capacity is already `usize::MAX` for SZTs, we can't go higher.
 130             return Err(AllocError);
 131         }
 132
 133         // We know cap is <= `isize::MAX` because `Layout::array` fails if the resulting byte size
 134         // is greater than `isize::MAX`. So the multiplication by two won't overflow.
 135         let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
 136         let layout = core::alloc::Layout::array::<T>(new_cap).map_err(|_| AllocError)?;
 137
 138         let (old_ptr, len, cap) = destructure(self);
 139
 140         // We need to make sure that `ptr` is either NULL or comes from a previous call to
 141         // `krealloc_aligned`. A `Vec<T>`'s `ptr` value is not guaranteed to be NULL and might be
 142         // dangling after being created with `Vec::new`. Instead, we can rely on `Vec<T>`'s capacity
 143         // to be zero if no memory has been allocated yet.
 144         let ptr = if cap == 0 { ptr::null_mut() } else { old_ptr };
 145
 146         // SAFETY: `ptr` is valid because it's either NULL or comes from a previous call to
 147         // `krealloc_aligned`. We also verified that the type is not a ZST.
 148         let new_ptr = unsafe { super::allocator::krealloc_aligned(ptr.cast(), layout, flags) };
 149         if new_ptr.is_null() {
 150             // SAFETY: We are just rebuilding the existing `Vec` with no changes.
 151             unsafe { rebuild(self, old_ptr, len, cap) };
 152             Err(AllocError)
 153         } else {
 154             // SAFETY: `ptr` has been reallocated with the layout for `new_cap` elements. New cap
 155             // is greater than `cap`, so it continues to be >= `len`.
 156             unsafe { rebuild(self, new_ptr.cast::<T>(), len, new_cap) };
 157             Ok(())
 158         }
 159     }
 160 }
 161
 162 #[cfg(not(any(test, testlib)))]
 163 fn destructure<T>(v: &mut Vec<T>) -> (*mut T, usize, usize) {
 164     let mut tmp = Vec::new();
 165     core::mem::swap(&mut tmp, v);
 166     let mut tmp = core::mem::ManuallyDrop::new(tmp);
 167     let len = tmp.len();
 168     let cap = tmp.capacity();
 169     (tmp.as_mut_ptr(), len, cap)
 170 }
 171
 172 /// Rebuilds a `Vec` from a pointer, length, and capacity.
 173 ///
 174 /// # Safety
 175 ///
 176 /// The same as [`Vec::from_raw_parts`].
 177 #[cfg(not(any(test, testlib)))]
 178 unsafe fn rebuild<T>(v: &mut Vec<T>, ptr: *mut T, len: usize, cap: usize) {
 179     // SAFETY: The safety requirements from this function satisfy those of `from_raw_parts`.
 180     let mut tmp = unsafe { Vec::from_raw_parts(ptr, len, cap) };
 181     core::mem::swap(&mut tmp, v);
 182 }