[PATCH v5 04/23] rust: add new `num` module with `PowerOfTwo` type

[Alexandre Courbot]

Introduce the `num` module, featuring the `PowerOfTwo` unsigned wrapper
that guarantees (at build-time or runtime) that a value is a power of
two.

Such a property is often useful to maintain. In the context of the
kernel, powers of two are often used to align addresses or sizes up and
down, or to create masks. These operations are provided by this type.

It is introduced to be first used by the nova-core driver.

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

diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 6b4774b2b1c37f4da1866e993be6230bc6715841..2955f65da1278dd4cba1e4272ff178b8211a892c 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -89,6 +89,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..ee0f67ad1a89e69f5f8d2077eba5541b472e7d8a
--- /dev/null
+++ b/rust/kernel/num.rs
@@ -0,0 +1,173 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Numerical and binary utilities for primitive types.
+
+use crate::build_assert;
+use core::borrow::Borrow;
+use core::fmt::Debug;
+use core::hash::Hash;
+use core::ops::Deref;
+
+/// An unsigned integer which is guaranteed to be a power of 2.
+#[derive(Debug, Clone, Copy)]
+#[repr(transparent)]
+pub struct PowerOfTwo<T>(T);
+
+macro_rules! power_of_two_impl {
+    ($($t:ty),+) => {
+        $(
+            impl PowerOfTwo<$t> {
+                /// Validates that `v` is a power of two at build-time, and returns it wrapped into
+                /// `PowerOfTwo`.
+                ///
+                /// A build error is triggered if `v` cannot be asserted to be a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                /// let v = PowerOfTwo::<u32>::new(256);
+                /// assert_eq!(v.value(), 256);
+                /// ```
+                #[inline(always)]
+                pub const fn new(v: $t) -> Self {
+                    build_assert!(v.count_ones() == 1);
+                    Self(v)
+                }
+
+                /// Validates that `v` is a power of two at runtime, and returns it wrapped into
+                /// `PowerOfTwo`.
+                ///
+                /// `None` is returned if `v` was not a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                /// assert_eq!(PowerOfTwo::<u32>::try_new(16).unwrap().value(), 16);
+                /// assert_eq!(PowerOfTwo::<u32>::try_new(15), None);
+                /// ```
+                #[inline(always)]
+                pub const fn try_new(v: $t) -> Option<Self> {
+                    match v.count_ones() {
+                        1 => Some(Self(v)),
+                        _ => None,
+                    }
+                }
+
+                /// Returns the value of this instance.
+                ///
+                /// It is guaranteed to be a power of two.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                /// let v = PowerOfTwo::<u32>::new(256);
+                /// assert_eq!(v.value(), 256);
+                /// ```
+                #[inline(always)]
+                pub const fn value(&self) -> $t {
+                    self.0
+                }
+
+                /// Returns the mask corresponding to `self.value() - 1`.
+                #[inline(always)]
+                pub const fn mask(&self) -> $t {
+                    self.0.wrapping_sub(1)
+                }
+
+                /// Aligns `self` down to `alignment`.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                /// assert_eq!(PowerOfTwo::<u32>::new(0x1000).align_down(0x4fff), 0x4000);
+                /// ```
+                #[inline(always)]
+                pub const fn align_down(self, value: $t) -> $t {
+                    value & !self.mask()
+                }
+
+                /// Aligns `value` up to `self`.
+                ///
+                /// Wraps around to `0` if the requested alignment pushes the result above the
+                /// type's limits.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                /// assert_eq!(PowerOfTwo::<u32>::new(0x1000).align_up(0x4fff), 0x5000);
+                /// assert_eq!(PowerOfTwo::<u32>::new(0x1000).align_up(0x4000), 0x4000);
+                /// assert_eq!(PowerOfTwo::<u32>::new(0x1000).align_up(0x0), 0x0);
+                /// assert_eq!(PowerOfTwo::<u16>::new(0x100).align_up(0xffff), 0x0);
+                /// ```
+                #[inline(always)]
+                pub const fn align_up(self, value: $t) -> $t {
+                    self.align_down(value.wrapping_add(self.mask()))
+                }
+            }
+        )+
+    };
+}
+
+power_of_two_impl!(usize, u8, u16, u32, u64, u128);
+
+impl<T> Deref for PowerOfTwo<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> PartialEq for PowerOfTwo<T>
+where
+    T: PartialEq,
+{
+    fn eq(&self, other: &Self) -> bool {
+        self.0 == other.0
+    }
+}
+
+impl<T> Eq for PowerOfTwo<T> where T: Eq {}
+
+impl<T> PartialOrd for PowerOfTwo<T>
+where
+    T: PartialOrd,
+{
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        self.0.partial_cmp(&other.0)
+    }
+}
+
+impl<T> Ord for PowerOfTwo<T>
+where
+    T: Ord,
+{
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        self.0.cmp(&other.0)
+    }
+}
+
+impl<T> Hash for PowerOfTwo<T>
+where
+    T: Hash,
+{
+    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+        self.0.hash(state);
+    }
+}
+
+impl<T> Borrow<T> for PowerOfTwo<T> {
+    fn borrow(&self) -> &T {
+        &self.0
+    }
+}

-- 
2.49.0