[RFC 2/2] pci: Add PCI capability infrastructure and SR-IOV capability support

[Zhi Wang <zhiw@nvidia.com>]

Rust drivers may need to access PCI capabilities (such as SR-IOV) to
configure firmware metadata.

Add a generic Capability<S, K> structure that:
- Wraps capability discovery via kernel's pci_find_ext_capability()
- Dynamically calculates capability size by reading next capability
  pointers (using 0xffc mask per PCI_EXT_CAP_NEXT for extended caps)
- Implements fallible I/O via the Io trait with runtime bounds checking

Add SR-IOV (Single Root I/O Virtualization) capability support:
- SriovCapability wraps Capability<Extended, Extended>
- read_vf_offset() reads PCI_SRIOV_VF_OFFSET register
- read_vf_bar(n) reads PCI_SRIOV_BAR + n*4 (32-bit BAR)
- read_vf_bar64(n) reads 64-bit BAR by combining two 32-bit reads

Signed-off-by: Zhi Wang <zhiw@nvidia.com>
---
 rust/kernel/pci.rs     |   9 ++
 rust/kernel/pci/cap.rs | 274 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 283 insertions(+)
 create mode 100644 rust/kernel/pci/cap.rs

diff --git a/rust/kernel/pci.rs b/rust/kernel/pci.rs
index cd46ac12812c..2deb26fb3775 100644
--- a/rust/kernel/pci.rs
+++ b/rust/kernel/pci.rs
@@ -31,6 +31,7 @@
     },
 };
 
+mod cap;
 mod id;
 mod io;
 mod irq;
@@ -42,6 +43,7 @@
 };
 pub use self::io::{
     Bar,
+    ConfigSpace,
     ConfigSpaceKind,
     ConfigSpaceSize,
     Extended,
@@ -52,6 +54,13 @@
     IrqTypes,
     IrqVector, //
 };
+pub use self::cap::{
+    Capability,
+    CapabilityId,
+    CapabilityKind,
+    ExtCapabilityId,
+    SriovCapability, //
+};
 
 /// An adapter for the registration of PCI drivers.
 pub struct Adapter<T: Driver>(T);
diff --git a/rust/kernel/pci/cap.rs b/rust/kernel/pci/cap.rs
new file mode 100644
index 000000000000..e3a903b79299
--- /dev/null
+++ b/rust/kernel/pci/cap.rs
@@ -0,0 +1,274 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! PCI Capability support.
+//!
+//! This module provides abstractions for discovering and accessing PCI capabilities.
+
+use super::{
+    ConfigSpace,
+    ConfigSpaceKind,
+    Extended,
+    Normal, //
+};
+use crate::{
+    bindings,
+    io::{
+        define_read,
+        define_write,
+        Io,
+        IoCapable, //
+    },
+    prelude::*,
+};
+
+/// Internal helper macro to call ConfigSpace fallible methods from Capability.
+macro_rules! call_cap_read {
+    (fallible, $method:ident, $self:ident, $ty:ty, $addr:expr) => {
+        $self.config_space.$method($self.offset() + $addr)
+    };
+}
+
+/// Internal helper macro to call ConfigSpace fallible write methods from Capability.
+macro_rules! call_cap_write {
+    (fallible, $method:ident, $self:ident, $ty:ty, $addr:expr, $value:expr) => {
+        $self.config_space.$method($value, $self.offset() + $addr)
+    };
+}
+
+/// PCI Capability IDs for normal capabilities (in 256-byte config space).
+///
+/// These are not currently implemented, but the enum is provided for API completeness.
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CapabilityId {
+    /// Power Management
+    PM = bindings::PCI_CAP_ID_PM as u8,
+    /// Message Signalled Interrupts
+    MSI = bindings::PCI_CAP_ID_MSI as u8,
+    /// MSI-X
+    MSIX = bindings::PCI_CAP_ID_MSIX as u8,
+    /// PCI Express
+    Express = bindings::PCI_CAP_ID_EXP as u8,
+    /// Vendor Specific
+    VendorSpecific = bindings::PCI_CAP_ID_VNDR as u8,
+}
+
+/// PCI Extended Capability IDs (in 4096-byte config space).
+///
+/// Currently only SR-IOV is implemented.
+#[repr(u16)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ExtCapabilityId {
+    /// Single Root I/O Virtualization
+    SRIOV = bindings::PCI_EXT_CAP_ID_SRIOV as u16,
+}
+
+/// Trait for capability kinds (Normal or Extended).
+pub trait CapabilityKind {
+    /// The capability ID type for this kind.
+    type IdType: Copy + PartialEq;
+
+    /// Start offset for capability scanning.
+    const START_OFFSET: usize;
+}
+
+/// Marker for normal (legacy) PCI capabilities.
+impl CapabilityKind for Normal {
+    type IdType = u8;
+    const START_OFFSET: usize = bindings::PCI_CAPABILITY_LIST as usize;
+}
+
+/// Marker for extended PCI capabilities.
+impl CapabilityKind for Extended {
+    type IdType = u16;
+    const START_OFFSET: usize = bindings::PCI_CFG_SPACE_SIZE as usize;
+}
+
+/// A PCI capability.
+///
+/// This type represents a discovered PCI capability and provides safe access
+/// to its registers. All I/O operations are relative to the capability's
+/// base offset in configuration space.
+pub struct Capability<'a, S: ConfigSpaceKind, K: CapabilityKind> {
+    config_space: &'a ConfigSpace<'a, S>,
+    offset: usize,
+    id: K::IdType,
+    size: usize,
+}
+
+impl<'a, S: ConfigSpaceKind, K: CapabilityKind> Capability<'a, S, K> {
+    /// Creates a new capability handle.
+    fn new(
+        config_space: &'a ConfigSpace<'a, S>,
+        offset: usize,
+        id: K::IdType,
+        size: usize,
+    ) -> Self {
+        Self {
+            config_space,
+            offset,
+            id,
+            size,
+        }
+    }
+
+    /// Returns the offset of this capability in configuration space.
+    pub fn offset(&self) -> usize {
+        self.offset
+    }
+
+    /// Returns the capability ID.
+    pub fn id(&self) -> K::IdType {
+        self.id
+    }
+
+    /// Returns the size of this capability in bytes.
+    pub fn size(&self) -> usize {
+        self.size
+    }
+}
+
+// Implement IoCapable for Capability
+impl<'a, S: ConfigSpaceKind, K: CapabilityKind> IoCapable<u8> for Capability<'a, S, K> {}
+impl<'a, S: ConfigSpaceKind, K: CapabilityKind> IoCapable<u16> for Capability<'a, S, K> {}
+impl<'a, S: ConfigSpaceKind, K: CapabilityKind> IoCapable<u32> for Capability<'a, S, K> {}
+
+// Implement Io trait for Capability using fallible methods (runtime checks)
+impl<'a, S: ConfigSpaceKind, K: CapabilityKind> Io for Capability<'a, S, K> {
+    const MIN_SIZE: usize = S::SIZE; // Use config space size for fallible bounds checking
+
+    #[inline]
+    fn addr(&self) -> usize {
+        0  // Offsets are relative to capability base, not absolute
+    }
+
+    #[inline]
+    fn maxsize(&self) -> usize {
+        self.size()
+    }
+
+    // Only implement fallible methods (no IoKnownSize, so infallible methods not available)
+    define_read!(fallible, try_read8, call_cap_read(try_read8) -> u8);
+    define_read!(fallible, try_read16, call_cap_read(try_read16) -> u16);
+    define_read!(fallible, try_read32, call_cap_read(try_read32) -> u32);
+
+    define_write!(fallible, try_write8, call_cap_write(try_write8) <- u8);
+    define_write!(fallible, try_write16, call_cap_write(try_write16) <- u16);
+    define_write!(fallible, try_write32, call_cap_write(try_write32) <- u32);
+}
+
+impl<'a> ConfigSpace<'a, Extended> {
+    /// Finds a specific extended capability by ID using the kernel's `pci_find_ext_capability`.
+    pub fn find_ext_capability(&self, id: u16) -> Option<Capability<'_, Extended, Extended>> {
+        // SAFETY: pdev is valid by ConfigSpace invariants
+        let offset = unsafe { bindings::pci_find_ext_capability(self.pdev.as_raw(), id as i32) };
+
+        if offset == 0 {
+            return None;
+        }
+
+        let size = self.calculate_ext_cap_size(offset as usize);
+        Some(Capability::new(self, offset as usize, id, size))
+    }
+
+    fn calculate_ext_cap_size(&self, offset: usize) -> usize {
+        // Extended capability header: [31:20] = next capability offset
+        // Use 0xffc mask (not 0xfff) to match kernel's PCI_EXT_CAP_NEXT macro
+        let header = self.try_read32(offset).unwrap_or(0);
+        let next_ptr = ((header >> 20) & 0xffc) as usize;
+
+        if next_ptr == 0 {
+            // Last capability, size goes to end of config space
+            self.pdev.cfg_size().into_raw() - offset
+        } else {
+            // Size is distance to next capability
+            next_ptr - offset
+        }
+    }
+
+    /// Finds the next occurrence of a specific extended capability starting from a given position.
+    pub fn find_next_ext_capability(
+        &self,
+        start_pos: u16,
+        id: u16,
+    ) -> Option<Capability<'_, Extended, Extended>> {
+        // SAFETY: pdev is valid by ConfigSpace invariants
+        let offset = unsafe {
+            bindings::pci_find_next_ext_capability(self.pdev.as_raw(), start_pos, id as i32)
+        };
+
+        if offset == 0 {
+            return None;
+        }
+
+        // Calculate real capability size
+        let size = self.calculate_ext_cap_size(offset as usize);
+        Some(Capability::new(self, offset as usize, id, size))
+    }
+}
+
+/// SR-IOV register offsets (relative to the capability base).
+mod sriov_offsets {
+    use crate::bindings;
+
+    /// First VF Offset register offset
+    pub(super) const VF_OFFSET: usize = bindings::PCI_SRIOV_VF_OFFSET as usize;
+    /// VF BAR0 register offset (first of 6 VF BARs)
+    pub(super) const VF_BAR0: usize = bindings::PCI_SRIOV_BAR as usize;
+}
+
+/// SR-IOV capability structure.
+///
+/// This structure provides typed access to the SR-IOV extended capability
+/// registers using the PCI configuration space backend.
+pub struct SriovCapability<'a> {
+    cap: Capability<'a, Extended, Extended>,
+}
+
+impl<'a> SriovCapability<'a> {
+    /// Creates a new SR-IOV capability from an extended capability.
+    pub fn new(cap: Capability<'a, Extended, Extended>) -> Result<Self> {
+        if cap.id() != ExtCapabilityId::SRIOV as u16 {
+            return Err(EINVAL);
+        }
+        Ok(Self { cap })
+    }
+
+    /// Tries to find and create an SR-IOV capability from a config space.
+    pub fn from_config_space(config_space: &'a ConfigSpace<'a, Extended>) -> Result<Self> {
+        let cap = config_space
+            .find_ext_capability(ExtCapabilityId::SRIOV as u16)
+            .ok_or(ENODEV)?;
+        Self::new(cap)
+    }
+
+    /// Returns the offset of this capability in configuration space.
+    pub fn offset(&self) -> usize {
+        self.cap.offset()
+    }
+
+    /// Reads the First VF Offset register.
+    pub fn read_vf_offset(&self) -> Result<u16> {
+        self.cap.try_read16(sriov_offsets::VF_OFFSET)
+    }
+
+    /// Reads a VF BAR register (32-bit).
+    /// Returns the 32-bit value of the specified VF BAR register.
+    pub fn read_vf_bar(&self, bar_index: usize) -> Result<u32> {
+        if bar_index >= 6 {
+            return Err(EINVAL);
+        }
+        self.cap.try_read32(sriov_offsets::VF_BAR0 + bar_index * 4)
+    }
+
+    /// Reads a 64-bit VF BAR register.
+    /// Returns the 64-bit address combining BAR[n] (low) and BAR[n+1] (high).
+    pub fn read_vf_bar64(&self, bar_index: usize) -> Result<u64> {
+        if bar_index >= 5 {
+            return Err(EINVAL);
+        }
+        let low = self.read_vf_bar(bar_index)?;
+        let high = self.read_vf_bar(bar_index + 1)?;
+        Ok((u64::from(high) << 32) | u64::from(low))
+    }
+}
-- 
2.51.0