Rollup merge of #132131 - celinval:smir-crate-defs, r=compiler-errors

[StableMIR] API to retrieve definitions from crates

Add functions to retrieve function definitions and static items from all crates (local and external).

For external crates, we're still missing items from trait implementation and primitives.

r? ````@compiler-errors:```` Do you know what is the best way to retrieve the associated items for primitives and trait implementations for external crates? Thanks!

7 files changed, 225 insertions, 4 deletions

diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
index f06f2fdc5e5..045fd0565ba 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -384,6 +384,7 @@ provide! { tcx, def_id, other, cdata,
     crate_hash => { cdata.root.header.hash }
     crate_host_hash => { cdata.host_hash }
     crate_name => { cdata.root.header.name }
+    num_extern_def_ids => { cdata.num_def_ids() }
 
     extra_filename => { cdata.root.extra_filename.clone() }
 
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index f58426601cd..8ae89ea6e2c 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1844,6 +1844,16 @@ rustc_queries! {
         desc { |tcx| "computing crate imported by `{}`", tcx.def_path_str(def_id) }
     }
 
+    /// Gets the number of definitions in a foreign crate.
+    ///
+    /// This allows external tools to iterate over all definitions in a foreign crate.
+    ///
+    /// This should never be used for the local crate, instead use `iter_local_def_id`.
+    query num_extern_def_ids(_: CrateNum) -> usize {
+        desc { "fetching the number of definitions in a crate" }
+        separate_provide_extern
+    }
+
     query lib_features(_: CrateNum) -> &'tcx LibFeatures {
         desc { "calculating the lib features defined in a crate" }
         separate_provide_extern
diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index 2d6ca3571fa..3db65692af7 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -34,7 +34,7 @@ use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, I
 
 use crate::rustc_internal::RustcInternal;
 use crate::rustc_smir::builder::BodyBuilder;
-use crate::rustc_smir::{Stable, Tables, alloc, new_item_kind, smir_crate};
+use crate::rustc_smir::{Stable, Tables, alloc, filter_def_ids, new_item_kind, smir_crate};
 
 impl<'tcx> Context for TablesWrapper<'tcx> {
     fn target_info(&self) -> MachineInfo {
@@ -80,6 +80,20 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
             .collect()
     }
 
+    fn crate_functions(&self, crate_num: CrateNum) -> Vec<FnDef> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let krate = crate_num.internal(&mut *tables, tcx);
+        filter_def_ids(tcx, krate, |def_id| tables.to_fn_def(def_id))
+    }
+
+    fn crate_statics(&self, crate_num: CrateNum) -> Vec<StaticDef> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let krate = crate_num.internal(&mut *tables, tcx);
+        filter_def_ids(tcx, krate, |def_id| tables.to_static(def_id))
+    }
+
     fn foreign_module(
         &self,
         mod_def: stable_mir::ty::ForeignModuleDef,
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 9032156b257..c5d33f090a0 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -15,8 +15,8 @@ use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use stable_mir::abi::Layout;
-use stable_mir::mir::mono::InstanceDef;
-use stable_mir::ty::{MirConstId, Span, TyConstId};
+use stable_mir::mir::mono::{InstanceDef, StaticDef};
+use stable_mir::ty::{FnDef, MirConstId, Span, TyConstId};
 use stable_mir::{CtorKind, ItemKind};
 use tracing::debug;
 
@@ -79,6 +79,36 @@ impl<'tcx> Tables<'tcx> {
         };
         !must_override && self.tcx.is_mir_available(def_id)
     }
+
+    fn to_fn_def(&mut self, def_id: DefId) -> Option<FnDef> {
+        if matches!(self.tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+            Some(self.fn_def(def_id))
+        } else {
+            None
+        }
+    }
+
+    fn to_static(&mut self, def_id: DefId) -> Option<StaticDef> {
+        matches!(self.tcx.def_kind(def_id), DefKind::Static { .. }).then(|| self.static_def(def_id))
+    }
+}
+
+/// Iterate over the definitions of the given crate.
+pub(crate) fn filter_def_ids<F, T>(tcx: TyCtxt<'_>, krate: CrateNum, mut func: F) -> Vec<T>
+where
+    F: FnMut(DefId) -> Option<T>,
+{
+    if krate == LOCAL_CRATE {
+        tcx.iter_local_def_id().filter_map(|did| func(did.to_def_id())).collect()
+    } else {
+        let num_definitions = tcx.num_extern_def_ids(krate);
+        (0..num_definitions)
+            .filter_map(move |i| {
+                let def_id = DefId { krate, index: rustc_span::def_id::DefIndex::from_usize(i) };
+                func(def_id)
+            })
+            .collect()
+    }
 }
 
 /// Build a stable mir crate from a given crate number.
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index 4b7707ebccf..a6f7c254583 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -34,6 +34,12 @@ pub trait Context {
     /// Check whether the body of a function is available.
     fn has_body(&self, item: DefId) -> bool;
     fn foreign_modules(&self, crate_num: CrateNum) -> Vec<ForeignModuleDef>;
+
+    /// Retrieve all functions defined in this crate.
+    fn crate_functions(&self, crate_num: CrateNum) -> Vec<FnDef>;
+
+    /// Retrieve all static items defined in this crate.
+    fn crate_statics(&self, crate_num: CrateNum) -> Vec<StaticDef>;
     fn foreign_module(&self, mod_def: ForeignModuleDef) -> ForeignModule;
     fn foreign_items(&self, mod_def: ForeignModuleDef) -> Vec<ForeignDef>;
     fn all_trait_decls(&self) -> TraitDecls;
diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs
index b523e949cde..0b4cebadad1 100644
--- a/compiler/stable_mir/src/lib.rs
+++ b/compiler/stable_mir/src/lib.rs
@@ -25,8 +25,9 @@ use serde::Serialize;
 use crate::compiler_interface::with;
 pub use crate::crate_def::{CrateDef, CrateDefType, DefId};
 pub use crate::error::*;
+use crate::mir::mono::StaticDef;
 use crate::mir::{Body, Mutability};
-use crate::ty::{ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
+use crate::ty::{FnDef, ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
 
 pub mod abi;
 #[macro_use]
@@ -96,6 +97,16 @@ impl Crate {
     pub fn trait_impls(&self) -> ImplTraitDecls {
         with(|cx| cx.trait_impls(self.id))
     }
+
+    /// Return a list of function definitions from this crate independent on their visibility.
+    pub fn fn_defs(&self) -> Vec<FnDef> {
+        with(|cx| cx.crate_functions(self.id))
+    }
+
+    /// Return a list of static items defined in this crate independent on their visibility.
+    pub fn statics(&self) -> Vec<StaticDef> {
+        with(|cx| cx.crate_statics(self.id))
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, Serialize)]
diff --git a/tests/ui-fulldeps/stable-mir/check_crate_defs.rs b/tests/ui-fulldeps/stable-mir/check_crate_defs.rs
new file mode 100644
index 00000000000..e039ca07dd4
--- /dev/null
+++ b/tests/ui-fulldeps/stable-mir/check_crate_defs.rs
@@ -0,0 +1,149 @@
+//@ run-pass
+//! Test information about crate definitions (local and external).
+
+//@ ignore-stage1
+//@ ignore-cross-compile
+//@ ignore-remote
+//@ ignore-windows-gnu mingw has troubles with linking https://github.com/rust-lang/rust/pull/116837
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+
+extern crate rustc_hir;
+#[macro_use]
+extern crate rustc_smir;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate stable_mir;
+
+use rustc_smir::rustc_internal;
+use stable_mir::CrateDef;
+use std::collections::HashSet;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "crate_defs";
+
+/// This function uses the Stable MIR APIs to get information about the test crate.
+fn test_stable_mir() -> ControlFlow<()> {
+    // Find items in the local crate.
+    let local = stable_mir::local_crate();
+    check_items(&local.statics(), &["PRIVATE_STATIC", "dummy::PUBLIC_STATIC"]);
+    check_items(
+        &local.fn_defs(),
+        &[
+            "top_level",
+            "dummy::public_fn",
+            "dummy::private_fn",
+            "dummy::PrivateStruct::new",
+            "<dummy::PrivateStruct as std::ops::Drop>::drop",
+            "DummyTrait::method",
+            "<T as DummyTrait>::method",
+        ],
+    );
+
+    // Find items inside core crate.
+    // FIXME: We are currently missing primitive type methods and trait implementations for external
+    // crates.
+    let core = stable_mir::find_crates("core").pop().expect("Cannot find `core` crate");
+    contains(
+        &core.fn_defs(),
+        &[
+            "std::fmt::Debug::fmt",
+            "std::option::Option::<T>::is_some",
+            "std::ptr::swap",
+            "<std::slice::Iter<'a, T> as std::iter::Iterator>::next",
+            "core::num::<impl u8>::abs_diff",
+        ],
+    );
+    // Ensure nothing crashes. There is no public static in core that we can test here.
+    let _ = core.statics();
+
+    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);
+}
+
+/// Check that the list contains the expected items.
+fn contains<T: CrateDef + std::fmt::Debug>(items: &[T], expected: &[&str]) {
+    let expected: HashSet<_> = expected.iter().map(|s| s.to_string()).collect();
+    let item_names = items.iter().map(|item| item.name()).collect();
+    let not_found: Vec<_> = expected.difference(&item_names).collect();
+    assert!(not_found.is_empty(), "Missing items: {:?}", not_found);
+}
+
+/// This test will generate and analyze a dummy crate using the stable mir.
+/// For that, it will first write the dummy crate into a file.
+/// Then it will create a `StableMir` using custom arguments and then
+/// it will run the compiler.
+fn main() {
+    let path = "crate_definitions.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_stable_mir).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)]
+        static PRIVATE_STATIC: u8 = 0;
+        fn top_level() -> &'static str {{
+            "hello"
+        }}
+
+        pub trait DummyTrait {{
+            fn method(&self) -> Self;
+        }}
+
+        impl<T: Copy> DummyTrait for T {{
+            fn method(&self) -> T {{
+                *self
+            }}
+        }}
+
+        pub mod dummy {{
+            pub static mut PUBLIC_STATIC: Option<char> = None;
+
+            pub fn public_fn(input: bool) -> bool {{
+                private_fn(!input)
+            }}
+
+            fn private_fn(input: bool) -> bool {{
+                todo!()
+            }}
+
+            struct PrivateStruct {{
+                field: u32,
+            }}
+
+            impl PrivateStruct {{
+                fn new() -> Self {{
+                    Self {{ field: 42 }}
+                }}
+            }}
+
+            impl Drop for PrivateStruct {{
+                fn drop(&mut self) {{
+                    println!("Dropping PrivateStruct");
+                }}
+            }}
+        }}
+        "#
+    )?;
+    Ok(())
+}