[PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

[Alexandre Courbot <acourbot@nvidia.com>]

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@nvidia.com>
---
 rust/kernel/lib.rs |   1 +
 rust/kernel/num.rs | 163 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 164 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..6ecff037893dd25420a6369ea0cd6adbe3460b29
--- /dev/null
+++ b/rust/kernel/num.rs
@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Numerical and binary utilities for primitive types.
+
+use crate::build_assert;
+use core::fmt::Debug;
+use core::hash::Hash;
+
+/// An unsigned integer which is guaranteed to be a power of 2.
+///
+/// # Invariants
+///
+/// The stored value is guaranteed to be a power of two.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[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;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
+                /// assert_eq!(v.value(), 16);
+                /// ```
+                #[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;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::try_new(16), Some(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(16)));"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::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;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(16);")]
+                /// assert_eq!(v.value(), 16);
+                /// ```
+                #[inline(always)]
+                pub const fn value(self) -> $t {
+                    self.0
+                }
+
+                /// Returns the mask corresponding to `self.value() - 1`.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!("let v = PowerOfTwo::<", stringify!($t), ">::new(0x10);")]
+                /// assert_eq!(v.mask(), 0xf);
+                /// ```
+                #[inline(always)]
+                pub const fn mask(self) -> $t {
+                    self.0.wrapping_sub(1)
+                }
+
+                /// Aligns `self` down to `alignment`.
+                ///
+                /// # Examples
+                ///
+                /// ```
+                /// use kernel::num::PowerOfTwo;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_down(0x2f), 0x20);"
+                )]
+                /// ```
+                #[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;
+                ///
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x4f), 0x50);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x40), 0x40);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(0x0), 0x0);"
+                )]
+                #[doc = concat!(
+                    "assert_eq!(PowerOfTwo::<",
+                    stringify!($t),
+                    ">::new(0x10).align_up(",
+                    stringify!($t), "::MAX), 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);

-- 
2.49.0