[PATCH v4 01/16] rust: io: register: reimplement as proc macro

[Gary Guo <gary@garyguo.net>]

The existing `register!` macro is implemented as a declarative macro.
Reimplement it as proc macro instead, with no functional changes intended.

The old implementation produces unhelpful diagnostics when things go wrong.
For example, for code like

    register! {
        pub(crate) TESTREG(u32) {
            31:0    data;
        }
    }

which misses out the "@ offset" part of the specification, and the
following error is produced:

error: no rules expected `{`
   --> test.rs:42:5
    |
 42 | /     register! {
 43 | |         pub(crate) TESTREG(u32) {
 44 | |             31:0    data;
...   |
100 | |     }
    | |_____^ no rules expected this token in macro call

which isn't very helpful. With the proc macro implementation, the following
error is produced:

error: expected `@` or `=>`
  --> tests.rs:43:33
   |
43 |         pub(crate) TESTREG(u32) {
   |                                 ^

which is much more helpful. Apart from diagnostics, proc macro also has a
benefit of not having follow-set restrictions, which makes syntax like

    register!(name: ty @ offset);

possible; declarative macro will reject this as `@` is not in the
follow-set of "ty" metavariable kind.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 MAINTAINERS                |   1 +
 rust/kernel/io/register.rs | 224 +-----------------------------------
 rust/macros/io/mod.rs      |   3 +
 rust/macros/io/register.rs | 276 +++++++++++++++++++++++++++++++++++++++++++++
 rust/macros/lib.rs         |   9 ++
 5 files changed, 291 insertions(+), 222 deletions(-)

diff --git a/MAINTAINERS b/MAINTAINERS
index 3a19da74d00c..d92c7d834693 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -7494,6 +7494,7 @@ F:	rust/kernel/io.rs
 F:	rust/kernel/io/
 F:	rust/kernel/irq.rs
 F:	rust/kernel/irq/
+F:	rust/macros/io/
 
 DEVICE RESOURCE MANAGEMENT HELPERS
 M:	Hans de Goede <hansg@kernel.org>
diff --git a/rust/kernel/io/register.rs b/rust/kernel/io/register.rs
index 03dfd2ff48c7..6a19552ffb95 100644
--- a/rust/kernel/io/register.rs
+++ b/rust/kernel/io/register.rs
@@ -801,227 +801,7 @@ fn into_io_op(self) -> (FixedRegisterLoc<T>, T) {
 /// ```
 #[macro_export]
 macro_rules! register {
-    // Entry point for the macro, allowing multiple registers to be defined in one call.
-    // It matches all possible register declaration patterns to dispatch them to corresponding
-    // `@reg` rule that defines a single register.
-    //
-    // TODO: change `alias:ident` to `alias:path` once relative registers are replaced by I/O
-    // projections.
-    (
-        $(
-            $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty)
-                $([ $size:expr $(, stride = $stride:expr)? ])?
-                $(@ $($base:ident +)? $offset:literal)?
-                $(=> $alias:ident $(+ $alias_offset:ident)? $([$alias_idx:expr])? )?
-            { $($fields:tt)* }
-        )*
-    ) => {
-        $(
-        $crate::register!(
-            @reg $(#[$attr])* $vis $name ($storage) $([$size $(, stride = $stride)?])?
-                $(@ $($base +)? $offset)?
-                $(=> $alias $(+ $alias_offset)? $([$alias_idx])? )?
-            { $($fields)* }
-        );
-        )*
-    };
-
-    // All the rules below are private helpers.
-
-    // Creates a register at a fixed offset of the MMIO space.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) @ $offset:literal
-            { $($fields:tt)* }
-    ) => {
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(@io_base $name($storage) @ $offset);
-        $crate::register!(@io_fixed $(#[$attr])* $vis $name);
-    };
-
-    // Creates an alias register of fixed offset register `alias` with its own fields.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) => $alias:path
-            { $($fields:tt)* }
-    ) => {
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(
-            @io_base $name($storage) @
-            <$alias as $crate::io::register::Register>::OFFSET
-        );
-        $crate::register!(@io_fixed $(#[$attr])* $vis $name);
-    };
-
-    // Creates a register at a relative offset from a base address provider.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) @ $base:ident + $offset:literal
-            { $($fields:tt)* }
-    ) => {
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(@io_base $name($storage) @ $offset);
-        $crate::register!(@io_relative $name @ $base);
-    };
-
-    // Creates an alias register of relative offset register `alias` with its own fields.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) => $base:ident + $alias:ident
-            { $($fields:tt)* }
-    ) => {
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(
-            @io_base $name($storage) @ <$alias as $crate::io::register::Register>::OFFSET
-        );
-        $crate::register!(@io_relative $name @ $base);
-    };
-
-    // Creates an array of registers at a fixed offset of the MMIO space.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty)
-            [ $size:expr, stride = $stride:expr ] @ $offset:literal { $($fields:tt)* }
-    ) => {
-        $crate::build_assert::static_assert!(::core::mem::size_of::<$storage>() <= $stride);
-
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(@io_base $name($storage) @ $offset);
-        $crate::register!(@io_array $name [ $size, stride = $stride ]);
-    };
-
-    // Shortcut for contiguous array of registers (stride == size of element).
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) [ $size:expr ] @ $offset:literal
-            { $($fields:tt)* }
-    ) => {
-        $crate::register!(
-            @reg $(#[$attr])* $vis $name($storage)
-                [ $size, stride = ::core::mem::size_of::<$storage>() ]
-                @ $offset { $($fields)* }
-        );
-    };
-
-    // Creates an alias of register `idx` of array of registers `alias` with its own fields.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) => $alias:path [ $idx:expr ]
-            { $($fields:tt)* }
-    ) => {
-        $crate::build_assert::static_assert!(
-            $idx < <$alias as $crate::io::register::RegisterArray>::SIZE
-        );
-
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(
-            @io_base $name($storage) @
-            <$alias as $crate::io::register::Register>::OFFSET
-                + $idx * <$alias as $crate::io::register::RegisterArray>::STRIDE
-        );
-        $crate::register!(@io_fixed $(#[$attr])* $vis $name);
-    };
-
-    // Creates an array of registers at a relative offset from a base address provider.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty)
-            [ $size:expr, stride = $stride:expr ]
-            @ $base:ident + $offset:literal { $($fields:tt)* }
-    ) => {
-        $crate::build_assert::static_assert!(::core::mem::size_of::<$storage>() <= $stride);
-
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(@io_base $name($storage) @ $offset);
-        $crate::register!(@io_relative_array $name [ $size, stride = $stride ] @ $base);
-    };
-
-    // Shortcut for contiguous array of relative registers (stride == size of element).
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty) [ $size:expr ]
-            @ $base:ident + $offset:literal { $($fields:tt)* }
-    ) => {
-        $crate::register!(
-            @reg $(#[$attr])* $vis $name($storage)
-                [ $size, stride = ::core::mem::size_of::<$storage>() ]
-                @ $base + $offset { $($fields)* }
-        );
-    };
-
-    // Creates an alias of register `idx` of relative array of registers `alias` with its own
-    // fields.
-    (
-        @reg $(#[$attr:meta])* $vis:vis $name:ident ($storage:ty)
-            => $base:ident + $alias:ident [ $idx:expr ] { $($fields:tt)* }
-    ) => {
-        $crate::build_assert::static_assert!(
-            $idx < <$alias as $crate::io::register::RegisterArray>::SIZE
-        );
-
-        $crate::register!(@bitfield $(#[$attr])* $vis struct $name($storage) { $($fields)* });
-        $crate::register!(
-            @io_base $name($storage) @
-                <$alias as $crate::io::register::Register>::OFFSET +
-                $idx * <$alias as $crate::io::register::RegisterArray>::STRIDE
-        );
-        $crate::register!(@io_relative $name @ $base);
-    };
-
-    // Generates the bitfield for the register.
-    //
-    // `#[allow(non_camel_case_types)]` is added since register names typically use
-    // `SCREAMING_CASE`.
-    (
-        @bitfield $(#[$attr:meta])* $vis:vis struct $name:ident($storage:ty) { $($fields:tt)* }
-    ) => {
-        $crate::bitfield!(
-            #[allow(non_camel_case_types)]
-            $(#[$attr])* $vis struct $name($storage) { $($fields)* }
-        );
-    };
-
-    // Implementations shared by all registers types.
-    (@io_base $name:ident($storage:ty) @ $offset:expr) => {
-        impl $crate::io::register::Register for $name {
-            type Storage = $storage;
-
-            const OFFSET: usize = $offset;
-        }
-    };
-
-    // Implementations of fixed registers.
-    (@io_fixed $(#[$attr:meta])* $vis:vis $name:ident) => {
-        impl $crate::io::register::FixedRegister for $name {}
-
-        $(#[$attr])*
-        $vis const $name: $crate::io::register::FixedRegisterLoc<$name> =
-            $crate::io::register::FixedRegisterLoc::<$name>::new();
-    };
-
-    // Implementations of relative registers.
-    (@io_relative $name:ident @ $base:ident) => {
-        impl $crate::io::register::WithBase for $name {
-            type BaseFamily = $base;
-        }
-
-        impl $crate::io::register::RelativeRegister for $name {}
-    };
-
-    // Implementations of register arrays.
-    (@io_array $name:ident [ $size:expr, stride = $stride:expr ]) => {
-        impl $crate::io::register::Array for $name {}
-
-        impl $crate::io::register::RegisterArray for $name {
-            const SIZE: usize = $size;
-            const STRIDE: usize = $stride;
-        }
-    };
-
-    // Implementations of relative array registers.
-    (
-        @io_relative_array $name:ident [ $size:expr, stride = $stride:expr ] @ $base:ident
-    ) => {
-        impl $crate::io::register::WithBase for $name {
-            type BaseFamily = $base;
-        }
-
-        impl $crate::io::register::RegisterArray for $name {
-            const SIZE: usize = $size;
-            const STRIDE: usize = $stride;
-        }
-
-        impl $crate::io::register::RelativeRegisterArray for $name {}
+    ($($tt:tt)*) => {
+        $crate::macros::register!($($tt)*);
     };
 }
diff --git a/rust/macros/io/mod.rs b/rust/macros/io/mod.rs
new file mode 100644
index 000000000000..87f7742f4619
--- /dev/null
+++ b/rust/macros/io/mod.rs
@@ -0,0 +1,3 @@
+// SPDX-License-Identifier: GPL-2.0
+
+pub(crate) mod register;
diff --git a/rust/macros/io/register.rs b/rust/macros/io/register.rs
new file mode 100644
index 000000000000..83ba4dd75ab3
--- /dev/null
+++ b/rust/macros/io/register.rs
@@ -0,0 +1,276 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Documention and usage example of the macro can be found at `rust/kernel/io/register.rs`.
+
+use proc_macro2::{
+    Group,
+    Literal,
+    Span,
+    TokenStream, //
+};
+use quote::{
+    quote,
+    quote_spanned, //
+};
+use syn::{
+    bracketed,
+    parenthesized,
+    parse::Parse,
+    spanned::Spanned,
+    token,
+    Attribute,
+    Expr,
+    Ident,
+    Path,
+    Result,
+    Token,
+    Type,
+    Visibility, //
+};
+
+mod kw {
+    syn::custom_keyword!(stride);
+}
+
+/// Definition of a register array.
+///
+/// Specify a size, and optionally a stride. Syntax is of form `[size: EXPR, $(stride = EXPR)?]`.
+struct RegArrayDef {
+    size: Expr,
+    stride: Option<Expr>,
+}
+
+/// Offset of a register.
+///
+/// Can be either of form
+/// * `@ offset` for fixed offset
+/// * `=> alias` for alias of register `alias`.
+/// * `=> alias[idx]` for alias of register array `alias`'s `idx`-th element.
+enum RegOffset {
+    /// Register is located at fixed address.
+    Fixed { offset: Literal },
+    /// Register is an alias of a fixed register.
+    Alias { alias: Path },
+    /// Register is an alias of an element of a register array.
+    ElementAlias { alias: Path, idx: Expr },
+}
+
+/// Definition of a single register.
+struct Reg {
+    attrs: Vec<Attribute>,
+    vis: Visibility,
+    name: Ident,
+    storage: Type,
+    array: Option<RegArrayDef>,
+    relative_base: Option<Path>,
+    offset: RegOffset,
+    bitfield_args: Group,
+}
+
+impl Parse for Reg {
+    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
+        let attrs = input.call(Attribute::parse_outer)?;
+        let vis = input.parse()?;
+        let name = input.parse()?;
+        let storage = {
+            let content;
+            parenthesized!(content in input);
+            content.parse()?
+        };
+        let array = if input.peek(token::Bracket) {
+            let content;
+            bracketed!(content in input);
+            let size = content.parse()?;
+            let stride = if content.peek(Token![,]) {
+                let _: Token![,] = content.parse()?;
+                let _: kw::stride = content.parse()?;
+                let _: Token![=] = content.parse()?;
+                Some(content.parse()?)
+            } else {
+                None
+            };
+            Some(RegArrayDef { size, stride })
+        } else {
+            None
+        };
+
+        // Parse offset and the base it's relative to.
+        let lh = input.lookahead1();
+        let mut relative_base = None;
+        let offset = if lh.peek(Token![@]) {
+            let _: Token![@] = input.parse()?;
+
+            if input.peek(Ident) {
+                relative_base = Some(input.parse()?);
+                let _: Token![+] = input.parse()?;
+            }
+
+            RegOffset::Fixed {
+                offset: input.parse()?,
+            }
+        } else if lh.peek(Token![=>]) {
+            let _: Token![=>] = input.parse()?;
+            let mut alias: Path = input.parse()?;
+            if input.peek(Token![+]) {
+                let _: Token![+] = input.parse()?;
+                relative_base = Some(alias);
+                alias = input.parse()?;
+            }
+
+            if input.peek(token::Bracket) {
+                let content;
+                bracketed!(content in input);
+                RegOffset::ElementAlias {
+                    alias,
+                    idx: content.parse()?,
+                }
+            } else {
+                RegOffset::Alias { alias }
+            }
+        } else {
+            Err(lh.error())?
+        };
+
+        let lh = input.lookahead1();
+        let bitfield_args = if lh.peek(token::Brace) {
+            input.parse()?
+        } else {
+            Err(lh.error())?
+        };
+
+        Ok(Self {
+            attrs,
+            vis,
+            name,
+            storage,
+            array,
+            relative_base,
+            offset,
+            bitfield_args,
+        })
+    }
+}
+
+pub(crate) struct RegDef {
+    regs: Vec<Reg>,
+}
+
+impl Parse for RegDef {
+    fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
+        let mut regs = Vec::new();
+        while !input.is_empty() {
+            regs.push(input.parse()?);
+        }
+        Ok(RegDef { regs })
+    }
+}
+
+pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
+    let mut outputs = TokenStream::new();
+
+    for reg in def.regs {
+        let Reg {
+            attrs,
+            vis,
+            name,
+            storage,
+            array,
+            relative_base,
+            offset,
+            bitfield_args,
+        } = reg;
+
+        // Use register name's span for generated code, so error messages (if any) can point to it
+        // instead of the entire register allocation.
+        let span = name.span().resolved_at(Span::mixed_site());
+
+        let offset = match offset {
+            RegOffset::Fixed { offset } => quote!(#offset),
+            RegOffset::Alias { alias } => {
+                quote_spanned!(alias.span().resolved_at(span) =>
+                    <#alias as ::kernel::io::register::Register>::OFFSET
+                )
+            }
+            RegOffset::ElementAlias { alias, idx } => {
+                outputs.extend(quote_spanned!(idx.span().resolved_at(span) =>
+                    ::kernel::build_assert::static_assert!(
+                        (#idx) < <#alias as ::kernel::io::register::RegisterArray>::SIZE
+                    );
+                ));
+                quote_spanned!(alias.span().resolved_at(span) =>
+                    <#alias as ::kernel::io::register::Register>::OFFSET
+                        + (#idx) * <#alias as ::kernel::io::register::RegisterArray>::STRIDE
+                )
+            }
+        };
+
+        outputs.extend(quote_spanned!(span =>
+            ::kernel::bitfield!(
+                // `#[allow(non_camel_case_types)]` is added since register names typically use
+                // `SCREAMING_CASE`.
+                #[allow(non_camel_case_types)]
+                #(#attrs)* #vis struct #name(#storage) #bitfield_args
+            );
+
+            impl ::kernel::io::register::Register for #name {
+                type Storage = #storage;
+
+                const OFFSET: usize = #offset;
+            }
+        ));
+
+        match array {
+            None => match relative_base {
+                None => outputs.extend(quote_spanned!(span =>
+                    impl ::kernel::io::register::FixedRegister for #name {}
+
+                    #(#attrs)* #vis const #name: ::kernel::io::register::FixedRegisterLoc<#name> =
+                        ::kernel::io::register::FixedRegisterLoc::<#name>::new();
+                )),
+                Some(relative_base) => outputs.extend(quote_spanned!(span =>
+                    impl ::kernel::io::register::WithBase for #name {
+                        type BaseFamily = #relative_base;
+                    }
+
+                    impl ::kernel::io::register::RelativeRegister for #name {}
+                )),
+            },
+
+            Some(def) => {
+                let size = &def.size;
+                let stride = if let Some(stride) = &def.stride {
+                    outputs.extend(quote_spanned!(stride.span().resolved_at(span) =>
+                        ::kernel::build_assert::static_assert!(
+                            ::core::mem::size_of::<#storage>() <= #stride
+                        );
+                    ));
+                    quote!(#stride)
+                } else {
+                    quote_spanned!(span => ::core::mem::size_of::<#storage>())
+                };
+
+                outputs.extend(quote_spanned!(span =>
+                    impl ::kernel::io::register::RegisterArray for #name {
+                        const SIZE: usize = #size;
+                        const STRIDE: usize = #stride;
+                    }
+                ));
+
+                match relative_base {
+                    None => outputs.extend(quote_spanned!(span =>
+                        impl ::kernel::io::register::Array for #name {}
+                    )),
+                    Some(relative_base) => outputs.extend(quote_spanned!(span =>
+                        impl ::kernel::io::register::WithBase for #name {
+                            type BaseFamily = #relative_base;
+                        }
+
+                        impl ::kernel::io::register::RelativeRegisterArray for #name {}
+                    )),
+                }
+            }
+        };
+    }
+
+    Ok(outputs)
+}
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
index 24f96feaeb34..9b76efe1476f 100644
--- a/rust/macros/lib.rs
+++ b/rust/macros/lib.rs
@@ -19,6 +19,7 @@
 mod fmt;
 mod for_lt;
 mod helpers;
+mod io;
 mod kunit;
 mod module;
 mod paste;
@@ -481,6 +482,14 @@ pub fn paste(input: TokenStream) -> TokenStream {
         .into()
 }
 
+#[doc(hidden)] // Documented in `kernel` crate.
+#[proc_macro]
+pub fn register(input: TokenStream) -> TokenStream {
+    io::register::register(parse_macro_input!(input))
+        .unwrap_or_else(|e| e.into_compile_error())
+        .into()
+}
+
 /// Registers a KUnit test suite and its test cases using a user-space like syntax.
 ///
 /// This macro should be used on modules. If `CONFIG_KUNIT` (in `.config`) is `n`, the target module

-- 
2.54.0