Introduce a PartitioningCx struct

rust-lang · Dylan-DPC-zz · Aug 29, 2020 · Aug 29, 2020 · Sep 2, 2020 · Sep 5, 2020
commit c9686cb31ad9dd9428825dbb3b362f5df72ab719
diff --git a/compiler/rustc_mir/src/monomorphize/partitioning/default.rs b/compiler/rustc_mir/src/monomorphize/partitioning/default.rs
@@ -11,6 +11,7 @@ use rustc_middle::ty::print::characteristic_def_id_of_type;
 use rustc_middle::ty::{self, DefIdTree, InstanceDef, TyCtxt};
 use rustc_span::symbol::Symbol;
 
+use super::PartitioningCx;
 use crate::monomorphize::collector::InliningMap;
 use crate::monomorphize::partitioning::merging;
 use crate::monomorphize::partitioning::{
@@ -22,35 +23,36 @@ pub struct DefaultPartitioning;
 impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
     fn place_root_mono_items(
         &mut self,
-        tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         mono_items: &mut dyn Iterator<Item = MonoItem<'tcx>>,
     ) -> PreInliningPartitioning<'tcx> {
         let mut roots = FxHashSet::default();
         let mut codegen_units = FxHashMap::default();
-        let is_incremental_build = tcx.sess.opts.incremental.is_some();
+        let is_incremental_build = cx.tcx.sess.opts.incremental.is_some();
         let mut internalization_candidates = FxHashSet::default();
 
         // Determine if monomorphizations instantiated in this crate will be made
         // available to downstream crates. This depends on whether we are in
         // share-generics mode and whether the current crate can even have
         // downstream crates.
-        let export_generics = tcx.sess.opts.share_generics() && tcx.local_crate_exports_generics();
+        let export_generics =
+            cx.tcx.sess.opts.share_generics() && cx.tcx.local_crate_exports_generics();
 
-        let cgu_name_builder = &mut CodegenUnitNameBuilder::new(tcx);
+        let cgu_name_builder = &mut CodegenUnitNameBuilder::new(cx.tcx);
         let cgu_name_cache = &mut FxHashMap::default();
 
         for mono_item in mono_items {
-            match mono_item.instantiation_mode(tcx) {
+            match mono_item.instantiation_mode(cx.tcx) {
                 InstantiationMode::GloballyShared { .. } => {}
                 InstantiationMode::LocalCopy => continue,
             }
 
-            let characteristic_def_id = characteristic_def_id_of_mono_item(tcx, mono_item);
+            let characteristic_def_id = characteristic_def_id_of_mono_item(cx.tcx, mono_item);
             let is_volatile = is_incremental_build && mono_item.is_generic_fn();
 
             let codegen_unit_name = match characteristic_def_id {
                 Some(def_id) => compute_codegen_unit_name(
-                    tcx,
+                    cx.tcx,
                     cgu_name_builder,
                     def_id,
                     is_volatile,
@@ -65,7 +67,7 @@ impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
 
             let mut can_be_internalized = true;
             let (linkage, visibility) = mono_item_linkage_and_visibility(
-                tcx,
+                cx.tcx,
                 &mono_item,
                 &mut can_be_internalized,
                 export_generics,
@@ -97,17 +99,16 @@ impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
 
     fn merge_codegen_units(
         &mut self,
-        tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         initial_partitioning: &mut PreInliningPartitioning<'tcx>,
-        target_cgu_count: usize,
     ) {
-        merging::merge_codegen_units(tcx, initial_partitioning, target_cgu_count);
+        merging::merge_codegen_units(cx, initial_partitioning);
     }
 
     fn place_inlined_mono_items(
         &mut self,
+        cx: &PartitioningCx<'_, 'tcx>,
         initial_partitioning: PreInliningPartitioning<'tcx>,
-        inlining_map: &InliningMap<'tcx>,
     ) -> PostInliningPartitioning<'tcx> {
         let mut new_partitioning = Vec::new();
         let mut mono_item_placements = FxHashMap::default();
@@ -124,7 +125,7 @@ impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
             // Collect all items that need to be available in this codegen unit.
             let mut reachable = FxHashSet::default();
             for root in old_codegen_unit.items().keys() {
-                follow_inlining(*root, inlining_map, &mut reachable);
+                follow_inlining(*root, cx.inlining_map, &mut reachable);
             }
 
             let mut new_codegen_unit = CodegenUnit::new(old_codegen_unit.name());
@@ -198,9 +199,8 @@ impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
 
     fn internalize_symbols(
         &mut self,
-        _tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         partitioning: &mut PostInliningPartitioning<'tcx>,
-        inlining_map: &InliningMap<'tcx>,
     ) {
         if partitioning.codegen_units.len() == 1 {
             // Fast path for when there is only one codegen unit. In this case we
@@ -218,7 +218,7 @@ impl<'tcx> Partitioner<'tcx> for DefaultPartitioning {
         // Build a map from every monomorphization to all the monomorphizations that
         // reference it.
         let mut accessor_map: FxHashMap<MonoItem<'tcx>, Vec<MonoItem<'tcx>>> = Default::default();
-        inlining_map.iter_accesses(|accessor, accessees| {
+        cx.inlining_map.iter_accesses(|accessor, accessees| {
             for accessee in accessees {
                 accessor_map.entry(*accessee).or_default().push(accessor);
             }

diff --git a/compiler/rustc_mir/src/monomorphize/partitioning/merging.rs b/compiler/rustc_mir/src/monomorphize/partitioning/merging.rs
@@ -3,17 +3,16 @@ use std::cmp;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder};
-use rustc_middle::ty::TyCtxt;
 use rustc_span::symbol::{Symbol, SymbolStr};
 
+use super::PartitioningCx;
 use crate::monomorphize::partitioning::PreInliningPartitioning;
 
 pub fn merge_codegen_units<'tcx>(
-    tcx: TyCtxt<'tcx>,
+    cx: &PartitioningCx<'_, 'tcx>,
     initial_partitioning: &mut PreInliningPartitioning<'tcx>,
-    target_cgu_count: usize,
 ) {
-    assert!(target_cgu_count >= 1);
+    assert!(cx.target_cgu_count >= 1);
     let codegen_units = &mut initial_partitioning.codegen_units;
 
     // Note that at this point in time the `codegen_units` here may not be in a
@@ -32,7 +31,7 @@ pub fn merge_codegen_units<'tcx>(
         codegen_units.iter().map(|cgu| (cgu.name(), vec![cgu.name().as_str()])).collect();
 
     // Merge the two smallest codegen units until the target size is reached.
-    while codegen_units.len() > target_cgu_count {
+    while codegen_units.len() > cx.target_cgu_count {
         // Sort small cgus to the back
         codegen_units.sort_by_cached_key(|cgu| cmp::Reverse(cgu.size_estimate()));
         let mut smallest = codegen_units.pop().unwrap();
@@ -56,9 +55,9 @@ pub fn merge_codegen_units<'tcx>(
         );
     }
 
-    let cgu_name_builder = &mut CodegenUnitNameBuilder::new(tcx);
+    let cgu_name_builder = &mut CodegenUnitNameBuilder::new(cx.tcx);
 
-    if tcx.sess.opts.incremental.is_some() {
+    if cx.tcx.sess.opts.incremental.is_some() {
         // If we are doing incremental compilation, we want CGU names to
         // reflect the path of the source level module they correspond to.
         // For CGUs that contain the code of multiple modules because of the
@@ -82,7 +81,7 @@ pub fn merge_codegen_units<'tcx>(
 
         for cgu in codegen_units.iter_mut() {
             if let Some(new_cgu_name) = new_cgu_names.get(&cgu.name()) {
-                if tcx.sess.opts.debugging_opts.human_readable_cgu_names {
+                if cx.tcx.sess.opts.debugging_opts.human_readable_cgu_names {
                     cgu.set_name(Symbol::intern(&new_cgu_name));
                 } else {
                     // If we don't require CGU names to be human-readable, we

diff --git a/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs b/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs
@@ -107,31 +107,35 @@ use rustc_span::symbol::Symbol;
 use crate::monomorphize::collector::InliningMap;
 use crate::monomorphize::collector::{self, MonoItemCollectionMode};
 
+pub struct PartitioningCx<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    target_cgu_count: usize,
+    inlining_map: &'a InliningMap<'tcx>,
+}
+
 trait Partitioner<'tcx> {
     fn place_root_mono_items(
         &mut self,
-        tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         mono_items: &mut dyn Iterator<Item = MonoItem<'tcx>>,
     ) -> PreInliningPartitioning<'tcx>;
 
     fn merge_codegen_units(
         &mut self,
-        tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         initial_partitioning: &mut PreInliningPartitioning<'tcx>,
-        target_cgu_count: usize,
     );
 
     fn place_inlined_mono_items(
         &mut self,
+        cx: &PartitioningCx<'_, 'tcx>,
         initial_partitioning: PreInliningPartitioning<'tcx>,
-        inlining_map: &InliningMap<'tcx>,
     ) -> PostInliningPartitioning<'tcx>;
 
     fn internalize_symbols(
         &mut self,
-        tcx: TyCtxt<'tcx>,
+        cx: &PartitioningCx<'_, 'tcx>,
         partitioning: &mut PostInliningPartitioning<'tcx>,
-        inlining_map: &InliningMap<'tcx>,
     );
 }
 
@@ -156,12 +160,13 @@ pub fn partition<'tcx>(
     let _prof_timer = tcx.prof.generic_activity("cgu_partitioning");
 
     let mut partitioner = get_partitioner(tcx);
+    let cx = &PartitioningCx { tcx, target_cgu_count: max_cgu_count, inlining_map };
     // In the first step, we place all regular monomorphizations into their
     // respective 'home' codegen unit. Regular monomorphizations are all
     // functions and statics defined in the local crate.
     let mut initial_partitioning = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_roots");
-        partitioner.place_root_mono_items(tcx, mono_items)
+        partitioner.place_root_mono_items(cx, mono_items)
     };
 
     initial_partitioning.codegen_units.iter_mut().for_each(|cgu| cgu.estimate_size(tcx));
@@ -171,7 +176,7 @@ pub fn partition<'tcx>(
     // Merge until we have at most `max_cgu_count` codegen units.
     {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_merge_cgus");
-        partitioner.merge_codegen_units(tcx, &mut initial_partitioning, max_cgu_count);
+        partitioner.merge_codegen_units(cx, &mut initial_partitioning);
         debug_dump(tcx, "POST MERGING:", initial_partitioning.codegen_units.iter());
     }
 
@@ -181,7 +186,7 @@ pub fn partition<'tcx>(
     // local functions the definition of which is marked with `#[inline]`.
     let mut post_inlining = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_inline_items");
-        partitioner.place_inlined_mono_items(initial_partitioning, inlining_map)
+        partitioner.place_inlined_mono_items(cx, initial_partitioning)
     };
 
     post_inlining.codegen_units.iter_mut().for_each(|cgu| cgu.estimate_size(tcx));
@@ -192,7 +197,7 @@ pub fn partition<'tcx>(
     // more freedom to optimize.
     if !tcx.sess.link_dead_code() {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_internalize_symbols");
-        partitioner.internalize_symbols(tcx, &mut post_inlining, inlining_map);
+        partitioner.internalize_symbols(cx, &mut post_inlining);
     }
 
     // Finally, sort by codegen unit name, so that we get deterministic results.