From 5a837515f28eed106dd031b0243d6476eba58dd9 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Sat, 29 Jun 2024 23:05:35 -0400
Subject: [PATCH] Make fn traits into first-class TraitSolverLangItems to avoid
 needing fn_trait_kind_from_def_id

---
 compiler/rustc_middle/src/ty/context.rs       | 16 +++++-----
 .../src/solve/assembly/mod.rs                 | 30 ++++++++++++++++---
 compiler/rustc_type_ir/src/interner.rs        |  4 ---
 compiler/rustc_type_ir/src/lang_items.rs      |  6 ++++
 4 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 155a76713cab..7c3d2df72faa 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -537,14 +537,6 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
         self.trait_def(trait_def_id).implement_via_object
     }
 
-    fn fn_trait_kind_from_def_id(self, trait_def_id: DefId) -> Option<ty::ClosureKind> {
-        self.fn_trait_kind_from_def_id(trait_def_id)
-    }
-
-    fn async_fn_trait_kind_from_def_id(self, trait_def_id: DefId) -> Option<ty::ClosureKind> {
-        self.async_fn_trait_kind_from_def_id(trait_def_id)
-    }
-
     fn supertrait_def_ids(self, trait_def_id: DefId) -> impl IntoIterator<Item = DefId> {
         self.supertrait_def_ids(trait_def_id)
     }
@@ -608,8 +600,11 @@ macro_rules! bidirectional_lang_item_map {
 bidirectional_lang_item_map! {
 // tidy-alphabetical-start
     AsyncDestruct,
+    AsyncFn,
     AsyncFnKindHelper,
     AsyncFnKindUpvars,
+    AsyncFnMut,
+    AsyncFnOnce,
     AsyncFnOnceOutput,
     AsyncIterator,
     CallOnceFuture,
@@ -622,11 +617,14 @@ bidirectional_lang_item_map! {
     Destruct,
     DiscriminantKind,
     DynMetadata,
-    EffectsMaybe,
     EffectsIntersection,
     EffectsIntersectionOutput,
+    EffectsMaybe,
     EffectsNoRuntime,
     EffectsRuntime,
+    Fn,
+    FnMut,
+    FnOnce,
     FnPtrTrait,
     FusedIterator,
     Future,
diff --git a/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs b/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
index 20a0da16c752..38a4f7dfe25c 100644
--- a/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/assembly/mod.rs
@@ -387,16 +387,38 @@ where
             G::consider_auto_trait_candidate(self, goal)
         } else if cx.trait_is_alias(trait_def_id) {
             G::consider_trait_alias_candidate(self, goal)
-        } else if let Some(kind) = self.cx().fn_trait_kind_from_def_id(trait_def_id) {
-            G::consider_builtin_fn_trait_candidates(self, goal, kind)
-        } else if let Some(kind) = self.cx().async_fn_trait_kind_from_def_id(trait_def_id) {
-            G::consider_builtin_async_fn_trait_candidates(self, goal, kind)
         } else {
             match cx.as_lang_item(trait_def_id) {
                 Some(TraitSolverLangItem::Sized) => G::consider_builtin_sized_candidate(self, goal),
                 Some(TraitSolverLangItem::Copy | TraitSolverLangItem::Clone) => {
                     G::consider_builtin_copy_clone_candidate(self, goal)
                 }
+                Some(TraitSolverLangItem::Fn) => {
+                    G::consider_builtin_fn_trait_candidates(self, goal, ty::ClosureKind::Fn)
+                }
+                Some(TraitSolverLangItem::FnMut) => {
+                    G::consider_builtin_fn_trait_candidates(self, goal, ty::ClosureKind::FnMut)
+                }
+                Some(TraitSolverLangItem::FnOnce) => {
+                    G::consider_builtin_fn_trait_candidates(self, goal, ty::ClosureKind::FnOnce)
+                }
+                Some(TraitSolverLangItem::AsyncFn) => {
+                    G::consider_builtin_async_fn_trait_candidates(self, goal, ty::ClosureKind::Fn)
+                }
+                Some(TraitSolverLangItem::AsyncFnMut) => {
+                    G::consider_builtin_async_fn_trait_candidates(
+                        self,
+                        goal,
+                        ty::ClosureKind::FnMut,
+                    )
+                }
+                Some(TraitSolverLangItem::AsyncFnOnce) => {
+                    G::consider_builtin_async_fn_trait_candidates(
+                        self,
+                        goal,
+                        ty::ClosureKind::FnOnce,
+                    )
+                }
                 Some(TraitSolverLangItem::PointerLike) => {
                     G::consider_builtin_pointer_like_candidate(self, goal)
                 }
diff --git a/compiler/rustc_type_ir/src/interner.rs b/compiler/rustc_type_ir/src/interner.rs
index 84c7e6e8f200..f76a2f55c482 100644
--- a/compiler/rustc_type_ir/src/interner.rs
+++ b/compiler/rustc_type_ir/src/interner.rs
@@ -254,10 +254,6 @@ pub trait Interner:
 
     fn trait_may_be_implemented_via_object(self, trait_def_id: Self::DefId) -> bool;
 
-    fn fn_trait_kind_from_def_id(self, trait_def_id: Self::DefId) -> Option<ty::ClosureKind>;
-
-    fn async_fn_trait_kind_from_def_id(self, trait_def_id: Self::DefId) -> Option<ty::ClosureKind>;
-
     fn supertrait_def_ids(self, trait_def_id: Self::DefId)
     -> impl IntoIterator<Item = Self::DefId>;
 
diff --git a/compiler/rustc_type_ir/src/lang_items.rs b/compiler/rustc_type_ir/src/lang_items.rs
index cf00c37caa28..265a41188273 100644
--- a/compiler/rustc_type_ir/src/lang_items.rs
+++ b/compiler/rustc_type_ir/src/lang_items.rs
@@ -3,8 +3,11 @@
 pub enum TraitSolverLangItem {
     // tidy-alphabetical-start
     AsyncDestruct,
+    AsyncFn,
     AsyncFnKindHelper,
     AsyncFnKindUpvars,
+    AsyncFnMut,
+    AsyncFnOnce,
     AsyncFnOnceOutput,
     AsyncIterator,
     CallOnceFuture,
@@ -22,6 +25,9 @@ pub enum TraitSolverLangItem {
     EffectsMaybe,
     EffectsNoRuntime,
     EffectsRuntime,
+    Fn,
+    FnMut,
+    FnOnce,
     FnPtrTrait,
     FusedIterator,
     Future,