Add associated_items api.

- Add `AssocItem` and related things to `stable_mir::ty`.
- Implement `Display` for `AssocKind`.

Author: makai410

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -147,6 +147,14 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::CoroutineWitnessDef(self.create_def_id(did))
     }
 
+    pub fn assoc_def(&mut self, did: DefId) -> stable_mir::ty::AssocDef {
+        stable_mir::ty::AssocDef(self.create_def_id(did))
+    }
+
+    pub fn opaque_def(&mut self, did: DefId) -> stable_mir::ty::OpaqueDef {
+        stable_mir::ty::OpaqueDef(self.create_def_id(did))
+    }
+
     pub fn prov(&mut self, aid: AllocId) -> stable_mir::ty::Prov {
         stable_mir::ty::Prov(self.create_alloc_id(aid))
     }

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -822,6 +822,27 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         let ty = un_op.internal(&mut *tables, tcx).ty(tcx, arg_internal);
         ty.stable(&mut *tables)
     }
+
+    fn associated_items(&self, def_id: stable_mir::DefId) -> stable_mir::ty::AssocItems {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let def_id = tables[def_id];
+        let assoc_items = if tcx.is_trait_alias(def_id) {
+            stable_mir::ty::AssocItems { items: Vec::new() }
+        } else {
+            let items: Vec<(Symbol, stable_mir::ty::AssocItem)> = tcx
+                .associated_item_def_ids(def_id)
+                .iter()
+                .map(|did| {
+                    let item = tcx.associated_item(*did).stable(&mut *tables);
+                    (item.name.clone(), item)
+                })
+                .collect();
+
+            stable_mir::ty::AssocItems { items }
+        };
+        assoc_items
+    }
 }
 
 pub(crate) struct TablesWrapper<'tcx>(pub RefCell<Tables<'tcx>>);

compiler/stable_mir/src/compiler_interface.rs

@@ -12,7 +12,7 @@ use crate::mir::mono::{Instance, InstanceDef, StaticDef};
 use crate::mir::{BinOp, Body, Place, UnOp};
 use crate::target::MachineInfo;
 use crate::ty::{
-    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef,
+    AdtDef, AdtKind, Allocation, AssocItems, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef,
     ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics,
     ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, TraitDecl,
     TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef,
@@ -251,6 +251,9 @@ pub trait Context {
 
     /// Get the resulting type of unary operation.
     fn unop_ty(&self, un_op: UnOp, arg: Ty) -> Ty;
+
+    /// Get all associated items of a definition.
+    fn associated_items(&self, def_id: DefId) -> AssocItems;
 }
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt

compiler/stable_mir/src/crate_def.rs

@@ -3,7 +3,7 @@
 
 use serde::Serialize;
 
-use crate::ty::{GenericArgs, Span, Ty};
+use crate::ty::{AssocItems, GenericArgs, Span, Ty};
 use crate::{Crate, Symbol, with};
 
 /// A unique identification number for each item accessible for the current compilation unit.
@@ -103,6 +103,14 @@ pub trait CrateDefType: CrateDef {
     }
 }
 
+/// A trait for retrieving all items from a definition within a crate.
+pub trait CrateDefItems: CrateDef {
+    /// Retrieve all associated items from a definition.
+    fn associated_items(&self) -> AssocItems {
+        with(|cx| cx.associated_items(self.def_id()))
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct Attribute {
     value: String,
@@ -158,3 +166,9 @@ macro_rules! crate_def_with_ty {
         impl CrateDefType for $name {}
     };
 }
+
+macro_rules! impl_crate_def_items {
+    ( $name:ident $(;)? ) => {
+        impl CrateDefItems for $name {}
+    };
+}

compiler/stable_mir/src/mir/pretty.rs

@@ -9,7 +9,7 @@ use super::{AggregateKind, AssertMessage, BinOp, BorrowKind, FakeBorrowKind, Ter
 use crate::mir::{
     Operand, Place, RawPtrKind, Rvalue, StatementKind, UnwindAction, VarDebugInfoContents,
 };
-use crate::ty::{AdtKind, IndexedVal, MirConst, Ty, TyConst};
+use crate::ty::{AdtKind, AssocKind, IndexedVal, MirConst, Ty, TyConst};
 use crate::{Body, CrateDef, Mutability, with};
 
 impl Display for Ty {
@@ -18,6 +18,16 @@ impl Display for Ty {
     }
 }
 
+impl Display for AssocKind {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            AssocKind::Fn => write!(f, "method"),
+            AssocKind::Const => write!(f, "associated const"),
+            AssocKind::Type => write!(f, "associated type"),
+        }
+    }
+}
+
 impl Debug for Place {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         with(|ctx| write!(f, "{}", ctx.place_pretty(self)))

compiler/stable_mir/src/ty.rs

@@ -6,7 +6,7 @@ use serde::Serialize;
 use super::mir::{Body, Mutability, Safety};
 use super::{DefId, Error, Symbol, with};
 use crate::abi::{FnAbi, Layout};
-use crate::crate_def::{CrateDef, CrateDefType};
+use crate::crate_def::{CrateDef, CrateDefItems, CrateDefType};
 use crate::mir::alloc::{AllocId, read_target_int, read_target_uint};
 use crate::mir::mono::StaticDef;
 use crate::target::MachineInfo;
@@ -773,6 +773,15 @@ pub enum AdtKind {
     Struct,
 }
 
+impl CrateDefItems for AdtDef {
+    fn associated_items(&self) -> AssocItems {
+        match &self.kind() {
+            AdtKind::Enum | AdtKind::Struct => with(|cx| cx.associated_items(self.def_id())),
+            AdtKind::Union => AssocItems { items: Vec::new() },
+        }
+    }
+}
+
 impl AdtDef {
     pub fn kind(&self) -> AdtKind {
         with(|cx| cx.adt_kind(*self))
@@ -910,6 +919,10 @@ crate_def! {
     pub TraitDef;
 }
 
+impl_crate_def_items! {
+    TraitDef;
+}
+
 impl TraitDef {
     pub fn declaration(trait_def: &TraitDef) -> TraitDecl {
         with(|cx| cx.trait_decl(trait_def))
@@ -932,6 +945,10 @@ crate_def! {
     pub ImplDef;
 }
 
+impl_crate_def_items! {
+    ImplDef;
+}
+
 impl ImplDef {
     /// Retrieve information about this implementation.
     pub fn trait_impl(&self) -> ImplTrait {
@@ -1555,3 +1572,65 @@ index_impl!(Span);
 pub struct VariantIdx(usize);
 
 index_impl!(VariantIdx);
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub struct AssocItems {
+    pub items: Vec<(Symbol, AssocItem)>,
+}
+
+crate_def! {
+    /// Hold infomation about an Opaque definition, particularly useful in `RPITIT`.
+    #[derive(Serialize)]
+    pub OpaqueDef;
+}
+
+crate_def! {
+    #[derive(Serialize)]
+    pub AssocDef;
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub struct AssocItem {
+    pub def_id: AssocDef,
+    pub name: Symbol,
+    pub kind: AssocKind,
+    pub container: AssocItemContainer,
+
+    /// If this is an item in an impl of a trait then this is the `DefId` of
+    /// the associated item on the trait that this implements.
+    pub trait_item_def_id: Option<AssocDef>,
+
+    /// Whether this is a method with an explicit self
+    /// as its first parameter, allowing method calls.
+    pub fn_has_self_parameter: bool,
+
+    /// `Some` if the associated item (an associated type) comes from the
+    /// return-position `impl Trait` in trait desugaring. The `ImplTraitInTraitData`
+    /// provides additional information about its source.
+    pub opt_rpitit_info: Option<ImplTraitInTraitData>,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub enum AssocKind {
+    Const,
+    Fn,
+    Type,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
+pub enum AssocItemContainer {
+    Trait,
+    Impl,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, Serialize)]
+pub enum ImplTraitInTraitData {
+    Trait { fn_def_id: FnDef, opaque_def_id: OpaqueDef },
+    Impl { fn_def_id: FnDef },
+}
+
+impl AssocItem {
+    pub fn is_impl_trait_in_trait(&self) -> bool {
+        self.opt_rpitit_info.is_some()
+    }
+}