From 2be012a0c664bf1ee90af0e5cac9a3e3e24f4666 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 8 May 2022 00:00:00 +0000
Subject: [PATCH] Use sparse representation of switch sources

to avoid quadratic space overhead
---
 compiler/rustc_middle/src/mir/mod.rs                 |  4 ++--
 compiler/rustc_middle/src/mir/switch_sources.rs      | 12 +++++-------
 .../rustc_mir_dataflow/src/framework/direction.rs    |  2 +-
 3 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index 4d26840fd62..2c97c7704e7 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -580,8 +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.
+    /// `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)
diff --git a/compiler/rustc_middle/src/mir/switch_sources.rs b/compiler/rustc_middle/src/mir/switch_sources.rs
index 7f62b4d0dba..adeeec70d1c 100644
--- a/compiler/rustc_middle/src/mir/switch_sources.rs
+++ b/compiler/rustc_middle/src/mir/switch_sources.rs
@@ -2,6 +2,7 @@
 //! `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;
@@ -9,7 +10,7 @@ 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);
                 }
             }
 
diff --git a/compiler/rustc_mir_dataflow/src/framework/direction.rs b/compiler/rustc_mir_dataflow/src/framework/direction.rs
index 327002219db..18795128b85 100644
--- a/compiler/rustc_mir_dataflow/src/framework/direction.rs
+++ b/compiler/rustc_mir_dataflow/src/framework/direction.rs
@@ -322,7 +322,7 @@ where
     fn apply(&mut self, mut apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget)) {
         assert!(!self.effects_applied);
 
-        let values = &self.body.switch_sources()[self.bb][self.pred];
+        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;