No need to collect test variants ahead of time

Nadrieril · Nadrieril · commit edea739292f6 · 2024-03-02T18:35:53.000+01:00
diff --git a/compiler/rustc_mir_build/src/build/matches/mod.rs b/compiler/rustc_mir_build/src/build/matches/mod.rs
@@ -14,7 +14,6 @@ use rustc_data_structures::{
     fx::{FxHashSet, FxIndexMap, FxIndexSet},
     stack::ensure_sufficient_stack,
 };
-use rustc_index::bit_set::BitSet;
 use rustc_middle::middle::region;
 use rustc_middle::mir::{self, *};
 use rustc_middle::thir::{self, *};
@@ -1116,19 +1115,10 @@ enum TestKind<'tcx> {
     Switch {
         /// The enum type being tested.
         adt_def: ty::AdtDef<'tcx>,
-        /// The set of variants that we should create a branch for. We also
-        /// create an additional "otherwise" case.
-        variants: BitSet<VariantIdx>,
     },
 
     /// Test what value an integer or `char` has.
-    SwitchInt {
-        /// The (ordered) set of values that we test for.
-        ///
-        /// We create a branch to each of the values in `options`, as well as an "otherwise" branch
-        /// for all other values, even in the (rare) case that `options` is exhaustive.
-        options: FxIndexMap<Const<'tcx>, u128>,
-    },
+    SwitchInt,
 
     /// Test what value a `bool` has.
     If,
@@ -1173,6 +1163,12 @@ enum TestBranch<'tcx> {
     Failure,
 }
 
+impl<'tcx> TestBranch<'tcx> {
+    fn as_constant(&self) -> Option<&Const<'tcx>> {
+        if let Self::Constant(v, _) = self { Some(v) } else { None }
+    }
+}
+
 /// `ArmHasGuard` is a wrapper around a boolean flag. It indicates whether
 /// a match arm has a guard expression attached to it.
 #[derive(Copy, Clone, Debug)]
@@ -1583,30 +1579,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     ) -> (PlaceBuilder<'tcx>, Test<'tcx>) {
         // Extract the match-pair from the highest priority candidate
         let match_pair = &candidates.first().unwrap().match_pairs[0];
-        let mut test = self.test(match_pair);
+        let test = self.test(match_pair);
         let match_place = match_pair.place.clone();
-
         debug!(?test, ?match_pair);
-        // Most of the time, the test to perform is simply a function of the main candidate; but for
-        // a test like SwitchInt, we may want to add cases based on the candidates that are
-        // available
-        match test.kind {
-            TestKind::SwitchInt { ref mut options } => {
-                for candidate in candidates.iter() {
-                    if !self.add_cases_to_switch(&match_place, candidate, options) {
-                        break;
-                    }
-                }
-            }
-            TestKind::Switch { adt_def: _, ref mut variants } => {
-                for candidate in candidates.iter() {
-                    if !self.add_variants_to_switch(&match_place, candidate, variants) {
-                        break;
-                    }
-                }
-            }
-            _ => {}
-        }
 
         (match_place, test)
     }
@@ -1663,7 +1638,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         // point we may encounter a candidate where the test is not relevant; at that point, we stop
         // sorting.
         while let Some(candidate) = candidates.first_mut() {
-            let Some(branch) = self.sort_candidate(&match_place, &test, candidate) else {
+            let Some(branch) =
+                self.sort_candidate(&match_place, test, candidate, &target_candidates)
+            else {
                 break;
             };
             let (candidate, rest) = candidates.split_first_mut().unwrap();
diff --git a/compiler/rustc_mir_build/src/build/matches/test.rs b/compiler/rustc_mir_build/src/build/matches/test.rs
@@ -10,17 +10,14 @@ use crate::build::matches::{Candidate, MatchPair, Test, TestBranch, TestCase, Te
 use crate::build::Builder;
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir::{LangItem, RangeEnd};
-use rustc_index::bit_set::BitSet;
 use rustc_middle::mir::*;
-use rustc_middle::thir::*;
 use rustc_middle::ty::util::IntTypeExt;
 use rustc_middle::ty::GenericArg;
 use rustc_middle::ty::{self, adjustment::PointerCoercion, Ty, TyCtxt};
 use rustc_span::def_id::DefId;
 use rustc_span::source_map::Spanned;
 use rustc_span::symbol::{sym, Symbol};
 use rustc_span::{Span, DUMMY_SP};
-use rustc_target::abi::VariantIdx;
 
 use std::cmp::Ordering;
 
@@ -30,22 +27,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// It is a bug to call this with a not-fully-simplified pattern.
     pub(super) fn test<'pat>(&mut self, match_pair: &MatchPair<'pat, 'tcx>) -> Test<'tcx> {
         let kind = match match_pair.test_case {
-            TestCase::Variant { adt_def, variant_index: _ } => {
-                TestKind::Switch { adt_def, variants: BitSet::new_empty(adt_def.variants().len()) }
-            }
+            TestCase::Variant { adt_def, variant_index: _ } => TestKind::Switch { adt_def },
 
             TestCase::Constant { .. } if match_pair.pattern.ty.is_bool() => TestKind::If,
-
-            TestCase::Constant { .. } if is_switch_ty(match_pair.pattern.ty) => {
-                // For integers, we use a `SwitchInt` match, which allows
-                // us to handle more cases.
-                TestKind::SwitchInt {
-                    // these maps are empty to start; cases are
-                    // added below in add_cases_to_switch
-                    options: Default::default(),
-                }
-            }
-
+            TestCase::Constant { .. } if is_switch_ty(match_pair.pattern.ty) => TestKind::SwitchInt,
             TestCase::Constant { value } => TestKind::Eq { value, ty: match_pair.pattern.ty },
 
             TestCase::Range(range) => {
@@ -70,57 +55,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         Test { span: match_pair.pattern.span, kind }
     }
 
-    pub(super) fn add_cases_to_switch<'pat>(
-        &mut self,
-        test_place: &PlaceBuilder<'tcx>,
-        candidate: &Candidate<'pat, 'tcx>,
-        options: &mut FxIndexMap<Const<'tcx>, u128>,
-    ) -> bool {
-        let Some(match_pair) = candidate.match_pairs.iter().find(|mp| mp.place == *test_place)
-        else {
-            return false;
-        };
-
-        match match_pair.test_case {
-            TestCase::Constant { value } => {
-                options.entry(value).or_insert_with(|| value.eval_bits(self.tcx, self.param_env));
-                true
-            }
-            TestCase::Variant { .. } => {
-                panic!("you should have called add_variants_to_switch instead!");
-            }
-            TestCase::Range(ref range) => {
-                // Check that none of the switch values are in the range.
-                self.values_not_contained_in_range(&*range, options).unwrap_or(false)
-            }
-            // don't know how to add these patterns to a switch
-            _ => false,
-        }
-    }
-
-    pub(super) fn add_variants_to_switch<'pat>(
-        &mut self,
-        test_place: &PlaceBuilder<'tcx>,
-        candidate: &Candidate<'pat, 'tcx>,
-        variants: &mut BitSet<VariantIdx>,
-    ) -> bool {
-        let Some(match_pair) = candidate.match_pairs.iter().find(|mp| mp.place == *test_place)
-        else {
-            return false;
-        };
-
-        match match_pair.test_case {
-            TestCase::Variant { variant_index, .. } => {
-                // We have a pattern testing for variant `variant_index`
-                // set the corresponding index to true
-                variants.insert(variant_index);
-                true
-            }
-            // don't know how to add these patterns to a switch
-            _ => false,
-        }
-    }
-
     #[instrument(skip(self, target_blocks, place_builder), level = "debug")]
     pub(super) fn perform_test(
         &mut self,
@@ -139,33 +73,20 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         let source_info = self.source_info(test.span);
         match test.kind {
-            TestKind::Switch { adt_def, ref variants } => {
-                // Variants is a BitVec of indexes into adt_def.variants.
-                let num_enum_variants = adt_def.variants().len();
+            TestKind::Switch { adt_def } => {
                 let otherwise_block = target_block(TestBranch::Failure);
-                let tcx = self.tcx;
                 let switch_targets = SwitchTargets::new(
-                    adt_def.discriminants(tcx).filter_map(|(idx, discr)| {
-                        if variants.contains(idx) {
-                            debug_assert_ne!(
-                                target_block(TestBranch::Variant(idx)),
-                                otherwise_block,
-                                "no candidates for tested discriminant: {discr:?}",
-                            );
-                            Some((discr.val, target_block(TestBranch::Variant(idx))))
+                    adt_def.discriminants(self.tcx).filter_map(|(idx, discr)| {
+                        if let Some(&block) = target_blocks.get(&TestBranch::Variant(idx)) {
+                            Some((discr.val, block))
                         } else {
-                            debug_assert_eq!(
-                                target_block(TestBranch::Variant(idx)),
-                                otherwise_block,
-                                "found candidates for untested discriminant: {discr:?}",
-                            );
                             None
                         }
                     }),
                     otherwise_block,
                 );
-                debug!("num_enum_variants: {}, variants: {:?}", num_enum_variants, variants);
-                let discr_ty = adt_def.repr().discr_type().to_ty(tcx);
+                debug!("num_enum_variants: {}", adt_def.variants().len());
+                let discr_ty = adt_def.repr().discr_type().to_ty(self.tcx);
                 let discr = self.temp(discr_ty, test.span);
                 self.cfg.push_assign(
                     block,
@@ -183,13 +104,17 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 );
             }
 
-            TestKind::SwitchInt { ref options } => {
+            TestKind::SwitchInt => {
                 // The switch may be inexhaustive so we have a catch-all block
                 let otherwise_block = target_block(TestBranch::Failure);
                 let switch_targets = SwitchTargets::new(
-                    options
-                        .iter()
-                        .map(|(&val, &bits)| (bits, target_block(TestBranch::Constant(val, bits)))),
+                    target_blocks.iter().filter_map(|(&branch, &block)| {
+                        if let TestBranch::Constant(_, bits) = branch {
+                            Some((bits, block))
+                        } else {
+                            None
+                        }
+                    }),
                     otherwise_block,
                 );
                 let terminator = TerminatorKind::SwitchInt {
@@ -548,11 +473,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// that it *doesn't* apply. For now, we return false, indicate that the
     /// test does not apply to this candidate, but it might be we can get
     /// tighter match code if we do something a bit different.
-    pub(super) fn sort_candidate<'pat>(
+    pub(super) fn sort_candidate(
         &mut self,
         test_place: &PlaceBuilder<'tcx>,
         test: &Test<'tcx>,
-        candidate: &mut Candidate<'pat, 'tcx>,
+        candidate: &mut Candidate<'_, 'tcx>,
+        sorted_candidates: &FxIndexMap<TestBranch<'tcx>, Vec<&mut Candidate<'_, 'tcx>>>,
     ) -> Option<TestBranch<'tcx>> {
         // Find the match_pair for this place (if any). At present,
         // afaik, there can be at most one. (In the future, if we
@@ -568,7 +494,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             // If we are performing a variant switch, then this
             // informs variant patterns, but nothing else.
             (
-                &TestKind::Switch { adt_def: tested_adt_def, .. },
+                &TestKind::Switch { adt_def: tested_adt_def },
                 &TestCase::Variant { adt_def, variant_index },
             ) => {
                 assert_eq!(adt_def, tested_adt_def);
@@ -581,17 +507,19 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             //
             // FIXME(#29623) we could use PatKind::Range to rule
             // things out here, in some cases.
-            (TestKind::SwitchInt { options }, &TestCase::Constant { value })
+            (TestKind::SwitchInt, &TestCase::Constant { value })
                 if is_switch_ty(match_pair.pattern.ty) =>
             {
                 fully_matched = true;
-                let bits = options.get(&value).unwrap();
-                Some(TestBranch::Constant(value, *bits))
+                let bits = value.eval_bits(self.tcx, self.param_env);
+                Some(TestBranch::Constant(value, bits))
             }
-            (TestKind::SwitchInt { options }, TestCase::Range(range)) => {
+            (TestKind::SwitchInt, TestCase::Range(range)) => {
                 fully_matched = false;
                 let not_contained =
-                    self.values_not_contained_in_range(&*range, options).unwrap_or(false);
+                    sorted_candidates.keys().filter_map(|br| br.as_constant()).copied().all(
+                        |val| matches!(range.contains(val, self.tcx, self.param_env), Some(false)),
+                    );
 
                 not_contained.then(|| {
                     // No switch values are contained in the pattern range,
@@ -732,20 +660,6 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
         ret
     }
-
-    fn values_not_contained_in_range(
-        &self,
-        range: &PatRange<'tcx>,
-        options: &FxIndexMap<Const<'tcx>, u128>,
-    ) -> Option<bool> {
-        for &val in options.keys() {
-            if range.contains(val, self.tcx, self.param_env)? {
-                return Some(false);
-            }
-        }
-
-        Some(true)
-    }
 }
 
 fn is_switch_ty(ty: Ty<'_>) -> bool {
