Nouveau - Jun 2025 - [PATCH 1/3] rust: add `num` module with `PowerOfTwo` type

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 | 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

On Fri, Jun 20, 2025 at 3:15?PM Alexandre Courbot <acourbot at nvidia.com>
wrote:
>
> 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.
Before I forget: the other day in a call we discussed powers of two
and I mentioned that there is `Alignment` in the standard library:

    https://doc.rust-lang.org/std/ptr/struct.Alignment.html

    "A type storing a `usize` which is a power of two"

So it would be nice to ask upstream the following if they have plans
to stabilize it, and whether they have considered a generic
`PowerOfTwo<T>` type like this one, rather than one just for alignment
purposes (possibly with an alias or newtype for `Alignment` if
needed).

Similarly, if they stabilize the `Alignment` one (only) and we end up
only using our `PowerOfTwo<T>` for `usize` and those use cases, then
we should consider using the upstream one (and adding any/all methods
that we need).

So I will ask them the next time we meet. I have added
`ptr_alignment_type` to our list (in the "nice to have" section).

(Apologies if this was already discussed!)

Cheers,
Miguel

On Fri Jun 20, 2025 at 10:35 PM JST, Miguel Ojeda wrote:
> On Fri, Jun 20, 2025 at 3:15?PM Alexandre Courbot <acourbot at
nvidia.com> wrote:
>>
>> 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.
>
> Before I forget: the other day in a call we discussed powers of two
> and I mentioned that there is `Alignment` in the standard library:
>
>     https://doc.rust-lang.org/std/ptr/struct.Alignment.html
>
>     "A type storing a `usize` which is a power of two"
Haha, I wasn't aware of this effort, and am quite amazed by how close
its API is to my own design. This is reassuring; maybe I am finally
starting to grok Rust after all. ;)
>
> So it would be nice to ask upstream the following if they have plans
> to stabilize it, and whether they have considered a generic
> `PowerOfTwo<T>` type like this one, rather than one just for
alignment
> purposes (possibly with an alias or newtype for `Alignment` if
> needed).
Mmm indeed I don't quite see the fundamental difference between
`Alignment` and `PowerOfTwo`, although `Alignment` might better capture
what we are doing with our type anyway.
>
> Similarly, if they stabilize the `Alignment` one (only) and we end up
> only using our `PowerOfTwo<T>` for `usize` and those use cases, then
> we should consider using the upstream one (and adding any/all methods
> that we need).
`Alignment` is very close to what we need, so I don't see a reason to
not adopt the same name at the very least.

This reminds me that I should also check whether upstream Rust would be
interested in `prev_multiple_of` and `last_set_bit`. The docs I've read
for contributing looked a bit intimidating, with RFCs to write and all.
Would you have a pointer for where I should start? Maybe a Zulip thread?
>
> So I will ask them the next time we meet. I have added
> `ptr_alignment_type` to our list (in the "nice to have" section).
>
> (Apologies if this was already discussed!)
I wasn't aware of this, so thanks for bringing it up!

On Fri Jun 20, 2025 at 3:14 PM CEST, Alexandre Courbot
wrote:
> +/// 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> {
I tried to use this type in a doctest like this:

    use kernel::num::PowerOfTwo;
   
    fn new(x: usize) -> PowerOfTwo<usize> {
        PowerOfTwo::new(1 << x)
    }

And it doesn't compile :(

    error[E0034]: multiple applicable items in scope
        --> rust/doctests_kernel_generated.rs:4930:17
         |
    4930 |     PowerOfTwo::new(1 << x)
         |                 ^^^ multiple `new` found
         |
         = note: candidate #1 is defined in an impl for the type
`PowerOfTwo<u128>`
         = note: candidate #2 is defined in an impl for the type
`PowerOfTwo<u16>`
         = note: candidate #3 is defined in an impl for the type
`PowerOfTwo<u32>`
         = note: candidate #4 is defined in an impl for the type
`PowerOfTwo<u64>`
         = note: and 2 others
    
    error: aborting due to 1 previous error

The problem is that the function `new` exists 6 times for each of the
integer types. You can write `PowerOfTwo::<usize>::new()` instead, but
that's annoying...

We probably need an `Integer` trait and then do

    impl<I: Integer> PowerOfTwo<I> {
        pub const fn new(value: I) -> Self;
    }
> +                /// 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);
Why not `v.is_power_of_two()`?
> +                    Self(v)
Missing `// INVARIANT` comment.
> +                }
> +
> +                /// 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> {
Maybe `new_checked` is a better name, since it doesn't return a result?
> +                    match v.count_ones() {
Why not `is_power_of_two()`?
> +                        1 => Some(Self(v)),
Missing `// INVARIANT` comment.
> +                        _ => 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
Let's add:

    if !self.0.is_power_of_two() {
        core::hint::unreachable_unchecked()
    }
    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)
Then use `self.value().wrapping_sub(1)` here instead to also propagate
the information.

---
Cheers,
Benno
> +                }