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

5 files changed, 66 insertions, 39 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()
     }};
 }
 
diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index fffc454804d..2a2626654a0 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -29,7 +29,7 @@ use stable_mir::{Crate, CrateItem, CrateNum, DefId, Error, Filename, ItemKind, S
 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 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
     }
 
     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 {
diff --git a/tests/ui-fulldeps/stable-mir/check_trait_queries.rs b/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
index fb1197e4ecc..c9fbe15ffb0 100644
--- a/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
+++ b/tests/ui-fulldeps/stable-mir/check_trait_queries.rs
@@ -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<()> {
diff --git a/tests/ui-fulldeps/stable-mir/compilation-result.rs b/tests/ui-fulldeps/stable-mir/compilation-result.rs
index c802316c1d8..e6dd9fa132d 100644
--- a/tests/ui-fulldeps/stable-mir/compilation-result.rs
+++ b/tests/ui-fulldeps/stable-mir/compilation-result.rs
@@ -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()));
 }
 
diff --git a/tests/ui-fulldeps/stable-mir/smir_internal.rs b/tests/ui-fulldeps/stable-mir/smir_internal.rs
index 63a6a5e47e8..b0811364c09 100644
--- a/tests/ui-fulldeps/stable-mir/smir_internal.rs
+++ b/tests/ui-fulldeps/stable-mir/smir_internal.rs
@@ -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<()> {