[PATCH v2 3/6] rust: ptr: add panicking index projection variant

[Gary Guo <gary@garyguo.net>]

There have been a few cases where the programmer knows that the indices are
in bounds but the compiler cannot deduce that. This is also
compiler-version-dependent, so using build indexing here can be
problematic. On the other hand, it is also not ideal to use the fallible
variant, as it adds an error handling path that is never hit.

Add a new panicking index projection for this scenario. Like all panicking
operations, this should be used carefully only in cases where the user
knows the index is going to be in bounds, and panicking would indicate
something is catastrophically wrong.

To signify this, require users to explicitly denote the type of index being
used. The existing two types of index projections also gain the keyworded
version, which will be the recommended way going forward.

The keyworded syntax also paves the way of perhaps adding more flavors in
the future, e.g. `unsafe` index projection. However, unless the code is
extremely performance sensitive and bounds checking cannot be tolerated,
the panicking variant is safer and should be preferred, so it will be left
to the future when demand arises.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/kernel/dma.rs            |  3 ++
 rust/kernel/ptr/projection.rs | 96 ++++++++++++++++++++++++++++++++++++-------
 2 files changed, 84 insertions(+), 15 deletions(-)

diff --git a/rust/kernel/dma.rs b/rust/kernel/dma.rs
index 8f97916e0688..ff69739eae35 100644
--- a/rust/kernel/dma.rs
+++ b/rust/kernel/dma.rs
@@ -1207,6 +1207,9 @@ macro_rules! dma_write {
     (@parse [$dma:expr] [$($proj:tt)*] [.$field:tt $($rest:tt)*]) => {
         $crate::dma_write!(@parse [$dma] [$($proj)* .$field] [$($rest)*])
     };
+    (@parse [$dma:expr] [$($proj:tt)*] [[$flavor:ident: $index:expr] $($rest:tt)*]) => {
+        $crate::dma_write!(@parse [$dma] [$($proj)* [$flavor: $index]] [$($rest)*])
+    };
     (@parse [$dma:expr] [$($proj:tt)*] [[$index:expr]? $($rest:tt)*]) => {
         $crate::dma_write!(@parse [$dma] [$($proj)* [$index]?] [$($rest)*])
     };
diff --git a/rust/kernel/ptr/projection.rs b/rust/kernel/ptr/projection.rs
index 1b54616c6cbb..e46877a208d9 100644
--- a/rust/kernel/ptr/projection.rs
+++ b/rust/kernel/ptr/projection.rs
@@ -26,14 +26,14 @@ fn from(_: OutOfBound) -> Self {
 ///
 /// # Safety
 ///
-/// For given input pointer `slice` and return value `output`, the implementation of `build_index`
-/// and `get` (if [`Some`] is returned) must ensure that:
+/// For given input pointer `slice` and return value `output`, the implementation of `index`,
+/// `build_index` and `get` (if [`Some`] is returned) must ensure that:
 /// - `output` has the same provenance as `slice`;
 /// - `output.byte_offset_from(slice)` is between 0 to
 ///   `KnownSize::size(slice) - KnownSize::size(output)`.
 ///
-/// This means that if the input pointer is valid, then pointer returned by `get` or `build_index`
-/// is also valid.
+/// This means that if the input pointer is valid, then pointer returned by `get`, `index` or
+/// `build_index` is also valid.
 #[diagnostic::on_unimplemented(message = "`{Self}` cannot be used to index `{T}`")]
 #[doc(hidden)]
 pub unsafe trait ProjectIndex<T: ?Sized>: Sized {
@@ -42,6 +42,9 @@ pub unsafe trait ProjectIndex<T: ?Sized>: Sized {
     /// Returns an index-projected pointer, if in bounds.
     fn get(self, slice: *mut T) -> Option<*mut Self::Output>;
 
+    /// Returns an index-projected pointer; panic if out of bounds.
+    fn index(self, slice: *mut T) -> *mut Self::Output;
+
     /// Returns an index-projected pointer; fail the build if it cannot be proved to be in bounds.
     #[inline(always)]
     fn build_index(self, slice: *mut T) -> *mut Self::Output {
@@ -66,6 +69,11 @@ fn build_index(self, slice: *mut T) -> *mut Self::Output {
         <I as ProjectIndex<[T]>>::get(self, slice)
     }
 
+    #[inline(always)]
+    fn index(self, slice: *mut [T; N]) -> *mut Self::Output {
+        <I as ProjectIndex<[T]>>::index(self, slice)
+    }
+
     #[inline(always)]
     fn build_index(self, slice: *mut [T; N]) -> *mut Self::Output {
         <I as ProjectIndex<[T]>>::build_index(self, slice)
@@ -85,6 +93,16 @@ fn get(self, slice: *mut [T]) -> Option<*mut T> {
             Some(slice.cast::<T>().wrapping_add(self))
         }
     }
+
+    #[inline(always)]
+    fn index(self, slice: *mut [T]) -> *mut T {
+        // Leverage Rust built-in operators for bounds checking.
+        // SAFETY: All non-null and aligned pointers are valid for ZST read.
+        let zst_slice =
+            unsafe { core::slice::from_raw_parts::<()>(core::ptr::dangling(), slice.len()) };
+        let () = zst_slice[self];
+        slice.cast::<T>().wrapping_add(self)
+    }
 }
 
 // SAFETY: `get`-returned pointer has the same provenance as `slice` and the offset is checked to
@@ -103,6 +121,18 @@ fn get(self, slice: *mut [T]) -> Option<*mut [T]> {
             new_len,
         ))
     }
+
+    #[inline(always)]
+    fn index(self, slice: *mut [T]) -> *mut [T] {
+        // Leverage Rust built-in operators for bounds checking.
+        // SAFETY: All non-null and aligned pointers are valid for ZST read.
+        let zst_slice =
+            unsafe { core::slice::from_raw_parts::<()>(core::ptr::dangling(), slice.len()) };
+        _ = zst_slice[self.clone()];
+
+        // SAFETY: bounds checked.
+        unsafe { self.get(slice).unwrap_unchecked() }
+    }
 }
 
 // SAFETY: Safety requirement guaranteed by the forwarded impl.
@@ -113,6 +143,11 @@ unsafe impl<T> ProjectIndex<[T]> for core::ops::RangeTo<usize> {
     fn get(self, slice: *mut [T]) -> Option<*mut [T]> {
         (0..self.end).get(slice)
     }
+
+    #[inline(always)]
+    fn index(self, slice: *mut [T]) -> *mut [T] {
+        (0..self.end).index(slice)
+    }
 }
 
 // SAFETY: Safety requirement guaranteed by the forwarded impl.
@@ -123,6 +158,11 @@ unsafe impl<T> ProjectIndex<[T]> for core::ops::RangeFrom<usize> {
     fn get(self, slice: *mut [T]) -> Option<*mut [T]> {
         (self.start..slice.len()).get(slice)
     }
+
+    #[inline(always)]
+    fn index(self, slice: *mut [T]) -> *mut [T] {
+        (self.start..slice.len()).index(slice)
+    }
 }
 
 // SAFETY: `get` returned the pointer as is, so it always has the same provenance and offset of 0.
@@ -133,6 +173,11 @@ unsafe impl<T> ProjectIndex<[T]> for core::ops::RangeFull {
     fn get(self, slice: *mut [T]) -> Option<*mut [T]> {
         Some(slice)
     }
+
+    #[inline(always)]
+    fn index(self, slice: *mut [T]) -> *mut [T] {
+        slice
+    }
 }
 
 /// A helper trait to perform field projection.
@@ -211,9 +256,12 @@ unsafe fn proj<F>(_: *mut Self, _: impl FnOnce(*mut Self) -> *mut F) -> *mut F {
 /// `kernel::ptr::project!(mut ptr, projection)`. By default, a const pointer is created.
 ///
 /// `ptr::project!` macro can perform both fallible indexing and build-time checked indexing.
-/// `[index]` form performs build-time bounds checking; if compiler fails to prove `[index]` is in
-/// bounds, compilation will fail. `[index]?` can be used to perform runtime bounds checking;
-/// `OutOfBound` error is raised via `?` if the index is out of bounds.
+/// The syntax is of form `[<flavor>: index]` where `flavor` indicates the way of handling index
+/// out-of-bound errors.
+/// - `try` will raise an `OutOfBound` error (which is convertible to `ERANGE`).
+/// - `build` will use [`build_assert!`](kernel::build_assert::build_assert) mechanism to have
+///   compiler validate the index is in bounds.
+/// - `panic` will cause a Rust [`panic!`] if index goes out of bound.
 ///
 /// # Examples
 ///
@@ -231,17 +279,21 @@ unsafe fn proj<F>(_: *mut Self, _: impl FnOnce(*mut Self) -> *mut F) -> *mut F {
 /// }
 /// ```
 ///
-/// Index projections are performed with `[index]`:
+/// Index projections are performed with `[<flavor>: index]`, where `flavor` is `try`, `build` or
+/// `panic`:
 ///
 /// ```
 /// fn proj(ptr: *const [u8; 32]) -> Result {
-///     let field_ptr: *const u8 = kernel::ptr::project!(ptr, [1]);
+///     let field_ptr: *const u8 = kernel::ptr::project!(ptr, [build: 1]);
 ///     // The following invocation, if uncommented, would fail the build.
 ///     //
-///     // kernel::ptr::project!(ptr, [128]);
+///     // kernel::ptr::project!(ptr, [build: 128]);
 ///
 ///     // This will raise an `OutOfBound` error (which is convertible to `ERANGE`).
-///     kernel::ptr::project!(ptr, [128]?);
+///     kernel::ptr::project!(ptr, [try: 128]);
+///
+///     // This will panic in runtime if executed.
+///     kernel::ptr::project!(ptr, [panic: 128]);
 ///     Ok(())
 /// }
 /// ```
@@ -251,7 +303,7 @@ unsafe fn proj<F>(_: *mut Self, _: impl FnOnce(*mut Self) -> *mut F) -> *mut F {
 /// ```
 /// let ptr: *const [u8; 32] = core::ptr::dangling();
 /// let field_ptr: Result<*const u8> = (|| -> Result<_> {
-///     Ok(kernel::ptr::project!(ptr, [128]?))
+///     Ok(kernel::ptr::project!(ptr, [try: 128]))
 /// })();
 /// assert!(field_ptr.is_err());
 /// ```
@@ -260,7 +312,7 @@ unsafe fn proj<F>(_: *mut Self, _: impl FnOnce(*mut Self) -> *mut F) -> *mut F {
 ///
 /// ```
 /// let ptr: *mut [(u8, u16); 32] = core::ptr::dangling_mut();
-/// let field_ptr: *mut u16 = kernel::ptr::project!(mut ptr, [1].1);
+/// let field_ptr: *mut u16 = kernel::ptr::project!(mut ptr, [build: 1].1);
 /// ```
 #[macro_export]
 macro_rules! project_pointer {
@@ -283,16 +335,30 @@ macro_rules! project_pointer {
         $crate::ptr::project!(@gen $ptr, $($rest)*)
     };
     // Fallible index projection.
-    (@gen $ptr:ident, [$index:expr]? $($rest:tt)*) => {
+    (@gen $ptr:ident, [try: $index:expr] $($rest:tt)*) => {
         let $ptr = $crate::ptr::projection::ProjectIndex::get($index, $ptr)
             .ok_or($crate::ptr::projection::OutOfBound)?;
         $crate::ptr::project!(@gen $ptr, $($rest)*)
     };
+    // Panicking index projection.
+    (@gen $ptr:ident, [panic: $index:expr] $($rest:tt)*) => {
+        let $ptr = $crate::ptr::projection::ProjectIndex::index($index, $ptr);
+        $crate::ptr::project!(@gen $ptr, $($rest)*)
+    };
     // Build-time checked index projection.
-    (@gen $ptr:ident, [$index:expr] $($rest:tt)*) => {
+    (@gen $ptr:ident, [build: $index:expr] $($rest:tt)*) => {
         let $ptr = $crate::ptr::projection::ProjectIndex::build_index($index, $ptr);
         $crate::ptr::project!(@gen $ptr, $($rest)*)
     };
+
+    // For compatibility
+    (@gen $ptr:ident, [$index:expr]? $($rest:tt)*) => {
+        $crate::ptr::project!(@gen $ptr, [try: $index] $($rest)*)
+    };
+    (@gen $ptr:ident, [$index:expr] $($rest:tt)*) => {
+        $crate::ptr::project!(@gen $ptr, [build: $index] $($rest)*)
+    };
+
     (mut $ptr:expr, $($proj:tt)*) => {{
         let ptr: *mut _ = $ptr;
         $crate::ptr::project!(@gen ptr, $($proj)*);

-- 
2.54.0