diff --git a/src/main.rs b/src/main.rs
index 0796ea50..efdea0c3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -172,19 +172,19 @@ pub extern "C" fn _start() -> ! {
 
     pci::print_bus();
 
-    let mut pci_transport;
-    let mut mmio_transport;
-
-    let mut device = if let Some(pci_device) =
-        pci::search_bus(VIRTIO_PCI_VENDOR_ID, VIRTIO_PCI_BLOCK_DEVICE_ID)
-    {
-        pci_transport = pci::VirtioPciTransport::new(pci_device);
-        block::VirtioBlockDevice::new(&mut pci_transport)
-    } else {
-        mmio_transport = mmio::VirtioMMIOTransport::new(0xd000_0000u64);
-        block::VirtioBlockDevice::new(&mut mmio_transport)
-    };
-
+    pci::with_devices(
+        VIRTIO_PCI_VENDOR_ID,
+        VIRTIO_PCI_BLOCK_DEVICE_ID,
+        |pci_device| {
+            let mut pci_transport = pci::VirtioPciTransport::new(pci_device);
+            block::VirtioBlockDevice::new(&mut pci_transport);
+            let mut device = block::VirtioBlockDevice::new(&mut pci_transport);
+            boot_from_device(&mut device)
+        },
+    );
+
+    let mut mmio_transport = mmio::VirtioMMIOTransport::new(0xd000_0000u64);
+    let mut device = block::VirtioBlockDevice::new(&mut mmio_transport);
     boot_from_device(&mut device);
 
     #[allow(clippy::empty_loop)]
diff --git a/src/pci.rs b/src/pci.rs
index ec63bd2e..09869ecc 100644
--- a/src/pci.rs
+++ b/src/pci.rs
@@ -90,14 +90,19 @@ pub fn print_bus() {
 }
 
 #[cfg(not(test))]
-pub fn search_bus(target_vendor_id: u16, target_device_id: u16) -> Option<PciDevice> {
+pub fn with_devices<F>(target_vendor_id: u16, target_device_id: u16, per_device: F)
+where
+    F: Fn(PciDevice) -> bool,
+{
     for device in 0..MAX_DEVICES {
         let (vendor_id, device_id) = get_device_details(0, device, 0);
-        if vendor_id == target_vendor_id && device_id == target_device_id {
-            return Some(PciDevice::new(0, device, 0));
+        if vendor_id == target_vendor_id
+            && device_id == target_device_id
+            && per_device(PciDevice::new(0, device, 0))
+        {
+            break;
         }
     }
-    None
 }
 
 #[cfg(not(test))]