rust: dma: add dma_handle_with_offset method to CoherentAllocation

Sometimes one may want to obtain a DMA handle starting at a given
offset. This can be done by adding said offset to the result of
`dma_handle()`, but doing so on the client side carries the risk that
the operation will go outside the bounds of the allocation.

Thus, add a `dma_handle_with_offset` method that adds the desired offset
after checking that it is still valid.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
Link: https://lore.kernel.org/r/20250619-nova-frts-v6-3-ecf41ef99252@nvidia.com
Signed-off-by: Danilo Krummrich <dakr@kernel.org>

1 files changed, 14 insertions, 0 deletions

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 34a406a0602e..25dfa0e6cc3c 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -238,6 +238,20 @@ impl<T: AsBytes + FromBytes> CoherentAllocation<T> {
         self.dma_handle
     }
 
+    /// 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.
+    ///
+    /// Returns `EINVAL` if `offset` is not within the bounds of the allocation.
+    pub fn dma_handle_with_offset(&self, offset: usize) -> Result<bindings::dma_addr_t> {
+        if offset >= self.count {
+            Err(EINVAL)
+        } else {
+            // INVARIANT: The type invariant of `Self` guarantees that `size_of::<T> * count` fits
+            // into a `usize`, and `offset` is inferior to `count`.
+            Ok(self.dma_handle + (offset * core::mem::size_of::<T>()) as bindings::dma_addr_t)
+        }
+    }
+
     /// Common helper to validate a range applied from the allocated region in the CPU's virtual
     /// address space.
     fn validate_range(&self, offset: usize, count: usize) -> Result {