From 593156a357994c2bd96805edfb65580456aa584c Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Mon, 4 Mar 2024 15:56:34 +0100
Subject: [PATCH] Re-use `InferenceTable` by snapshotting in method resolution

---
 crates/hir-ty/src/autoderef.rs         |  6 +-
 crates/hir-ty/src/infer/expr.rs        |  2 +-
 crates/hir-ty/src/infer/unify.rs       | 16 ++----
 crates/hir-ty/src/method_resolution.rs | 80 ++++++++++++--------------
 4 files changed, 48 insertions(+), 56 deletions(-)

diff --git a/crates/hir-ty/src/autoderef.rs b/crates/hir-ty/src/autoderef.rs
index 8d819e41aa..e2446c3425 100644
--- a/crates/hir-ty/src/autoderef.rs
+++ b/crates/hir-ty/src/autoderef.rs
@@ -113,7 +113,7 @@ pub(crate) fn autoderef_step(
     ty: Ty,
     explicit: bool,
 ) -> Option<(AutoderefKind, Ty)> {
-    if let Some(derefed) = builtin_deref(table, &ty, explicit) {
+    if let Some(derefed) = builtin_deref(table.db, &ty, explicit) {
         Some((AutoderefKind::Builtin, table.resolve_ty_shallow(derefed)))
     } else {
         Some((AutoderefKind::Overloaded, deref_by_trait(table, ty)?))
@@ -121,7 +121,7 @@ pub(crate) fn autoderef_step(
 }
 
 pub(crate) fn builtin_deref<'ty>(
-    table: &mut InferenceTable<'_>,
+    db: &dyn HirDatabase,
     ty: &'ty Ty,
     explicit: bool,
 ) -> Option<&'ty Ty> {
@@ -129,7 +129,7 @@ pub(crate) fn builtin_deref<'ty>(
         TyKind::Ref(.., ty) => Some(ty),
         TyKind::Raw(.., ty) if explicit => Some(ty),
         &TyKind::Adt(chalk_ir::AdtId(adt), ref substs) => {
-            if crate::lang_items::is_box(table.db, adt) {
+            if crate::lang_items::is_box(db, adt) {
                 substs.at(Interner, 0).ty(Interner)
             } else {
                 None
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index c377a51e7d..5c50f42d56 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -657,7 +657,7 @@ impl InferenceContext<'_> {
                                 );
                             }
                         }
-                        if let Some(derefed) = builtin_deref(&mut self.table, &inner_ty, true) {
+                        if let Some(derefed) = builtin_deref(self.table.db, &inner_ty, true) {
                             self.resolve_ty_shallow(derefed)
                         } else {
                             deref_by_trait(&mut self.table, inner_ty)
diff --git a/crates/hir-ty/src/infer/unify.rs b/crates/hir-ty/src/infer/unify.rs
index 1d0150d850..00c9246d43 100644
--- a/crates/hir-ty/src/infer/unify.rs
+++ b/crates/hir-ty/src/infer/unify.rs
@@ -243,7 +243,7 @@ pub(crate) struct InferenceTable<'a> {
     pub(crate) db: &'a dyn HirDatabase,
     pub(crate) trait_env: Arc<TraitEnvironment>,
     var_unification_table: ChalkInferenceTable,
-    type_variable_table: Vec<TypeVariableFlags>,
+    type_variable_table: SmallVec<[TypeVariableFlags; 16]>,
     pending_obligations: Vec<Canonicalized<InEnvironment<Goal>>>,
     /// Double buffer used in [`Self::resolve_obligations_as_possible`] to cut down on
     /// temporary allocations.
@@ -252,8 +252,8 @@ pub(crate) struct InferenceTable<'a> {
 
 pub(crate) struct InferenceTableSnapshot {
     var_table_snapshot: chalk_solve::infer::InferenceSnapshot<Interner>,
+    type_variable_table: SmallVec<[TypeVariableFlags; 16]>,
     pending_obligations: Vec<Canonicalized<InEnvironment<Goal>>>,
-    type_variable_table_snapshot: Vec<TypeVariableFlags>,
 }
 
 impl<'a> InferenceTable<'a> {
@@ -262,7 +262,7 @@ impl<'a> InferenceTable<'a> {
             db,
             trait_env,
             var_unification_table: ChalkInferenceTable::new(),
-            type_variable_table: Vec::new(),
+            type_variable_table: SmallVec::new(),
             pending_obligations: Vec::new(),
             resolve_obligations_buffer: Vec::new(),
         }
@@ -575,19 +575,15 @@ impl<'a> InferenceTable<'a> {
 
     pub(crate) fn snapshot(&mut self) -> InferenceTableSnapshot {
         let var_table_snapshot = self.var_unification_table.snapshot();
-        let type_variable_table_snapshot = self.type_variable_table.clone();
+        let type_variable_table = self.type_variable_table.clone();
         let pending_obligations = self.pending_obligations.clone();
-        InferenceTableSnapshot {
-            var_table_snapshot,
-            pending_obligations,
-            type_variable_table_snapshot,
-        }
+        InferenceTableSnapshot { var_table_snapshot, pending_obligations, type_variable_table }
     }
 
     #[tracing::instrument(skip_all)]
     pub(crate) fn rollback_to(&mut self, snapshot: InferenceTableSnapshot) {
         self.var_unification_table.rollback_to(snapshot.var_table_snapshot);
-        self.type_variable_table = snapshot.type_variable_table_snapshot;
+        self.type_variable_table = snapshot.type_variable_table;
         self.pending_obligations = snapshot.pending_obligations;
     }
 
diff --git a/crates/hir-ty/src/method_resolution.rs b/crates/hir-ty/src/method_resolution.rs
index e68dbe7b02..bb2c436a99 100644
--- a/crates/hir-ty/src/method_resolution.rs
+++ b/crates/hir-ty/src/method_resolution.rs
@@ -972,10 +972,9 @@ pub fn iterate_method_candidates_dyn(
 
             deref_chain.into_iter().try_for_each(|(receiver_ty, adj)| {
                 iterate_method_candidates_with_autoref(
+                    &mut table,
                     &receiver_ty,
                     adj,
-                    db,
-                    env.clone(),
                     traits_in_scope,
                     visible_from_module,
                     name,
@@ -1000,10 +999,9 @@ pub fn iterate_method_candidates_dyn(
 
 #[tracing::instrument(skip_all, fields(name = ?name))]
 fn iterate_method_candidates_with_autoref(
+    table: &mut InferenceTable<'_>,
     receiver_ty: &Canonical<Ty>,
     first_adjustment: ReceiverAdjustments,
-    db: &dyn HirDatabase,
-    env: Arc<TraitEnvironment>,
     traits_in_scope: &FxHashSet<TraitId>,
     visible_from_module: VisibleFromModule,
     name: Option<&Name>,
@@ -1016,10 +1014,9 @@ fn iterate_method_candidates_with_autoref(
 
     let mut iterate_method_candidates_by_receiver = move |receiver_ty, first_adjustment| {
         iterate_method_candidates_by_receiver(
+            table,
             receiver_ty,
             first_adjustment,
-            db,
-            env.clone(),
             traits_in_scope,
             visible_from_module,
             name,
@@ -1058,50 +1055,49 @@ fn iterate_method_candidates_with_autoref(
 
 #[tracing::instrument(skip_all, fields(name = ?name))]
 fn iterate_method_candidates_by_receiver(
+    table: &mut InferenceTable<'_>,
     receiver_ty: &Canonical<Ty>,
     receiver_adjustments: ReceiverAdjustments,
-    db: &dyn HirDatabase,
-    env: Arc<TraitEnvironment>,
     traits_in_scope: &FxHashSet<TraitId>,
     visible_from_module: VisibleFromModule,
     name: Option<&Name>,
     mut callback: &mut dyn FnMut(ReceiverAdjustments, AssocItemId, bool) -> ControlFlow<()>,
 ) -> ControlFlow<()> {
-    let mut table = InferenceTable::new(db, env);
-    let receiver_ty = table.instantiate_canonical(receiver_ty.clone());
-    let snapshot = table.snapshot();
-    // We're looking for methods with *receiver* type receiver_ty. These could
-    // be found in any of the derefs of receiver_ty, so we have to go through
-    // that, including raw derefs.
-    let mut autoderef = autoderef::Autoderef::new(&mut table, receiver_ty.clone(), true);
-    while let Some((self_ty, _)) = autoderef.next() {
-        iterate_inherent_methods(
-            &self_ty,
-            autoderef.table,
-            name,
-            Some(&receiver_ty),
-            Some(receiver_adjustments.clone()),
-            visible_from_module,
-            &mut callback,
-        )?
-    }
+    table.run_in_snapshot(|table| {
+        let receiver_ty = table.instantiate_canonical(receiver_ty.clone());
+        // We're looking for methods with *receiver* type receiver_ty. These could
+        // be found in any of the derefs of receiver_ty, so we have to go through
+        // that, including raw derefs.
+        table.run_in_snapshot(|table| {
+            let mut autoderef = autoderef::Autoderef::new(table, receiver_ty.clone(), true);
+            while let Some((self_ty, _)) = autoderef.next() {
+                iterate_inherent_methods(
+                    &self_ty,
+                    autoderef.table,
+                    name,
+                    Some(&receiver_ty),
+                    Some(receiver_adjustments.clone()),
+                    visible_from_module,
+                    &mut callback,
+                )?
+            }
+            ControlFlow::Continue(())
+        })?;
 
-    table.rollback_to(snapshot);
-
-    let mut autoderef = autoderef::Autoderef::new(&mut table, receiver_ty.clone(), true);
-    while let Some((self_ty, _)) = autoderef.next() {
-        iterate_trait_method_candidates(
-            &self_ty,
-            autoderef.table,
-            traits_in_scope,
-            name,
-            Some(&receiver_ty),
-            Some(receiver_adjustments.clone()),
-            &mut callback,
-        )?
-    }
-
-    ControlFlow::Continue(())
+        let mut autoderef = autoderef::Autoderef::new(table, receiver_ty.clone(), true);
+        while let Some((self_ty, _)) = autoderef.next() {
+            iterate_trait_method_candidates(
+                &self_ty,
+                autoderef.table,
+                traits_in_scope,
+                name,
+                Some(&receiver_ty),
+                Some(receiver_adjustments.clone()),
+                &mut callback,
+            )?
+        }
+        ControlFlow::Continue(())
+    })
 }
 
 #[tracing::instrument(skip_all, fields(name = ?name))]