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

[Alexandre Courbot]

On Sat May 3, 2025 at 4:59 AM JST, Joel Fernandes wrote:
> 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 at 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 at nvidia.com: handled conflicts due to removal of patch creating
>>> num.rs]
>>>
>>> Signed-off-by: Joel Fernandes <joelagnelf at nvidia.com>
>>> Signed-off-by: Alexandre Courbot <acourbot at 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!

No, I think this is great - monomorphization only happens as the code is
used, so it won't be a problem. And the compiler should just inline that
anyway (let's add an `#[inline]` to be sure although I'm not convinced
it is entirely necessary).

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

This looks very rust-y and I like it. I guess Alignable (maybe with a
more generic name because I suspect these properties can be used for
other things than aligning as well) could be in the `num` module without
any problem.

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

I think that's the correct way to do it and am fully on board with this
approach.

The only thing this doesn't solve is that it doesn't provide `const`
functions. But maybe for that purpose we can use a single macro that
nicely panics at build-time should an overflow occur.

Cheers,
Alex.