From d4fa5648c39ec50821eb7e0ef36c62adbbe7cd2b Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 8 Feb 2024 21:21:30 +0000
Subject: [PATCH] Move query providers

---
 compiler/rustc_metadata/src/rmeta/encoder.rs | 25 -------------
 compiler/rustc_middle/src/ty/mod.rs          |  2 +
 compiler/rustc_middle/src/ty/trait_def.rs    | 39 +++++++++++++++++---
 3 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index 2ffbcc182f5..8f037d25163 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -13,7 +13,6 @@ use rustc_hir_pretty::id_to_string;
 use rustc_middle::middle::dependency_format::Linkage;
 use rustc_middle::middle::exported_symbols::metadata_symbol_name;
 use rustc_middle::mir::interpret;
-use rustc_middle::query::LocalCrate;
 use rustc_middle::query::Providers;
 use rustc_middle::traits::specialization_graph;
 use rustc_middle::ty::codec::TyEncoder;
@@ -2298,30 +2297,6 @@ pub fn provide(providers: &mut Providers) {
             })
         },
 
-        // TODO: Uplift these into
-        traits: |tcx, LocalCrate| {
-            let mut traits = Vec::new();
-            for id in tcx.hir().items() {
-                if matches!(tcx.def_kind(id.owner_id), DefKind::Trait | DefKind::TraitAlias) {
-                    traits.push(id.owner_id.to_def_id())
-                }
-            }
-
-            tcx.arena.alloc_slice(&traits)
-        },
-        trait_impls_in_crate: |tcx, LocalCrate| {
-            let mut trait_impls = Vec::new();
-            for id in tcx.hir().items() {
-                if matches!(tcx.def_kind(id.owner_id), DefKind::Impl { .. })
-                    && tcx.impl_trait_ref(id.owner_id).is_some()
-                {
-                    trait_impls.push(id.owner_id.to_def_id())
-                }
-            }
-
-            tcx.arena.alloc_slice(&trait_impls)
-        },
-
         ..*providers
     }
 }
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index fdf6e3f423f..0451be420cf 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -2082,6 +2082,8 @@ pub fn provide(providers: &mut Providers) {
     *providers = Providers {
         trait_impls_of: trait_def::trait_impls_of_provider,
         incoherent_impls: trait_def::incoherent_impls_provider,
+        trait_impls_in_crate: trait_def::trait_impls_in_crate_provider,
+        traits: trait_def::traits_provider,
         const_param_default: consts::const_param_default,
         vtable_allocation: vtable::vtable_allocation_provider,
         ..*providers
diff --git a/compiler/rustc_middle/src/ty/trait_def.rs b/compiler/rustc_middle/src/ty/trait_def.rs
index da5860043c9..3bd9f6ad11b 100644
--- a/compiler/rustc_middle/src/ty/trait_def.rs
+++ b/compiler/rustc_middle/src/ty/trait_def.rs
@@ -1,16 +1,19 @@
-use crate::traits::specialization_graph;
-use crate::ty::fast_reject::{self, SimplifiedType, TreatParams};
-use crate::ty::{Ident, Ty, TyCtxt};
-use hir::def_id::LOCAL_CRATE;
-use rustc_hir as hir;
-use rustc_hir::def_id::DefId;
 use std::iter;
 use tracing::debug;
 
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_errors::ErrorGuaranteed;
+use rustc_hir as hir;
+use rustc_hir::def::DefKind;
+use rustc_hir::def_id::DefId;
+use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_macros::{Decodable, Encodable, HashStable};
 
+use crate::query::LocalCrate;
+use crate::traits::specialization_graph;
+use crate::ty::fast_reject::{self, SimplifiedType, TreatParams};
+use crate::ty::{Ident, Ty, TyCtxt};
+
 /// A trait's definition with type information.
 #[derive(HashStable, Encodable, Decodable)]
 pub struct TraitDef {
@@ -274,3 +277,27 @@ pub(super) fn incoherent_impls_provider(
 
     Ok(tcx.arena.alloc_slice(&impls))
 }
+
+pub(super) fn traits_provider(tcx: TyCtxt<'_>, _: LocalCrate) -> &[DefId] {
+    let mut traits = Vec::new();
+    for id in tcx.hir().items() {
+        if matches!(tcx.def_kind(id.owner_id), DefKind::Trait | DefKind::TraitAlias) {
+            traits.push(id.owner_id.to_def_id())
+        }
+    }
+
+    tcx.arena.alloc_slice(&traits)
+}
+
+pub(super) fn trait_impls_in_crate_provider(tcx: TyCtxt<'_>, _: LocalCrate) -> &[DefId] {
+    let mut trait_impls = Vec::new();
+    for id in tcx.hir().items() {
+        if matches!(tcx.def_kind(id.owner_id), DefKind::Impl { .. })
+            && tcx.impl_trait_ref(id.owner_id).is_some()
+        {
+            trait_impls.push(id.owner_id.to_def_id())
+        }
+    }
+
+    tcx.arena.alloc_slice(&trait_impls)
+}