44
55use core:: mem;
66
7- use alloc:: collections:: btree_set:: BTreeSet ;
7+ use alloc:: collections:: btree_map:: BTreeMap ;
8+ use alloc:: collections:: btree_set:: { self , BTreeSet } ;
89
10+ use super :: PciIoAddress ;
911use super :: root_bridge:: PciRootBridgeIo ;
10- use super :: { FullPciIoAddress , PciIoAddress } ;
1112
1213#[ allow( unused) ]
1314#[ derive( Clone , Copy , Debug ) ]
@@ -55,6 +56,93 @@ fn read_device_register_u32<T: Sized + Copy>(
5556 }
5657}
5758
59+ // ##########################################################################################
60+
61+ /// Struct representing the tree structure of PCI devices.
62+ ///
63+ /// This allows iterating over all valid PCI device addresses in a tree, as well as querying
64+ /// the tree topology.
65+ #[ derive( Debug ) ]
66+ pub struct PciTree {
67+ segment : u32 ,
68+ devices : BTreeSet < PciIoAddress > ,
69+ bus_anchors : BTreeMap < u8 /* bus */ , PciIoAddress > ,
70+ }
71+ impl PciTree {
72+ pub ( crate ) const fn new ( segment : u32 ) -> Self {
73+ Self {
74+ segment,
75+ devices : BTreeSet :: new ( ) ,
76+ bus_anchors : BTreeMap :: new ( ) ,
77+ }
78+ }
79+
80+ pub ( crate ) fn should_visit_bus ( & self , bus : u8 ) -> bool {
81+ !self . bus_anchors . contains_key ( & bus)
82+ }
83+
84+ pub ( crate ) fn push_device ( & mut self , addr : PciIoAddress ) {
85+ self . devices . insert ( addr) ;
86+ }
87+
88+ /// Pushes a new bridge into the topology.
89+ ///
90+ /// Returns `false` if the bus is already in the topology and `true`
91+ /// if the bridge was added to the topology.
92+ pub ( crate ) fn push_bridge ( & mut self , addr : PciIoAddress , child_bus : u8 ) -> bool {
93+ match self . bus_anchors . contains_key ( & child_bus) {
94+ true => false ,
95+ false => {
96+ self . bus_anchors . insert ( child_bus, addr) ;
97+ true
98+ }
99+ }
100+ }
101+
102+ /// Iterate over all valid PCI device addresses in this tree structure.
103+ pub fn iter ( & self ) -> btree_set:: Iter < ' _ , PciIoAddress > {
104+ self . devices . iter ( )
105+ }
106+
107+ /// Get the segment number of this PCI tree.
108+ #[ must_use]
109+ pub const fn segment_nr ( & self ) -> u32 {
110+ self . segment
111+ }
112+
113+ /// Query the address of the parent PCI bridge this `addr`'s bus is subordinate to.
114+ #[ must_use]
115+ pub fn parent_for ( & self , addr : PciIoAddress ) -> Option < PciIoAddress > {
116+ self . bus_anchors . get ( & addr. bus ) . cloned ( )
117+ }
118+
119+ /// Iterate over all subsequent busses below the given `addr`.
120+ /// This yields 0 results if `addr` doesn't point to a PCI bridge.
121+ pub fn child_bus_of_iter ( & self , addr : PciIoAddress ) -> impl Iterator < Item = u8 > {
122+ self . bus_anchors
123+ . iter ( )
124+ . filter ( move |& ( _, parent) | * parent == addr)
125+ . map ( |( bus, _) | bus)
126+ . cloned ( )
127+ }
128+ }
129+ impl IntoIterator for PciTree {
130+ type Item = PciIoAddress ;
131+ type IntoIter = btree_set:: IntoIter < PciIoAddress > ;
132+
133+ fn into_iter ( self ) -> Self :: IntoIter {
134+ self . devices . into_iter ( )
135+ }
136+ }
137+ impl < ' a > IntoIterator for & ' a PciTree {
138+ type Item = & ' a PciIoAddress ;
139+ type IntoIter = btree_set:: Iter < ' a , PciIoAddress > ;
140+
141+ fn into_iter ( self ) -> Self :: IntoIter {
142+ self . devices . iter ( )
143+ }
144+ }
145+
58146// ##########################################################################################
59147// # Query Helpers (read from a device's configuration registers)
60148
@@ -86,12 +174,12 @@ fn get_secondary_bus_range(
86174fn visit_function (
87175 proto : & mut PciRootBridgeIo ,
88176 addr : PciIoAddress ,
89- queue : & mut BTreeSet < FullPciIoAddress > ,
177+ tree : & mut PciTree ,
90178) -> uefi:: Result < ( ) > {
91179 if get_vendor_id ( proto, addr) ? == 0xFFFF {
92180 return Ok ( ( ) ) ; // function doesn't exist - bail instantly
93181 }
94- queue . insert ( FullPciIoAddress :: new ( proto . segment_nr ( ) , addr) ) ;
182+ tree . push_device ( addr) ;
95183 let ( base_class, sub_class) = get_classes ( proto, addr) ?;
96184 let header_type = get_header_type ( proto, addr) ? & 0b01111111 ;
97185 if base_class == 0x6 && sub_class == 0x4 && header_type == 0x01 {
@@ -106,8 +194,11 @@ fn visit_function(
106194 return Ok ( ( ) ) ;
107195 }
108196 for bus in secondary_bus_nr..=subordinate_bus_nr {
109- // Recurse into the bus namespaces on the other side of the bridge
110- visit_bus ( proto, PciIoAddress :: new ( bus, 0 , 0 ) , queue) ?;
197+ // Recurse into the bus namespaces on the other side of the bridge, if we haven't visited
198+ // the subordinate bus through a more direct path already
199+ if tree. push_bridge ( addr, bus) {
200+ visit_bus ( proto, PciIoAddress :: new ( bus, 0 , 0 ) , tree) ?;
201+ }
111202 }
112203 }
113204 Ok ( ( ) )
@@ -116,17 +207,17 @@ fn visit_function(
116207fn visit_device (
117208 proto : & mut PciRootBridgeIo ,
118209 addr : PciIoAddress ,
119- queue : & mut BTreeSet < FullPciIoAddress > ,
210+ tree : & mut PciTree ,
120211) -> uefi:: Result < ( ) > {
121212 if get_vendor_id ( proto, addr) ? == 0xFFFF {
122213 return Ok ( ( ) ) ; // device doesn't exist
123214 }
124- visit_function ( proto, addr. with_function ( 0 ) , queue ) ?;
215+ visit_function ( proto, addr. with_function ( 0 ) , tree ) ?;
125216 if get_header_type ( proto, addr. with_function ( 0 ) ) ? & 0x80 != 0 {
126217 // This is a multi-function device - also try the remaining functions [1;7]
127218 // These remaining functions can be sparsely populated - as long as function 0 exists.
128219 for fun in 1 ..=7 {
129- visit_function ( proto, addr. with_function ( fun) , queue ) ?;
220+ visit_function ( proto, addr. with_function ( fun) , tree ) ?;
130221 }
131222 }
132223
@@ -136,11 +227,11 @@ fn visit_device(
136227pub ( crate ) fn visit_bus (
137228 proto : & mut PciRootBridgeIo ,
138229 addr : PciIoAddress ,
139- queue : & mut BTreeSet < FullPciIoAddress > ,
230+ tree : & mut PciTree ,
140231) -> uefi:: Result < ( ) > {
141232 // Given a valid bus entry point - simply try all possible devices addresses
142233 for dev in 0 ..32 {
143- visit_device ( proto, addr. with_device ( dev) , queue ) ?;
234+ visit_device ( proto, addr. with_device ( dev) , tree ) ?;
144235 }
145236 Ok ( ( ) )
146237}
0 commit comments