diff options
Diffstat (limited to 'rust/kernel/alloc')
| -rw-r--r-- | rust/kernel/alloc/arrayvec.rs | 351 | ||||
| -rw-r--r-- | rust/kernel/alloc/kvec.rs | 86 | ||||
| -rw-r--r-- | rust/kernel/alloc/kvec/errors.rs | 32 |
3 files changed, 466 insertions, 3 deletions
diff --git a/rust/kernel/alloc/arrayvec.rs b/rust/kernel/alloc/arrayvec.rs new file mode 100644 index 000000000000..44cc09ee1e51 --- /dev/null +++ b/rust/kernel/alloc/arrayvec.rs @@ -0,0 +1,351 @@ +// SPDX-License-Identifier: GPL-2.0 + +//! Implementation of [`ArrayVec`]. + +use crate::{ + alloc::kvec::{ + impl_slice_eq, + PushError, // + }, + const_assert, + error::{ + code::EINVAL, + Error, + Result, // + }, + fmt, // +}; + +use core::{ + borrow::{ + Borrow, + BorrowMut, // + }, + mem::MaybeUninit, + ops::{ + Deref, + DerefMut, // + }, + ptr, + slice, // +}; + +use pin_init::{ + init_from_closure, + Init, // +}; + +/// A fixed capacity vector that holds at most `N` elements. +/// +/// # Invariants +/// +/// - `len` is at most `N`. +/// - The first `len` elements of `data` are initialized. +/// +/// # Examples +/// +/// ``` +/// use kernel::alloc::ArrayVec; +/// +/// let mut v = ArrayVec::<u8, 4>::new(); +/// v.extend_from_slice(b"abc")?; +/// assert_eq!(*v, *b"abc"); +/// +/// assert!(v.extend_from_slice(b"ab").is_err()); +/// +/// v.push(4u8)?; +/// assert_eq!(*v, *b"abc\x04"); +/// assert!(v.push(5u8).is_err()); +/// +/// v.clear(); +/// assert!(v.is_empty()); +/// # Ok::<(), Error>(()) +/// ``` +pub struct ArrayVec<T, const N: usize> { + data: [MaybeUninit<T>; N], + len: usize, +} + +/// The maximum size in bytes of an [`ArrayVec`] returned by value from [`ArrayVec::new`]. +const MAX_STACK_SIZE: usize = 512; + +impl<T, const N: usize> ArrayVec<T, N> { + /// Creates an empty [`ArrayVec`]. + #[inline] + pub const fn new() -> Self { + const_assert!( + size_of::<Self>() <= MAX_STACK_SIZE, + "use `init_empty` instead of constructing a large ArrayVec on the stack" + ); + + // INVARIANT: An empty ArrayVec trivially has all its elements initialized. + Self { + data: [const { MaybeUninit::uninit() }; N], + len: 0, + } + } + + /// Creates an initializer for an empty [`ArrayVec`]. + /// + /// # Examples + /// + /// ``` + /// use kernel::alloc::ArrayVec; + /// + /// let v = KBox::init(ArrayVec::<u8, 4096>::init_empty(), GFP_KERNEL)?; + /// assert!(v.is_empty()); + /// # Ok::<(), Error>(()) + /// ``` + #[inline] + pub fn init_empty() -> impl Init<Self> { + Self::init_with(|_| Ok(())) + } + + /// Creates an initializer for an [`ArrayVec`] populated by `f`. + /// + /// `f` gets an empty [`ArrayVec`] and can fill it in place. + /// + /// # Examples + /// + /// ``` + /// use kernel::alloc::ArrayVec; + /// + /// let v = KBox::init( + /// ArrayVec::<u8, 4096>::init_with(|v| v.extend_from_slice(b"abc")), + /// GFP_KERNEL, + /// )?; + /// assert_eq!(**v, *b"abc"); + /// # Ok::<(), Error>(()) + /// ``` + #[inline] + pub fn init_with<E>(f: impl FnOnce(&mut Self) -> Result<(), E>) -> impl Init<Self, E> { + let init_empty = |slot: *mut Self| { + // SAFETY: By the initializer contract `slot` is valid for writes. Once `len` is zero + // the slot holds a valid empty ArrayVec, since `data` requires no initialization. + // INVARIANT: An empty ArrayVec trivially has all its elements initialized. + unsafe { (&raw mut (*slot).len).write(0) }; + Ok(()) + }; + + // SAFETY: `init_empty` fully initializes the slot and never fails. + unsafe { init_from_closure(init_empty) }.chain(f) + } + + /// Appends an element to the back of the [`ArrayVec`]. + /// + /// Fails when the [`ArrayVec`] is full, handing the element back in [`PushError`]. + #[inline] + pub fn push(&mut self, v: T) -> Result<(), PushError<T>> { + self.push_init(v) + } + + /// Appends an element to the back of the [`ArrayVec`] by initializing it in place. + /// + /// Fails when the [`ArrayVec`] is full, handing the initializer back in [`PushError`]. + #[inline] + pub fn push_init<I: Init<T>>(&mut self, init: I) -> Result<(), PushError<I>> { + self.try_push_init(init) + .map_err(|TryPushInitError::Full(init)| PushError(init)) + } + + /// Appends an element to the back of the [`ArrayVec`] by initializing it in place. + /// + /// Unlike [`ArrayVec::push_init`], the initializer may be fallible. If the [`ArrayVec`] is + /// full, the original initializer `init` is handed back in [`TryPushInitError::Full`]. If the + /// initializer itself fails, its error is returned in [`TryPushInitError::InitError`]. + pub fn try_push_init<I, E>(&mut self, init: I) -> Result<(), TryPushInitError<I, E>> + where + I: Init<T, E>, + { + let Some(slot) = self.spare_capacity_mut().first_mut() else { + return Err(TryPushInitError::Full(init)); + }; + + // SAFETY: + // - `slot.as_mut_ptr()` points to a spare slot, which is valid, uninitialized memory. + // - If `raw_try_init` fails, this returns without increasing `len` or touching `slot`. + // - `init` implements `Init<T, E>`, which removes the pinning requirement. + unsafe { pin_init::raw_try_init(slot.as_mut_ptr(), init) } + .map_err(TryPushInitError::InitError)?; + + // INVARIANT: The element at index `len` was just initialized, and the new `len` does not + // exceed `N` because a spare slot existed. + self.len += 1; + + Ok(()) + } + + /// Appends a clone of each element in `slice` to the back of the [`ArrayVec`]. + /// + /// Fails with [`EINVAL`] if `slice` is longer than the remaining capacity. + pub fn extend_from_slice(&mut self, slice: &[T]) -> Result + where + T: Clone, + { + let Some(dst) = self.spare_capacity_mut().get_mut(..slice.len()) else { + return Err(EINVAL); + }; + + for (d, s) in dst.iter_mut().zip(slice) { + d.write(s.clone()); + } + // INVARIANT: The next `slice.len()` elements after `len` were just initialized, and the + // new `len` does not exceed `N` because the spare capacity was enough. + self.len += slice.len(); + + Ok(()) + } + + /// Removes all elements. + #[inline] + pub fn clear(&mut self) { + let elems: *mut [T] = self.as_mut_slice(); + // INVARIANT: An empty ArrayVec trivially has all its elements initialized. + self.len = 0; + // SAFETY: There are no references to the elements since we hold `&mut self`. The elements + // can't be dropped again because `len` is already 0. + unsafe { ptr::drop_in_place(elems) }; + } + + /// Returns the initialized elements as a slice. + #[inline] + pub fn as_slice(&self) -> &[T] { + let ptr = self.data.as_ptr().cast::<T>(); + // SAFETY: `MaybeUninit<T>` has the same layout as `T`, and by the type invariants the first + // `len` elements of `data` are initialized. + unsafe { slice::from_raw_parts(ptr, self.len) } + } + + /// Returns the initialized elements as a mutable slice. + #[inline] + pub fn as_mut_slice(&mut self) -> &mut [T] { + let ptr = self.data.as_mut_ptr().cast::<T>(); + // SAFETY: `MaybeUninit<T>` has the same layout as `T`, and by the type invariants the first + // `len` elements of `data` are initialized. + unsafe { slice::from_raw_parts_mut(ptr, self.len) } + } + + /// Returns a slice of `MaybeUninit<T>` for the remaining spare capacity of the [`ArrayVec`]. + #[inline] + fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] { + // PANIC: `len` never exceeds `N` by the type invariants. + &mut self.data[self.len..] + } +} + +/// Error type for [`ArrayVec::try_push_init`]. +pub enum TryPushInitError<I, E> { + /// The [`ArrayVec`] is full. Hand the initializer back. + Full(I), + /// The initializer failed. + InitError(E), +} + +impl<I, E> fmt::Debug for TryPushInitError<I, E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TryPushInitError::Full(_) => write!(f, "Not enough capacity"), + TryPushInitError::InitError(_) => write!(f, "Initializer failed"), + } + } +} + +impl<I, E> From<TryPushInitError<I, E>> for Error +where + Error: From<E>, +{ + #[inline] + fn from(e: TryPushInitError<I, E>) -> Error { + match e { + TryPushInitError::Full(_) => EINVAL, + TryPushInitError::InitError(e) => Error::from(e), + } + } +} + +impl<T, const N: usize> Default for ArrayVec<T, N> { + #[inline] + fn default() -> Self { + Self::new() + } +} + +impl<T, const N: usize> Drop for ArrayVec<T, N> { + #[inline] + fn drop(&mut self) { + // SAFETY: The slice holds initialized elements that are never accessed again after this + // point. + unsafe { ptr::drop_in_place(self.as_mut_slice()) }; + } +} + +impl<T, const N: usize> Deref for ArrayVec<T, N> { + type Target = [T]; + + #[inline] + fn deref(&self) -> &Self::Target { + self.as_slice() + } +} + +impl<T, const N: usize> DerefMut for ArrayVec<T, N> { + #[inline] + fn deref_mut(&mut self) -> &mut Self::Target { + self.as_mut_slice() + } +} + +impl<T, const N: usize> Borrow<[T]> for ArrayVec<T, N> { + #[inline] + fn borrow(&self) -> &[T] { + self.as_slice() + } +} + +impl<T, const N: usize> BorrowMut<[T]> for ArrayVec<T, N> { + #[inline] + fn borrow_mut(&mut self) -> &mut [T] { + self.as_mut_slice() + } +} + +impl<T: Eq, const N: usize> Eq for ArrayVec<T, N> {} + +impl_slice_eq! { + [const N: usize, const M: usize] ArrayVec<T, N>, ArrayVec<U, M>, + [const N: usize] ArrayVec<T, N>, &[U], + [const N: usize] ArrayVec<T, N>, &mut [U], + [const N: usize] &[T], ArrayVec<U, N>, + [const N: usize] &mut [T], ArrayVec<U, N>, + [const N: usize] ArrayVec<T, N>, [U], + [const N: usize] [T], ArrayVec<U, N>, + [const N: usize, const M: usize] ArrayVec<T, N>, [U; M], + [const N: usize, const M: usize] ArrayVec<T, N>, &[U; M], +} + +impl<'a, T, const N: usize> IntoIterator for &'a ArrayVec<T, N> { + type Item = &'a T; + type IntoIter = slice::Iter<'a, T>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'a, T, const N: usize> IntoIterator for &'a mut ArrayVec<T, N> { + type Item = &'a mut T; + type IntoIter = slice::IterMut<'a, T>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.iter_mut() + } +} + +impl<T: fmt::Debug, const N: usize> fmt::Debug for ArrayVec<T, N> { + #[inline] + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(self.as_slice(), f) + } +} diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs index c7546b9da4fa..cd0c57a9ad90 100644 --- a/rust/kernel/alloc/kvec.rs +++ b/rust/kernel/alloc/kvec.rs @@ -52,10 +52,18 @@ use core::{ }, // }; -use pin_init::Zeroable; +use pin_init::{ + Init, + Zeroable, // +}; mod errors; -pub use self::errors::{InsertError, PushError, RemoveError}; +pub use self::errors::{ + InsertError, + PushError, + RemoveError, + TryPushInitError, // +}; /// Create a [`KVec`] containing the arguments. /// @@ -359,6 +367,79 @@ where Ok(()) } + /// Appends an element to the back of the [`Vec`] instance by initializing it in place. + /// + /// # Examples + /// + /// ``` + /// use pin_init::init_zeroed; + /// + /// let mut v = KVec::<[u8; 200]>::new(); + /// v.push_init(init_zeroed(), GFP_KERNEL)?; + /// assert_eq!(v[0], [0; 200]); + /// # Ok::<(), Error>(()) + /// ``` + #[inline] + pub fn push_init(&mut self, init: impl Init<T>, flags: Flags) -> Result<(), AllocError> { + self.try_push_init(init, flags) + .map_err(|TryPushInitError::AllocError(_)| AllocError) + } + + /// Appends an element to the back of the [`Vec`] instance by initializing it in place. + /// + /// Unlike [`Vec::push_init`], the initializer may be fallible. If the allocation fails, the + /// original initializer `init` is handed back in [`TryPushInitError::AllocError`]. If the + /// initializer itself fails, its error is returned in [`TryPushInitError::InitError`]. + /// + /// # Examples + /// + /// ``` + /// struct Element { + /// value: u8, + /// } + /// + /// impl Element { + /// fn new(value: u32) -> impl Init<Self, Error> { + /// try_init!(Element { + /// value: u8::try_from(value)?, + /// }? Error) + /// } + /// } + /// + /// let mut v: KVec<Element> = KVec::new(); + /// v.try_push_init(Element::new(16), GFP_KERNEL)?; + /// assert_eq!(v[0].value, 16); + /// + /// // The initializer fails because 256 does not fit in a `u8`. + /// assert!(v.try_push_init(Element::new(256), GFP_KERNEL).is_err()); + /// assert_eq!(v.len(), 1); + /// # Ok::<(), Error>(()) + /// ``` + pub fn try_push_init<I, E>( + &mut self, + init: I, + flags: Flags, + ) -> Result<(), TryPushInitError<I, E>> + where + I: Init<T, E>, + { + if self.reserve(1, flags).is_err() { + return Err(TryPushInitError::AllocError(init)); + } + // SAFETY: The call to `reserve` was successful, so there is at least one spare slot. + let slot = unsafe { self.spare_capacity_mut().get_unchecked_mut(0) }.as_mut_ptr(); + // SAFETY: + // - `slot` points to a spare slot, which is valid, uninitialized memory. + // - If `raw_try_init` fails, this returns without increasing the length or touching `slot`. + // - `init` implements `Init<T, E>`, which removes the pinning requirement. + unsafe { pin_init::raw_try_init(slot, init) }.map_err(TryPushInitError::InitError)?; + // SAFETY: The call to `raw_try_init` returned `Ok`, so the first spare slot now holds an + // initialized `T`. The new length does not exceed the capacity because `reserve` ensured + // the capacity is greater than the length by at least one. + unsafe { self.inc_len(1) }; + Ok(()) + } + /// Appends an element to the back of the [`Vec`] instance without reallocating. /// /// Fails if the vector does not have capacity for the new element. @@ -1174,6 +1255,7 @@ macro_rules! impl_slice_eq { )* } } +pub(super) use impl_slice_eq; impl_slice_eq! { [A1: Allocator, A2: Allocator] Vec<T, A1>, Vec<U, A2>, diff --git a/rust/kernel/alloc/kvec/errors.rs b/rust/kernel/alloc/kvec/errors.rs index aaca6446516a..f5c44fdc7016 100644 --- a/rust/kernel/alloc/kvec/errors.rs +++ b/rust/kernel/alloc/kvec/errors.rs @@ -7,7 +7,7 @@ use crate::{ prelude::*, // }; -/// Error type for [`Vec::push_within_capacity`]. +/// Error type for a push that fails because there is no spare capacity. pub struct PushError<T>(pub T); impl<T> fmt::Debug for PushError<T> { @@ -25,6 +25,36 @@ impl<T> From<PushError<T>> for Error { } } +/// Error type for [`Vec::try_push_init`]. +pub enum TryPushInitError<I, E> { + /// The allocation failed. Hand the initializer back. + AllocError(I), + /// The initializer failed. + InitError(E), +} + +impl<I, E> fmt::Debug for TryPushInitError<I, E> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + TryPushInitError::AllocError(_) => write!(f, "Failed to allocate"), + TryPushInitError::InitError(_) => write!(f, "Initializer failed"), + } + } +} + +impl<I, E> From<TryPushInitError<I, E>> for Error +where + Error: From<E>, +{ + #[inline] + fn from(e: TryPushInitError<I, E>) -> Error { + match e { + TryPushInitError::AllocError(_) => ENOMEM, + TryPushInitError::InitError(e) => Error::from(e), + } + } +} + /// Error type for [`Vec::remove`]. pub struct RemoveError; |