99use std:: { fmt, mem, ops, str:: FromStr } ;
1010
1111use cfg:: CfgOptions ;
12- use la_arena:: { Arena , Idx } ;
12+ use la_arena:: { Arena , Idx , RawIdx } ;
1313use rustc_hash:: { FxHashMap , FxHashSet } ;
1414use semver:: Version ;
1515use syntax:: SmolStr ;
@@ -157,6 +157,10 @@ impl CrateOrigin {
157157 pub fn is_lib ( & self ) -> bool {
158158 matches ! ( self , CrateOrigin :: Library { .. } )
159159 }
160+
161+ pub fn is_lang ( & self ) -> bool {
162+ matches ! ( self , CrateOrigin :: Lang { .. } )
163+ }
160164}
161165
162166#[ derive( Debug , Clone , Copy , PartialEq , Eq , Hash ) ]
@@ -174,7 +178,7 @@ impl From<&str> for LangCrateOrigin {
174178 match s {
175179 "alloc" => LangCrateOrigin :: Alloc ,
176180 "core" => LangCrateOrigin :: Core ,
177- "proc-macro" => LangCrateOrigin :: ProcMacro ,
181+ "proc-macro" | "proc_macro" => LangCrateOrigin :: ProcMacro ,
178182 "std" => LangCrateOrigin :: Std ,
179183 "test" => LangCrateOrigin :: Test ,
180184 _ => LangCrateOrigin :: Other ,
@@ -522,7 +526,7 @@ impl CrateGraph {
522526 self . arena . iter ( ) . map ( |( idx, _) | idx)
523527 }
524528
525- // FIXME: used for `handle_hack_cargo_workspace` , should be removed later
529+ // FIXME: used for fixing up the toolchain sysroot , should be removed and done differently
526530 #[ doc( hidden) ]
527531 pub fn iter_mut ( & mut self ) -> impl Iterator < Item = ( CrateId , & mut CrateData ) > + ' _ {
528532 self . arena . iter_mut ( )
@@ -619,7 +623,12 @@ impl CrateGraph {
619623 /// This will deduplicate the crates of the graph where possible.
620624 /// Note that for deduplication to fully work, `self`'s crate dependencies must be sorted by crate id.
621625 /// If the crate dependencies were sorted, the resulting graph from this `extend` call will also have the crate dependencies sorted.
622- pub fn extend ( & mut self , mut other : CrateGraph , proc_macros : & mut ProcMacroPaths ) {
626+ pub fn extend (
627+ & mut self ,
628+ mut other : CrateGraph ,
629+ proc_macros : & mut ProcMacroPaths ,
630+ on_finished : impl FnOnce ( & FxHashMap < CrateId , CrateId > ) ,
631+ ) {
623632 let topo = other. crates_in_topological_order ( ) ;
624633 let mut id_map: FxHashMap < CrateId , CrateId > = FxHashMap :: default ( ) ;
625634 for topo in topo {
@@ -670,6 +679,8 @@ impl CrateGraph {
670679
671680 * proc_macros =
672681 mem:: take ( proc_macros) . into_iter ( ) . map ( |( id, macros) | ( id_map[ & id] , macros) ) . collect ( ) ;
682+
683+ on_finished ( & id_map) ;
673684 }
674685
675686 fn find_path (
@@ -721,6 +732,29 @@ impl CrateGraph {
721732 fn hacky_find_crate < ' a > ( & ' a self , display_name : & ' a str ) -> impl Iterator < Item = CrateId > + ' a {
722733 self . iter ( ) . filter ( move |it| self [ * it] . display_name . as_deref ( ) == Some ( display_name) )
723734 }
735+
736+ /// Removes all crates from this crate graph except for the ones in `to_keep` and fixes up the dependencies.
737+ /// Returns a mapping from old crate ids to new crate ids.
738+ pub fn remove_crates_except ( & mut self , to_keep : & [ CrateId ] ) -> Vec < Option < CrateId > > {
739+ let mut id_map = vec ! [ None ; self . arena. len( ) ] ;
740+ self . arena = std:: mem:: take ( & mut self . arena )
741+ . into_iter ( )
742+ . filter_map ( |( id, data) | if to_keep. contains ( & id) { Some ( ( id, data) ) } else { None } )
743+ . enumerate ( )
744+ . map ( |( new_id, ( id, data) ) | {
745+ id_map[ id. into_raw ( ) . into_u32 ( ) as usize ] =
746+ Some ( CrateId :: from_raw ( RawIdx :: from_u32 ( new_id as u32 ) ) ) ;
747+ data
748+ } )
749+ . collect ( ) ;
750+ for ( _, data) in self . arena . iter_mut ( ) {
751+ data. dependencies . iter_mut ( ) . for_each ( |dep| {
752+ dep. crate_id =
753+ id_map[ dep. crate_id . into_raw ( ) . into_u32 ( ) as usize ] . expect ( "crate was filtered" )
754+ } ) ;
755+ }
756+ id_map
757+ }
724758}
725759
726760impl ops:: Index < CrateId > for CrateGraph {
0 commit comments