uefi: Refactor PciRootBridgeIo::enumerate() with tree-topology information

- Refactored return type from standard BTreeSet to custom PciTree struct
- Removed special FullPciIoAddress type, since segment number is PciRoot dependent
- During enumeration, skip branches we have already seen
- During enumeration, collect tree topology information (which child bus linked from where)
- Add complicated pci structure in integration test vm
- Print child busses for every device entry in integration test

Author: seijikun

uefi-test-runner/src/proto/pci/root_bridge.rs

@@ -21,9 +21,8 @@ pub fn test() {
     for pci_handle in pci_handles {
         let mut pci_proto = get_open_protocol::<PciRootBridgeIo>(pci_handle);
 
-        let devices = pci_proto.enumerate().unwrap();
-        for fqaddr in devices {
-            let addr = fqaddr.addr();
+        let pci_tree = pci_proto.enumerate().unwrap();
+        for addr in pci_tree.iter().cloned() {
             let Ok(reg0) = pci_proto.pci().read_one::<u32>(addr.with_register(0)) else {
                 continue;
             };
@@ -53,8 +52,11 @@ pub fn test() {
 
             let (bus, dev, fun) = (addr.bus, addr.dev, addr.fun);
             log::info!(
-                "PCI Device: [{bus}, {dev}, {fun}]: vendor={vendor_id:04X}, device={device_id:04X}, class={class_code:02X}, subclass={subclass_code:02X}"
+                "PCI Device: [{bus:02x}, {dev:02x}, {fun:02x}]: vendor={vendor_id:04X}, device={device_id:04X}, class={class_code:02X}, subclass={subclass_code:02X}"
             );
+            for child_bus in pci_tree.iter_subsequent_busses_for(addr) {
+                log::info!(" |- Bus: {child_bus:02x}");
+            }
         }
     }
 

uefi-test-runner/src/proto/scsi/pass_thru.rs

@@ -16,9 +16,9 @@ fn test_allocating_api() {
     // by default respectively. We manually configure an additional SCSI controller.
     // Thus, we should see two controllers with support for EXT_SCSI_PASS_THRU on this platform
     #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-    assert_eq!(scsi_ctrl_handles.len(), 2);
+    assert_eq!(scsi_ctrl_handles.len(), 4);
     #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
-    assert_eq!(scsi_ctrl_handles.len(), 1);
+    assert_eq!(scsi_ctrl_handles.len(), 3);
 
     let mut found_drive = false;
     for handle in scsi_ctrl_handles {

uefi/src/proto/pci/enumeration.rs

@@ -4,10 +4,11 @@
 
 use core::mem;
 
-use alloc::collections::btree_set::BTreeSet;
+use alloc::collections::btree_map::BTreeMap;
+use alloc::collections::btree_set::{self, BTreeSet};
 
+use super::PciIoAddress;
 use super::root_bridge::PciRootBridgeIo;
-use super::{FullPciIoAddress, PciIoAddress};
 
 #[allow(unused)]
 #[derive(Clone, Copy, Debug)]
@@ -55,6 +56,86 @@ fn read_device_register_u32<T: Sized + Copy>(
     }
 }
 
+// ##########################################################################################
+
+/// Struct representing the tree structure of PCI devices.
+/// This allows iterating over all valid PCI device addresses in a tree, as well as query
+/// the tree topology.
+#[derive(Debug)]
+pub struct PciTree {
+    segment: u32,
+    devices: BTreeSet<PciIoAddress>,
+    bus_anchors: BTreeMap<u8, PciIoAddress>,
+}
+impl PciTree {
+    pub(crate) fn new(segment: u32) -> Self {
+        Self {
+            segment,
+            devices: BTreeSet::new(),
+            bus_anchors: BTreeMap::new(),
+        }
+    }
+
+    pub(crate) fn push_bus(&self, bus: u8) -> bool {
+        !self.bus_anchors.contains_key(&bus)
+    }
+
+    pub(crate) fn push_device(&mut self, addr: PciIoAddress) {
+        self.devices.insert(addr);
+    }
+
+    pub(crate) fn push_bridge(&mut self, addr: PciIoAddress, child_bus: u8) -> bool {
+        match self.bus_anchors.contains_key(&child_bus) {
+            true => false,
+            false => {
+                self.bus_anchors.insert(child_bus, addr);
+                true
+            }
+        }
+    }
+
+    /// Iterate over all valid PCI device addresses in this tree structure.
+    pub fn iter(&self) -> btree_set::Iter<'_, PciIoAddress> {
+        self.devices.iter()
+    }
+
+    /// Get the segment number of this PCI tree.
+    pub fn segment_nr(&self) -> u32 {
+        self.segment
+    }
+
+    /// Query the address of the parent PCI bridge this `addr`'s bus is subordinate to.
+    pub fn parent_for(&self, addr: PciIoAddress) -> Option<PciIoAddress> {
+        self.bus_anchors.get(&addr.bus).cloned()
+    }
+
+    /// Iterate over all subsequent busses below the given `addr`.
+    /// This yields 0 results if `addr` doesn't point to a PCI bridge.
+    pub fn iter_subsequent_busses_for(&self, addr: PciIoAddress) -> impl Iterator<Item = u8> {
+        self.bus_anchors
+            .iter()
+            .filter(move |&(_, parent)| *parent == addr)
+            .map(|(bus, _)| bus)
+            .cloned()
+    }
+}
+impl IntoIterator for PciTree {
+    type Item = PciIoAddress;
+    type IntoIter = btree_set::IntoIter<PciIoAddress>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.devices.into_iter()
+    }
+}
+impl<'a> IntoIterator for &'a PciTree {
+    type Item = &'a PciIoAddress;
+    type IntoIter = btree_set::Iter<'a, PciIoAddress>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.devices.iter()
+    }
+}
+
 // ##########################################################################################
 // # Query Helpers (read from a device's configuration registers)
 
@@ -86,12 +167,12 @@ fn get_secondary_bus_range(
 fn visit_function(
     proto: &mut PciRootBridgeIo,
     addr: PciIoAddress,
-    queue: &mut BTreeSet<FullPciIoAddress>,
+    tree: &mut PciTree,
 ) -> uefi::Result<()> {
     if get_vendor_id(proto, addr)? == 0xFFFF {
         return Ok(()); // function doesn't exist - bail instantly
     }
-    queue.insert(FullPciIoAddress::new(proto.segment_nr(), addr));
+    tree.push_device(addr);
     let (base_class, sub_class) = get_classes(proto, addr)?;
     let header_type = get_header_type(proto, addr)? & 0b01111111;
     if base_class == 0x6 && sub_class == 0x4 && header_type == 0x01 {
@@ -106,8 +187,11 @@ fn visit_function(
             return Ok(());
         }
         for bus in secondary_bus_nr..=subordinate_bus_nr {
-            // Recurse into the bus namespaces on the other side of the bridge
-            visit_bus(proto, PciIoAddress::new(bus, 0, 0), queue)?;
+            // Recurse into the bus namespaces on the other side of the bridge, if we haven't visited
+            // the subordinate bus through a more direct path already
+            if tree.push_bridge(addr, bus) {
+                visit_bus(proto, PciIoAddress::new(bus, 0, 0), tree)?;
+            }
         }
     }
     Ok(())
@@ -116,17 +200,17 @@ fn visit_function(
 fn visit_device(
     proto: &mut PciRootBridgeIo,
     addr: PciIoAddress,
-    queue: &mut BTreeSet<FullPciIoAddress>,
+    tree: &mut PciTree,
 ) -> uefi::Result<()> {
     if get_vendor_id(proto, addr)? == 0xFFFF {
         return Ok(()); // device doesn't exist
     }
-    visit_function(proto, addr.with_function(0), queue)?;
+    visit_function(proto, addr.with_function(0), tree)?;
     if get_header_type(proto, addr.with_function(0))? & 0x80 != 0 {
         // This is a multi-function device - also try the remaining functions [1;7]
         // These remaining functions can be sparsely populated - as long as function 0 exists.
         for fun in 1..=7 {
-            visit_function(proto, addr.with_function(fun), queue)?;
+            visit_function(proto, addr.with_function(fun), tree)?;
         }
     }
 
@@ -136,11 +220,11 @@ fn visit_device(
 pub(crate) fn visit_bus(
     proto: &mut PciRootBridgeIo,
     addr: PciIoAddress,
-    queue: &mut BTreeSet<FullPciIoAddress>,
+    tree: &mut PciTree,
 ) -> uefi::Result<()> {
     // Given a valid bus entry point - simply try all possible devices addresses
     for dev in 0..32 {
-        visit_device(proto, addr.with_device(dev), queue)?;
+        visit_device(proto, addr.with_device(dev), tree)?;
     }
     Ok(())
 }

uefi/src/proto/pci/mod.rs

@@ -8,7 +8,7 @@ use uefi_raw::protocol::pci::root_bridge::PciRootBridgeIoProtocolWidth;
 
 pub mod configuration;
 #[cfg(feature = "alloc")]
-mod enumeration;
+pub mod enumeration;
 pub mod root_bridge;
 
 /// IO Address for PCI/register IO operations
@@ -124,36 +124,6 @@ impl Ord for PciIoAddress {
     }
 }
 
-// --------------------------------------------------------------------------------------------
-
-/// Fully qualified pci address. This address is valid across root bridges.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
-pub struct FullPciIoAddress {
-    /// PCI segment number
-    segment: u32,
-    /// Subsequent PCI address
-    addr: PciIoAddress,
-}
-impl FullPciIoAddress {
-    /// Construct a new fully qualified pci address.
-    #[must_use]
-    pub const fn new(segment: u32, addr: PciIoAddress) -> Self {
-        Self { segment, addr }
-    }
-
-    /// Get the segment number this address belongs to.
-    #[must_use]
-    pub const fn segment(&self) -> u32 {
-        self.segment
-    }
-
-    /// Get the internal RootBridge-specific portion of the address.
-    #[must_use]
-    pub const fn addr(&self) -> PciIoAddress {
-        self.addr
-    }
-}
-
 // ############################################################################################
 
 /// Trait implemented by all data types that can natively be read from a PCI device.

uefi/src/proto/pci/root_bridge.rs

@@ -7,17 +7,13 @@ use crate::StatusExt;
 #[cfg(feature = "alloc")]
 use crate::proto::pci::configuration::QwordAddressSpaceDescriptor;
 #[cfg(feature = "alloc")]
-use alloc::collections::btree_set::BTreeSet;
-#[cfg(feature = "alloc")]
 use alloc::vec::Vec;
 #[cfg(feature = "alloc")]
 use core::ffi::c_void;
 use core::ptr;
 use uefi_macros::unsafe_protocol;
 use uefi_raw::protocol::pci::root_bridge::{PciRootBridgeIoAccess, PciRootBridgeIoProtocol};
 
-#[cfg(doc)]
-use super::FullPciIoAddress;
 #[cfg(doc)]
 use crate::Status;
 
@@ -84,10 +80,10 @@ impl PciRootBridgeIo {
     // ###################################################
     // # Convenience functionality
 
-    /// Recursively enumerate all devices, device functions and pci-to-pci bridges on this root bridge.
+    /// Recursively enumerate all devices, device functions and pci(e)-to-pci(e) bridges, starting from this pci root.
     ///
     /// The returned addresses might overlap with the addresses returned by another [`PciRootBridgeIo`] instance.
-    /// Make sure to perform some form of cross-[`PciRootBridgeIo`] deduplication on the returned [`FullPciIoAddress`]es.
+    /// Make sure to perform some form of cross-[`PciRootBridgeIo`] deduplication on the discovered addresses.
     /// **WARNING:** Only use the returned addresses with the respective [`PciRootBridgeIo`] instance that returned them.
     ///
     /// # Returns
@@ -96,24 +92,27 @@ impl PciRootBridgeIo {
     /// # Errors
     /// This can basically fail with all the IO errors found in [`PciIoAccessPci`] methods.
     #[cfg(feature = "alloc")]
-    pub fn enumerate(&mut self) -> crate::Result<BTreeSet<super::FullPciIoAddress>> {
+    pub fn enumerate(&mut self) -> crate::Result<super::enumeration::PciTree> {
+        use super::enumeration::{self, PciTree};
         use crate::proto::pci::configuration::ResourceRangeType;
-        use crate::proto::pci::enumeration;
 
-        let mut devices = BTreeSet::new();
-        // In the descriptors, the entry with range_type bus specifies the bus numbers that were
-        // allocated to devices below this root bridge. The first bus number in this range is
-        // the starting point. All subsequent numbers are reached via PCI bridge recursion during enumeration.
-        if let Some(descriptor) = self
-            .configuration()?
-            .iter()
-            .find(|d| d.resource_range_type == ResourceRangeType::Bus)
-        {
-            let addr = PciIoAddress::new(descriptor.address_min as u8, 0, 0);
-            enumeration::visit_bus(self, addr, &mut devices)?;
+        let mut tree = PciTree::new(self.segment_nr());
+        for descriptor in self.configuration()? {
+            // In the descriptors we can query for the current root bridge, Bus entries contain ranges of valid
+            // bus addresses. These are all bus addresses found below ourselves. These are not only the busses
+            // linked to **directly** from ourselves, but also recursively. Thus we use PciTree::push_bus() to
+            // determine whether we have already visited a given bus number.
+            if descriptor.resource_range_type == ResourceRangeType::Bus {
+                for bus in (descriptor.address_min as u8)..=(descriptor.address_max as u8) {
+                    if tree.push_bus(bus) {
+                        let addr = PciIoAddress::new(bus, 0, 0);
+                        enumeration::visit_bus(self, addr, &mut tree)?;
+                    }
+                }
+            }
         }
 
-        Ok(devices)
+        Ok(tree)
     }
 }
 

xtask/src/qemu.rs

@@ -506,6 +506,23 @@ pub fn run_qemu(arch: UefiArch, opt: &QemuOpt) -> Result<()> {
         None
     };
 
+    // Make PCI tree a little more complicated so the PCI enumeration integration
+    // test is more interesting.
+    cmd.args([
+        "-device",
+        "ioh3420,id=root_port1,bus=pcie.0",
+        "-device",
+        "x3130-upstream,id=upstream1,bus=root_port1",
+        "-device",
+        "xio3130-downstream,id=downstream1,bus=upstream1,chassis=9",
+        "-device",
+        "virtio-scsi-pci,bus=downstream1",
+        "-device",
+        "xio3130-downstream,id=downstream2,bus=upstream1,chassis=10",
+        "-device",
+        "virtio-scsi-pci,bus=downstream2",
+    ]);
+
     // Pass CA certificate database to the edk2 firmware, for TLS support.
     cmd.args([
         "-fw_cfg",