[PATCH RFC 1/2] rust: kernel: add bounded integer types

[Alexandre Courbot <acourbot@nvidia.com>]

Add the BoundedInt type, which restricts the number of bits allowed to
be used in a given integer value. This is useful to carry guarantees
when setting bitfields.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 rust/kernel/lib.rs |   1 +
 rust/kernel/num.rs | 120 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index fcffc3988a90392f1d5fc19f15c75d9ba7104f9a..21c1f452ee6a30d46d7ed7f0847488fac068042a 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -101,6 +101,7 @@
 pub mod mm;
 #[cfg(CONFIG_NET)]
 pub mod net;
+pub mod num;
 pub mod of;
 #[cfg(CONFIG_PM_OPP)]
 pub mod opp;
diff --git a/rust/kernel/num.rs b/rust/kernel/num.rs
new file mode 100644
index 0000000000000000000000000000000000000000..f452a4e229ca962ae0c2382d0cda55b5ee335973
--- /dev/null
+++ b/rust/kernel/num.rs
@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Numerical types for the kernel.
+
+/// Integer type for which only the bits `0..NUM_BITS` are valid.
+///
+/// TODO: Find a better name? Bounded sounds like we also have a lower bound...
+///
+/// # Invariants
+///
+/// Only bits `0..NUM_BITS` can be set on this type.
+#[repr(transparent)]
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub struct BoundedInt<T, const NUM_BITS: u32>(T);
+
+// TODO: this should be implemented by a macro for all integer types.
+impl<const NUM_BITS: u32> BoundedInt<u32, NUM_BITS> {
+    /// Mask of the valid bits for this type.
+    pub const MASK: u32 = crate::bits::genmask_u32(0..=(NUM_BITS - 1));
+
+    /// Validates that `value` is within the bounds supported by this type.
+    const fn is_in_bounds(value: u32) -> bool {
+        value & !Self::MASK == 0
+    }
+
+    /// Checks that `value` is valid for this type at compile-time and build a new value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::num::BoundedInt;
+    ///
+    /// assert_eq!(BoundedInt::<u32, 1>::new(1).get(), 1);
+    /// assert_eq!(BoundedInt::<u32, 8>::new(0xff).get(), 0xff);
+    /// assert_eq!(BoundedInt::<u32, { 10 - 3 }>::new(1).get(), 1);
+    pub const fn new(value: u32) -> Self {
+        crate::build_assert!(
+            Self::is_in_bounds(value),
+            "Provided parameter is larger than maximal supported value"
+        );
+
+        Self(value)
+    }
+
+    /// Returns the contained value as a primitive type.
+    pub const fn get(self) -> u32 {
+        if !Self::is_in_bounds(self.0) {
+            // SAFETY: Per the invariants, `self.0` cannot have bits set outside of `MASK`, so
+            // this block will
+            // never be reached.
+            unsafe { core::hint::unreachable_unchecked() }
+        }
+        self.0
+    }
+}
+
+impl<const NUM_BITS: u32> From<BoundedInt<u32, NUM_BITS>> for u32 {
+    fn from(value: BoundedInt<u32, NUM_BITS>) -> Self {
+        value.get()
+    }
+}
+
+/// Attempt to convert a non-bounded integer to a bounded equivalent.
+impl<const NUM_BITS: u32> TryFrom<u32> for BoundedInt<u32, NUM_BITS> {
+    type Error = crate::error::Error;
+
+    fn try_from(value: u32) -> Result<Self, Self::Error> {
+        if !Self::is_in_bounds(value) {
+            Err(crate::prelude::EINVAL)
+        } else {
+            Ok(Self(value))
+        }
+    }
+}
+
+/// Allow comparison with non-bounded values.
+impl<const NUM_BITS: u32, T> PartialEq<T> for BoundedInt<T, NUM_BITS>
+where
+    T: PartialEq,
+{
+    fn eq(&self, other: &T) -> bool {
+        self.0 == *other
+    }
+}
+
+impl<const NUM_BITS: u32, T> core::fmt::Debug for BoundedInt<T, NUM_BITS>
+where
+    T: core::fmt::Debug,
+{
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<const NUM_BITS: u32, T> core::fmt::Display for BoundedInt<T, NUM_BITS>
+where
+    T: core::fmt::Display,
+{
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<const NUM_BITS: u32, T> core::fmt::LowerHex for BoundedInt<T, NUM_BITS>
+where
+    T: core::fmt::LowerHex,
+{
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<const NUM_BITS: u32, T> core::fmt::UpperHex for BoundedInt<T, NUM_BITS>
+where
+    T: core::fmt::UpperHex,
+{
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        self.0.fmt(f)
+    }
+}

-- 
2.51.0