[PATCH 2/4] rust: add pinned wrapper of Vec

[Markus Probst <markus.probst@posteo.de>]

Implement a wrapper of Vec that guarantees that its content will never be
moved, unless the item implements Unpin, allowing PinInit to be
initialized on the Vec.

Signed-off-by: Markus Probst <markus.probst@posteo.de>
---
 rust/kernel/alloc/kvec.rs | 86 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 3c72e0bdddb8..3576929b2e12 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -21,6 +21,7 @@
     slice,
     slice::SliceIndex,
 };
+use pin_init::PinInit;
 
 mod errors;
 pub use self::errors::{InsertError, PushError, RemoveError};
@@ -109,6 +110,11 @@ pub struct Vec<T, A: Allocator> {
     _p: PhantomData<A>,
 }
 
+/// A pinned wrapper of the [`Vec`] type.
+///
+/// It is guaranteed that the contents will never be moved, unless T implements Unpin.
+pub struct PinnedVec<T, A: Allocator>(Vec<T, A>);
+
 /// Type alias for [`Vec`] with a [`Kmalloc`] allocator.
 ///
 /// # Examples
@@ -121,6 +127,8 @@ pub struct Vec<T, A: Allocator> {
 /// # Ok::<(), Error>(())
 /// ```
 pub type KVec<T> = Vec<T, Kmalloc>;
+/// Type alias for [`PinnedVec`] with a [`Kmalloc`] allocator.
+pub type KPinnedVec<T> = PinnedVec<T, Kmalloc>;
 
 /// Type alias for [`Vec`] with a [`Vmalloc`] allocator.
 ///
@@ -134,6 +142,8 @@ pub struct Vec<T, A: Allocator> {
 /// # Ok::<(), Error>(())
 /// ```
 pub type VVec<T> = Vec<T, Vmalloc>;
+/// Type alias for [`PinnedVec`] with a [`Vmalloc`] allocator.
+pub type VPinnedVec<T> = PinnedVec<T, Vmalloc>;
 
 /// Type alias for [`Vec`] with a [`KVmalloc`] allocator.
 ///
@@ -147,6 +157,8 @@ pub struct Vec<T, A: Allocator> {
 /// # Ok::<(), Error>(())
 /// ```
 pub type KVVec<T> = Vec<T, KVmalloc>;
+/// Type alias for [`PinnedVec`] with a [`KVmalloc`] allocator.
+pub type KVPinnedVec<T> = PinnedVec<T, KVmalloc>;
 
 // SAFETY: `Vec` is `Send` if `T` is `Send` because `Vec` owns its elements.
 unsafe impl<T, A> Send for Vec<T, A>
@@ -1294,6 +1306,80 @@ fn drop(&mut self) {
     }
 }
 
+impl<T, A: Allocator> PinnedVec<T, A> {
+    /// Creates a new [`PinnedVec`] instance with at least the given capacity.
+    pub fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
+        Vec::with_capacity(capacity, flags).map(Self)
+    }
+
+    /// Shortens the vector, setting the length to `len` and drops the removed values.
+    /// If `len` is greater than or equal to the current length, this does nothing.
+    ///
+    /// This has no effect on the capacity and will not allocate.
+    pub fn truncate(&mut self, len: usize) {
+        self.0.truncate(len);
+    }
+
+    /// Pin-initializes P and appends it to the back of the [`Vec`] instance without reallocating.
+    pub fn push_pin_init<E, P: PinInit<T, E>>(&mut self, init: P) -> Result<(), E>
+    where
+        E: From<PushError<P>>,
+    {
+        if self.0.len() < self.0.capacity() {
+            let spare = self.0.spare_capacity_mut();
+            // SAFETY: the length is less than the capacity, so `spare` is non-empty.
+            unsafe { init.__pinned_init(spare.get_unchecked_mut(0).as_mut_ptr())? };
+            // SAFETY: We just initialised the first spare entry, so it is safe to
+            // increase the length by 1. We also know that the new length is <= capacity.
+            unsafe { self.0.inc_len(1) };
+            Ok(())
+        } else {
+            Err(E::from(PushError(init)))
+        }
+    }
+
+    /// Removes the last element from a vector and drops it returning true, or false if it is empty.
+    pub fn pop(&mut self) -> bool {
+        if self.is_empty() {
+            return false;
+        }
+
+        // SAFETY: We just checked that the length is at least one.
+        let ptr: *mut [T] = unsafe { self.0.dec_len(1) };
+
+        // SAFETY: the contract of `dec_len` guarantees that the elements in `ptr` are
+        // valid elements whose ownership has been transferred to the caller.
+        unsafe { ptr::drop_in_place(ptr) };
+        true
+    }
+}
+
+impl<T, A: Allocator> Deref for PinnedVec<T, A> {
+    type Target = Vec<T, A>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T: Unpin, A: Allocator> DerefMut for PinnedVec<T, A> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl<T, A: Allocator> From<Vec<T, A>> for PinnedVec<T, A> {
+    fn from(value: Vec<T, A>) -> Self {
+        Self(value)
+    }
+}
+
+impl<T: Unpin, A: Allocator> From<PinnedVec<T, A>> for Vec<T, A> {
+    fn from(value: PinnedVec<T, A>) -> Self {
+        value.0
+    }
+}
+
 #[macros::kunit_tests(rust_kvec_kunit)]
 mod tests {
     use super::*;
-- 
2.49.1