Skip to content

Commit bd53d1e

Browse files
authored
Rollup merge of #122405 - celinval:smir-new-const, r=oli-obk
Add methods to create StableMIR constant I've been experimenting with transforming the StableMIR to instrument the code with potential UB checks. The modified body will only be used by our analysis tool, however, constants in StableMIR must be backed by rustc constants. Thus, I'm adding a few functions to build constants, such as building string and other primitives. One question I have is whether we should create a global allocation instead for strings. r? ``````@oli-obk``````
2 parents 7a744af + 893a910 commit bd53d1e

File tree

4 files changed

+237
-12
lines changed

4 files changed

+237
-12
lines changed

compiler/rustc_smir/src/rustc_smir/context.rs

+51-7
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ use stable_mir::mir::Body;
2323
use stable_mir::target::{MachineInfo, MachineSize};
2424
use stable_mir::ty::{
2525
AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, ForeignDef,
26-
ForeignItemKind, GenericArgs, LineInfo, PolyFnSig, RigidTy, Span, Ty, TyKind, VariantDef,
26+
ForeignItemKind, GenericArgs, LineInfo, PolyFnSig, RigidTy, Span, Ty, TyKind, UintTy,
27+
VariantDef,
2728
};
2829
use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, ItemKind, Symbol};
2930
use std::cell::RefCell;
@@ -341,15 +342,56 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
341342
.ok_or_else(|| Error::new(format!("Const `{cnst:?}` cannot be encoded as u64")))
342343
}
343344

344-
fn usize_to_const(&self, val: u64) -> Result<Const, Error> {
345+
fn try_new_const_zst(&self, ty: Ty) -> Result<Const, Error> {
345346
let mut tables = self.0.borrow_mut();
346-
let ty = tables.tcx.types.usize;
347+
let tcx = tables.tcx;
348+
let ty_internal = ty.internal(&mut *tables, tcx);
349+
let size = tables
350+
.tcx
351+
.layout_of(ParamEnv::empty().and(ty_internal))
352+
.map_err(|err| {
353+
Error::new(format!(
354+
"Cannot create a zero-sized constant for type `{ty_internal}`: {err}"
355+
))
356+
})?
357+
.size;
358+
if size.bytes() != 0 {
359+
return Err(Error::new(format!(
360+
"Cannot create a zero-sized constant for type `{ty_internal}`: \
361+
Type `{ty_internal}` has {} bytes",
362+
size.bytes()
363+
)));
364+
}
365+
366+
Ok(ty::Const::zero_sized(tables.tcx, ty_internal).stable(&mut *tables))
367+
}
368+
369+
fn new_const_str(&self, value: &str) -> Const {
370+
let mut tables = self.0.borrow_mut();
371+
let tcx = tables.tcx;
372+
let ty = ty::Ty::new_static_str(tcx);
373+
let bytes = value.as_bytes();
374+
let val_tree = ty::ValTree::from_raw_bytes(tcx, bytes);
375+
376+
ty::Const::new_value(tcx, val_tree, ty).stable(&mut *tables)
377+
}
378+
379+
fn new_const_bool(&self, value: bool) -> Const {
380+
let mut tables = self.0.borrow_mut();
381+
ty::Const::from_bool(tables.tcx, value).stable(&mut *tables)
382+
}
383+
384+
fn try_new_const_uint(&self, value: u128, uint_ty: UintTy) -> Result<Const, Error> {
385+
let mut tables = self.0.borrow_mut();
386+
let tcx = tables.tcx;
387+
let ty = ty::Ty::new_uint(tcx, uint_ty.internal(&mut *tables, tcx));
347388
let size = tables.tcx.layout_of(ParamEnv::empty().and(ty)).unwrap().size;
348389

349-
let scalar = ScalarInt::try_from_uint(val, size).ok_or_else(|| {
350-
Error::new(format!("Value overflow: cannot convert `{val}` to usize."))
390+
// We don't use Const::from_bits since it doesn't have any error checking.
391+
let scalar = ScalarInt::try_from_uint(value, size).ok_or_else(|| {
392+
Error::new(format!("Value overflow: cannot convert `{value}` to `{ty}`."))
351393
})?;
352-
Ok(rustc_middle::ty::Const::new_value(tables.tcx, ValTree::from_scalar_int(scalar), ty)
394+
Ok(ty::Const::new_value(tables.tcx, ValTree::from_scalar_int(scalar), ty)
353395
.stable(&mut *tables))
354396
}
355397

@@ -556,7 +598,9 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
556598
global_alloc: &GlobalAlloc,
557599
) -> Option<stable_mir::mir::alloc::AllocId> {
558600
let mut tables = self.0.borrow_mut();
559-
let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else { return None };
601+
let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else {
602+
return None;
603+
};
560604
let tcx = tables.tcx;
561605
let alloc_id = tables.tcx.vtable_allocation((
562606
ty.internal(&mut *tables, tcx),

compiler/stable_mir/src/compiler_interface.rs

+13-4
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::ty::{
1414
AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, ForeignDef,
1515
ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics,
1616
ImplDef, ImplTrait, LineInfo, PolyFnSig, RigidTy, Span, TraitDecl, TraitDef, Ty, TyKind,
17-
VariantDef,
17+
UintTy, VariantDef,
1818
};
1919
use crate::{
2020
mir, Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls, ItemKind,
@@ -101,8 +101,17 @@ pub trait Context {
101101
/// Evaluate constant as a target usize.
102102
fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error>;
103103

104-
/// Create a target usize constant for the given value.
105-
fn usize_to_const(&self, val: u64) -> Result<Const, Error>;
104+
/// Create a new zero-sized constant.
105+
fn try_new_const_zst(&self, ty: Ty) -> Result<Const, Error>;
106+
107+
/// Create a new constant that represents the given string value.
108+
fn new_const_str(&self, value: &str) -> Const;
109+
110+
/// Create a new constant that represents the given boolean value.
111+
fn new_const_bool(&self, value: bool) -> Const;
112+
113+
/// Create a new constant that represents the given value.
114+
fn try_new_const_uint(&self, value: u128, uint_ty: UintTy) -> Result<Const, Error>;
106115

107116
/// Create a new type from the given kind.
108117
fn new_rigid_ty(&self, kind: RigidTy) -> Ty;
@@ -200,7 +209,7 @@ pub trait Context {
200209

201210
// A thread local variable that stores a pointer to the tables mapping between TyCtxt
202211
// datastructures and stable MIR datastructures
203-
scoped_thread_local! (static TLV: Cell<*const ()>);
212+
scoped_thread_local!(static TLV: Cell<*const ()>);
204213

205214
pub fn run<F, T>(context: &dyn Context, f: F) -> Result<T, Error>
206215
where

compiler/stable_mir/src/ty.rs

+26-1
Original file line numberDiff line numberDiff line change
@@ -128,13 +128,38 @@ impl Const {
128128

129129
/// Creates an interned usize constant.
130130
fn try_from_target_usize(val: u64) -> Result<Self, Error> {
131-
with(|cx| cx.usize_to_const(val))
131+
with(|cx| cx.try_new_const_uint(val.into(), UintTy::Usize))
132132
}
133133

134134
/// Try to evaluate to a target `usize`.
135135
pub fn eval_target_usize(&self) -> Result<u64, Error> {
136136
with(|cx| cx.eval_target_usize(self))
137137
}
138+
139+
/// Create a constant that represents a new zero-sized constant of type T.
140+
/// Fails if the type is not a ZST or if it doesn't have a known size.
141+
pub fn try_new_zero_sized(ty: Ty) -> Result<Const, Error> {
142+
with(|cx| cx.try_new_const_zst(ty))
143+
}
144+
145+
/// Build a new constant that represents the given string.
146+
///
147+
/// Note that there is no guarantee today about duplication of the same constant.
148+
/// I.e.: Calling this function multiple times with the same argument may or may not return
149+
/// the same allocation.
150+
pub fn from_str(value: &str) -> Const {
151+
with(|cx| cx.new_const_str(value))
152+
}
153+
154+
/// Build a new constant that represents the given boolean value.
155+
pub fn from_bool(value: bool) -> Const {
156+
with(|cx| cx.new_const_bool(value))
157+
}
158+
159+
/// Build a new constant that represents the given unsigned integer.
160+
pub fn try_from_uint(value: u128, uint_ty: UintTy) -> Result<Const, Error> {
161+
with(|cx| cx.try_new_const_uint(value, uint_ty))
162+
}
138163
}
139164

140165
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
//@ run-pass
2+
//! Test a few methods to transform StableMIR.
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+
#![feature(control_flow_enum)]
12+
#![feature(ascii_char, ascii_char_variants)]
13+
14+
extern crate rustc_hir;
15+
#[macro_use]
16+
extern crate rustc_smir;
17+
extern crate rustc_driver;
18+
extern crate rustc_interface;
19+
extern crate stable_mir;
20+
21+
use rustc_smir::rustc_internal;
22+
use stable_mir::mir::alloc::GlobalAlloc;
23+
use stable_mir::mir::mono::Instance;
24+
use stable_mir::mir::{Body, Constant, Operand, Rvalue, StatementKind, TerminatorKind};
25+
use stable_mir::ty::{Const, ConstantKind};
26+
use stable_mir::{CrateDef, CrateItems, ItemKind};
27+
use std::convert::TryFrom;
28+
use std::io::Write;
29+
use std::ops::ControlFlow;
30+
31+
const CRATE_NAME: &str = "input";
32+
33+
/// This function uses the Stable MIR APIs to transform the MIR.
34+
fn test_transform() -> ControlFlow<()> {
35+
// Find items in the local crate.
36+
let items = stable_mir::all_local_items();
37+
38+
// Test fn_abi
39+
let target_fn = *get_item(&items, (ItemKind::Fn, "dummy")).unwrap();
40+
let instance = Instance::try_from(target_fn).unwrap();
41+
let body = instance.body().unwrap();
42+
check_msg(&body, "oops");
43+
44+
let new_msg = "new panic message";
45+
let new_body = change_panic_msg(body, new_msg);
46+
check_msg(&new_body, new_msg);
47+
48+
ControlFlow::Continue(())
49+
}
50+
51+
/// Check that the body panic message matches the given message.
52+
fn check_msg(body: &Body, expected: &str) {
53+
let msg = body
54+
.blocks
55+
.iter()
56+
.find_map(|bb| match &bb.terminator.kind {
57+
TerminatorKind::Call { args, .. } => {
58+
assert_eq!(args.len(), 1, "Expected panic message, but found {args:?}");
59+
let msg_const = match &args[0] {
60+
Operand::Constant(msg_const) => msg_const,
61+
Operand::Copy(place) | Operand::Move(place) => {
62+
assert!(place.projection.is_empty());
63+
bb.statements
64+
.iter()
65+
.find_map(|stmt| match &stmt.kind {
66+
StatementKind::Assign(
67+
destination,
68+
Rvalue::Use(Operand::Constant(msg_const)),
69+
) if destination == place => Some(msg_const),
70+
_ => None,
71+
})
72+
.unwrap()
73+
}
74+
};
75+
let ConstantKind::Allocated(alloc) = msg_const.literal.kind() else {
76+
unreachable!()
77+
};
78+
assert_eq!(alloc.provenance.ptrs.len(), 1);
79+
80+
let alloc_prov_id = alloc.provenance.ptrs[0].1 .0;
81+
let GlobalAlloc::Memory(val) = GlobalAlloc::from(alloc_prov_id) else {
82+
unreachable!()
83+
};
84+
let bytes = val.raw_bytes().unwrap();
85+
Some(std::str::from_utf8(&bytes).unwrap().to_string())
86+
}
87+
_ => None,
88+
})
89+
.expect("Failed to find panic message");
90+
assert_eq!(&msg, expected);
91+
}
92+
93+
/// Modify body to use a different panic message.
94+
fn change_panic_msg(mut body: Body, new_msg: &str) -> Body {
95+
for bb in &mut body.blocks {
96+
match &mut bb.terminator.kind {
97+
TerminatorKind::Call { args, .. } => {
98+
let new_const = Const::from_str(new_msg);
99+
args[0] = Operand::Constant(Constant {
100+
literal: new_const,
101+
span: bb.terminator.span,
102+
user_ty: None,
103+
});
104+
}
105+
_ => {}
106+
}
107+
}
108+
body
109+
}
110+
111+
fn get_item<'a>(
112+
items: &'a CrateItems,
113+
item: (ItemKind, &str),
114+
) -> Option<&'a stable_mir::CrateItem> {
115+
items.iter().find(|crate_item| (item.0 == crate_item.kind()) && crate_item.name() == item.1)
116+
}
117+
118+
/// This test will generate and analyze a dummy crate using the stable mir.
119+
/// For that, it will first write the dummy crate into a file.
120+
/// Then it will create a `StableMir` using custom arguments and then
121+
/// it will run the compiler.
122+
fn main() {
123+
let path = "transform_input.rs";
124+
generate_input(&path).unwrap();
125+
let args = vec![
126+
"rustc".to_string(),
127+
"--crate-type=lib".to_string(),
128+
"--crate-name".to_string(),
129+
CRATE_NAME.to_string(),
130+
path.to_string(),
131+
];
132+
run!(args, test_transform).unwrap();
133+
}
134+
135+
fn generate_input(path: &str) -> std::io::Result<()> {
136+
let mut file = std::fs::File::create(path)?;
137+
write!(
138+
file,
139+
r#"
140+
#![feature(panic_internals)]
141+
pub fn dummy() {{
142+
core::panicking::panic_str("oops");
143+
}}
144+
"#
145+
)?;
146+
Ok(())
147+
}

0 commit comments

Comments
 (0)