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

[Boqun Feng <boqun.feng@gmail.com>]

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

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

Regards,
Boqun

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