Remove {Pre,Post}InliningPartitioning.

nnethercote · nnethercote · commit 20de2ba75963 · 2023-05-24T12:33:35.000+10:00
I find that these structs obfuscate the code. Removing them and just
passing the individual fields around makes the `Partition` method
signatures a little longer, but makes the data flow much clearer. E.g.

- `codegen_units` is mutable all the way through.
- `codegen_units`'s length is changed by `merge_codegen_units`, but only
  the individual elements are changed by `place_inlined_mono_items` and
  `internalize_symbols`.
- `roots`, `internalization_candidates`, and `mono_item_placements` are
  all immutable after creation, and all used by just one of the four
  methods.
diff --git a/compiler/rustc_monomorphize/src/partitioning/default.rs b/compiler/rustc_monomorphize/src/partitioning/default.rs
@@ -15,9 +15,7 @@ use rustc_span::symbol::Symbol;
 
 use super::PartitioningCx;
 use crate::collector::InliningMap;
-use crate::partitioning::{
-    MonoItemPlacement, Partition, PostInliningPartitioning, PreInliningPartitioning,
-};
+use crate::partitioning::{MonoItemPlacement, Partition};
 
 pub struct DefaultPartitioning;
 
@@ -26,7 +24,7 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
         mono_items: &mut I,
-    ) -> PreInliningPartitioning<'tcx>
+    ) -> (Vec<CodegenUnit<'tcx>>, FxHashSet<MonoItem<'tcx>>, FxHashSet<MonoItem<'tcx>>)
     where
         I: Iterator<Item = MonoItem<'tcx>>,
     {
@@ -91,20 +89,15 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
             codegen_units.insert(codegen_unit_name, CodegenUnit::new(codegen_unit_name));
         }
 
-        PreInliningPartitioning {
-            codegen_units: codegen_units.into_values().collect(),
-            roots,
-            internalization_candidates,
-        }
+        (codegen_units.into_values().collect(), roots, internalization_candidates)
     }
 
     fn merge_codegen_units(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: &mut PreInliningPartitioning<'tcx>,
+        codegen_units: &mut Vec<CodegenUnit<'tcx>>,
     ) {
         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 deterministic order (but we know they're deterministically the
@@ -201,20 +194,14 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
     fn place_inlined_mono_items(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: PreInliningPartitioning<'tcx>,
-    ) -> PostInliningPartitioning<'tcx> {
-        let mut new_partitioning = Vec::new();
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        roots: FxHashSet<MonoItem<'tcx>>,
+    ) -> FxHashMap<MonoItem<'tcx>, MonoItemPlacement> {
         let mut mono_item_placements = FxHashMap::default();
 
-        let PreInliningPartitioning {
-            codegen_units: initial_cgus,
-            roots,
-            internalization_candidates,
-        } = initial_partitioning;
-
-        let single_codegen_unit = initial_cgus.len() == 1;
+        let single_codegen_unit = codegen_units.len() == 1;
 
-        for old_codegen_unit in initial_cgus {
+        for old_codegen_unit in codegen_units.iter_mut() {
             // 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() {
@@ -266,14 +253,10 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
                 }
             }
 
-            new_partitioning.push(new_codegen_unit);
+            *old_codegen_unit = new_codegen_unit;
         }
 
-        return PostInliningPartitioning {
-            codegen_units: new_partitioning,
-            mono_item_placements,
-            internalization_candidates,
-        };
+        return mono_item_placements;
 
         fn follow_inlining<'tcx>(
             mono_item: MonoItem<'tcx>,
@@ -293,14 +276,16 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
     fn internalize_symbols(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        partitioning: &mut PostInliningPartitioning<'tcx>,
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        mono_item_placements: FxHashMap<MonoItem<'tcx>, MonoItemPlacement>,
+        internalization_candidates: FxHashSet<MonoItem<'tcx>>,
     ) {
-        if partitioning.codegen_units.len() == 1 {
+        if codegen_units.len() == 1 {
             // Fast path for when there is only one codegen unit. In this case we
             // can internalize all candidates, since there is nowhere else they
             // could be accessed from.
-            for cgu in &mut partitioning.codegen_units {
-                for candidate in &partitioning.internalization_candidates {
+            for cgu in codegen_units {
+                for candidate in &internalization_candidates {
                     cgu.items_mut().insert(*candidate, (Linkage::Internal, Visibility::Default));
                 }
             }
@@ -317,15 +302,13 @@ impl<'tcx> Partition<'tcx> for DefaultPartitioning {
             }
         });
 
-        let mono_item_placements = &partitioning.mono_item_placements;
-
         // For each internalization candidates in each codegen unit, check if it is
         // accessed from outside its defining codegen unit.
-        for cgu in &mut partitioning.codegen_units {
+        for cgu in codegen_units {
             let home_cgu = MonoItemPlacement::SingleCgu { cgu_name: cgu.name() };
 
             for (accessee, linkage_and_visibility) in cgu.items_mut() {
-                if !partitioning.internalization_candidates.contains(accessee) {
+                if !internalization_candidates.contains(accessee) {
                     // This item is no candidate for internalizing, so skip it.
                     continue;
                 }
diff --git a/compiler/rustc_monomorphize/src/partitioning/mod.rs b/compiler/rustc_monomorphize/src/partitioning/mod.rs
@@ -128,7 +128,7 @@ impl<'tcx> Partition<'tcx> for Partitioner {
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
         mono_items: &mut I,
-    ) -> PreInliningPartitioning<'tcx>
+    ) -> (Vec<CodegenUnit<'tcx>>, FxHashSet<MonoItem<'tcx>>, FxHashSet<MonoItem<'tcx>>)
     where
         I: Iterator<Item = MonoItem<'tcx>>,
     {
@@ -141,24 +141,23 @@ impl<'tcx> Partition<'tcx> for Partitioner {
     fn merge_codegen_units(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: &mut PreInliningPartitioning<'tcx>,
+        codegen_units: &mut Vec<CodegenUnit<'tcx>>,
     ) {
         match self {
-            Partitioner::Default(partitioner) => {
-                partitioner.merge_codegen_units(cx, initial_partitioning)
-            }
+            Partitioner::Default(partitioner) => partitioner.merge_codegen_units(cx, codegen_units),
             Partitioner::Unknown => cx.tcx.sess.emit_fatal(UnknownPartitionStrategy),
         }
     }
 
     fn place_inlined_mono_items(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: PreInliningPartitioning<'tcx>,
-    ) -> PostInliningPartitioning<'tcx> {
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        roots: FxHashSet<MonoItem<'tcx>>,
+    ) -> FxHashMap<MonoItem<'tcx>, MonoItemPlacement> {
         match self {
             Partitioner::Default(partitioner) => {
-                partitioner.place_inlined_mono_items(cx, initial_partitioning)
+                partitioner.place_inlined_mono_items(cx, codegen_units, roots)
             }
             Partitioner::Unknown => cx.tcx.sess.emit_fatal(UnknownPartitionStrategy),
         }
@@ -167,12 +166,17 @@ impl<'tcx> Partition<'tcx> for Partitioner {
     fn internalize_symbols(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        post_inlining_partitioning: &mut PostInliningPartitioning<'tcx>,
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        mono_item_placements: FxHashMap<MonoItem<'tcx>, MonoItemPlacement>,
+        internalization_candidates: FxHashSet<MonoItem<'tcx>>,
     ) {
         match self {
-            Partitioner::Default(partitioner) => {
-                partitioner.internalize_symbols(cx, post_inlining_partitioning)
-            }
+            Partitioner::Default(partitioner) => partitioner.internalize_symbols(
+                cx,
+                codegen_units,
+                mono_item_placements,
+                internalization_candidates,
+            ),
             Partitioner::Unknown => cx.tcx.sess.emit_fatal(UnknownPartitionStrategy),
         }
     }
@@ -189,26 +193,29 @@ trait Partition<'tcx> {
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
         mono_items: &mut I,
-    ) -> PreInliningPartitioning<'tcx>
+    ) -> (Vec<CodegenUnit<'tcx>>, FxHashSet<MonoItem<'tcx>>, FxHashSet<MonoItem<'tcx>>)
     where
         I: Iterator<Item = MonoItem<'tcx>>;
 
     fn merge_codegen_units(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: &mut PreInliningPartitioning<'tcx>,
+        codegen_units: &mut Vec<CodegenUnit<'tcx>>,
     );
 
     fn place_inlined_mono_items(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        initial_partitioning: PreInliningPartitioning<'tcx>,
-    ) -> PostInliningPartitioning<'tcx>;
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        roots: FxHashSet<MonoItem<'tcx>>,
+    ) -> FxHashMap<MonoItem<'tcx>, MonoItemPlacement>;
 
     fn internalize_symbols(
         &mut self,
         cx: &PartitioningCx<'_, 'tcx>,
-        partitioning: &mut PostInliningPartitioning<'tcx>,
+        codegen_units: &mut [CodegenUnit<'tcx>],
+        mono_item_placements: FxHashMap<MonoItem<'tcx>, MonoItemPlacement>,
+        internalization_candidates: FxHashSet<MonoItem<'tcx>>,
     );
 }
 
@@ -240,52 +247,57 @@ where
     // 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 (mut codegen_units, roots, internalization_candidates) = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_roots");
         partitioner.place_root_mono_items(cx, mono_items)
     };
 
-    for cgu in &mut initial_partitioning.codegen_units {
+    for cgu in &mut codegen_units {
         cgu.create_size_estimate(tcx);
     }
 
-    debug_dump(tcx, "INITIAL PARTITIONING", &initial_partitioning.codegen_units);
+    debug_dump(tcx, "INITIAL PARTITIONING", &codegen_units);
 
     // 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(cx, &mut initial_partitioning);
-        debug_dump(tcx, "POST MERGING", &initial_partitioning.codegen_units);
+        partitioner.merge_codegen_units(cx, &mut codegen_units);
+        debug_dump(tcx, "POST MERGING", &codegen_units);
     }
 
     // In the next step, we use the inlining map to determine which additional
     // monomorphizations have to go into each codegen unit. These additional
     // monomorphizations can be drop-glue, functions from external crates, and
     // local functions the definition of which is marked with `#[inline]`.
-    let mut post_inlining = {
+    let mono_item_placements = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_inline_items");
-        partitioner.place_inlined_mono_items(cx, initial_partitioning)
+        partitioner.place_inlined_mono_items(cx, &mut codegen_units, roots)
     };
 
-    for cgu in &mut post_inlining.codegen_units {
+    for cgu in &mut codegen_units {
         cgu.create_size_estimate(tcx);
     }
 
-    debug_dump(tcx, "POST INLINING", &post_inlining.codegen_units);
+    debug_dump(tcx, "POST INLINING", &codegen_units);
 
     // Next we try to make as many symbols "internal" as possible, so LLVM has
     // more freedom to optimize.
     if !tcx.sess.link_dead_code() {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_internalize_symbols");
-        partitioner.internalize_symbols(cx, &mut post_inlining);
+        partitioner.internalize_symbols(
+            cx,
+            &mut codegen_units,
+            mono_item_placements,
+            internalization_candidates,
+        );
     }
 
     let instrument_dead_code =
         tcx.sess.instrument_coverage() && !tcx.sess.instrument_coverage_except_unused_functions();
 
     if instrument_dead_code {
         assert!(
-            post_inlining.codegen_units.len() > 0,
+            codegen_units.len() > 0,
             "There must be at least one CGU that code coverage data can be generated in."
         );
 
@@ -296,7 +308,7 @@ where
         // the object file (CGU) containing the dead function stubs is included
         // in the final binary. This will probably require forcing these
         // function symbols to be included via `-u` or `/include` linker args.
-        let mut cgus: Vec<_> = post_inlining.codegen_units.iter_mut().collect();
+        let mut cgus: Vec<_> = codegen_units.iter_mut().collect();
         cgus.sort_by_key(|cgu| cgu.size_estimate());
 
         let dead_code_cgu =
@@ -307,29 +319,17 @@ where
             } else {
                 // If there are no CGUs that have externally linked items,
                 // then we just pick the first CGU as a fallback.
-                &mut post_inlining.codegen_units[0]
+                &mut codegen_units[0]
             };
         dead_code_cgu.make_code_coverage_dead_code_cgu();
     }
 
     // Finally, sort by codegen unit name, so that we get deterministic results.
-    let PostInliningPartitioning {
-        codegen_units: mut result,
-        mono_item_placements: _,
-        internalization_candidates: _,
-    } = post_inlining;
+    codegen_units.sort_by(|a, b| a.name().as_str().cmp(b.name().as_str()));
 
-    result.sort_by(|a, b| a.name().as_str().cmp(b.name().as_str()));
+    debug_dump(tcx, "FINAL", &codegen_units);
 
-    debug_dump(tcx, "FINAL", &result);
-
-    result
-}
-
-pub struct PreInliningPartitioning<'tcx> {
-    codegen_units: Vec<CodegenUnit<'tcx>>,
-    roots: FxHashSet<MonoItem<'tcx>>,
-    internalization_candidates: FxHashSet<MonoItem<'tcx>>,
+    codegen_units
 }
 
 /// For symbol internalization, we need to know whether a symbol/mono-item is
@@ -341,12 +341,6 @@ enum MonoItemPlacement {
     MultipleCgus,
 }
 
-struct PostInliningPartitioning<'tcx> {
-    codegen_units: Vec<CodegenUnit<'tcx>>,
-    mono_item_placements: FxHashMap<MonoItem<'tcx>, MonoItemPlacement>,
-    internalization_candidates: FxHashSet<MonoItem<'tcx>>,
-}
-
 fn debug_dump<'a, 'tcx: 'a>(tcx: TyCtxt<'tcx>, label: &str, cgus: &[CodegenUnit<'tcx>]) {
     let dump = move || {
         use std::fmt::Write;
