[PATCH v2 3/8] rust: alloc: add ArrayVec

[Eliot Courtney <ecourtney@nvidia.com>]

Add a fixed capacity vector backed by [MaybeUninit<T>; N]. The ArrayVec
is also initializable with a closure, returning an Init instance, to
avoid constructing it on the stack. ArrayVec is useful for small but
varying size arrays stored on the stack, to avoid a heap allocation, or,
for larger varying size arrays initialized into caller provided memory
but not wanting to provide an allocator.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 rust/kernel/alloc.rs          |   3 +
 rust/kernel/alloc/arrayvec.rs | 347 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 350 insertions(+)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 21067bde6860..510e2c7f9f72 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -3,10 +3,13 @@
 //! Implementation of the kernel's memory allocation infrastructure.
 
 pub mod allocator;
+pub mod arrayvec;
 pub mod kbox;
 pub mod kvec;
 pub mod layout;
 
+pub use self::arrayvec::ArrayVec;
+
 pub use self::kbox::Box;
 pub use self::kbox::KBox;
 pub use self::kbox::KVBox;
diff --git a/rust/kernel/alloc/arrayvec.rs b/rust/kernel/alloc/arrayvec.rs
new file mode 100644
index 000000000000..4172a982e477
--- /dev/null
+++ b/rust/kernel/alloc/arrayvec.rs
@@ -0,0 +1,347 @@
+// 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,
+    Zeroable, //
+};
+
+/// 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>(())
+/// ```
+#[derive(Zeroable)]
+pub struct ArrayVec<T, const N: usize> {
+    data: [MaybeUninit<T>; N],
+    len: usize,
+}
+
+impl<T, const N: usize> ArrayVec<T, N> {
+    /// Creates an empty [`ArrayVec`].
+    #[inline]
+    pub const fn new() -> Self {
+        // Clippy triggers this even if the enclosing function is never called, so skip if clippy is
+        // on.
+        const_assert!(
+            cfg!(clippy) || size_of::<Self>() <= 512,
+            "use `init_with` 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 [`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>(())
+    /// ```
+    pub fn init_with<E>(f: impl FnOnce(&mut Self) -> Result<(), E>) -> impl Init<Self, E> {
+        let init = move |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 { ptr::addr_of_mut!((*slot).len).write(0) };
+
+            // SAFETY: `slot` holds a valid ArrayVec and no other reference to it exists.
+            let v = unsafe { &mut *slot };
+            f(v).inspect_err(|_| {
+                // SAFETY: `slot` holds a valid ArrayVec, and on failure the slot is never accessed
+                // again, so the elements can't be dropped twice.
+                unsafe { ptr::drop_in_place(slot) }
+            })
+        };
+
+        // SAFETY: `init` fully initializes the slot on success and drops the potentially filled
+        // ArrayVec on failure.
+        unsafe { init_from_closure(init) }
+    }
+
+    /// Appends an element to the back of the [`ArrayVec`].
+    ///
+    /// Fails when the [`ArrayVec`] is full, handing the element back in [`PushError`].
+    pub fn push(&mut self, v: T) -> Result<(), PushError<T>> {
+        self.try_push_init(v)
+            .map_err(|PushInitError::Full(v)| PushError(v))
+    }
+
+    /// Appends an element to the back of the [`ArrayVec`] by initializing it in place.
+    ///
+    /// Fails with [`FullError`] when the [`ArrayVec`] is full.
+    pub fn push_init(&mut self, init: impl Init<T>) -> Result<(), FullError> {
+        self.try_push_init(init)
+            .map_err(|PushInitError::Full(_)| FullError)
+    }
+
+    /// 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 [`PushInitError::Full`]. If the
+    /// initializer itself fails, its error is returned in [`PushInitError::InitError`].
+    pub fn try_push_init<I, E>(&mut self, init: I) -> Result<(), PushInitError<I, E>>
+    where
+        I: Init<T, E>,
+    {
+        let Some(slot) = self.spare_capacity_mut().first_mut() else {
+            return Err(PushInitError::Full(init));
+        };
+
+        // SAFETY: `slot` refers to allocated, aligned memory valid for a write of one `T`.
+        unsafe { init.__init(slot.as_mut_ptr()) }.map_err(PushInitError::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`].
+    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 PushInitError<I, E> {
+    /// The [`ArrayVec`] is full. Hand the initializer back.
+    Full(I),
+    /// The initializer failed.
+    InitError(E),
+}
+
+impl<I, E> fmt::Debug for PushInitError<I, E> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            PushInitError::Full(_) => write!(f, "Not enough capacity"),
+            PushInitError::InitError(_) => write!(f, "Initializer failed"),
+        }
+    }
+}
+
+impl<I, E> From<PushInitError<I, E>> for Error
+where
+    Error: From<E>,
+{
+    #[inline]
+    fn from(e: PushInitError<I, E>) -> Error {
+        match e {
+            PushInitError::Full(_) => EINVAL,
+            PushInitError::InitError(e) => Error::from(e),
+        }
+    }
+}
+
+/// Error type for [`ArrayVec::push_init`].
+pub struct FullError;
+
+impl fmt::Debug for FullError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Not enough capacity")
+    }
+}
+
+impl From<FullError> for Error {
+    #[inline]
+    fn from(_: FullError) -> Error {
+        EINVAL
+    }
+}
+
+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> {
+    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> {
+    fn borrow(&self) -> &[T] {
+        self.as_slice()
+    }
+}
+
+impl<T, const N: usize> BorrowMut<[T]> for ArrayVec<T, N> {
+    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>;
+
+    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>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter_mut()
+    }
+}
+
+impl<T: fmt::Debug, const N: usize> fmt::Debug for ArrayVec<T, N> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(self.as_slice(), f)
+    }
+}

-- 
2.55.0