[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, add a query to retrieve the number of defs in a
foreign crate.

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -387,6 +387,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() }
 

compiler/rustc_middle/src/query/mod.rs

@@ -1812,6 +1812,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

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,

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::mir::mono::{InstanceDef, StaticDef};
-use stable_mir::ty::{MirConstId, Span, TyConstId};
+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.

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;

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)]

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(())
+}