From 9aed7d560e8f760452f93141c217408dc0f4c682 Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Sun, 9 Jun 2024 18:21:50 +0200
Subject: [PATCH] pcid: Merge CapabilityOffsetsIter into CapabilitiesIter

---
 pcid/src/main.rs    |  2 +-
 pcid/src/pci/cap.rs | 36 +++++++++++++-----------------------
 2 files changed, 14 insertions(+), 24 deletions(-)

diff --git a/pcid/src/main.rs b/pcid/src/main.rs
index 08f8eb09a6..7e852bb982 100644
--- a/pcid/src/main.rs
+++ b/pcid/src/main.rs
@@ -267,7 +267,7 @@ fn handle_parsed_header(state: Arc<State>, config: &Config, header: PciEndpointH
                 pci: &state.pcie,
                 addr: header.address(),
             };
-            crate::pci::cap::CapabilitiesIter { inner: crate::pci::cap::CapabilityOffsetsIter::new(header.cap_pointer(), &func) }.collect::<Vec<_>>()
+            crate::pci::cap::CapabilitiesIter::new(header.cap_pointer(), &func).collect::<Vec<_>>()
         } else {
             Vec::new()
         };
diff --git a/pcid/src/pci/cap.rs b/pcid/src/pci/cap.rs
index 5e852a528f..41e480a71c 100644
--- a/pcid/src/pci/cap.rs
+++ b/pcid/src/pci/cap.rs
@@ -1,11 +1,11 @@
 use super::func::ConfigReader;
 use serde::{Serialize, Deserialize};
 
-pub struct CapabilityOffsetsIter<'a, R> {
+pub struct CapabilitiesIter<'a, R> {
     offset: u8,
     reader: &'a R,
 }
-impl<'a, R> CapabilityOffsetsIter<'a, R> {
+impl<'a, R> CapabilitiesIter<'a, R> {
     pub fn new(offset: u8, reader: &'a R) -> Self {
         Self {
             offset,
@@ -13,14 +13,14 @@ impl<'a, R> CapabilityOffsetsIter<'a, R> {
         }
     }
 }
-impl<'a, R> Iterator for CapabilityOffsetsIter<'a, R>
+impl<'a, R> Iterator for CapabilitiesIter<'a, R>
 where
     R: ConfigReader
 {
-    type Item = u8;
+    type Item = (u8, Capability);
 
     fn next(&mut self) -> Option<Self::Item> {
-        unsafe {
+        let offset = unsafe {
             // mask RsvdP bits
             self.offset = self.offset & 0xFC;
 
@@ -32,8 +32,14 @@ where
             let offset = self.offset;
             self.offset = next;
 
-            Some(offset)
-        }
+            offset
+        };
+
+        let cap = unsafe {
+            Capability::parse(self.reader, offset)
+        };
+
+        Some((offset, cap))
     }
 }
 
@@ -229,19 +235,3 @@ impl Capability {
         }
     }
 }
-
-pub struct CapabilitiesIter<'a, R> {
-    pub inner: CapabilityOffsetsIter<'a, R>,
-}
-
-impl<'a, R> Iterator for CapabilitiesIter<'a, R>
-where
-    R: ConfigReader
-{
-    type Item = (u8, Capability);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let offset = self.inner.next()?;
-        Some((offset, unsafe { Capability::parse(self.inner.reader, offset) }))
-    }
-}