Implement `associated_items` api.

9 files changed, 330 insertions, 7 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index ad38ea228bf..30e9b2e72f7 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/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))
     }
diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index aa1921fc8e7..322e86147c1 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -822,6 +822,21 @@ 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::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) {
+            Vec::new()
+        } else {
+            tcx.associated_item_def_ids(def_id)
+                .iter()
+                .map(|did| tcx.associated_item(*did).stable(&mut *tables))
+                .collect()
+        };
+        assoc_items
+    }
 }
 
 pub(crate) struct TablesWrapper<'tcx>(pub RefCell<Tables<'tcx>>);
diff --git a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
index aa0eac628dd..8309809d7b5 100644
--- a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
+++ b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
@@ -890,3 +890,63 @@ impl<'tcx> Stable<'tcx> for rustc_session::cstore::ForeignModule {
         }
     }
 }
+
+impl<'tcx> Stable<'tcx> for ty::AssocKind {
+    type T = stable_mir::ty::AssocKind;
+
+    fn stable(&self, _tables: &mut Tables<'_>) -> Self::T {
+        use stable_mir::ty::AssocKind;
+        match self {
+            ty::AssocKind::Const => AssocKind::Const,
+            ty::AssocKind::Fn => AssocKind::Fn,
+            ty::AssocKind::Type => AssocKind::Type,
+        }
+    }
+}
+
+impl<'tcx> Stable<'tcx> for ty::AssocItemContainer {
+    type T = stable_mir::ty::AssocItemContainer;
+
+    fn stable(&self, _tables: &mut Tables<'_>) -> Self::T {
+        use stable_mir::ty::AssocItemContainer;
+        match self {
+            ty::AssocItemContainer::Trait => AssocItemContainer::Trait,
+            ty::AssocItemContainer::Impl => AssocItemContainer::Impl,
+        }
+    }
+}
+
+impl<'tcx> Stable<'tcx> for ty::AssocItem {
+    type T = stable_mir::ty::AssocItem;
+
+    fn stable(&self, tables: &mut Tables<'_>) -> Self::T {
+        stable_mir::ty::AssocItem {
+            def_id: tables.assoc_def(self.def_id),
+            name: self.name.to_string(),
+            kind: self.kind.stable(tables),
+            container: self.container.stable(tables),
+            trait_item_def_id: self.trait_item_def_id.map(|did| tables.assoc_def(did)),
+            fn_has_self_parameter: self.fn_has_self_parameter,
+            opt_rpitit_info: self.opt_rpitit_info.map(|rpitit| rpitit.stable(tables)),
+        }
+    }
+}
+
+impl<'tcx> Stable<'tcx> for ty::ImplTraitInTraitData {
+    type T = stable_mir::ty::ImplTraitInTraitData;
+
+    fn stable(&self, tables: &mut Tables<'_>) -> Self::T {
+        use stable_mir::ty::ImplTraitInTraitData;
+        match self {
+            ty::ImplTraitInTraitData::Trait { fn_def_id, opaque_def_id } => {
+                ImplTraitInTraitData::Trait {
+                    fn_def_id: tables.fn_def(*fn_def_id),
+                    opaque_def_id: tables.opaque_def(*opaque_def_id),
+                }
+            }
+            ty::ImplTraitInTraitData::Impl { fn_def_id } => {
+                ImplTraitInTraitData::Impl { fn_def_id: tables.fn_def(*fn_def_id) }
+            }
+        }
+    }
+}
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index e82c957c34e..46154da36ca 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -18,8 +18,8 @@ use crate::ty::{
     TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef,
 };
 use crate::{
-    Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls, ItemKind,
-    Symbol, TraitDecls, mir,
+    AssocItems, Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls,
+    ItemKind, Symbol, TraitDecls, mir,
 };
 
 /// This trait defines the interface between stable_mir and the Rust compiler.
@@ -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
diff --git a/compiler/stable_mir/src/crate_def.rs b/compiler/stable_mir/src/crate_def.rs
index 2577c281ca4..75228135e4c 100644
--- a/compiler/stable_mir/src/crate_def.rs
+++ b/compiler/stable_mir/src/crate_def.rs
@@ -4,7 +4,7 @@
 use serde::Serialize;
 
 use crate::ty::{GenericArgs, Span, Ty};
-use crate::{Crate, Symbol, with};
+use crate::{AssocItems, Crate, Symbol, with};
 
 /// A unique identification number for each item accessible for the current compilation unit.
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize)]
@@ -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 {}
+    };
+}
diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs
index 70d42dfbfcb..df90d3e5a08 100644
--- a/compiler/stable_mir/src/lib.rs
+++ b/compiler/stable_mir/src/lib.rs
@@ -23,11 +23,11 @@ use std::{fmt, io};
 use serde::Serialize;
 
 use crate::compiler_interface::with;
-pub use crate::crate_def::{CrateDef, CrateDefType, DefId};
+pub use crate::crate_def::{CrateDef, CrateDefItems, CrateDefType, DefId};
 pub use crate::error::*;
 use crate::mir::mono::StaticDef;
 use crate::mir::{Body, Mutability};
-use crate::ty::{FnDef, ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
+use crate::ty::{AssocItem, FnDef, ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
 
 pub mod abi;
 #[macro_use]
@@ -71,6 +71,9 @@ pub type TraitDecls = Vec<TraitDef>;
 /// A list of impl trait decls.
 pub type ImplTraitDecls = Vec<ImplDef>;
 
+/// A list of associated items.
+pub type AssocItems = Vec<AssocItem>;
+
 /// Holds information about a crate.
 #[derive(Clone, PartialEq, Eq, Debug, Serialize)]
 pub struct Crate {
diff --git a/compiler/stable_mir/src/mir/pretty.rs b/compiler/stable_mir/src/mir/pretty.rs
index 8278afb7a2f..65d9f20f0a3 100644
--- a/compiler/stable_mir/src/mir/pretty.rs
+++ b/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)))
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index b857a735b72..25ec4a440d6 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/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;
@@ -910,6 +910,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 +936,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 +1563,60 @@ index_impl!(Span);
 pub struct VariantIdx(usize);
 
 index_impl!(VariantIdx);
+
+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()
+    }
+}
diff --git a/tests/ui-fulldeps/stable-mir/check_assoc_items.rs b/tests/ui-fulldeps/stable-mir/check_assoc_items.rs
new file mode 100644
index 00000000000..f6f895588f2
--- /dev/null
+++ b/tests/ui-fulldeps/stable-mir/check_assoc_items.rs
@@ -0,0 +1,145 @@
+//@ run-pass
+//! Test that users are able to retrieve all associated items from a definition.
+//! definition.
+
+//@ ignore-stage1
+//@ ignore-cross-compile
+//@ ignore-remote
+//@ edition: 2021
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+
+extern crate rustc_middle;
+#[macro_use]
+extern crate rustc_smir;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate stable_mir;
+
+use rustc_smir::rustc_internal;
+use std::io::Write;
+use std::collections::HashSet;
+use stable_mir::CrateDef;
+use stable_mir::*;
+use stable_mir::ty::*;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "crate_assoc_items";
+
+/// This function uses the Stable MIR APIs to get information about the test crate.
+fn test_assoc_items() -> ControlFlow<()> {
+    let local_crate = stable_mir::local_crate();
+    check_items(
+        &local_crate.fn_defs(),
+        &[
+            "AStruct::new",
+            "<AStruct as ATrait>::assoc_fn_no_self",
+            "<AStruct as ATrait>::assoc_fn_has_self",
+            "ATrait::rpitit",
+            "ATrait::assoc_fn_has_self",
+            "ATrait::assoc_fn_no_self",
+            "<AStruct as ATrait>::rpitit",
+        ],
+    );
+
+    let local_impls = local_crate.trait_impls();
+    let local_traits = local_crate.trait_decls();
+
+    let trait_assoc_item_defs: Vec<AssocDef> = local_traits[0].associated_items()
+        .iter().map(|assoc_item| assoc_item.def_id).collect();
+    check_items(
+        &trait_assoc_item_defs,
+        &[
+            "ATrait::{synthetic#0}",
+            "ATrait::rpitit",
+            "ATrait::Assoc",
+            "ATrait::assoc_fn_no_self",
+            "ATrait::assoc_fn_has_self",
+        ]
+    );
+
+    let impl_assoc_item_defs: Vec<AssocDef> = local_impls[0].associated_items()
+        .iter().map(|assoc_item| assoc_item.def_id).collect();
+    check_items(
+        &impl_assoc_item_defs,
+        &[
+            "<AStruct as ATrait>::{synthetic#0}",
+            "<AStruct as ATrait>::rpitit",
+            "<AStruct as ATrait>::Assoc",
+            "<AStruct as ATrait>::assoc_fn_no_self",
+            "<AStruct as ATrait>::assoc_fn_has_self",
+        ]
+    );
+
+    ControlFlow::Continue(())
+}
+
+/// Check if the list of definitions matches the expected list.
+/// Note that order doesn't matter.
+fn check_items<T: CrateDef>(items: &[T], expected: &[&str]) {
+    let expected: HashSet<_> = expected.iter().map(|s| s.to_string()).collect();
+    let item_names: HashSet<_> = items.iter().map(|item| item.name()).collect();
+    assert_eq!(item_names, expected);
+}
+
+fn main() {
+    let path = "assoc_items.rs";
+    generate_input(&path).unwrap();
+    let args = vec![
+        "rustc".to_string(),
+        "--crate-type=lib".to_string(),
+        "--crate-name".to_string(),
+        CRATE_NAME.to_string(),
+        path.to_string(),
+    ];
+    run!(args, test_assoc_items).unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+        #![allow(dead_code, unused_variables)]
+        struct AStruct;
+
+        impl AStruct {{
+            const ASSOC_CONST: &str = "Nina";
+
+            fn new() -> Self {{
+                AStruct{{}}
+            }}
+        }}
+
+        trait ATrait {{
+            type Assoc;
+
+            fn assoc_fn_no_self() {{
+            }}
+
+            fn assoc_fn_has_self(&self) {{
+            }}
+
+            fn rpitit(&self) -> impl std::fmt::Debug {{
+                "ciallo"
+            }}
+        }}
+
+        impl ATrait for AStruct {{
+            type Assoc = u32;
+
+            fn assoc_fn_no_self() {{
+            }}
+
+            fn assoc_fn_has_self(&self) {{
+            }}
+
+            fn rpitit(&self) -> impl std::fmt::Debug {{
+                "ciallo~"
+            }}
+        }}
+    "#
+    )?;
+    Ok(())
+}