From 9920baba5f459668b70131df64bb85bd82c3b149 Mon Sep 17 00:00:00 2001
From: Andrew Xie <ndrew.xie@gmail.com>
Date: Fri, 7 Apr 2023 19:41:08 -0400
Subject: [PATCH] Switched provisional evaluation cache map to FxIndexMap, and
 replaced map.drain_filter with map.retain

---
 .../src/traits/select/mod.rs                  | 24 +++++++------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs
index d2d75d2aa4c..6bb53418bea 100644
--- a/compiler/rustc_trait_selection/src/traits/select/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs
@@ -2,12 +2,6 @@
 //!
 //! [rustc dev guide]: https://rustc-dev-guide.rust-lang.org/traits/resolution.html#selection
 
-// FIXME: The `map` field in ProvisionalEvaluationCache should be changed to
-// a `FxIndexMap` to avoid query instability, but right now it causes a perf regression. This would be
-// fixed or at least lightened by the addition of the `drain_filter` method to `FxIndexMap`
-// Relevant: https://github.com/rust-lang/rust/pull/103723 and https://github.com/bluss/indexmap/issues/242
-#![allow(rustc::potential_query_instability)]
-
 use self::EvaluationResult::*;
 use self::SelectionCandidate::*;
 
@@ -32,8 +26,7 @@ use crate::traits::project::ProjectAndUnifyResult;
 use crate::traits::project::ProjectionCacheKeyExt;
 use crate::traits::ProjectionCacheKey;
 use crate::traits::Unimplemented;
-use rustc_data_structures::fx::FxHashMap;
-use rustc_data_structures::fx::{FxHashSet, FxIndexSet};
+use rustc_data_structures::fx::{FxHashSet, FxIndexMap, FxIndexSet};
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_errors::Diagnostic;
 use rustc_hir as hir;
@@ -2782,7 +2775,7 @@ struct ProvisionalEvaluationCache<'tcx> {
     /// - then we determine that `E` is in error -- we will then clear
     ///   all cache values whose DFN is >= 4 -- in this case, that
     ///   means the cached value for `F`.
-    map: RefCell<FxHashMap<ty::PolyTraitPredicate<'tcx>, ProvisionalEvaluation>>,
+    map: RefCell<FxIndexMap<ty::PolyTraitPredicate<'tcx>, ProvisionalEvaluation>>,
 
     /// The stack of args that we assume to be true because a `WF(arg)` predicate
     /// is on the stack above (and because of wellformedness is coinductive).
@@ -2930,12 +2923,13 @@ impl<'tcx> ProvisionalEvaluationCache<'tcx> {
     /// have a performance impact in practice.
     fn on_completion(&self, dfn: usize) {
         debug!(?dfn, "on_completion");
-
-        for (fresh_trait_pred, eval) in
-            self.map.borrow_mut().drain_filter(|_k, eval| eval.from_dfn >= dfn)
-        {
-            debug!(?fresh_trait_pred, ?eval, "on_completion");
-        }
+        self.map.borrow_mut().retain(|fresh_trait_pred, eval| {
+            if eval.from_dfn >= dfn {
+                debug!(?fresh_trait_pred, ?eval, "on_completion");
+                return false;
+            }
+            true
+        });
     }
 }