rust: pci: support driver data for `DeviceId`

Add support for PCI drivers to link driver specific data to the device
`IdTable`.

This can easily be generalized by substituting the
`bindings::pci_device_id` type with a generic type.

Signed-off-by: Danilo Krummrich <dakr@kernel.org>

2 files changed, 92 insertions, 22 deletions

diff --git a/rust/kernel/pci.rs b/rust/kernel/pci.rs
index 22eddec25437..3c4ea5e2b4af 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -47,7 +47,9 @@ impl DeviceId {
     }
 
     /// Convert `DeviceId` to raw `bindings::pci_device_id`.
-    pub const fn to_rawid(&self) -> bindings::pci_device_id {
+    ///
+    /// `offset` is the offset of `ids[i]` to `id_infos[i]` within `IdArray`.
+    pub const fn to_rawid(&self, offset: usize) -> bindings::pci_device_id {
         let mut raw = Self::ZERO;
 
         raw.vendor = self.vendor;
@@ -56,6 +58,7 @@ impl DeviceId {
         raw.subdevice = self.subdevice;
         raw.class = self.class;
         raw.class_mask = self.class_mask;
+        raw.driver_data = offset as _;
 
         raw
     }
@@ -63,47 +66,85 @@ impl DeviceId {
 
 /// A zero-terminated PCI device ID array.
 #[repr(C)]
-pub struct IdArray<const N: usize> {
+pub struct IdArray<U, const N: usize> {
     ids: [bindings::pci_device_id; N],
     sentinel: bindings::pci_device_id,
+    id_infos: [Option<U>; N],
 }
 
-impl<const N: usize> IdArray<N> {
+impl<U, const N: usize> IdArray<U, N> {
     /// Creates a new instance of the ID array.
     ///
     /// The contents are derived from the given identifiers.
     #[doc(hidden)]
-    pub const fn new(ids: [bindings::pci_device_id; N]) -> Self {
+    pub const fn new(ids: [bindings::pci_device_id; N], infos: [Option<U>; N]) -> Self {
         Self {
             ids,
             sentinel: DeviceId::ZERO,
+            id_infos: infos,
         }
     }
 
     /// Returns an `IdTable` backed by `self`.
     ///
     /// This is used to essentially erase the array size.
-    pub const fn as_table(&self) -> IdTable<'_> {
+    pub const fn as_table(&self) -> IdTable<'_, U> {
         IdTable {
             first: &self.ids[0],
+            _p: PhantomData,
         }
     }
+
+    /// Returns the offset of `ids[i]` to `id_infos[i]` within `IdArray`.
+    #[doc(hidden)]
+    pub const fn get_offset(index: usize) -> usize {
+        let id_size = core::mem::size_of::<bindings::pci_device_id>();
+        let info_size = core::mem::size_of::<Option<U>>();
+
+        id_size * (N - index + 1) + info_size * index
+    }
 }
 
 /// A device ID table.
 ///
 /// The table is guaranteed to be zero-terminated.
 #[repr(C)]
-pub struct IdTable<'a> {
+pub struct IdTable<'a, U> {
     first: &'a bindings::pci_device_id,
+    _p: PhantomData<&'a U>,
 }
 
-impl AsRef<bindings::pci_device_id> for IdTable<'_> {
+impl<U> AsRef<bindings::pci_device_id> for IdTable<'_, U> {
     fn as_ref(&self) -> &bindings::pci_device_id {
         self.first
     }
 }
 
+/// Converts a comma-separated list of pairs into an array with the first element. That is, it
+/// discards the second element of the pair.
+///
+/// Additionally, it automatically introduces a type if the first element is warpped in curly
+/// braces, for example, if it's `{v: 10}`, it becomes `X { v: 10 }`; this is to avoid repeating
+/// the type.
+#[macro_export]
+macro_rules! first_item {
+    ($id_type:ty, $(({$($first:tt)*}, $second:expr)),* $(,)?) => {
+        {
+            type IdType = $id_type;
+            [$(IdType{$($first)*},)*]
+        }
+    };
+    ($id_type:ty, $(($first:expr, $second:expr)),* $(,)?) => { [$($first,)*] };
+}
+
+/// Converts a comma-separated list of pairs into an array with the second element. That is, it
+/// discards the first element of the pair.
+#[macro_export]
+macro_rules! second_item {
+    ($(({$($first:tt)*}, $second:expr)),* $(,)?) => { [$($second,)*] };
+    ($(($first:expr, $second:expr)),* $(,)?) => { [$($second,)*] };
+}
+
 /// Counts the number of parenthesis-delimited, comma-separated items.
 #[macro_export]
 macro_rules! count_paren_items {
@@ -115,30 +156,38 @@ macro_rules! count_paren_items {
 #[macro_export]
 #[doc(hidden)]
 macro_rules! define_pci_id_array {
-    ($($args:tt)*) => {{
-        const fn new<const N: usize>(ids: [$crate::pci::DeviceId; N]) -> $crate::pci::IdArray<N> {
+    ($id_type:ty, $($args:tt)*) => {{
+        const fn new<U, const N: usize>(
+            ids: [$crate::pci::DeviceId; N],
+            infos: [Option<U>; N]) -> $crate::pci::IdArray<U, N>
+        {
             let mut raw_ids = [$crate::pci::DeviceId::ZERO; N];
 
+            // The offset of `ids[i]` to `id_infos[i]` within `IdArray`.
+
             let mut i = 0usize;
             while i < N {
-                raw_ids[i] = ids[i].to_rawid();
+                let offset = $crate::pci::IdArray::<U, N>::get_offset(i);
+
+                raw_ids[i] = ids[i].to_rawid(offset);
                 i += 1;
             }
 
-            $crate::pci::IdArray::<N>::new(raw_ids)
+            $crate::pci::IdArray::<U, N>::new(raw_ids, infos)
         }
 
-        new([ $($args)* ])
+        new($crate::first_item!($id_type, $($args)*), $crate::second_item!($($args)*))
     }}
 }
 
 /// Define a const PCI device ID table.
 #[macro_export]
 macro_rules! define_pci_id_table {
-    ([ $($args:tt)* ]) => {
-        const ID_TABLE: $crate::pci::IdTable<'static> = {
-            const ARRAY: $crate::pci::IdArray<{ $crate::count_paren_items!($($args)*) }> =
-                $crate::define_pci_id_array!($($args)*);
+    ($id_type:ty, [ $($args:tt)* ]) => {
+        type IdInfo = $id_type;
+        const ID_TABLE: $crate::pci::IdTable<'static, $id_type> = {
+            const ARRAY: $crate::pci::IdArray<$id_type, { $crate::count_paren_items!($($args)*) }> =
+                $crate::define_pci_id_array!($id_type, $($args)*);
             ARRAY.as_table()
         };
     };
@@ -147,13 +196,16 @@ pub use define_pci_id_table;
 
 /// Drivers must implement this trait to register a PCI driver.
 pub trait Driver {
+    /// The type holding information about each device ID supported by the drivID.
+    type IdInfo: 'static;
+
     /// The table of device IDs supported by this driver.
-    const ID_TABLE: IdTable<'static>;
+    const ID_TABLE: IdTable<'static, Self::IdInfo>;
 
     /// PCI driver probe.
     ///
     /// Called when a PCI device is matched against a PCI driver.
-    fn probe(pdev: *mut bindings::pci_dev) -> Result;
+    fn probe(pdev: *mut bindings::pci_dev, id: Option<&Self::IdInfo>) -> Result;
 
     /// PCI driver remove.
     ///
@@ -170,10 +222,26 @@ where
 {
     extern "C" fn probe(
         pdev: *mut bindings::pci_dev,
-        _ent: *const bindings::pci_device_id,
+        id: *const bindings::pci_device_id,
     ) -> core::ffi::c_int {
+        // SAFETY: `id` is a pointer within the static table, so it's always valid.
+        let offset = unsafe { (*id).driver_data };
+
+        let info = {
+            // SAFETY: The offset comes from a previous call to `offset_from` in `IdArray::new`,
+            // which guarantees that the resulting pointer is within the table.
+            let ptr = unsafe {
+                id.cast::<u8>()
+                    .offset(offset as _)
+                    .cast::<Option<T::IdInfo>>()
+            };
+
+            // SAFETY: Guaranteed by the preceding safety requirement.
+            unsafe { (*ptr).as_ref() }
+        };
+
         from_result(|| {
-            T::probe(pdev)?;
+            T::probe(pdev, info)?;
             Ok(0)
         })
     }
diff --git a/samples/rust/rust_pci_driver/driver.rs b/samples/rust/rust_pci_driver/driver.rs
index 64a3189d8689..7bd5c36228ae 100644
--- a/samples/rust/rust_pci_driver/driver.rs
+++ b/samples/rust/rust_pci_driver/driver.rs
@@ -10,12 +10,14 @@ pub(crate) struct Driver;
 
 impl pci::Driver for Driver {
     define_pci_id_table! {
+        (),
         [ (pci::DeviceId::new(bindings::PCI_VENDOR_ID_REDHAT,
-                             PCI_DEVICE_ID_REDHAT_QEMU_PCI_TESTDEV)) ]
+                             PCI_DEVICE_ID_REDHAT_QEMU_PCI_TESTDEV), None) ]
     }
 
-    fn probe(_pdev: *mut bindings::pci_dev) -> Result {
+    fn probe(_pdev: *mut bindings::pci_dev, id: Option<&Self::IdInfo>) -> Result {
         pr_info!("Probe Rust PCI driver sample.\n");
+        pr_info!("Info: {:?}\n", id);
 
         Ok(())
     }