[PATCH v6 5/9] rust: io: add IoRef and IoWrite types

[Alexandre Courbot <acourbot@nvidia.com>]

I/O accesses are defined by the following properties:

- For reads, a start address, a width, and a type to interpret the read
  value as,
- For writes, the same as above, and a value to write.

Introduce the `IoRef` trait, which allows implementing types to specify
the address a type expects to be accessed at, as well as the width of
the access, and the user-facing type used to perform the access.

This allows read operations to be made generic with the `read` method
over an `IoRef` argument.

Write operations need a value to write on top of the `IoRef`: fulfill
that purpose with the `IoWrite`, which is the combination of an `IoRef`
and a value of the type it expects. This allows write operations to be
made generic with the `write` method over a single `IoWrite` argument.

The main purpose of these new entities is to allow register types to be
written using these generic `read` and `write` methods of `Io`.

Co-developed-by: Gary Guo <gary@garyguo.net>
Signed-off-by: Gary Guo <gary@garyguo.net>
Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
 rust/kernel/io.rs | 243 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 243 insertions(+)

diff --git a/rust/kernel/io.rs b/rust/kernel/io.rs
index b150743ffa4f..6da8593f7858 100644
--- a/rust/kernel/io.rs
+++ b/rust/kernel/io.rs
@@ -173,6 +173,160 @@ pub trait IoCapable<T> {
     unsafe fn io_write(&self, value: T, address: usize);
 }
 
+/// Reference to an I/O location, describing the offset, width, and return type of an access.
+///
+/// This trait is the key abstraction allowing [`Io::read`], [`Io::write`], and [`Io::update`]
+/// to work uniformly with both raw `usize` offsets (for primitive types like `u32`) and typed
+/// references.
+///
+/// An `IoRef<T>` carries three pieces of information:
+///
+/// - The offset to access (returned by [`IoRef::offset`]),
+/// - The width of the access (determined by [`IoRef::IoType`]),
+/// - The type `T` in which data is returned or provided.
+///
+/// `T` and `IoType` may differ: for instance, a typed register has `T` = the register type with
+/// its bitfields, and `IoType` = its backing primitive (e.g. `u32`), with `From`/`Into`
+/// conversions between them.
+///
+/// An `IoRef` can be passed directly to [`Io::read`] or [`Io::try_read`] to obtain a value, or
+/// turned into an [`IoWrite`] via [`IoRef::set`] to be passed to [`Io::write`] or
+/// [`Io::try_write`].
+pub trait IoRef<T>: Copy
+where
+    T: From<Self::IoType> + Into<Self::IoType>,
+{
+    /// Size (`u8`, `u16`, etc) of the I/O performed on this reference.
+    type IoType;
+
+    /// Returns the offset of this reference.
+    fn offset(self) -> usize;
+
+    /// Turns this reference into an [`IoWrite`] with the initial `value`.
+    fn set(self, value: T) -> IoWrite<T, Self> {
+        IoWrite {
+            value,
+            reference: self,
+        }
+    }
+
+    /// Turns this reference into an [`IoWrite`] with the initial value `0`.
+    fn zeroed(self) -> IoWrite<T, Self>
+    where
+        T: Zeroable,
+    {
+        self.set(pin_init::zeroed())
+    }
+
+    /// Turns this reference into an [`IoWrite`] with the initial [`Default`] value of `T`.
+    fn default(self) -> IoWrite<T, Self>
+    where
+        T: Default,
+    {
+        self.set(Default::default())
+    }
+
+    /// Turns this reference into an [`IoWrite`] initialized from `0` and transformed by `f`.
+    ///
+    /// This is a shortcut for `self.zeroed().update(f)`.
+    fn init<F>(self, f: F) -> IoWrite<T, Self>
+    where
+        T: Zeroable,
+        F: FnOnce(T) -> T,
+    {
+        self.zeroed().update(f)
+    }
+
+    /// Turns this reference into an [`IoWrite`] initialized from `0` and transformed by `f`.
+    ///
+    /// `f` is expected to return a [`Result<T>`].
+    ///
+    /// This is a shortcut for `self.zeroed().try_update(f)`.
+    fn try_init<F, E>(self, f: F) -> Result<IoWrite<T, Self>, E>
+    where
+        T: Zeroable,
+        F: FnOnce(T) -> Result<T, E>,
+    {
+        self.zeroed().try_update(f)
+    }
+
+    /// Turns this reference into an [`IoWrite`] initialized from [`Default`] and transformed
+    /// by `f`.
+    ///
+    /// This is a shortcut for `self.default().update(f)`.
+    fn init_default<F>(self, f: F) -> IoWrite<T, Self>
+    where
+        T: Default,
+        F: FnOnce(T) -> T,
+    {
+        self.default().update(f)
+    }
+
+    /// Turns this reference into an [`IoWrite`] initialized from [`Default`] and transformed by
+    /// `f`.
+    ///
+    /// `f` is expected to return a [`Result<T>`].
+    ///
+    /// This is a shortcut for `self.default().try_update(f)`.
+    fn try_init_default<F, E>(self, f: F) -> Result<IoWrite<T, Self>, E>
+    where
+        T: Default,
+        F: FnOnce(T) -> Result<T, E>,
+    {
+        self.default().try_update(f)
+    }
+}
+
+/// A pending I/O write operation, bundling a value with the [`IoRef`] it should be written to.
+///
+/// Created by [`IoRef::set`], [`IoRef::zeroed`], [`IoRef::default`], [`IoRef::init`], or
+/// [`IoRef::init_default`], and consumed by [`Io::write`] or [`Io::try_write`] to perform the
+/// actual write.
+///
+/// The value can be modified before writing using [`IoWrite::update`] or [`IoWrite::try_update`],
+/// enabling a builder pattern:
+///
+/// ```ignore
+/// io.write(REGISTER.init(|v| v.with_field(x)));
+/// ```
+pub struct IoWrite<T, R>
+where
+    R: IoRef<T>,
+    T: From<R::IoType> + Into<R::IoType>,
+{
+    value: T,
+    reference: R,
+}
+
+impl<R, T> IoWrite<T, R>
+where
+    R: IoRef<T>,
+    T: From<R::IoType> + Into<R::IoType>,
+{
+    /// Transforms the value to be written by applying `f`, returning the modified [`IoWrite`].
+    #[inline(always)]
+    pub fn update<F>(mut self, f: F) -> Self
+    where
+        F: FnOnce(T) -> T,
+    {
+        self.value = f(self.value);
+
+        self
+    }
+
+    /// Transforms the value to be written by applying `f`, returning the modified [`IoWrite`] on
+    /// success.
+    #[inline(always)]
+    pub fn try_update<F, E>(mut self, f: F) -> Result<Self, E>
+    where
+        F: FnOnce(T) -> Result<T, E>,
+    {
+        self.value = f(self.value)?;
+
+        Ok(self)
+    }
+}
+
 /// Types implementing this trait (e.g. MMIO BARs or PCI config regions)
 /// can perform I/O operations on regions of memory.
 ///
@@ -406,6 +560,95 @@ fn write64(&self, value: u64, offset: usize)
         // SAFETY: `address` has been validated by `io_addr_assert`.
         unsafe { self.io_write(value, address) }
     }
+
+    /// Generic fallible read with runtime bounds check.
+    #[inline(always)]
+    fn try_read<T, R>(&self, r: R) -> Result<T>
+    where
+        T: From<R::IoType> + Into<R::IoType>,
+        R: IoRef<T>,
+        Self: IoCapable<R::IoType>,
+    {
+        let address = self.io_addr::<R::IoType>(r.offset())?;
+
+        // SAFETY: `address` has been validated by `io_addr`.
+        Ok(T::from(unsafe { self.io_read(address) }))
+    }
+
+    /// Generic fallible write with runtime bounds check.
+    #[inline(always)]
+    fn try_write<T, R>(&self, access: IoWrite<T, R>) -> Result
+    where
+        T: From<R::IoType> + Into<R::IoType>,
+        R: IoRef<T>,
+        Self: IoCapable<R::IoType>,
+    {
+        let address = self.io_addr::<R::IoType>(access.reference.offset())?;
+
+        // SAFETY: `address` has been validated by `io_addr`.
+        unsafe { self.io_write(access.value.into(), address) };
+        Ok(())
+    }
+
+    /// Generic fallible update with runtime bounds check.
+    ///
+    /// Caution: this does not perform any synchronization. Race conditions can occur in case of
+    /// concurrent access.
+    #[inline(always)]
+    fn try_update<T, R, F>(&self, r: R, f: F) -> Result
+    where
+        T: From<R::IoType> + Into<R::IoType>,
+        R: IoRef<T>,
+        Self: IoCapable<R::IoType>,
+        F: FnOnce(T) -> T,
+    {
+        let v = self.try_read(r)?;
+        self.try_write(r.set(f(v)))
+    }
+
+    /// Generic infallible read with compile-time bounds check.
+    #[inline(always)]
+    fn read<T, R>(&self, r: R) -> T
+    where
+        T: From<R::IoType> + Into<R::IoType>,
+        R: IoRef<T>,
+        Self: IoKnownSize + IoCapable<R::IoType>,
+    {
+        let address = self.io_addr_assert::<R::IoType>(r.offset());
+
+        // SAFETY: `address` has been validated by `io_addr_assert`.
+        T::from(unsafe { self.io_read(address) })
+    }
+
+    /// Generic infallible write with compile-time bounds check.
+    #[inline(always)]
+    fn write<T, R>(&self, access: IoWrite<T, R>)
+    where
+        R: IoRef<T>,
+        T: From<R::IoType> + Into<R::IoType>,
+        Self: IoKnownSize + IoCapable<R::IoType>,
+    {
+        let address = self.io_addr_assert::<R::IoType>(access.reference.offset());
+
+        // SAFETY: `address` has been validated by `io_addr_assert`.
+        unsafe { self.io_write(access.value.into(), address) }
+    }
+
+    /// Generic infallible update with compile-time bounds check.
+    ///
+    /// Caution: this does not perform any synchronization. Race conditions can occur in case of
+    /// concurrent access.
+    #[inline(always)]
+    fn update<T, R, F>(&self, r: R, f: F)
+    where
+        T: From<R::IoType> + Into<R::IoType>,
+        R: IoRef<T>,
+        Self: IoKnownSize + IoCapable<R::IoType> + Sized,
+        F: FnOnce(T) -> T,
+    {
+        let v = self.read(r);
+        self.write(r.set(f(v)));
+    }
 }
 
 /// Trait for types with a known size at compile time.

-- 
2.53.0