Auto merge of #96815 - SparrowLii:promote_const, r=oli-obk

optimize `promote_consts` by caching the results of `validate_local`

From the FIXME in the impl of `promote_consts`. Early return the `validate_local` should save some compile time.

`qualif_local` is similar to this, but requires futher changing because there are different types of qualif checks. If this PR is effective, I will do it as well.

Author: bors

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: Result<(), ()> },
     /// Any other combination of assignments/uses.
     Unpromotable,
     /// This temp was part of an rvalue which got extracted
@@ -133,7 +133,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: Err(()) };
                     return;
                 }
                 _ => { /* mark as unpromotable below */ }
@@ -188,7 +188,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 +202,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 +234,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,40 +272,50 @@ 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] {
+            valid.or_else(|_| {
+                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] = match ok {
+                    Ok(()) => TempState::Defined { location: loc, uses, valid: Ok(()) },
+                    Err(_) => TempState::Unpromotable,
+                };
+                ok
+            })
         } else {
             Err(Unpromotable)
         }
     }
 
-    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 +427,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 +457,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 +490,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 +633,7 @@ impl<'tcx> Validator<'_, 'tcx> {
     }
 
     fn validate_call(
-        &self,
+        &mut self,
         callee: &Operand<'tcx>,
         args: &[Operand<'tcx>],
     ) -> Result<(), Unpromotable> {
@@ -665,10 +675,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 +730,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;
                 }