Auto merge of #96838 - tmiasko:lazy-switch-sources, r=oli-obk

Optimize switch sources representation and usage

* Avoid constructing switch sources unless necessary - switch sources are used by backward analysis with a custom switch int edge effects, but are otherwise unnecessarily computed.
* Use sparse representation of switch sources to avoid quadratic space overhead.

Author: bors

compiler/rustc_middle/src/mir/mod.rs

@@ -580,6 +580,8 @@ impl<'tcx> Body<'tcx> {
         self.predecessor_cache.compute(&self.basic_blocks)
     }
 
+    /// `body.switch_sources()[&(target, switch)]` returns a list of switch
+    /// values that lead to a `target` block from a `switch` block.
     #[inline]
     pub fn switch_sources(&self) -> &SwitchSources {
         self.switch_source_cache.compute(&self.basic_blocks)

compiler/rustc_middle/src/mir/switch_sources.rs

@@ -2,14 +2,15 @@
 //! `Predecessors`/`PredecessorCache`.
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
+use rustc_data_structures::stable_map::FxHashMap;
 use rustc_data_structures::sync::OnceCell;
 use rustc_index::vec::IndexVec;
 use rustc_serialize as serialize;
 use smallvec::SmallVec;
 
 use crate::mir::{BasicBlock, BasicBlockData, Terminator, TerminatorKind};
 
-pub type SwitchSources = IndexVec<BasicBlock, IndexVec<BasicBlock, SmallVec<[Option<u128>; 1]>>>;
+pub type SwitchSources = FxHashMap<(BasicBlock, BasicBlock), SmallVec<[Option<u128>; 1]>>;
 
 #[derive(Clone, Debug)]
 pub(super) struct SwitchSourceCache {
@@ -35,19 +36,16 @@ impl SwitchSourceCache {
         basic_blocks: &IndexVec<BasicBlock, BasicBlockData<'_>>,
     ) -> &SwitchSources {
         self.cache.get_or_init(|| {
-            let mut switch_sources = IndexVec::from_elem(
-                IndexVec::from_elem(SmallVec::new(), basic_blocks),
-                basic_blocks,
-            );
+            let mut switch_sources: SwitchSources = FxHashMap::default();
             for (bb, data) in basic_blocks.iter_enumerated() {
                 if let Some(Terminator {
                     kind: TerminatorKind::SwitchInt { targets, .. }, ..
                 }) = &data.terminator
                 {
                     for (value, target) in targets.iter() {
-                        switch_sources[target][bb].push(Some(value));
+                        switch_sources.entry((target, bb)).or_default().push(Some(value));
                     }
-                    switch_sources[targets.otherwise()][bb].push(None);
+                    switch_sources.entry((targets.otherwise(), bb)).or_default().push(None);
                 }
             }
 

compiler/rustc_mir_dataflow/src/framework/direction.rs

@@ -269,9 +269,9 @@ impl Direction for Backward {
 
                 mir::TerminatorKind::SwitchInt { targets: _, ref discr, switch_ty: _ } => {
                     let mut applier = BackwardSwitchIntEdgeEffectsApplier {
+                        body,
                         pred,
                         exit_state,
-                        values: &body.switch_sources()[bb][pred],
                         bb,
                         propagate: &mut propagate,
                         effects_applied: false,
@@ -305,9 +305,9 @@ impl Direction for Backward {
 }
 
 struct BackwardSwitchIntEdgeEffectsApplier<'a, D, F> {
+    body: &'a mir::Body<'a>,
     pred: BasicBlock,
     exit_state: &'a mut D,
-    values: &'a [Option<u128>],
     bb: BasicBlock,
     propagate: &'a mut F,
 
@@ -322,7 +322,8 @@ where
     fn apply(&mut self, mut apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget)) {
         assert!(!self.effects_applied);
 
-        let targets = self.values.iter().map(|&value| SwitchIntTarget { value, target: self.bb });
+        let values = &self.body.switch_sources()[&(self.bb, self.pred)];
+        let targets = values.iter().map(|&value| SwitchIntTarget { value, target: self.bb });
 
         let mut tmp = None;
         for target in targets {