Skip to content

Commit e4a47ad

Browse files
committed
[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.
1 parent f61306d commit e4a47ad

File tree

5 files changed

+251
-4
lines changed

5 files changed

+251
-4
lines changed

compiler/rustc_smir/src/rustc_smir/context.rs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,34 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
8080
.collect()
8181
}
8282

83+
fn crate_functions(&self, crate_num: CrateNum) -> Vec<FnDef> {
84+
let mut tables = self.0.borrow_mut();
85+
let tcx = tables.tcx;
86+
let krate = crate_num.internal(&mut *tables, tcx);
87+
if krate == LOCAL_CRATE {
88+
tcx.hir_crate_items(())
89+
.definitions()
90+
.filter_map(|local_id| tables.to_fn_def(tcx, local_id.to_def_id()))
91+
.collect()
92+
} else {
93+
tables.filter_map_items(tcx, krate, |tables, tcx, def_id| tables.to_fn_def(tcx, def_id))
94+
}
95+
}
96+
97+
fn crate_statics(&self, crate_num: CrateNum) -> Vec<StaticDef> {
98+
let mut tables = self.0.borrow_mut();
99+
let tcx = tables.tcx;
100+
let krate = crate_num.internal(&mut *tables, tcx);
101+
if krate == LOCAL_CRATE {
102+
tcx.hir_crate_items(())
103+
.definitions()
104+
.filter_map(|local_id| tables.to_static(tcx, local_id.to_def_id()))
105+
.collect()
106+
} else {
107+
tables.filter_map_items(tcx, krate, |tables, tcx, def_id| tables.to_static(tcx, def_id))
108+
}
109+
}
110+
83111
fn foreign_module(
84112
&self,
85113
mod_def: stable_mir::ty::ForeignModuleDef,

compiler/rustc_smir/src/rustc_smir/mod.rs

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,12 @@ use std::ops::RangeInclusive;
1212
use rustc_hir::def::DefKind;
1313
use rustc_middle::mir;
1414
use rustc_middle::mir::interpret::AllocId;
15+
use rustc_middle::ty::data_structures::HashSet;
1516
use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
16-
use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
17+
use rustc_span::def_id::{CRATE_DEF_INDEX, CrateNum, DefId, LOCAL_CRATE};
1718
use stable_mir::abi::Layout;
18-
use stable_mir::mir::mono::InstanceDef;
19-
use stable_mir::ty::{MirConstId, Span, TyConstId};
19+
use stable_mir::mir::mono::{InstanceDef, StaticDef};
20+
use stable_mir::ty::{FnDef, MirConstId, Span, TyConstId};
2021
use stable_mir::{CtorKind, ItemKind};
2122
use tracing::debug;
2223

@@ -79,6 +80,57 @@ impl<'tcx> Tables<'tcx> {
7980
};
8081
!must_override && self.tcx.is_mir_available(def_id)
8182
}
83+
84+
fn to_fn_def(&mut self, tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<FnDef> {
85+
if matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
86+
Some(self.fn_def(def_id))
87+
} else {
88+
None
89+
}
90+
}
91+
92+
fn to_static(&mut self, tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<StaticDef> {
93+
matches!(tcx.def_kind(def_id), DefKind::Static { .. }).then(|| self.static_def(def_id))
94+
}
95+
96+
fn filter_map_items<T, F>(&mut self, tcx: TyCtxt<'tcx>, krate: CrateNum, func: F) -> Vec<T>
97+
where
98+
F: Fn(&mut Tables<'tcx>, TyCtxt<'tcx>, DefId) -> Option<T>,
99+
{
100+
let crate_def = rustc_span::def_id::DefId { index: CRATE_DEF_INDEX, krate };
101+
let mut queue = vec![crate_def];
102+
let mut visited = HashSet::default();
103+
let mut result = vec![];
104+
while let Some(next) = queue.pop() {
105+
for def_id in tcx.module_children(next).iter().filter_map(|item| item.res.opt_def_id())
106+
{
107+
if !visited.contains(&def_id) {
108+
visited.insert(def_id);
109+
result.extend(func(self, tcx, def_id));
110+
match tcx.def_kind(def_id) {
111+
DefKind::Mod | DefKind::ForeignMod => queue.push(def_id),
112+
DefKind::Struct | DefKind::Enum | DefKind::Union => {
113+
for associated_item in tcx
114+
.inherent_impls(def_id)
115+
.iter()
116+
.flat_map(|impl_id| tcx.associated_item_def_ids(impl_id))
117+
{
118+
result.extend(func(self, tcx, *associated_item));
119+
}
120+
}
121+
DefKind::Trait => {
122+
for associated_item in tcx.associated_item_def_ids(def_id) {
123+
result.extend(func(self, tcx, *associated_item));
124+
}
125+
}
126+
_ => {}
127+
}
128+
}
129+
}
130+
}
131+
132+
result
133+
}
82134
}
83135

84136
/// Build a stable mir crate from a given crate number.

compiler/stable_mir/src/compiler_interface.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,12 @@ pub trait Context {
3434
/// Check whether the body of a function is available.
3535
fn has_body(&self, item: DefId) -> bool;
3636
fn foreign_modules(&self, crate_num: CrateNum) -> Vec<ForeignModuleDef>;
37+
38+
/// Retrieve all functions defined in this crate.
39+
fn crate_functions(&self, crate_num: CrateNum) -> Vec<FnDef>;
40+
41+
/// Retrieve all static items defined in this crate.
42+
fn crate_statics(&self, crate_num: CrateNum) -> Vec<StaticDef>;
3743
fn foreign_module(&self, mod_def: ForeignModuleDef) -> ForeignModule;
3844
fn foreign_items(&self, mod_def: ForeignModuleDef) -> Vec<ForeignDef>;
3945
fn all_trait_decls(&self) -> TraitDecls;

compiler/stable_mir/src/lib.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@ use serde::Serialize;
2525
use crate::compiler_interface::with;
2626
pub use crate::crate_def::{CrateDef, CrateDefType, DefId};
2727
pub use crate::error::*;
28+
use crate::mir::mono::StaticDef;
2829
use crate::mir::{Body, Mutability};
29-
use crate::ty::{ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
30+
use crate::ty::{FnDef, ForeignModuleDef, ImplDef, IndexedVal, Span, TraitDef, Ty};
3031

3132
pub mod abi;
3233
#[macro_use]
@@ -96,6 +97,16 @@ impl Crate {
9697
pub fn trait_impls(&self) -> ImplTraitDecls {
9798
with(|cx| cx.trait_impls(self.id))
9899
}
100+
101+
/// Return a list of function definitions from this crate independent on their visibility.
102+
pub fn fn_defs(&self) -> Vec<FnDef> {
103+
with(|cx| cx.crate_functions(self.id))
104+
}
105+
106+
/// Return a list of static items defined in this crate independent on their visibility.
107+
pub fn statics(&self) -> Vec<StaticDef> {
108+
with(|cx| cx.crate_statics(self.id))
109+
}
99110
}
100111

101112
#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, Serialize)]
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
//@ run-pass
2+
//! Test information about crate definitions (local and external).
3+
4+
//@ ignore-stage1
5+
//@ ignore-cross-compile
6+
//@ ignore-remote
7+
//@ ignore-windows-gnu mingw has troubles with linking https://github.com/rust-lang/rust/pull/116837
8+
9+
#![feature(rustc_private)]
10+
#![feature(assert_matches)]
11+
12+
extern crate rustc_hir;
13+
#[macro_use]
14+
extern crate rustc_smir;
15+
extern crate rustc_driver;
16+
extern crate rustc_interface;
17+
extern crate stable_mir;
18+
19+
use rustc_smir::rustc_internal;
20+
use stable_mir::CrateDef;
21+
use std::io::Write;
22+
use std::ops::ControlFlow;
23+
24+
const CRATE_NAME: &str = "crate_defs";
25+
26+
/// This function uses the Stable MIR APIs to get information about the test crate.
27+
fn test_stable_mir() -> ControlFlow<()> {
28+
// Find items in the local crate.
29+
let local = stable_mir::local_crate();
30+
check_items(&local.statics(), &["PRIVATE_STATIC", "dummy::PUBLIC_STATIC"]);
31+
check_items(
32+
&local.fn_defs(),
33+
&[
34+
"top_level",
35+
"dummy::public_fn",
36+
"dummy::private_fn",
37+
"dummy::PrivateStruct::new",
38+
"<dummy::PrivateStruct as std::ops::Drop>::drop",
39+
"DummyTrait::method",
40+
"<T as DummyTrait>::method",
41+
],
42+
);
43+
44+
// Find items inside core crate.
45+
// FIXME: We are currently missing primitive type methods and trait implementations for external
46+
// crates.
47+
let core = stable_mir::find_crates("core").pop().expect("Cannot find `core` crate");
48+
contains(
49+
&core.fn_defs(),
50+
&["std::fmt::Debug::fmt", "std::option::Option::<T>::is_some", "std::ptr::swap"],
51+
);
52+
// Ensure nothing crashes. There is no public static in core that we can test here.
53+
let _ = core.statics();
54+
55+
ControlFlow::Continue(())
56+
}
57+
58+
/// Check if the list of definitions matches the expected list.
59+
/// Note that order doesn't matter.
60+
fn check_items<T: CrateDef>(items: &[T], expected: &[&str]) {
61+
let item_names: Vec<_> = items.iter().map(|item| item.name()).collect();
62+
let unexpected: Vec<_> =
63+
item_names.iter().filter(|item| !expected.contains(&item.as_str())).collect();
64+
assert!(unexpected.is_empty(), "Unexpected items: {:?}", unexpected);
65+
assert_eq!(items.len(), expected.len(), "Expected items: {expected:?}\nFound: {item_names:?}");
66+
}
67+
68+
/// Check that the list contains the expected items.
69+
fn contains<T: CrateDef>(items: &[T], expected: &[&str]) {
70+
let item_names: Vec<_> = items.iter().map(|item| item.name()).collect();
71+
let found: Vec<_> =
72+
item_names.iter().filter(|item| expected.contains(&item.as_str())).collect();
73+
assert!(items.len() >= expected.len(), "Missing items: {:?}", item_names);
74+
assert_eq!(
75+
found.len(),
76+
expected.len(),
77+
"Expected items: {:?}\nFound: {found:?}",
78+
&item_names[0..10]
79+
);
80+
}
81+
82+
/// This test will generate and analyze a dummy crate using the stable mir.
83+
/// For that, it will first write the dummy crate into a file.
84+
/// Then it will create a `StableMir` using custom arguments and then
85+
/// it will run the compiler.
86+
fn main() {
87+
let path = "crate_definitions.rs";
88+
generate_input(&path).unwrap();
89+
let args = vec![
90+
"rustc".to_string(),
91+
"--crate-type=lib".to_string(),
92+
"--crate-name".to_string(),
93+
CRATE_NAME.to_string(),
94+
path.to_string(),
95+
];
96+
run!(args, test_stable_mir).unwrap();
97+
}
98+
99+
fn generate_input(path: &str) -> std::io::Result<()> {
100+
let mut file = std::fs::File::create(path)?;
101+
write!(
102+
file,
103+
r#"
104+
#![allow(dead_code, unused_variables)]
105+
static PRIVATE_STATIC: u8 = 0;
106+
fn top_level() -> &'static str {{
107+
"hello"
108+
}}
109+
110+
pub trait DummyTrait {{
111+
fn method(&self) -> Self;
112+
}}
113+
114+
impl<T: Copy> DummyTrait for T {{
115+
fn method(&self) -> T {{
116+
*self
117+
}}
118+
}}
119+
120+
pub mod dummy {{
121+
pub static mut PUBLIC_STATIC: Option<char> = None;
122+
123+
pub fn public_fn(input: bool) -> bool {{
124+
private_fn(!input)
125+
}}
126+
127+
fn private_fn(input: bool) -> bool {{
128+
todo!()
129+
}}
130+
131+
struct PrivateStruct {{
132+
field: u32,
133+
}}
134+
135+
impl PrivateStruct {{
136+
fn new() -> Self {{
137+
Self {{ field: 42 }}
138+
}}
139+
}}
140+
141+
impl Drop for PrivateStruct {{
142+
fn drop(&mut self) {{
143+
println!("Dropping PrivateStruct");
144+
}}
145+
}}
146+
}}
147+
"#
148+
)?;
149+
Ok(())
150+
}

0 commit comments

Comments
 (0)