Re: [PATCH v2 17/21] rust: num: Add an upward alignment helper for usize

[Joel Fernandes <joelagnelf@nvidia.com>]

Hello, Alex,

On 5/2/2025 12:57 AM, Alexandre Courbot wrote:
> On Thu May 1, 2025 at 9:58 PM JST, Alexandre Courbot wrote:
>> From: Joel Fernandes <joelagnelf@nvidia.com>
>>
>> This will be used in the nova-core driver where we need to upward-align
>> the image size to get to the next image in the VBIOS ROM.
>>
>> [acourbot@nvidia.com: handled conflicts due to removal of patch creating
>> num.rs]
>>
>> Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
>> Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
>> ---
>>  rust/kernel/lib.rs |  1 +
>>  rust/kernel/num.rs | 21 +++++++++++++++++++++
>>  2 files changed, 22 insertions(+)
>>
>> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
>> index ab0286857061d2de1be0279cbd2cd3490e5a48c3..be75b196aa7a29cf3eed7c902ed8fb98689bbb50 100644
>> --- a/rust/kernel/lib.rs
>> +++ b/rust/kernel/lib.rs
>> @@ -67,6 +67,7 @@
>>  pub mod miscdevice;
>>  #[cfg(CONFIG_NET)]
>>  pub mod net;
>> +pub mod num;
>>  pub mod of;
>>  pub mod page;
>>  #[cfg(CONFIG_PCI)]
>> diff --git a/rust/kernel/num.rs b/rust/kernel/num.rs
>> new file mode 100644
>> index 0000000000000000000000000000000000000000..953c6ab012601efb3c9387b4b299e22233670c4b
>> --- /dev/null
>> +++ b/rust/kernel/num.rs
>> @@ -0,0 +1,21 @@
>> +// SPDX-License-Identifier: GPL-2.0
>> +
>> +//! Numerical and binary utilities for primitive types.
>> +
>> +/// A trait providing alignment operations for `usize`.
>> +pub trait UsizeAlign {
>> +    /// Aligns `self` upwards to the nearest multiple of `align`.
>> +    fn align_up(self, align: usize) -> usize;
>> +}
> 
> As it turns out I will also need the same functionality for u64 in a
> future patch. :) Could we turn this into a more generic trait? E.g:
> 
>     /// A trait providing alignment operations for `usize`.
>     pub trait Align {
>         /// Aligns `self` upwards to the nearest multiple of `align`.
>         fn align_up(self, align: Self) -> Self;
>     }
> 
> That way it can be implemented for all basic types. I'd suggest having a local
> macro that takes an arbitrary number of arguments and generates the impl blocks
> for each of them, which would be invoked as:
> 
>     impl_align!(i8, u8, i16, u16, ...);

I actually tried this initially before settling for the simpler solution.

We can do this in 3 ways I think, generics, macros or both.

Unlike the last attempt, this time I did get generics to work. So how about this?

use core::ops::{Add, Sub, BitAnd, Not};

pub trait AlignUp {
    fn align_up(self, alignment: Self) -> Self;
}

impl<T> AlignUp for T
where
    T: Copy
        + Add<Output = T>
        + Sub<Output = T>
        + BitAnd<Output = T>
        + Not<Output = T>
        + From<u8>,
{
    fn align_up(self, alignment: Self) -> Self {
        let one = T::from(1u8);
        (self + alignment - one) & !(alignment - one)
    }
}

I know you must be screaming the word monomorphization, but it is clean code and
I'm willing to see just how much code the compiler generates and does not
requiring specifying any specific types at all!

This also could be simpler if we had num_traits in r4L like userspace, because
num_trait's PrimInt encapsulates all the needed ops.

use num_traits::PrimInt;

pub trait RoundUp {
    fn round_up(self, alignment: Self) -> Self;
}

impl<T> RoundUp for T
where
    T: PrimInt,
{
    fn round_up(self, alignment: Self) -> Self {
        (self + alignment - T::one()) & !(alignment - T::one())
    }
}

fn main() {
    let x: u32 = 1234;
    let y: usize = 1234;

    // Output 1536
    println!("u32: {}", x.round_up(512));
    println!("usize: {}", y.round_up(512));
}

For the monomorphization issues, I do wish Rust in general supported restricting
types further so if we could say "where T is u32, usize etc." but it does not
afaics, so maybe we can do this then?

/// This bit can go into separate module if we want to call it
/// 'num_traits' something.

ppub trait Alignable:
    Copy
    + Add<Output = Self>
    + Sub<Output = Self>
    + BitAnd<Output = Self>
    + Not<Output = Self>
    + From<u8>
{
}

impl Alignable for usize {}
impl Alignable for u8 {}
impl Alignable for u16 {}
impl Alignable for u32 {}
impl Alignable for u64 {}
impl Alignable for u128 {}

//////////////////////

And then num.rs remains simple but restricted to those types:

pub trait AlignUp {
    fn align_up(self, alignment: Self) -> Self;
}

impl<T> AlignUp for T
where
    T: Alignable ,
{
    fn align_up(self, alignment: Self) -> Self {
        let one = T::from(1u8);
        (self + alignment - one) & !(alignment - one)
    }
}

If we dislike that, we could go with the pure macro implementation as well. But
I do like the 'num_traits' approach better, since it keeps the align_up
implementation quite clean.

pub trait AlignUp {
    fn align_up(self, alignment: Self) -> Self;
}

macro_rules! align_up_impl {
    ($($t:ty),+) => {
        $(
            impl AlignUp for $t {
                fn align_up(self, alignment: Self) -> Self {
                    (self + alignment - 1) & !(alignment - 1)
                }
            }
        )+
    }
}

align_up_impl!(usize, u8, u16, u32, u64, u128);

Or, we can even combine the 2 approaches. Use macros for the "impl Alignable"
and use generics on the Alignable trait.

macro_rules! impl_alignable {
    ($($t:ty),+) => {
        $(
            impl Alignable for $t {}
        )+
    };
}

impl_alignable!(usize, u8, u16, u32, u64, u128);

pub trait AlignUp {
    fn align_up(self, alignment: Self) -> Self;
}

impl<T> AlignUp for T
where
    T: Alignable,
{
    fn align_up(self, alignment: Self) -> Self {
        let one = T::from(1u8);
        (self + alignment - one) & !(alignment - one)
    }
}

Thoughts?

>
>> +impl UsizeAlign for usize {
>> +    fn align_up(mut self, align: usize) -> usize {
>> +        self = (self + align - 1) & !(align - 1);
>> +        self
>> +    }
>> +}
> 
> Note that there is no need to mutate `self`, the body can just be:
> 
>     (self + align - 1) & !(align - 1)

Sure, that's fair enough. I am Ok with either.

> 
> This operation can trigger overflows/underflows, so I guess for safety it
> should return a `Result` and use `checked_add` and `checked_sub`, after moving
> `align - 1` into a local variable to avoid checking it twice.
> 
> There is also the interesting question of, what if `align` is not a power of 2.
> We should probably document the fact that it is expected to be.

Good idea, makes sense to return Result and also check for power-of-2.

> 
>> +/// Aligns `val` upwards to the nearest multiple of `align`.
>> +pub const fn usize_align_up(val: usize, align: usize) -> usize {
>> +    (val + align - 1) & !(align - 1)
>> +}
> 
> Let's add a statement in the documentation saying that this exists for use in
> `const` contexts. The `impl` version can maybe call this one directly to avoid
> repeating the same operation twice.

Sure.

> Actually I'm not sure whether we want the `impl` block at all since this>
provides the same functionality, albeit in a slightly less cosmetic way. Can
> core R4L folks provide guidance about that?

Yeah I am Ok with this as well, however as you noted it is less cosmetic.

To clarify for r4l folks, Alex is asking do we add an impl on the types for
aligning them up as done by this patch, or do we provide a bunch of helpers
(possibly using generics or macros to keep the code clean)?

thanks,

 - Joel