Rollup merge of #115300 - spastorino:smir-tweaks, r=oli-obk

Tweaks and improvements on SMIR around generics_of and predicates_of

r? `@oli-obk`

This allows an API like the following ...

```rust
    let trait_decls = stable_mir::all_trait_decls().iter().map(|trait_def| {
        let trait_decl = stable_mir::trait_decl(trait_def);
        let generics = trait_decl.generics_of();
        let predicates = trait_decl.predicates_of().predicates;
```

I didn't like that much `trait_def.trait_decl()` which is it possible but adding a method to a def_id that loads up a whole trait definition looks backwards to me.

4 files changed, 41 insertions, 71 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 4385d3f51ed..5334a75dc06 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -4,6 +4,7 @@
 //! until stable MIR is complete.
 
 use std::fmt::Debug;
+use std::ops::Index;
 use std::string::ToString;
 
 use crate::rustc_internal;
@@ -24,7 +25,7 @@ fn with_tables<R>(mut f: impl FnMut(&mut Tables<'_>) -> R) -> R {
 }
 
 pub fn item_def_id(item: &stable_mir::CrateItem) -> DefId {
-    with_tables(|t| t.item_def_id(item))
+    with_tables(|t| t[item.0])
 }
 
 pub fn crate_item(did: DefId) -> stable_mir::CrateItem {
@@ -71,23 +72,16 @@ pub fn impl_def(did: DefId) -> stable_mir::ty::ImplDef {
     with_tables(|t| t.impl_def(did))
 }
 
-impl<'tcx> Tables<'tcx> {
-    pub fn item_def_id(&self, item: &stable_mir::CrateItem) -> DefId {
-        self.def_ids[item.0]
-    }
-
-    pub fn trait_def_id(&self, trait_def: &stable_mir::ty::TraitDef) -> DefId {
-        self.def_ids[trait_def.0]
-    }
+impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
+    type Output = DefId;
 
-    pub fn impl_trait_def_id(&self, impl_def: &stable_mir::ty::ImplDef) -> DefId {
-        self.def_ids[impl_def.0]
-    }
-
-    pub fn generic_def_id(&self, generic_def: &stable_mir::ty::GenericDef) -> DefId {
-        self.def_ids[generic_def.0]
+    #[inline(always)]
+    fn index(&self, index: stable_mir::DefId) -> &Self::Output {
+        &self.def_ids[index.0]
     }
+}
 
+impl<'tcx> Tables<'tcx> {
     pub fn crate_item(&mut self, did: DefId) -> stable_mir::CrateItem {
         stable_mir::CrateItem(self.create_def_id(did))
     }
@@ -144,12 +138,12 @@ impl<'tcx> Tables<'tcx> {
         // FIXME: this becomes inefficient when we have too many ids
         for (i, &d) in self.def_ids.iter().enumerate() {
             if d == did {
-                return i;
+                return stable_mir::DefId(i);
             }
         }
         let id = self.def_ids.len();
         self.def_ids.push(did);
-        id
+        stable_mir::DefId(id)
     }
 }
 
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 925d4a3bdd8..7877c91c2cf 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -10,8 +10,8 @@
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
 use crate::stable_mir::ty::{
-    allocation_filter, new_allocation, Const, FloatTy, GenericDef, GenericParamDef, IntTy,
-    Movability, RigidTy, TyKind, UintTy,
+    allocation_filter, new_allocation, Const, FloatTy, GenericParamDef, IntTy, Movability, RigidTy,
+    TyKind, UintTy,
 };
 use crate::stable_mir::{self, Context};
 use rustc_hir as hir;
@@ -54,7 +54,7 @@ impl<'tcx> Context for Tables<'tcx> {
     }
 
     fn trait_decl(&mut self, trait_def: &stable_mir::ty::TraitDef) -> stable_mir::ty::TraitDecl {
-        let def_id = self.trait_def_id(trait_def);
+        let def_id = self[trait_def.0];
         let trait_def = self.tcx.trait_def(def_id);
         trait_def.stable(self)
     }
@@ -68,13 +68,13 @@ impl<'tcx> Context for Tables<'tcx> {
     }
 
     fn trait_impl(&mut self, impl_def: &stable_mir::ty::ImplDef) -> stable_mir::ty::ImplTrait {
-        let def_id = self.impl_trait_def_id(impl_def);
+        let def_id = self[impl_def.0];
         let impl_trait = self.tcx.impl_trait_ref(def_id).unwrap();
         impl_trait.stable(self)
     }
 
     fn mir_body(&mut self, item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
-        let def_id = self.item_def_id(item);
+        let def_id = self[item.0];
         let mir = self.tcx.optimized_mir(def_id);
         stable_mir::mir::Body {
             blocks: mir
@@ -102,19 +102,16 @@ impl<'tcx> Context for Tables<'tcx> {
         ty.stable(self)
     }
 
-    fn generics_of(&mut self, generic_def: &GenericDef) -> stable_mir::ty::Generics {
-        let def_id = self.generic_def_id(generic_def);
-        let generic_def = self.tcx.generics_of(def_id);
-        generic_def.stable(self)
+    fn generics_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::Generics {
+        let def_id = self[def_id];
+        let generics = self.tcx.generics_of(def_id);
+        generics.stable(self)
     }
 
-    fn predicates_of(
-        &mut self,
-        trait_def: &stable_mir::ty::TraitDef,
-    ) -> stable_mir::GenericPredicates {
-        let trait_def_id = self.trait_def_id(trait_def);
-        let ty::GenericPredicates { parent, predicates } = self.tcx.predicates_of(trait_def_id);
-        stable_mir::GenericPredicates {
+    fn predicates_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::GenericPredicates {
+        let def_id = self[def_id];
+        let ty::GenericPredicates { parent, predicates } = self.tcx.predicates_of(def_id);
+        stable_mir::ty::GenericPredicates {
             parent: parent.map(|did| self.trait_def(did)),
             predicates: predicates
                 .iter()
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 9a2c5b2c7c1..1b834628175 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -14,7 +14,7 @@
 use std::cell::Cell;
 
 use self::ty::{
-    GenericDef, Generics, ImplDef, ImplTrait, PredicateKind, Span, TraitDecl, TraitDef, Ty, TyKind,
+    GenericPredicates, Generics, ImplDef, ImplTrait, Span, TraitDecl, TraitDef, Ty, TyKind,
 };
 use crate::rustc_smir::Tables;
 
@@ -28,7 +28,8 @@ pub type Symbol = String;
 pub type CrateNum = usize;
 
 /// A unique identification number for each item accessible for the current compilation unit.
-pub type DefId = usize;
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct DefId(pub(crate) usize);
 
 /// A list of crate items.
 pub type CrateItems = Vec<CrateItem>;
@@ -39,12 +40,6 @@ pub type TraitDecls = Vec<TraitDef>;
 /// A list of impl trait decls.
 pub type ImplTraitDecls = Vec<ImplDef>;
 
-/// A list of predicates.
-pub struct GenericPredicates {
-    pub parent: Option<TraitDef>,
-    pub predicates: Vec<(PredicateKind, Span)>,
-}
-
 /// Holds information about a crate.
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Crate {
@@ -108,14 +103,6 @@ pub fn trait_impl(trait_impl: &ImplDef) -> ImplTrait {
     with(|cx| cx.trait_impl(trait_impl))
 }
 
-pub fn generics_of(generic_def: &GenericDef) -> Generics {
-    with(|cx| cx.generics_of(generic_def))
-}
-
-pub fn predicates_of(trait_def: &TraitDef) -> GenericPredicates {
-    with(|cx| cx.predicates_of(trait_def))
-}
-
 pub trait Context {
     fn entry_fn(&mut self) -> Option<CrateItem>;
     /// Retrieve all items of the local crate that have a MIR associated with them.
@@ -125,8 +112,8 @@ pub trait Context {
     fn trait_decl(&mut self, trait_def: &TraitDef) -> TraitDecl;
     fn all_trait_impls(&mut self) -> ImplTraitDecls;
     fn trait_impl(&mut self, trait_impl: &ImplDef) -> ImplTrait;
-    fn generics_of(&mut self, generic_def: &GenericDef) -> Generics;
-    fn predicates_of(&mut self, trait_def: &TraitDef) -> GenericPredicates;
+    fn generics_of(&mut self, def_id: DefId) -> Generics;
+    fn predicates_of(&mut self, def_id: DefId) -> GenericPredicates;
     /// Get information about the local crate.
     fn local_crate(&self) -> Crate;
     /// Retrieve a list of all external crates.
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index 5929823b1bb..fce5375f848 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -120,27 +120,9 @@ pub struct GenericDef(pub(crate) DefId);
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct ConstDef(pub(crate) DefId);
 
-impl TraitDef {
-    pub fn trait_decl(&self) -> TraitDecl {
-        with(|cx| cx.trait_decl(self))
-    }
-}
-
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct ImplDef(pub(crate) DefId);
 
-impl ImplDef {
-    pub fn trait_impl(&self) -> ImplTrait {
-        with(|cx| cx.trait_impl(self))
-    }
-}
-
-impl GenericDef {
-    pub fn generics_of(&self) -> Generics {
-        with(|tcx| tcx.generics_of(self))
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct GenericArgs(pub Vec<GenericArgKind>);
 
@@ -463,6 +445,16 @@ pub struct TraitDecl {
     pub deny_explicit_impl: bool,
 }
 
+impl TraitDecl {
+    pub fn generics_of(&self) -> Generics {
+        with(|cx| cx.generics_of(self.def_id.0))
+    }
+
+    pub fn predicates_of(&self) -> GenericPredicates {
+        with(|cx| cx.predicates_of(self.def_id.0))
+    }
+}
+
 pub type ImplTrait = EarlyBinder<TraitRef>;
 
 #[derive(Clone, Debug)]
@@ -499,8 +491,8 @@ pub struct GenericParamDef {
 }
 
 pub struct GenericPredicates {
-    pub parent: Option<DefId>,
-    pub predicates: Vec<PredicateKind>,
+    pub parent: Option<TraitDef>,
+    pub predicates: Vec<(PredicateKind, Span)>,
 }
 
 #[derive(Clone, Debug)]