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`.

1 files changed, 52 insertions, 26 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 7633739f289..873bd94e88f 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -28,7 +28,7 @@ pub mod pretty;
 ///
 /// # 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 @@ where
 
 /// 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 @@ where
 /// #   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 @@ where
 ///     }
 /// #   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) => {
+        run_driver!($args, $callback)
     };
-    ($args:expr, $callback:expr) => {{
+}
+
+/// 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 @@ macro_rules! run {
         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 @@ macro_rules! run {
         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 @@ macro_rules! run {
                 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 @@ macro_rules! run {
             }
         }
 
-        StableMir::new($args, || $callback).run()
+        StableMir::new($args, $callback).run()
     }};
 }