[PATCH v3 16/16] rust: io: register: unify handling of register with/without bitfields

Move the `FixedRegister` to be a property of register to become a property
of type. Name the new trait `FixedIoLoc` indicating if I/O location of a
type is unique for a specific base. Thus, bitfields become just a special
case of this (where type is unique because we're generating it in the
register macro), and expose feature to registers without inline bitfield
definition with the `#[unique]` attribute.

Signed-off-by: Gary Guo <[email protected]>
---
 rust/kernel/io/register.rs | 110 ++++++++++++++++++++-------------------------
 rust/macros/io/register.rs |  58 +++++++++++++++---------
 2 files changed, 85 insertions(+), 83 deletions(-)

diff --git a/rust/kernel/io/register.rs b/rust/kernel/io/register.rs
index 6aeee6c0a5f2..6d4ffe6e7b5f 100644
--- a/rust/kernel/io/register.rs
+++ b/rust/kernel/io/register.rs
@@ -121,61 +121,15 @@
     io::IoLoc, //
 };
 
-/// Trait implemented by registers with a fixed offset.
-pub trait FixedRegister: Sized {
-    /// Base type for this register.
-    type Base: ?Sized;
-
-    /// Start offset of the register.
-    ///
-    /// The interpretation of this offset depends on the type of the register.
-    const OFFSET: usize;
-}
-
 /// Allows `()` to be used as the `location` parameter of [`Io::write`](super::Io::write) when
-/// passing a [`FixedRegister`] value.
+/// passing a [`FixedIoLoc`] value.
 impl<Base: ?Sized, T> IoLoc<Base, T> for ()
 where
-    T: FixedRegister<Base = Base>,
-{
-    #[inline(always)]
-    fn offset(self) -> usize {
-        T::OFFSET
-    }
-}
-
-/// A [`FixedRegister`] carries its location in its type. Thus `FixedRegister` values can be used
-/// as an [`IoLoc`].
-impl<Base: ?Sized, T> IoLoc<Base, T> for T
-where
-    T: FixedRegister<Base = Base>,
-{
-    #[inline(always)]
-    fn offset(self) -> usize {
-        T::OFFSET
-    }
-}
-
-/// Location of a fixed register.
-pub struct FixedRegisterLoc<T: FixedRegister>(PhantomData<T>);
-
-impl<T: FixedRegister> FixedRegisterLoc<T> {
-    /// Returns the location of `T`.
-    #[inline(always)]
-    // We do not implement `Default` so we can be const.
-    #[expect(clippy::new_without_default)]
-    pub const fn new() -> Self {
-        Self(PhantomData)
-    }
-}
-
-impl<Base: ?Sized, T> IoLoc<Base, T> for FixedRegisterLoc<T>
-where
-    T: FixedRegister<Base = Base>,
+    T: FixedIoLoc<Base>,
 {
     #[inline(always)]
     fn offset(self) -> usize {
-        T::OFFSET
+        T::LOCATION.offset()
     }
 }
 
@@ -187,6 +141,11 @@ impl<Base: ?Sized, T> OffsetLoc<Base, T> {
     pub const fn new(offset: usize) -> Self {
         Self(offset, PhantomData)
     }
+
+    #[inline]
+    pub const fn const_offset(self) -> usize {
+        self.0
+    }
 }
 
 impl<Base: ?Sized, T> IoLoc<Base, T> for OffsetLoc<Base, T> {
@@ -265,6 +224,17 @@ fn try_at(idx: usize) -> Option<RegisterArrayLoc<Self>>
     }
 }
 
+/// Trait implemented by types that indicate there is a fixed I/O location for this given type.
+///
+/// Implementors can be used with [`Io::write_reg`](super::Io::write_reg).
+pub trait FixedIoLoc<Base: ?Sized>: Sized {
+    /// Type of [`FixedIoLoc::location`].
+    type Location: IoLoc<Base, Self>;
+
+    /// Location of this type within given base.
+    const LOCATION: Self::Location;
+}
+
 /// Trait implemented by items that contain both a register value and the absolute I/O location at
 /// which to write it.
 ///
@@ -282,26 +252,17 @@ pub trait LocatedRegister<Base: ?Sized> {
 
 impl<Base: ?Sized, T> LocatedRegister<Base> for T
 where
-    T: FixedRegister<Base = Base>,
+    T: FixedIoLoc<Base>,
 {
-    type Location = FixedRegisterLoc<Self::Value>;
+    type Location = T::Location;
     type Value = T;
 
     #[inline(always)]
-    fn into_io_op(self) -> (FixedRegisterLoc<T>, T) {
-        (FixedRegisterLoc::new(), self)
+    fn into_io_op(self) -> (T::Location, T) {
+        (T::LOCATION, self)
     }
 }
 
-/// Helper function for register alias implementation.
-///
-/// This is used to enforce base matching. Only called during const eval.
-#[doc(hidden)]
-#[inline(always)]
-pub const fn alias_offset<Base: ?Sized, Alias: FixedRegister<Base = Base>>() -> usize {
-    Alias::OFFSET
-}
-
 /// Helper function for register element alias implementation.
 ///
 /// This is used to enforce base matching and provide bounds checking. Only called during const
@@ -419,6 +380,31 @@ pub const fn element_alias_offset<Base: ?Sized, Alias: RegisterArray<Base = Base
 /// }
 /// ```
 ///
+/// In case there is a fixed register associated with a specific type in the base, you can apply
+/// `#[unique]` attribute which enables `write_reg` shorthand. This is automatically applied to
+/// bitfields instantiated via the `register!` macro.
+///
+/// ```no_run
+/// # use kernel::{bitfield, io::*};
+///
+/// bitfield! {
+///     pub struct Reset(u32) {
+///         0:0 reset;
+///     }
+/// }
+///
+/// register! {
+///     base: Region<0x1000>;
+///
+///     pub RESET: #[unique] Reset @ 0x100;
+/// }
+///
+/// # fn test(mmio: Mmio<'_, Region<0x1000>>) {
+/// // let mmio: Mmio<'_, Region<0x1000>>;
+/// mmio.write_reg(Reset::zeroed().with_const_reset::<1>());
+/// # }
+/// ```
+///
 /// ## Arrays of registers
 ///
 /// Some I/O areas contain consecutive registers that share the same field layout. These areas can
diff --git a/rust/macros/io/register.rs b/rust/macros/io/register.rs
index e0bba6b71ebc..84be1a593959 100644
--- a/rust/macros/io/register.rs
+++ b/rust/macros/io/register.rs
@@ -51,6 +51,7 @@ struct Reg {
     attrs: Vec<Attribute>,
     vis: Visibility,
     name: Ident,
+    unique: bool,
     ty: Type,
     array: Option<RegArrayDef>,
     offset: RegOffset,
@@ -65,16 +66,26 @@ fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
 
         let lh = input.lookahead1();
         let mut bitfield_storage = None;
-        let ty = if lh.peek(Token![:]) {
+        let (unique, ty) = if lh.peek(Token![:]) {
             let _: Token![:] = input.parse()?;
-            input.parse()?
+
+            let mut attrs = input.call(Attribute::parse_outer)?;
+            let unique = attrs
+                .extract_if(.., |attr| attr.path().is_ident("unique"))
+                .count()
+                != 0;
+            if !attrs.is_empty() {
+                Err(Error::new_spanned(&attrs[0], "unexpected attributes"))?
+            }
+
+            (unique, input.parse()?)
         } else if lh.peek(token::Paren) {
             let content;
             parenthesized!(content in input);
             bitfield_storage = Some(content.parse()?);
 
             // For bitfields, bitfield macro will generate a type with the same name as `name`.
-            parse_quote!(#name)
+            (true, parse_quote!(#name))
         } else {
             Err(lh.error())?
         };
@@ -138,6 +149,7 @@ fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
             attrs,
             vis,
             name,
+            unique,
             ty,
             array,
             offset,
@@ -181,6 +193,7 @@ pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
             attrs,
             vis,
             name,
+            unique,
             ty,
             array,
             offset,
@@ -195,7 +208,7 @@ pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
             RegOffset::Fixed { offset } => quote!(#offset),
             RegOffset::Alias { alias } => {
                 quote_spanned!(alias.span().resolved_at(span) =>
-                    ::kernel::io::register::alias_offset::<#base, #alias>()
+                    ::kernel::io::register::OffsetLoc::<#base, _>::const_offset(#alias)
                 )
             }
             RegOffset::ElementAlias { alias, idx } => {
@@ -217,27 +230,30 @@ pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
         }
 
         match array {
-            None if bitfield.is_none() => outputs.extend(quote!(
-                #(#attrs)* #vis const #name: ::kernel::io::register::OffsetLoc<#base, #ty> =
-                    ::kernel::io::register::OffsetLoc::new(#offset);
-            )),
-
-            _ if bitfield.is_none() => Err(Error::new_spanned(
-                ty,
-                "defining without bitfield is not yet supported for this type of register",
-            ))?,
-
-            None => outputs.extend(quote_spanned!(span =>
-                impl ::kernel::io::register::FixedRegister for #name {
-                    type Base = #base;
-                    const OFFSET: usize = #offset;
+            None => {
+                if unique {
+                    outputs.extend(quote!(
+                        impl ::kernel::io::register::FixedIoLoc<#base> for #ty {
+                            type Location = ::kernel::io::register::OffsetLoc<#base, #ty>;
+                            const LOCATION: Self::Location = #name;
+                        }
+                    ))
                 }
 
-                #(#attrs)* #vis const #name: ::kernel::io::register::FixedRegisterLoc<#name> =
-                    ::kernel::io::register::FixedRegisterLoc::<#name>::new();
-            )),
+                outputs.extend(quote_spanned!(span =>
+                    #(#attrs)* #vis const #name: ::kernel::io::register::OffsetLoc<#base, #ty> =
+                        ::kernel::io::register::OffsetLoc::new(#offset);
+                ));
+            }
 
             Some(def) => {
+                if !unique {
+                    Err(Error::new_spanned(
+                        &ty,
+                        "defining without bitfield is not yet supported for this type of register",
+                    ))?
+                }
+
                 let size = &def.size;
                 let stride = if let Some(stride) = &def.stride {
                     outputs.extend(quote_spanned!(stride.span().resolved_at(span) =>

-- 
2.54.0