optimize `promote_consts` by cache the validate check

1 files changed, 70 insertions, 41 deletions

diff --git a/compiler/rustc_const_eval/src/transform/promote_consts.rs b/compiler/rustc_const_eval/src/transform/promote_consts.rs
index 1052d588fad..daba60fca38 100644
--- a/compiler/rustc_const_eval/src/transform/promote_consts.rs
+++ b/compiler/rustc_const_eval/src/transform/promote_consts.rs
@@ -60,9 +60,9 @@ impl<'tcx> MirPass<'tcx> for PromoteTemps<'tcx> {
 
         let mut rpo = traversal::reverse_postorder(body);
         let ccx = ConstCx::new(tcx, body);
-        let (temps, all_candidates) = collect_temps_and_candidates(&ccx, &mut rpo);
+        let (mut temps, all_candidates) = collect_temps_and_candidates(&ccx, &mut rpo);
 
-        let promotable_candidates = validate_candidates(&ccx, &temps, &all_candidates);
+        let promotable_candidates = validate_candidates(&ccx, &mut temps, &all_candidates);
 
         let promoted = promote_candidates(body, tcx, temps, promotable_candidates);
         self.promoted_fragments.set(promoted);
@@ -77,7 +77,7 @@ pub enum TempState {
     /// One direct assignment and any number of direct uses.
     /// A borrow of this temp is promotable if the assigned
     /// value is qualified as constant.
-    Defined { location: Location, uses: usize },
+    Defined { location: Location, uses: usize, valid: Valid },
     /// Any other combination of assignments/uses.
     Unpromotable,
     /// This temp was part of an rvalue which got extracted
@@ -85,6 +85,13 @@ pub enum TempState {
     PromotedOut,
 }
 
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub enum Valid {
+    Unknown,
+    InValid,
+    Validated,
+}
+
 impl TempState {
     pub fn is_promotable(&self) -> bool {
         debug!("is_promotable: self={:?}", self);
@@ -133,7 +140,7 @@ impl<'tcx> Visitor<'tcx> for Collector<'_, 'tcx> {
             match context {
                 PlaceContext::MutatingUse(MutatingUseContext::Store)
                 | PlaceContext::MutatingUse(MutatingUseContext::Call) => {
-                    *temp = TempState::Defined { location, uses: 0 };
+                    *temp = TempState::Defined { location, uses: 0, valid: Valid::Unknown };
                     return;
                 }
                 _ => { /* mark as unpromotable below */ }
@@ -188,7 +195,7 @@ pub fn collect_temps_and_candidates<'tcx>(
 /// This wraps an `Item`, and has access to all fields of that `Item` via `Deref` coercion.
 struct Validator<'a, 'tcx> {
     ccx: &'a ConstCx<'a, 'tcx>,
-    temps: &'a IndexVec<Local, TempState>,
+    temps: &'a mut IndexVec<Local, TempState>,
 }
 
 impl<'a, 'tcx> std::ops::Deref for Validator<'a, 'tcx> {
@@ -202,7 +209,7 @@ impl<'a, 'tcx> std::ops::Deref for Validator<'a, 'tcx> {
 struct Unpromotable;
 
 impl<'tcx> Validator<'_, 'tcx> {
-    fn validate_candidate(&self, candidate: Candidate) -> Result<(), Unpromotable> {
+    fn validate_candidate(&mut self, candidate: Candidate) -> Result<(), Unpromotable> {
         let loc = candidate.location;
         let statement = &self.body[loc.block].statements[loc.statement_index];
         match &statement.kind {
@@ -234,7 +241,7 @@ impl<'tcx> Validator<'_, 'tcx> {
     }
 
     // FIXME(eddyb) maybe cache this?
-    fn qualif_local<Q: qualifs::Qualif>(&self, local: Local) -> bool {
+    fn qualif_local<Q: qualifs::Qualif>(&mut self, local: Local) -> bool {
         if let TempState::Defined { location: loc, .. } = self.temps[local] {
             let num_stmts = self.body[loc.block].statements.len();
 
@@ -272,32 +279,54 @@ impl<'tcx> Validator<'_, 'tcx> {
         }
     }
 
-    // FIXME(eddyb) maybe cache this?
-    fn validate_local(&self, local: Local) -> Result<(), Unpromotable> {
-        if let TempState::Defined { location: loc, .. } = self.temps[local] {
-            let block = &self.body[loc.block];
-            let num_stmts = block.statements.len();
-
-            if loc.statement_index < num_stmts {
-                let statement = &block.statements[loc.statement_index];
-                match &statement.kind {
-                    StatementKind::Assign(box (_, rhs)) => self.validate_rvalue(rhs),
-                    _ => {
-                        span_bug!(
-                            statement.source_info.span,
-                            "{:?} is not an assignment",
-                            statement
-                        );
-                    }
-                }
-            } else {
-                let terminator = block.terminator();
-                match &terminator.kind {
-                    TerminatorKind::Call { func, args, .. } => self.validate_call(func, args),
-                    TerminatorKind::Yield { .. } => Err(Unpromotable),
-                    kind => {
-                        span_bug!(terminator.source_info.span, "{:?} not promotable", kind);
-                    }
+    fn validate_local(&mut self, local: Local) -> Result<(), Unpromotable> {
+        if let TempState::Defined { location: loc, uses, valid } = self.temps[local] {
+            match valid {
+                Valid::InValid => Err(Unpromotable),
+                Valid::Validated => Ok(()),
+                Valid::Unknown => {
+                    let ok = {
+                        let block = &self.body[loc.block];
+                        let num_stmts = block.statements.len();
+
+                        if loc.statement_index < num_stmts {
+                            let statement = &block.statements[loc.statement_index];
+                            match &statement.kind {
+                                StatementKind::Assign(box (_, rhs)) => self.validate_rvalue(rhs),
+                                _ => {
+                                    span_bug!(
+                                        statement.source_info.span,
+                                        "{:?} is not an assignment",
+                                        statement
+                                    );
+                                }
+                            }
+                        } else {
+                            let terminator = block.terminator();
+                            match &terminator.kind {
+                                TerminatorKind::Call { func, args, .. } => {
+                                    self.validate_call(func, args)
+                                }
+                                TerminatorKind::Yield { .. } => Err(Unpromotable),
+                                kind => {
+                                    span_bug!(
+                                        terminator.source_info.span,
+                                        "{:?} not promotable",
+                                        kind
+                                    );
+                                }
+                            }
+                        }
+                    };
+                    self.temps[local] = TempState::Defined {
+                        location: loc,
+                        uses,
+                        valid: match ok {
+                            Ok(()) => Valid::Validated,
+                            Err(_) => Valid::InValid,
+                        },
+                    };
+                    ok
                 }
             }
         } else {
@@ -305,7 +334,7 @@ impl<'tcx> Validator<'_, 'tcx> {
         }
     }
 
-    fn validate_place(&self, place: PlaceRef<'tcx>) -> Result<(), Unpromotable> {
+    fn validate_place(&mut self, place: PlaceRef<'tcx>) -> Result<(), Unpromotable> {
         match place.last_projection() {
             None => self.validate_local(place.local),
             Some((place_base, elem)) => {
@@ -417,7 +446,7 @@ impl<'tcx> Validator<'_, 'tcx> {
         }
     }
 
-    fn validate_operand(&self, operand: &Operand<'tcx>) -> Result<(), Unpromotable> {
+    fn validate_operand(&mut self, operand: &Operand<'tcx>) -> Result<(), Unpromotable> {
         match operand {
             Operand::Copy(place) | Operand::Move(place) => self.validate_place(place.as_ref()),
 
@@ -447,7 +476,7 @@ impl<'tcx> Validator<'_, 'tcx> {
         }
     }
 
-    fn validate_ref(&self, kind: BorrowKind, place: &Place<'tcx>) -> Result<(), Unpromotable> {
+    fn validate_ref(&mut self, kind: BorrowKind, place: &Place<'tcx>) -> Result<(), Unpromotable> {
         match kind {
             // Reject these borrow types just to be safe.
             // FIXME(RalfJung): could we allow them? Should we? No point in it until we have a usecase.
@@ -480,7 +509,7 @@ impl<'tcx> Validator<'_, 'tcx> {
         Ok(())
     }
 
-    fn validate_rvalue(&self, rvalue: &Rvalue<'tcx>) -> Result<(), Unpromotable> {
+    fn validate_rvalue(&mut self, rvalue: &Rvalue<'tcx>) -> Result<(), Unpromotable> {
         match rvalue {
             Rvalue::Use(operand) | Rvalue::Repeat(operand, _) => {
                 self.validate_operand(operand)?;
@@ -623,7 +652,7 @@ impl<'tcx> Validator<'_, 'tcx> {
     }
 
     fn validate_call(
-        &self,
+        &mut self,
         callee: &Operand<'tcx>,
         args: &[Operand<'tcx>],
     ) -> Result<(), Unpromotable> {
@@ -665,10 +694,10 @@ impl<'tcx> Validator<'_, 'tcx> {
 // FIXME(eddyb) remove the differences for promotability in `static`, `const`, `const fn`.
 pub fn validate_candidates(
     ccx: &ConstCx<'_, '_>,
-    temps: &IndexVec<Local, TempState>,
+    temps: &mut IndexVec<Local, TempState>,
     candidates: &[Candidate],
 ) -> Vec<Candidate> {
-    let validator = Validator { ccx, temps };
+    let mut validator = Validator { ccx, temps };
 
     candidates
         .iter()
@@ -720,7 +749,7 @@ impl<'a, 'tcx> Promoter<'a, 'tcx> {
     fn promote_temp(&mut self, temp: Local) -> Local {
         let old_keep_original = self.keep_original;
         let loc = match self.temps[temp] {
-            TempState::Defined { location, uses } if uses > 0 => {
+            TempState::Defined { location, uses, .. } if uses > 0 => {
                 if uses > 1 {
                     self.keep_original = true;
                 }