support duplicates in the opaque_type_storage

lcnr · lcnr · commit 5d4793bb7764 · 2025-05-03T03:40:23.000Z
diff --git a/compiler/rustc_hir_typeck/src/_match.rs b/compiler/rustc_hir_typeck/src/_match.rs
@@ -601,7 +601,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             // FIXME(-Znext-solver): Remove this branch once `replace_opaque_types_with_infer` is gone.
             ty::Infer(ty::TyVar(_)) => self
                 .inner
-                .borrow()
+                .borrow_mut()
+                .opaque_types()
                 .iter_opaque_types()
                 .find(|(_, v)| v.ty == expected_ty)
                 .map(|(k, _)| (k.def_id, k.args))?,
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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),
diff --git a/compiler/rustc_next_trait_solver/src/delegate.rs b/compiler/rustc_next_trait_solver/src/delegate.rs
@@ -39,7 +39,10 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         term: <Self::Interner as Interner>::Term,
     ) -> Option<Vec<Goal<Self::Interner, <Self::Interner as Interner>::Predicate>>>;
 
-    fn clone_opaque_types_for_query_response(
+    fn clone_opaque_types_lookup_table(
+        &self,
+    ) -> Vec<(ty::OpaqueTypeKey<Self::Interner>, <Self::Interner as Interner>::Ty)>;
+    fn clone_duplicate_opaque_types(
         &self,
     ) -> Vec<(ty::OpaqueTypeKey<Self::Interner>, <Self::Interner as Interner>::Ty)>;
 
@@ -68,6 +71,12 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         hidden_ty: <Self::Interner as Interner>::Ty,
         span: <Self::Interner as Interner>::Span,
     ) -> Option<<Self::Interner as Interner>::Ty>;
+    fn add_duplicate_opaque_type(
+        &self,
+        opaque_type_key: ty::OpaqueTypeKey<Self::Interner>,
+        hidden_ty: <Self::Interner as Interner>::Ty,
+        span: <Self::Interner as Interner>::Span,
+    );
 
     fn add_item_bounds_for_hidden_type(
         &self,
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
@@ -56,7 +56,10 @@ where
         &self,
         goal: Goal<I, T>,
     ) -> (Vec<I::GenericArg>, CanonicalInput<I, T>) {
-        let opaque_types = self.delegate.clone_opaque_types_for_query_response();
+        // We only care about one entry per `OpaqueTypeKey` here,
+        // so we only canonicalize the lookup table and ignore
+        // duplicate entries.
+        let opaque_types = self.delegate.clone_opaque_types_lookup_table();
         let (goal, opaque_types) =
             (goal, opaque_types).fold_with(&mut EagerResolver::new(self.delegate));
 
@@ -241,19 +244,21 @@ where
             Default::default()
         };
 
-        ExternalConstraintsData {
-            region_constraints,
-            opaque_types: self
-                .delegate
-                .clone_opaque_types_for_query_response()
-                .into_iter()
-                // Only return *newly defined* opaque types.
-                .filter(|(a, _)| {
-                    self.predefined_opaques_in_body.opaque_types.iter().all(|(pa, _)| pa != a)
-                })
-                .collect(),
-            normalization_nested_goals,
-        }
+        // We only return *newly defined* opaque types from canonical queries.
+        //
+        // Constraints for any existing opaque types are already tracked by changes
+        // to the `var_values`.
+        let opaque_types = self
+            .delegate
+            .clone_opaque_types_lookup_table()
+            .into_iter()
+            .filter(|(a, _)| {
+                self.predefined_opaques_in_body.opaque_types.iter().all(|(pa, _)| pa != a)
+            })
+            .chain(self.delegate.clone_duplicate_opaque_types())
+            .collect();
+
+        ExternalConstraintsData { region_constraints, opaque_types, normalization_nested_goals }
     }
 
     /// After calling a canonical query, we apply the constraints returned
@@ -432,7 +437,16 @@ where
     fn register_new_opaque_types(&mut self, opaque_types: &[(ty::OpaqueTypeKey<I>, I::Ty)]) {
         for &(key, ty) in opaque_types {
             let prev = self.delegate.register_hidden_type_in_storage(key, ty, self.origin_span);
-            assert_eq!(prev, None);
+            // 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);
+            }
         }
     }
 }
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -14,7 +14,7 @@ use rustc_type_ir::{
     TypeSuperFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
     TypingMode,
 };
-use tracing::{instrument, trace};
+use tracing::{debug, instrument, trace};
 
 use super::has_only_region_constraints;
 use crate::coherence;
@@ -361,7 +361,20 @@ where
 
         for &(key, ty) in &input.predefined_opaques_in_body.opaque_types {
             let prev = ecx.delegate.register_hidden_type_in_storage(key, ty, ecx.origin_span);
-            assert_eq!(prev, None);
+            // It may be possible that two entries in the opaque type storage end up
+            // with the same key after resolving contained inference variables.
+            //
+            // We could put them in the duplicate list but don't have to. The opaques we
+            // encounter here are already tracked in the caller, so there's no need to
+            // also store them here. We'd take them out when computing the query response
+            // and then discard them, as they're already present in the input.
+            //
+            // Ideally we'd drop duplicate opaque type definitions when computing
+            // the canonical input. This is more annoying to implement and may cause a
+            // perf regression, so we do it inside of the query for now.
+            if let Some(prev) = prev {
+                debug!(?key, ?ty, ?prev, "ignore duplicate in `opaque_type_storage`");
+            }
         }
 
         if !ecx.nested_goals.is_empty() {
@@ -1065,14 +1078,13 @@ where
         &mut self,
         key: ty::OpaqueTypeKey<I>,
     ) -> Option<(ty::OpaqueTypeKey<I>, I::Ty)> {
-        let mut matching =
-            self.delegate.clone_opaque_types_for_query_response().into_iter().filter(
-                |(candidate_key, _)| {
-                    candidate_key.def_id == key.def_id
-                        && DeepRejectCtxt::relate_rigid_rigid(self.cx())
-                            .args_may_unify(candidate_key.args, key.args)
-                },
-            );
+        let mut matching = self.delegate.clone_opaque_types_lookup_table().into_iter().filter(
+            |(candidate_key, _)| {
+                candidate_key.def_id == key.def_id
+                    && DeepRejectCtxt::relate_rigid_rigid(self.cx())
+                        .args_may_unify(candidate_key.args, key.args)
+            },
+        );
         let first = matching.next();
         let second = matching.next();
         assert_eq!(second, None);
diff --git a/compiler/rustc_trait_selection/src/solve/delegate.rs b/compiler/rustc_trait_selection/src/solve/delegate.rs
diff --git a/tests/ui/traits/next-solver/opaques/duplicate-opaque-type-entries.rs b/tests/ui/traits/next-solver/opaques/duplicate-opaque-type-entries.rs
