Remove tcx function and make internal fn safer

I added `tcx` argument to `internal` to force 'tcx to be the same
lifetime as TyCtxt. The only other solution I could think is to change
this function to be `unsafe`.
This commit is contained in:
Celina G. Val 2024-01-16 14:31:25 -08:00
parent f91ccf9ace
commit 2564811e7b
5 changed files with 66 additions and 39 deletions

View File

@ -28,7 +28,7 @@
///
/// # Warning
///
/// This function is unstable, and it's behavior may change at any point.
/// This function is unstable, and its behavior may change at any point.
/// E.g.: Items that were previously supported, may no longer be supported, or its translation may
/// change.
///
@ -50,23 +50,12 @@ pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
/// # Panics
///
/// This function will panic if StableMIR has not been properly initialized.
pub fn internal<'tcx, S: RustcInternal<'tcx>>(item: S) -> S::T {
pub fn internal<'tcx, S: RustcInternal<'tcx>>(tcx: TyCtxt<'tcx>, item: S) -> S::T {
// The tcx argument ensures that the item won't outlive the type context.
let _ = tcx;
with_tables(|tables| item.internal(tables))
}
/// Retrieve the internal Rust compiler type context.
///
/// # Warning
///
/// This function is unstable, and it's behavior may change at any point.
///
/// # Panics
///
/// This function will panic if StableMIR has not been properly initialized.
pub fn tcx<'tcx>() -> TyCtxt<'tcx> {
with_tables(|tables| tables.tcx)
}
impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
type Output = DefId;
@ -227,7 +216,7 @@ pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
/// Instantiate and run the compiler with the provided arguments and callback.
///
/// The callback will be invoked after the compiler ran all its analysis, but before code generation.
/// The callback will be invoked after the compiler ran all its analyses, but before code generation.
/// Note that this macro accepts two different formats for the callback:
/// 1. An ident that resolves to a function that accepts no argument and returns `ControlFlow<B, C>`
/// ```ignore(needs-extern-crate)
@ -249,7 +238,7 @@ pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
/// # assert_eq!(result, Err(CompilerError::Skipped))
/// # }
/// ```
/// 2. An expression that represents the body of a closure:
/// 2. A closure expression:
/// ```ignore(needs-extern-crate)
/// # extern crate rustc_driver;
/// # extern crate rustc_interface;
@ -267,26 +256,63 @@ pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
/// }
/// # let args = vec!["--verbose".to_string()];
/// # let extra_args = vec![];
/// let result = run!(args, analyze_code(extra_args));
/// let result = run!(args, || analyze_code(extra_args));
/// # assert_eq!(result, Err(CompilerError::Skipped))
/// # }
/// ```
#[macro_export]
macro_rules! run {
($args:expr, $callback_fn:ident) => {
run!($args, $callback_fn())
run_driver!($args, || $callback_fn())
};
($args:expr, $callback:expr) => {{
($args:expr, $callback:expr) => {
run_driver!($args, $callback)
};
}
/// Instantiate and run the compiler with the provided arguments and callback.
///
/// This is similar to `run` but it invokes the callback with the compiler's `TyCtxt`,
/// which can be used to invoke internal APIs.
#[macro_export]
macro_rules! run_with_tcx {
($args:expr, $callback_fn:ident) => {
run_driver!($args, |tcx| $callback_fn(tcx), with_tcx)
};
($args:expr, $callback:expr) => {
run_driver!($args, $callback, with_tcx)
};
}
/// Optionally include an ident. This is needed due to macro hygiene.
#[macro_export]
#[doc(hidden)]
macro_rules! optional {
(with_tcx $ident:ident) => {
$ident
};
}
/// Prefer using [run] and [run_with_tcx] instead.
///
/// This macro implements the instantiation of a StableMIR driver, and it will invoke
/// the given callback after the compiler analyses.
///
/// The third argument determines whether the callback requires `tcx` as an argument.
#[macro_export]
#[doc(hidden)]
macro_rules! run_driver {
($args:expr, $callback:expr $(, $with_tcx:ident)?) => {{
use rustc_driver::{Callbacks, Compilation, RunCompiler};
use rustc_interface::{interface, Queries};
use stable_mir::CompilerError;
use std::ops::ControlFlow;
pub struct StableMir<B = (), C = (), F = fn() -> ControlFlow<B, C>>
pub struct StableMir<B = (), C = (), F = fn($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C>>
where
B: Send,
C: Send,
F: FnOnce() -> ControlFlow<B, C> + Send,
F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
{
args: Vec<String>,
callback: Option<F>,
@ -297,7 +323,7 @@ impl<B, C, F> StableMir<B, C, F>
where
B: Send,
C: Send,
F: FnOnce() -> ControlFlow<B, C> + Send,
F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
{
/// Creates a new `StableMir` instance, with given test_function and arguments.
pub fn new(args: Vec<String>, callback: F) -> Self {
@ -325,7 +351,7 @@ impl<B, C, F> Callbacks for StableMir<B, C, F>
where
B: Send,
C: Send,
F: FnOnce() -> ControlFlow<B, C> + Send,
F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
{
/// Called after analysis. Return value instructs the compiler whether to
/// continue the compilation afterwards (defaults to `Compilation::Continue`)
@ -337,7 +363,7 @@ fn after_analysis<'tcx>(
queries.global_ctxt().unwrap().enter(|tcx| {
if let Some(callback) = self.callback.take() {
rustc_internal::run(tcx, || {
self.result = Some((callback)());
self.result = Some(callback($(optional!($with_tcx tcx))?));
})
.unwrap();
if self.result.as_ref().is_some_and(|val| val.is_continue()) {
@ -352,7 +378,7 @@ fn after_analysis<'tcx>(
}
}
StableMir::new($args, || $callback).run()
StableMir::new($args, $callback).run()
}};
}

View File

@ -29,7 +29,7 @@
use std::cell::RefCell;
use std::iter;
use crate::rustc_internal::{internal, RustcInternal};
use crate::rustc_internal::RustcInternal;
use crate::rustc_smir::builder::BodyBuilder;
use crate::rustc_smir::{alloc, new_item_kind, smir_crate, Stable, Tables};
@ -322,7 +322,8 @@ fn def_ty_with_args(&self, item: stable_mir::DefId, args: &GenericArgs) -> stabl
}
fn const_literal(&self, cnst: &stable_mir::ty::Const) -> String {
internal(cnst).to_string()
let mut tables = self.0.borrow_mut();
cnst.internal(&mut *tables).to_string()
}
fn span_of_an_item(&self, def_id: stable_mir::DefId) -> Span {

View File

@ -11,14 +11,12 @@
#![feature(assert_matches)]
#![feature(control_flow_enum)]
extern crate rustc_middle;
#[macro_use]
extern crate rustc_smir;
extern crate rustc_driver;
extern crate rustc_interface;
extern crate stable_mir;
use rustc_middle::ty::TyCtxt;
use rustc_smir::rustc_internal;
use stable_mir::CrateDef;
use std::collections::HashSet;
@ -83,7 +81,7 @@ fn main() {
CRATE_NAME.to_string(),
path.to_string(),
];
run!(args, test_traits()).unwrap();
run!(args, test_traits).unwrap();
}
fn generate_input(path: &str) -> std::io::Result<()> {

View File

@ -35,33 +35,33 @@ fn main() {
}
fn test_continue(args: Vec<String>) {
let result = run!(args, ControlFlow::Continue::<(), bool>(true));
let result = run!(args, || ControlFlow::Continue::<(), bool>(true));
assert_eq!(result, Ok(true));
}
fn test_break(args: Vec<String>) {
let result = run!(args, ControlFlow::Break::<bool, i32>(false));
let result = run!(args, || ControlFlow::Break::<bool, i32>(false));
assert_eq!(result, Err(stable_mir::CompilerError::Interrupted(false)));
}
#[allow(unreachable_code)]
fn test_skipped(mut args: Vec<String>) {
args.push("--version".to_string());
let result = run!(args, unreachable!() as ControlFlow<()>);
let result = run!(args, || unreachable!() as ControlFlow<()>);
assert_eq!(result, Err(stable_mir::CompilerError::Skipped));
}
#[allow(unreachable_code)]
fn test_failed(mut args: Vec<String>) {
args.push("--cfg=broken".to_string());
let result = run!(args, unreachable!() as ControlFlow<()>);
let result = run!(args, || unreachable!() as ControlFlow<()>);
assert_eq!(result, Err(stable_mir::CompilerError::CompilationFailed));
}
/// Test that we are able to pass a closure and set the return according to the captured value.
fn test_captured(args: Vec<String>) {
let captured = "10".to_string();
let result = run!(args, ControlFlow::Continue::<(), usize>(captured.len()));
let result = run!(args, || ControlFlow::Continue::<(), usize>(captured.len()));
assert_eq!(result, Ok(captured.len()));
}

View File

@ -16,19 +16,21 @@
extern crate rustc_smir;
extern crate rustc_driver;
extern crate rustc_interface;
extern crate rustc_middle;
extern crate stable_mir;
use rustc_middle::ty::TyCtxt;
use rustc_smir::rustc_internal;
use std::io::Write;
use std::ops::ControlFlow;
const CRATE_NAME: &str = "input";
fn test_translation() -> ControlFlow<()> {
fn test_translation(tcx: TyCtxt) -> ControlFlow<()> {
let main_fn = stable_mir::entry_fn().unwrap();
let body = main_fn.body();
let orig_ty = body.locals()[0].ty;
let rustc_ty = rustc_internal::internal(&orig_ty);
let rustc_ty = rustc_internal::internal(tcx, &orig_ty);
assert!(rustc_ty.is_unit());
ControlFlow::Continue(())
}
@ -46,7 +48,7 @@ fn main() {
CRATE_NAME.to_string(),
path.to_string(),
];
run!(args, test_translation).unwrap();
run_with_tcx!(args, test_translation).unwrap();
}
fn generate_input(path: &str) -> std::io::Result<()> {