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

Alexandre Courbot
2025-05-03T01:59:46Z

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@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
> + Sub
> + BitAnd
> + Not
> + 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
> + Sub
> + BitAnd
> + Not
> + 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.