[PATCH 3/9] rust: dma: add CoherentArray for compile-time sized allocations

[Eliot Courtney <ecourtney@nvidia.com>]

Add a CoherentArray type alias which takes the size parameter directly,
without using the StaticSize<N> marker type. This makes it a bit nicer
to use.

Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 drivers/gpu/nova-core/dma.rs |   8 +--
 rust/kernel/dma.rs           | 127 ++++++++++++++++++++++++++++++++++++++++++-
 samples/rust/rust_dma.rs     |   8 +--
 3 files changed, 132 insertions(+), 11 deletions(-)

diff --git a/drivers/gpu/nova-core/dma.rs b/drivers/gpu/nova-core/dma.rs
index f77754f12f02..c217cdb14223 100644
--- a/drivers/gpu/nova-core/dma.rs
+++ b/drivers/gpu/nova-core/dma.rs
@@ -9,13 +9,13 @@
 
 use kernel::{
     device,
-    dma::CoherentAllocation,
+    dma::CoherentSlice,
     page::PAGE_SIZE,
     prelude::*, //
 };
 
 pub(crate) struct DmaObject {
-    dma: CoherentAllocation<u8>,
+    dma: CoherentSlice<u8>,
 }
 
 impl DmaObject {
@@ -24,7 +24,7 @@ pub(crate) fn new(dev: &device::Device<device::Bound>, len: usize) -> Result<Sel
             .map_err(|_| EINVAL)?
             .pad_to_align()
             .size();
-        let dma = CoherentAllocation::alloc_coherent(dev, len, GFP_KERNEL | __GFP_ZERO)?;
+        let dma = CoherentSlice::alloc_coherent(dev, len, GFP_KERNEL | __GFP_ZERO)?;
 
         Ok(Self { dma })
     }
@@ -40,7 +40,7 @@ pub(crate) fn from_data(dev: &device::Device<device::Bound>, data: &[u8]) -> Res
 }
 
 impl Deref for DmaObject {
-    type Target = CoherentAllocation<u8>;
+    type Target = CoherentSlice<u8>;
 
     fn deref(&self) -> &Self::Target {
         &self.dma
diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 6e6d91a9cd62..43ed0dfdbc08 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -194,12 +194,12 @@ pub const fn value(&self) -> u64 {
 ///
 /// ```
 /// # use kernel::device::{Bound, Device};
-/// use kernel::dma::{attrs::*, CoherentAllocation};
+/// use kernel::dma::{attrs::*, CoherentArray};
 ///
 /// # fn test(dev: &Device<Bound>) -> Result {
 /// let attribs = DMA_ATTR_FORCE_CONTIGUOUS | DMA_ATTR_NO_WARN;
-/// let c: CoherentAllocation<u64> =
-///     CoherentAllocation::alloc_attrs(dev, 4, GFP_KERNEL, attribs)?;
+/// let c: CoherentArray<u64, 4> =
+///     CoherentArray::alloc_attrs(dev, GFP_KERNEL, attribs)?;
 /// # Ok::<(), Error>(()) }
 /// ```
 #[derive(Clone, Copy, PartialEq)]
@@ -414,6 +414,9 @@ pub struct CoherentAllocation<T: AsBytes + FromBytes, Size: AllocationSize = Run
 /// A coherent DMA allocation with a runtime-determined size.
 pub type CoherentSlice<T> = CoherentAllocation<T, RuntimeSize>;
 
+/// A coherent DMA allocation for an array of `N` elements.
+pub type CoherentArray<T, const N: usize> = CoherentAllocation<T, StaticSize<N>>;
+
 impl<T: AsBytes + FromBytes, Size: AllocationSize> CoherentAllocation<T, Size> {
     /// Returns the number of elements `T` in this allocation.
     ///
@@ -692,6 +695,124 @@ pub fn alloc_coherent(
     }
 }
 
+impl<T: AsBytes + FromBytes, const N: usize> CoherentArray<T, N> {
+    /// Allocates a region of `size_of::<T> * N` of coherent memory.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use kernel::device::{Bound, Device};
+    /// use kernel::dma::{attrs::*, CoherentArray};
+    ///
+    /// # fn test(dev: &Device<Bound>) -> Result {
+    /// let c: CoherentArray<u64, 4> =
+    ///     CoherentArray::alloc_attrs(dev, GFP_KERNEL, DMA_ATTR_NO_WARN)?;
+    /// # Ok::<(), Error>(()) }
+    /// ```
+    pub fn alloc_attrs(
+        dev: &device::Device<Bound>,
+        gfp_flags: kernel::alloc::Flags,
+        dma_attrs: Attrs,
+    ) -> Result<Self> {
+        Self::alloc_impl(dev, N, gfp_flags, dma_attrs)
+    }
+
+    /// Performs the same functionality as [`CoherentArray::alloc_attrs`], except the
+    /// `dma_attrs` is 0 by default.
+    pub fn alloc_coherent(
+        dev: &device::Device<Bound>,
+        gfp_flags: kernel::alloc::Flags,
+    ) -> Result<Self> {
+        Self::alloc_attrs(dev, gfp_flags, Attrs(0))
+    }
+
+    /// Returns a DMA handle starting at `OFFSET` (in units of `T`) which may be given to the
+    /// device as the DMA address base of the region.
+    pub fn dma_handle_with_offset<const OFFSET: usize>(&self) -> DmaAddress {
+        build_assert!(OFFSET < N, "Offset is out of bounds for the allocation.");
+
+        // INVARIANT: The type invariant of `Self` guarantees that `size_of::<T> * N` fits
+        // into a `usize`, and `OFFSET` is inferior to `N`.
+        self.dma_handle + (OFFSET * core::mem::size_of::<T>()) as DmaAddress
+    }
+
+    /// Returns the data from the region starting from `OFFSET` as a slice.
+    /// `OFFSET` and `COUNT` are in units of `T`, not the number of bytes.
+    ///
+    /// For ringbuffer type of r/w access or use-cases where the pointer to the live data is needed,
+    /// [`CoherentAllocation::start_ptr`] or [`CoherentAllocation::start_ptr_mut`] could be used
+    /// instead.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that the device does not read/write to/from memory while the returned
+    ///   slice is live.
+    /// * Callers must ensure that this call does not race with a write to the same region while
+    ///   the returned slice is live.
+    pub unsafe fn as_slice<const OFFSET: usize, const COUNT: usize>(&self) -> &[T] {
+        build_assert!(
+            OFFSET + COUNT <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`,
+        //   we've just checked that the range and index is within bounds. The immutability of the
+        //   data is also guaranteed by the safety requirements of the function.
+        // - `OFFSET + COUNT` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe { core::slice::from_raw_parts(self.start_ptr().add(OFFSET), COUNT) }
+    }
+
+    /// Performs the same functionality as [`CoherentArray::as_slice`], except that a mutable
+    /// slice is returned.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that the device does not read/write to/from memory while the returned
+    ///   slice is live.
+    /// * Callers must ensure that this call does not race with a read or write to the same region
+    ///   while the returned slice is live.
+    pub unsafe fn as_slice_mut<const OFFSET: usize, const COUNT: usize>(&mut self) -> &mut [T] {
+        build_assert!(
+            OFFSET + COUNT <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`,
+        //   we've just checked that the range and index is within bounds. The immutability of the
+        //   data is also guaranteed by the safety requirements of the function.
+        // - `OFFSET + COUNT` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe { core::slice::from_raw_parts_mut(self.start_ptr_mut().add(OFFSET), COUNT) }
+    }
+
+    /// Writes data to the region starting from `OFFSET`. `OFFSET` is in units of `T`, not the
+    /// number of bytes.
+    ///
+    /// # Safety
+    ///
+    /// * Callers must ensure that this call does not race with a read or write to the same region
+    ///   that overlaps with this write.
+    pub unsafe fn write<const OFFSET: usize, const SIZE: usize>(&mut self, src: &[T; SIZE]) {
+        build_assert!(
+            OFFSET + SIZE <= N,
+            "Range is out of bounds for the allocation."
+        );
+        // SAFETY:
+        // - The pointer is valid due to type invariant on `CoherentAllocation`
+        //   and we've just checked that the range and index is within bounds.
+        // - `OFFSET + SIZE` can't overflow since it is smaller than `N` and we've checked
+        //   that `N` won't overflow early in the constructor.
+        unsafe {
+            core::ptr::copy_nonoverlapping(
+                src.as_ptr(),
+                self.start_ptr_mut().add(OFFSET),
+                src.len(),
+            )
+        };
+    }
+}
+
 /// Note that the device configured to do DMA must be halted before this object is dropped.
 impl<T: AsBytes + FromBytes, Size: AllocationSize> Drop for CoherentAllocation<T, Size> {
     fn drop(&mut self) {
diff --git a/samples/rust/rust_dma.rs b/samples/rust/rust_dma.rs
index 7a87048575df..97711a99ac8b 100644
--- a/samples/rust/rust_dma.rs
+++ b/samples/rust/rust_dma.rs
@@ -6,7 +6,7 @@
 
 use kernel::{
     device::Core,
-    dma::{CoherentAllocation, DataDirection, Device, DmaMask},
+    dma::{CoherentSlice, DataDirection, Device, DmaMask},
     page, pci,
     prelude::*,
     scatterlist::{Owned, SGTable},
@@ -16,7 +16,7 @@
 #[pin_data(PinnedDrop)]
 struct DmaSampleDriver {
     pdev: ARef<pci::Device>,
-    ca: CoherentAllocation<MyStruct>,
+    ca: CoherentSlice<MyStruct>,
     #[pin]
     sgt: SGTable<Owned<VVec<u8>>>,
 }
@@ -64,8 +64,8 @@ fn probe(pdev: &pci::Device<Core>, _info: &Self::IdInfo) -> impl PinInit<Self, E
             // SAFETY: There are no concurrent calls to DMA allocation and mapping primitives.
             unsafe { pdev.dma_set_mask_and_coherent(mask)? };
 
-            let ca: CoherentAllocation<MyStruct> =
-                CoherentAllocation::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
+            let ca: CoherentSlice<MyStruct> =
+                CoherentSlice::alloc_coherent(pdev.as_ref(), TEST_VALUES.len(), GFP_KERNEL)?;
 
             for (i, value) in TEST_VALUES.into_iter().enumerate() {
                 kernel::try_dma_write!(ca[i] = MyStruct::new(value.0, value.1))?;

-- 
2.52.0