1
1
use self :: collector:: NodeCollector ;
2
2
3
- use crate :: hir:: { AttributeMap , IndexedHir , Owner } ;
3
+ use crate :: hir:: { AttributeMap , IndexedHir , ModuleItems , Owner } ;
4
4
use crate :: ty:: TyCtxt ;
5
5
use rustc_ast as ast;
6
6
use rustc_data_structures:: fingerprint:: Fingerprint ;
@@ -19,6 +19,7 @@ use rustc_span::source_map::Spanned;
19
19
use rustc_span:: symbol:: { kw, sym, Ident , Symbol } ;
20
20
use rustc_span:: Span ;
21
21
use rustc_target:: spec:: abi:: Abi ;
22
+ use std:: collections:: VecDeque ;
22
23
23
24
pub mod blocks;
24
25
mod collector;
@@ -536,23 +537,55 @@ impl<'hir> Map<'hir> {
536
537
{
537
538
let module = self . tcx . hir_module_items ( module) ;
538
539
539
- for id in & module. items {
540
+ for id in module. items . iter ( ) {
540
541
visitor. visit_item ( self . item ( * id) ) ;
541
542
}
542
543
543
- for id in & module. trait_items {
544
+ for id in module. trait_items . iter ( ) {
544
545
visitor. visit_trait_item ( self . trait_item ( * id) ) ;
545
546
}
546
547
547
- for id in & module. impl_items {
548
+ for id in module. impl_items . iter ( ) {
548
549
visitor. visit_impl_item ( self . impl_item ( * id) ) ;
549
550
}
550
551
551
- for id in & module. foreign_items {
552
+ for id in module. foreign_items . iter ( ) {
552
553
visitor. visit_foreign_item ( self . foreign_item ( * id) ) ;
553
554
}
554
555
}
555
556
557
+ pub fn for_each_module ( & self , f : impl Fn ( LocalDefId ) ) {
558
+ let mut queue = VecDeque :: new ( ) ;
559
+ queue. push_back ( CRATE_DEF_ID ) ;
560
+
561
+ while let Some ( id) = queue. pop_front ( ) {
562
+ f ( id) ;
563
+ let items = self . tcx . hir_module_items ( id) ;
564
+ queue. extend ( items. submodules . iter ( ) . copied ( ) )
565
+ }
566
+ }
567
+
568
+ #[ cfg( not( parallel_compiler) ) ]
569
+ #[ inline]
570
+ pub fn par_for_each_module ( & self , f : impl Fn ( LocalDefId ) ) {
571
+ self . for_each_module ( f)
572
+ }
573
+
574
+ #[ cfg( parallel_compiler) ]
575
+ pub fn par_for_each_module ( & self , f : impl Fn ( LocalDefId ) + Sync ) {
576
+ use rustc_data_structures:: sync:: { par_iter, ParallelIterator } ;
577
+ par_iter_submodules ( self . tcx , CRATE_DEF_ID , & f) ;
578
+
579
+ fn par_iter_submodules < F > ( tcx : TyCtxt < ' _ > , module : LocalDefId , f : & F )
580
+ where
581
+ F : Fn ( LocalDefId ) + Sync ,
582
+ {
583
+ ( * f) ( module) ;
584
+ let items = tcx. hir_module_items ( module) ;
585
+ par_iter ( & items. submodules [ ..] ) . for_each ( |& sm| par_iter_submodules ( tcx, sm, f) ) ;
586
+ }
587
+ }
588
+
556
589
/// Returns an iterator for the nodes in the ancestor tree of the `current_id`
557
590
/// until the crate root is reached. Prefer this over your own loop using `get_parent_node`.
558
591
pub fn parent_iter ( & self , current_id : HirId ) -> ParentHirIterator < ' _ , ' hir > {
@@ -1113,3 +1146,69 @@ fn hir_id_to_string(map: &Map<'_>, id: HirId) -> String {
1113
1146
None => format ! ( "unknown node{}" , id_str) ,
1114
1147
}
1115
1148
}
1149
+
1150
+ pub ( super ) fn hir_module_items ( tcx : TyCtxt < ' _ > , module_id : LocalDefId ) -> ModuleItems {
1151
+ let mut collector = ModuleCollector {
1152
+ tcx,
1153
+ submodules : Vec :: default ( ) ,
1154
+ items : Vec :: default ( ) ,
1155
+ trait_items : Vec :: default ( ) ,
1156
+ impl_items : Vec :: default ( ) ,
1157
+ foreign_items : Vec :: default ( ) ,
1158
+ } ;
1159
+
1160
+ let ( hir_mod, span, hir_id) = tcx. hir ( ) . get_module ( module_id) ;
1161
+ collector. visit_mod ( hir_mod, span, hir_id) ;
1162
+
1163
+ let ModuleCollector { submodules, items, trait_items, impl_items, foreign_items, .. } =
1164
+ collector;
1165
+ return ModuleItems {
1166
+ submodules : submodules. into_boxed_slice ( ) ,
1167
+ items : items. into_boxed_slice ( ) ,
1168
+ trait_items : trait_items. into_boxed_slice ( ) ,
1169
+ impl_items : impl_items. into_boxed_slice ( ) ,
1170
+ foreign_items : foreign_items. into_boxed_slice ( ) ,
1171
+ } ;
1172
+
1173
+ struct ModuleCollector < ' tcx > {
1174
+ tcx : TyCtxt < ' tcx > ,
1175
+ submodules : Vec < LocalDefId > ,
1176
+ items : Vec < ItemId > ,
1177
+ trait_items : Vec < TraitItemId > ,
1178
+ impl_items : Vec < ImplItemId > ,
1179
+ foreign_items : Vec < ForeignItemId > ,
1180
+ }
1181
+
1182
+ impl < ' hir > Visitor < ' hir > for ModuleCollector < ' hir > {
1183
+ type Map = Map < ' hir > ;
1184
+
1185
+ fn nested_visit_map ( & mut self ) -> intravisit:: NestedVisitorMap < Self :: Map > {
1186
+ intravisit:: NestedVisitorMap :: All ( self . tcx . hir ( ) )
1187
+ }
1188
+
1189
+ fn visit_item ( & mut self , item : & ' hir Item < ' hir > ) {
1190
+ self . items . push ( item. item_id ( ) ) ;
1191
+ if let ItemKind :: Mod ( ..) = item. kind {
1192
+ // If this declares another module, do not recurse inside it.
1193
+ self . submodules . push ( item. def_id ) ;
1194
+ } else {
1195
+ intravisit:: walk_item ( self , item)
1196
+ }
1197
+ }
1198
+
1199
+ fn visit_trait_item ( & mut self , item : & ' hir TraitItem < ' hir > ) {
1200
+ self . trait_items . push ( item. trait_item_id ( ) ) ;
1201
+ intravisit:: walk_trait_item ( self , item)
1202
+ }
1203
+
1204
+ fn visit_impl_item ( & mut self , item : & ' hir ImplItem < ' hir > ) {
1205
+ self . impl_items . push ( item. impl_item_id ( ) ) ;
1206
+ intravisit:: walk_impl_item ( self , item)
1207
+ }
1208
+
1209
+ fn visit_foreign_item ( & mut self , item : & ' hir ForeignItem < ' hir > ) {
1210
+ self . foreign_items . push ( item. foreign_item_id ( ) ) ;
1211
+ intravisit:: walk_foreign_item ( self , item)
1212
+ }
1213
+ }
1214
+ }
0 commit comments