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

[Alexandre Courbot]

On Fri Jun 13, 2025 at 12:07 AM JST, Boqun Feng wrote:
> On Thu, Jun 12, 2025 at 11:01:32PM +0900, Alexandre Courbot 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.
>> 
>> 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 {
>
> Then this function should be unsafe, because an invalid `v` can create
> an invalid PowerOfTwo.

Doesn't the `build_assert` below allow us to keep this method safe,
since it will fail at build-time if it cannot be asserted that `v` is a
power of two?

>
>> +                    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 {
>
> I'm late to party, but could we instead implement:
>
>     pub const fn round_down<i32>(value: i32, shift: i32) -> i32 {
>         value & !((1 << shift) - 1)
>     }
>
>     pub const fn round_up<i32>(value: i32, shift: i32) -> i32 {
>         let mask = (1 << shift) - 1;
>         value.wrapping_add(mask) & !mask
>     }
>
> ? It's much harder to pass an invalid alignment with this.

It also forces you to think in terms of shifts instead of values - i.e.
you cannot round to `0x1000` as it commonly done in the kernel, now you
need to do some mental gymnastics to know it is actually a shift of `12`.
Being able to use the actual value to round to is more familiar (and
natural) to me.