11use self :: collector:: NodeCollector ;
22
3- use crate :: hir:: { AttributeMap , IndexedHir , Owner } ;
3+ use crate :: hir:: { AttributeMap , IndexedHir , ModuleItems , Owner } ;
44use crate :: ty:: TyCtxt ;
55use rustc_ast as ast;
66use rustc_data_structures:: fingerprint:: Fingerprint ;
@@ -19,6 +19,7 @@ use rustc_span::source_map::Spanned;
1919use rustc_span:: symbol:: { kw, sym, Ident , Symbol } ;
2020use rustc_span:: Span ;
2121use rustc_target:: spec:: abi:: Abi ;
22+ use std:: collections:: VecDeque ;
2223
2324pub mod blocks;
2425mod collector;
@@ -536,23 +537,55 @@ impl<'hir> Map<'hir> {
536537 {
537538 let module = self . tcx . hir_module_items ( module) ;
538539
539- for id in & module. items {
540+ for id in module. items . iter ( ) {
540541 visitor. visit_item ( self . item ( * id) ) ;
541542 }
542543
543- for id in & module. trait_items {
544+ for id in module. trait_items . iter ( ) {
544545 visitor. visit_trait_item ( self . trait_item ( * id) ) ;
545546 }
546547
547- for id in & module. impl_items {
548+ for id in module. impl_items . iter ( ) {
548549 visitor. visit_impl_item ( self . impl_item ( * id) ) ;
549550 }
550551
551- for id in & module. foreign_items {
552+ for id in module. foreign_items . iter ( ) {
552553 visitor. visit_foreign_item ( self . foreign_item ( * id) ) ;
553554 }
554555 }
555556
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+
556589 /// Returns an iterator for the nodes in the ancestor tree of the `current_id`
557590 /// until the crate root is reached. Prefer this over your own loop using `get_parent_node`.
558591 pub fn parent_iter ( & self , current_id : HirId ) -> ParentHirIterator < ' _ , ' hir > {
@@ -1113,3 +1146,69 @@ fn hir_id_to_string(map: &Map<'_>, id: HirId) -> String {
11131146 None => format ! ( "unknown node{}" , id_str) ,
11141147 }
11151148}
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