From 470e9fa1af3e0ec36768e57cdb63bde46af8ee6d Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Sat, 12 Oct 2024 09:10:04 -0400
Subject: [PATCH] Shallowly match opaque key in storage

---
 .../src/solve/eval_ctxt/mod.rs                | 51 +++++++------------
 .../src/solve/normalizes_to/opaque_types.rs   | 34 +++++++++----
 2 files changed, 42 insertions(+), 43 deletions(-)

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
index 458edf920ab..daacc669118 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -4,6 +4,7 @@ use derive_where::derive_where;
 #[cfg(feature = "nightly")]
 use rustc_macros::{HashStable_NoContext, TyDecodable, TyEncodable};
 use rustc_type_ir::data_structures::{HashMap, HashSet, ensure_sufficient_stack};
+use rustc_type_ir::fast_reject::DeepRejectCtxt;
 use rustc_type_ir::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::Relate;
@@ -18,9 +19,9 @@ use crate::delegate::SolverDelegate;
 use crate::solve::inspect::{self, ProofTreeBuilder};
 use crate::solve::search_graph::SearchGraph;
 use crate::solve::{
-    CanonicalInput, CanonicalResponse, Certainty, FIXPOINT_STEP_LIMIT, Goal, GoalEvaluationKind,
-    GoalSource, HasChanged, NestedNormalizationGoals, NoSolution, PredefinedOpaquesData,
-    QueryResult, SolverMode,
+    CanonicalInput, Certainty, FIXPOINT_STEP_LIMIT, Goal, GoalEvaluationKind, GoalSource,
+    HasChanged, NestedNormalizationGoals, NoSolution, PredefinedOpaquesData, QueryResult,
+    SolverMode,
 };
 
 pub(super) mod canonical;
@@ -987,40 +988,22 @@ where
 
     // Do something for each opaque/hidden pair defined with `def_id` in the
     // current inference context.
-    pub(super) fn unify_existing_opaque_tys(
+    pub(super) fn probe_existing_opaque_ty(
         &mut self,
-        param_env: I::ParamEnv,
         key: ty::OpaqueTypeKey<I>,
-        ty: I::Ty,
-    ) -> Vec<CanonicalResponse<I>> {
-        // FIXME: Super inefficient to be cloning this...
-        let opaques = self.delegate.clone_opaque_types_for_query_response();
-
-        let mut values = vec![];
-        for (candidate_key, candidate_ty) in opaques {
-            if candidate_key.def_id != key.def_id {
-                continue;
-            }
-            values.extend(
-                self.probe(|result| inspect::ProbeKind::OpaqueTypeStorageLookup {
-                    result: *result,
-                })
-                .enter(|ecx| {
-                    for (a, b) in std::iter::zip(candidate_key.args.iter(), key.args.iter()) {
-                        ecx.eq(param_env, a, b)?;
-                    }
-                    ecx.eq(param_env, candidate_ty, ty)?;
-                    ecx.add_item_bounds_for_hidden_type(
-                        candidate_key.def_id.into(),
-                        candidate_key.args,
-                        param_env,
-                        candidate_ty,
-                    );
-                    ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
-                }),
+    ) -> 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)
+                },
             );
-        }
-        values
+        let first = matching.next();
+        let second = matching.next();
+        assert_eq!(second, None);
+        first
     }
 
     // Try to evaluate a const, or return `None` if the const is too generic.
diff --git a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
index 6a0703c5313..f8d51f304f3 100644
--- a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
@@ -7,7 +7,9 @@ use rustc_type_ir::inherent::*;
 use rustc_type_ir::{self as ty, Interner};
 
 use crate::delegate::SolverDelegate;
-use crate::solve::{Certainty, EvalCtxt, Goal, NoSolution, QueryResult, Reveal, SolverMode};
+use crate::solve::{
+    Certainty, EvalCtxt, Goal, NoSolution, QueryResult, Reveal, SolverMode, inspect,
+};
 
 impl<D, I> EvalCtxt<'_, D>
 where
@@ -52,14 +54,28 @@ where
                 //
                 // If that fails, we insert `expected` as a new hidden type instead of
                 // eagerly emitting an error.
-                let matches =
-                    self.unify_existing_opaque_tys(goal.param_env, opaque_type_key, expected);
-                if !matches.is_empty() {
-                    if let Some(response) = self.try_merge_responses(&matches) {
-                        return Ok(response);
-                    } else {
-                        return self.flounder(&matches);
-                    }
+                let existing = self.probe_existing_opaque_ty(opaque_type_key);
+                if let Some((candidate_key, candidate_ty)) = existing {
+                    return self
+                        .probe(|result| inspect::ProbeKind::OpaqueTypeStorageLookup {
+                            result: *result,
+                        })
+                        .enter(|ecx| {
+                            for (a, b) in std::iter::zip(
+                                candidate_key.args.iter(),
+                                opaque_type_key.args.iter(),
+                            ) {
+                                ecx.eq(goal.param_env, a, b)?;
+                            }
+                            ecx.eq(goal.param_env, candidate_ty, expected)?;
+                            ecx.add_item_bounds_for_hidden_type(
+                                candidate_key.def_id.into(),
+                                candidate_key.args,
+                                goal.param_env,
+                                candidate_ty,
+                            );
+                            ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+                        });
                 }
 
                 // Otherwise, define a new opaque type