[PATCH v3 10/16] rust: io: register: make register have a typed base

Previously `register!` defined registers can be used on any untyped I/O
regions. With all users specifying their desired register type now,
propagate the specified type and restrict I/O access only when type
matches.

Also, add an `io_project!` example which is enabled by this change.

Signed-off-by: Gary Guo <[email protected]>
---
 rust/kernel/io.rs          | 13 +++++++++++
 rust/kernel/io/register.rs | 55 ++++++++++++++++++++++++++++++++--------------
 rust/macros/io/register.rs | 34 ++++++++++++++--------------
 3 files changed, 68 insertions(+), 34 deletions(-)

diff --git a/rust/kernel/io.rs b/rust/kernel/io.rs
index 84dd876b3407..4542187d6b91 100644
--- a/rust/kernel/io.rs
+++ b/rust/kernel/io.rs
@@ -1692,21 +1692,34 @@ pub fn project_loc<U, L>(self, location: L) -> <T::Backend as IoBackend>::View<'
 /// The syntax is of form `io_project!(io, proj)` where `io` is an expression to a type that
 /// implements [`Io`] and `proj` is a [projection specification](kernel::ptr::project!).
 ///
+/// `io_project!` can also project to subview of registers defined with [`register!`] macro.
+/// Register projection has syntax `io_project!(io, try: REGISTER)` for fallible projection and
+/// `io_project!(io, build: REGISTER)` for infallible projection.
+///
 /// # Examples
 ///
 /// ```
 /// use kernel::io::{
 ///     io_project,
+///     register,
 ///     Mmio,
 /// };
 /// #[repr(C)]
 /// struct MyStruct { field: u32, }
 ///
+/// register! {
+///     base: MyStruct;
+///     FIELD(u32) @ 0 {
+///         31:0 val;
+///     }
+/// }
+///
 /// # fn test(mmio: Mmio<'_, [MyStruct]>) -> Result {
 /// // let mmio: Mmio<[MyStruct]>;
 /// let field: Mmio<'_, u32> = io_project!(mmio, [try: 1].field);
 /// let whole: Mmio<'_, MyStruct> = io_project!(mmio, [try: 2]);
 /// let nested: Mmio<'_, u32> = io_project!(whole, .field);
+/// let reg: Mmio<'_, FIELD> = io_project!(whole, build: FIELD);
 /// # Ok::<(), Error>(()) }
 /// ```
 #[macro_export]
diff --git a/rust/kernel/io/register.rs b/rust/kernel/io/register.rs
index 49a5fe069e47..f622057f7346 100644
--- a/rust/kernel/io/register.rs
+++ b/rust/kernel/io/register.rs
@@ -121,10 +121,11 @@
     io::IoLoc, //
 };
 
-use super::Region;
-
 /// Trait implemented by all registers.
 pub trait Register: 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.
@@ -136,9 +137,9 @@ pub trait FixedRegister: Register {}
 
 /// Allows `()` to be used as the `location` parameter of [`Io::write`](super::Io::write) when
 /// passing a [`FixedRegister`] value.
-impl<const SIZE: usize, T> IoLoc<Region<SIZE>, T> for ()
+impl<Base: ?Sized, T> IoLoc<Base, T> for ()
 where
-    T: FixedRegister,
+    T: FixedRegister<Base = Base>,
 {
     #[inline(always)]
     fn offset(self) -> usize {
@@ -148,9 +149,9 @@ fn offset(self) -> usize {
 
 /// A [`FixedRegister`] carries its location in its type. Thus `FixedRegister` values can be used
 /// as an [`IoLoc`].
-impl<const SIZE: usize, T> IoLoc<Region<SIZE>, T> for T
+impl<Base: ?Sized, T> IoLoc<Base, T> for T
 where
-    T: FixedRegister,
+    T: FixedRegister<Base = Base>,
 {
     #[inline(always)]
     fn offset(self) -> usize {
@@ -171,9 +172,9 @@ pub const fn new() -> Self {
     }
 }
 
-impl<const SIZE: usize, T> IoLoc<Region<SIZE>, T> for FixedRegisterLoc<T>
+impl<Base: ?Sized, T> IoLoc<Base, T> for FixedRegisterLoc<T>
 where
-    T: FixedRegister,
+    T: FixedRegister<Base = Base>,
 {
     #[inline(always)]
     fn offset(self) -> usize {
@@ -240,9 +241,9 @@ const fn offset(self) -> usize {
     }
 }
 
-impl<const SIZE: usize, T, B> IoLoc<Region<SIZE>, T> for RelativeRegisterLoc<T, B>
+impl<SuperBase: ?Sized, T, B> IoLoc<SuperBase, T> for RelativeRegisterLoc<T, B>
 where
-    T: RelativeRegister,
+    T: RelativeRegister<Base = SuperBase>,
     B: RegisterBase<T::BaseFamily> + ?Sized,
 {
     #[inline(always)]
@@ -282,9 +283,9 @@ pub fn try_new(idx: usize) -> Option<Self> {
     }
 }
 
-impl<const SIZE: usize, T> IoLoc<Region<SIZE>, T> for RegisterArrayLoc<T>
+impl<Base: ?Sized, T> IoLoc<Base, T> for RegisterArrayLoc<T>
 where
-    T: RegisterArray,
+    T: RegisterArray<Base = Base>,
 {
     #[inline(always)]
     fn offset(self) -> usize {
@@ -367,9 +368,9 @@ pub fn try_at(self, idx: usize) -> Option<RelativeRegisterArrayLoc<T, B>> {
     }
 }
 
-impl<const SIZE: usize, T, B> IoLoc<Region<SIZE>, T> for RelativeRegisterArrayLoc<T, B>
+impl<SuperBase: ?Sized, T, B> IoLoc<SuperBase, T> for RelativeRegisterArrayLoc<T, B>
 where
-    T: RelativeRegisterArray,
+    T: RelativeRegisterArray<Base = SuperBase>,
     B: RegisterBase<T::BaseFamily> + ?Sized,
 {
     #[inline(always)]
@@ -393,9 +394,9 @@ pub trait LocatedRegister<Base: ?Sized> {
     fn into_io_op(self) -> (Self::Location, Self::Value);
 }
 
-impl<const SIZE: usize, T> LocatedRegister<Region<SIZE>> for T
+impl<Base: ?Sized, T> LocatedRegister<Base> for T
 where
-    T: FixedRegister,
+    T: FixedRegister<Base = Base>,
 {
     type Location = FixedRegisterLoc<Self::Value>;
     type Value = T;
@@ -406,6 +407,28 @@ fn into_io_op(self) -> (FixedRegisterLoc<T>, T) {
     }
 }
 
+/// 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: Register<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
+/// eval.
+#[doc(hidden)]
+#[inline(always)]
+pub const fn element_alias_offset<Base: ?Sized, Alias: RegisterArray<Base = Base>>(
+    idx: usize,
+) -> usize {
+    build_assert!(idx < Alias::SIZE);
+    Alias::OFFSET + idx * Alias::STRIDE
+}
+
 /// Defines a dedicated type for a register, including getter and setter methods for its fields and
 /// methods to read and write it from an [`Io`](kernel::io::Io) region.
 ///
diff --git a/rust/macros/io/register.rs b/rust/macros/io/register.rs
index 2fb48e1be82d..cb02e850b23f 100644
--- a/rust/macros/io/register.rs
+++ b/rust/macros/io/register.rs
@@ -140,21 +140,23 @@ fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
 }
 
 pub(crate) struct RegDef {
-    base: Option<Type>,
+    base: Type,
     regs: Vec<Reg>,
 }
 
 impl Parse for RegDef {
     fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
-        let base = if input.peek(kw::base) {
-            let _: kw::base = input.parse()?;
-            let _: Token![:] = input.parse()?;
-            let base = input.parse()?;
-            let _: Token![;] = input.parse()?;
-            Some(base)
-        } else {
-            None
-        };
+        if !input.peek(kw::base) {
+            Err(input.error(
+                "a base type needs to be specified for `register!` invocation with `base: ty;`",
+            ))?;
+        }
+
+        let _: kw::base = input.parse()?;
+        let _: Token![:] = input.parse()?;
+        let base = input.parse()?;
+        let _: Token![;] = input.parse()?;
+
         let mut regs = Vec::new();
         while !input.is_empty() {
             regs.push(input.parse()?);
@@ -166,6 +168,7 @@ fn parse(input: syn::parse::ParseStream<'_>) -> Result<Self> {
 pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
     let mut outputs = TokenStream::new();
 
+    let base = &def.base;
     for reg in def.regs {
         let Reg {
             attrs,
@@ -186,18 +189,12 @@ pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
             RegOffset::Fixed { offset } => quote!(#offset),
             RegOffset::Alias { alias } => {
                 quote_spanned!(alias.span().resolved_at(span) =>
-                    <#alias as ::kernel::io::register::Register>::OFFSET
+                    ::kernel::io::register::alias_offset::<#base, #alias>()
                 )
             }
             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
+                    ::kernel::io::register::element_alias_offset::<#base, #alias>(#idx)
                 )
             }
         };
@@ -211,6 +208,7 @@ pub(crate) fn register(def: RegDef) -> Result<TokenStream> {
             );
 
             impl ::kernel::io::register::Register for #name {
+                type Base = #base;
                 const OFFSET: usize = #offset;
             }
         ));

-- 
2.54.0