Rollup merge of #140607 - lcnr:opaque-type-storage, r=compiler-errors

support duplicate entries in the opaque_type_storage

Necessary for the new solver as we may unify keys when eagerly resolving for canonical queries. See the relevant comment when instantiating query responses:
```rust
            // We eagerly resolve inference variables when computing the query response.
            // This can cause previously distinct opaque type keys to now be structurally equal.
            //
            // To handle this, we store any duplicate entries in a separate list to check them
            // at the end of typeck/borrowck. We could alternatively eagerly equate the hidden
            // types here. However, doing so is difficult as it may result in nested goals and
            // any errors may make it harder to track the control flow for diagnostics.
            if let Some(prev) = prev {
                self.delegate.add_duplicate_opaque_type(key, prev, self.origin_span);
            }
```

This will be far more relevant with #140497.

r? `@compiler-errors`

5 files changed, 83 insertions, 48 deletions

diff --git a/compiler/rustc_infer/src/infer/canonical/query_response.rs b/compiler/rustc_infer/src/infer/canonical/query_response.rs
index 5220071c500..1ae864c454f 100644
--- a/compiler/rustc_infer/src/infer/canonical/query_response.rs
+++ b/compiler/rustc_infer/src/infer/canonical/query_response.rs
@@ -132,7 +132,13 @@ impl<'tcx> InferCtxt<'tcx> {
 
         let certainty = if errors.is_empty() { Certainty::Proven } else { Certainty::Ambiguous };
 
-        let opaque_types = self.take_opaque_types_for_query_response();
+        let opaque_types = self
+            .inner
+            .borrow_mut()
+            .opaque_type_storage
+            .take_opaque_types()
+            .map(|(k, v)| (k, v.ty))
+            .collect();
 
         Ok(QueryResponse {
             var_values: inference_vars,
@@ -143,24 +149,6 @@ impl<'tcx> InferCtxt<'tcx> {
         })
     }
 
-    /// Used by the new solver as that one takes the opaque types at the end of a probe
-    /// to deal with multiple candidates without having to recompute them.
-    pub fn clone_opaque_types_for_query_response(
-        &self,
-    ) -> Vec<(ty::OpaqueTypeKey<'tcx>, Ty<'tcx>)> {
-        self.inner
-            .borrow()
-            .opaque_type_storage
-            .opaque_types
-            .iter()
-            .map(|(k, v)| (*k, v.ty))
-            .collect()
-    }
-
-    fn take_opaque_types_for_query_response(&self) -> Vec<(ty::OpaqueTypeKey<'tcx>, Ty<'tcx>)> {
-        self.take_opaque_types().into_iter().map(|(k, v)| (k, v.ty)).collect()
-    }
-
     /// Given the (canonicalized) result to a canonical query,
     /// instantiates the result so it can be used, plugging in the
     /// values from the canonical query. (Note that the result may
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index d25542dadd5..070d285b5a6 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -31,9 +31,9 @@ use rustc_middle::traits::solve::Goal;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
 use rustc_middle::ty::{
     self, BoundVarReplacerDelegate, ConstVid, FloatVid, GenericArg, GenericArgKind, GenericArgs,
-    GenericArgsRef, GenericParamDefKind, InferConst, IntVid, PseudoCanonicalInput, Term, TermKind,
-    Ty, TyCtxt, TyVid, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitable,
-    TypeVisitableExt, TypingEnv, TypingMode, fold_regions,
+    GenericArgsRef, GenericParamDefKind, InferConst, IntVid, OpaqueHiddenType, OpaqueTypeKey,
+    PseudoCanonicalInput, Term, TermKind, Ty, TyCtxt, TyVid, TypeFoldable, TypeFolder,
+    TypeSuperFoldable, TypeVisitable, TypeVisitableExt, TypingEnv, TypingMode, fold_regions,
 };
 use rustc_span::{Span, Symbol};
 use snapshot::undo_log::InferCtxtUndoLogs;
@@ -198,7 +198,7 @@ impl<'tcx> InferCtxtInner<'tcx> {
     }
 
     #[inline]
-    fn opaque_types(&mut self) -> opaque_types::OpaqueTypeTable<'_, 'tcx> {
+    pub fn opaque_types(&mut self) -> opaque_types::OpaqueTypeTable<'_, 'tcx> {
         self.opaque_type_storage.with_log(&mut self.undo_log)
     }
 
@@ -224,15 +224,6 @@ impl<'tcx> InferCtxtInner<'tcx> {
             .expect("region constraints already solved")
             .with_log(&mut self.undo_log)
     }
-
-    // Iterates through the opaque type definitions without taking them; this holds the
-    // `InferCtxtInner` lock, so make sure to not do anything with `InferCtxt` side-effects
-    // while looping through this.
-    pub fn iter_opaque_types(
-        &self,
-    ) -> impl Iterator<Item = (ty::OpaqueTypeKey<'tcx>, ty::OpaqueHiddenType<'tcx>)> {
-        self.opaque_type_storage.opaque_types.iter().map(|(&k, &v)| (k, v))
-    }
 }
 
 pub struct InferCtxt<'tcx> {
@@ -954,13 +945,13 @@ impl<'tcx> InferCtxt<'tcx> {
     }
 
     #[instrument(level = "debug", skip(self), ret)]
-    pub fn take_opaque_types(&self) -> opaque_types::OpaqueTypeMap<'tcx> {
-        std::mem::take(&mut self.inner.borrow_mut().opaque_type_storage.opaque_types)
+    pub fn take_opaque_types(&self) -> Vec<(OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        self.inner.borrow_mut().opaque_type_storage.take_opaque_types().collect()
     }
 
     #[instrument(level = "debug", skip(self), ret)]
-    pub fn clone_opaque_types(&self) -> opaque_types::OpaqueTypeMap<'tcx> {
-        self.inner.borrow().opaque_type_storage.opaque_types.clone()
+    pub fn clone_opaque_types(&self) -> Vec<(OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        self.inner.borrow_mut().opaque_type_storage.iter_opaque_types().collect()
     }
 
     #[inline(always)]
diff --git a/compiler/rustc_infer/src/infer/opaque_types/mod.rs b/compiler/rustc_infer/src/infer/opaque_types/mod.rs
index ce5d2e6e17a..df7144c31da 100644
--- a/compiler/rustc_infer/src/infer/opaque_types/mod.rs
+++ b/compiler/rustc_infer/src/infer/opaque_types/mod.rs
@@ -1,5 +1,4 @@
 use hir::def_id::{DefId, LocalDefId};
-use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir as hir;
 use rustc_middle::bug;
 use rustc_middle::traits::ObligationCause;
@@ -19,7 +18,6 @@ use crate::traits::{self, Obligation, PredicateObligations};
 
 mod table;
 
-pub(crate) type OpaqueTypeMap<'tcx> = FxIndexMap<OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>>;
 pub(crate) use table::{OpaqueTypeStorage, OpaqueTypeTable};
 
 impl<'tcx> InferCtxt<'tcx> {
diff --git a/compiler/rustc_infer/src/infer/opaque_types/table.rs b/compiler/rustc_infer/src/infer/opaque_types/table.rs
index ba6cc0d783d..3c5bf9d722b 100644
--- a/compiler/rustc_infer/src/infer/opaque_types/table.rs
+++ b/compiler/rustc_infer/src/infer/opaque_types/table.rs
@@ -1,18 +1,17 @@
+use std::ops::Deref;
+
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::undo_log::UndoLogs;
 use rustc_middle::bug;
 use rustc_middle::ty::{self, OpaqueHiddenType, OpaqueTypeKey, Ty};
 use tracing::instrument;
 
-use super::OpaqueTypeMap;
 use crate::infer::snapshot::undo_log::{InferCtxtUndoLogs, UndoLog};
 
 #[derive(Default, Debug, Clone)]
-pub(crate) struct OpaqueTypeStorage<'tcx> {
-    /// Opaque types found in explicit return types and their
-    /// associated fresh inference variable. Writeback resolves these
-    /// variables to get the concrete type, which can be used to
-    /// 'de-opaque' OpaqueHiddenType, after typeck is done with all functions.
-    pub opaque_types: OpaqueTypeMap<'tcx>,
+pub struct OpaqueTypeStorage<'tcx> {
+    opaque_types: FxIndexMap<OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>>,
+    duplicate_entries: Vec<(OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)>,
 }
 
 impl<'tcx> OpaqueTypeStorage<'tcx> {
@@ -33,6 +32,52 @@ impl<'tcx> OpaqueTypeStorage<'tcx> {
         }
     }
 
+    pub(crate) fn pop_duplicate_entry(&mut self) {
+        let entry = self.duplicate_entries.pop();
+        assert!(entry.is_some());
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        let OpaqueTypeStorage { opaque_types, duplicate_entries } = self;
+        opaque_types.is_empty() && duplicate_entries.is_empty()
+    }
+
+    pub(crate) fn take_opaque_types(
+        &mut self,
+    ) -> impl Iterator<Item = (OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        let OpaqueTypeStorage { opaque_types, duplicate_entries } = self;
+        std::mem::take(opaque_types).into_iter().chain(std::mem::take(duplicate_entries))
+    }
+
+    /// Only returns the opaque types from the lookup table. These are used
+    /// when normalizing opaque types and have a unique key.
+    ///
+    /// Outside of canonicalization one should generally use `iter_opaque_types`
+    /// to also consider duplicate entries.
+    pub fn iter_lookup_table(
+        &self,
+    ) -> impl Iterator<Item = (OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        self.opaque_types.iter().map(|(k, v)| (*k, *v))
+    }
+
+    /// Only returns the opaque types which are stored in `duplicate_entries`.
+    ///
+    /// These have to considered when checking all opaque type uses but are e.g.
+    /// irrelevant for canonical inputs as nested queries never meaningfully
+    /// accesses them.
+    pub fn iter_duplicate_entries(
+        &self,
+    ) -> impl Iterator<Item = (OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        self.duplicate_entries.iter().copied()
+    }
+
+    pub fn iter_opaque_types(
+        &self,
+    ) -> impl Iterator<Item = (OpaqueTypeKey<'tcx>, OpaqueHiddenType<'tcx>)> {
+        let OpaqueTypeStorage { opaque_types, duplicate_entries } = self;
+        opaque_types.iter().map(|(k, v)| (*k, *v)).chain(duplicate_entries.iter().copied())
+    }
+
     #[inline]
     pub(crate) fn with_log<'a>(
         &'a mut self,
@@ -44,21 +89,27 @@ impl<'tcx> OpaqueTypeStorage<'tcx> {
 
 impl<'tcx> Drop for OpaqueTypeStorage<'tcx> {
     fn drop(&mut self) {
-        if !self.opaque_types.is_empty() {
+        if !self.is_empty() {
             ty::tls::with(|tcx| tcx.dcx().delayed_bug(format!("{:?}", self.opaque_types)));
         }
     }
 }
 
-pub(crate) struct OpaqueTypeTable<'a, 'tcx> {
+pub struct OpaqueTypeTable<'a, 'tcx> {
     storage: &'a mut OpaqueTypeStorage<'tcx>,
 
     undo_log: &'a mut InferCtxtUndoLogs<'tcx>,
 }
+impl<'tcx> Deref for OpaqueTypeTable<'_, 'tcx> {
+    type Target = OpaqueTypeStorage<'tcx>;
+    fn deref(&self) -> &Self::Target {
+        self.storage
+    }
+}
 
 impl<'a, 'tcx> OpaqueTypeTable<'a, 'tcx> {
     #[instrument(skip(self), level = "debug")]
-    pub(crate) fn register(
+    pub fn register(
         &mut self,
         key: OpaqueTypeKey<'tcx>,
         hidden_type: OpaqueHiddenType<'tcx>,
@@ -72,4 +123,9 @@ impl<'a, 'tcx> OpaqueTypeTable<'a, 'tcx> {
         self.undo_log.push(UndoLog::OpaqueTypes(key, None));
         None
     }
+
+    pub fn add_duplicate(&mut self, key: OpaqueTypeKey<'tcx>, hidden_type: OpaqueHiddenType<'tcx>) {
+        self.storage.duplicate_entries.push((key, hidden_type));
+        self.undo_log.push(UndoLog::DuplicateOpaqueType);
+    }
 }
diff --git a/compiler/rustc_infer/src/infer/snapshot/undo_log.rs b/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
index ba7d8f588e6..b7412d3d6a6 100644
--- a/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
+++ b/compiler/rustc_infer/src/infer/snapshot/undo_log.rs
@@ -17,6 +17,7 @@ pub struct Snapshot<'tcx> {
 /// Records the "undo" data for a single operation that affects some form of inference variable.
 #[derive(Clone)]
 pub(crate) enum UndoLog<'tcx> {
+    DuplicateOpaqueType,
     OpaqueTypes(OpaqueTypeKey<'tcx>, Option<OpaqueHiddenType<'tcx>>),
     TypeVariables(sv::UndoLog<ut::Delegate<type_variable::TyVidEqKey<'tcx>>>),
     ConstUnificationTable(sv::UndoLog<ut::Delegate<ConstVidKey<'tcx>>>),
@@ -58,6 +59,7 @@ impl_from! {
 impl<'tcx> Rollback<UndoLog<'tcx>> for InferCtxtInner<'tcx> {
     fn reverse(&mut self, undo: UndoLog<'tcx>) {
         match undo {
+            UndoLog::DuplicateOpaqueType => self.opaque_type_storage.pop_duplicate_entry(),
             UndoLog::OpaqueTypes(key, idx) => self.opaque_type_storage.remove(key, idx),
             UndoLog::TypeVariables(undo) => self.type_variable_storage.reverse(undo),
             UndoLog::ConstUnificationTable(undo) => self.const_unification_storage.reverse(undo),