Rollup merge of #119833 - celinval:smir-accept-closures, r=oli-obk

Make tcx optional from StableMIR run macro and extend it to accept closures

Change `run` macro to avoid sometimes unnecessary dependency on `TyCtxt`, and introduce `run_with_tcx` to capture use cases where `tcx` is required. Additionally, extend both macros to accept closures that may capture variables.

I've also modified the `internal()` method to make it safer, by accepting the type context to force the `'tcx` lifetime to match the context lifetime.

These are non-backward compatible changes, but they only affect internal APIs which are provided today as helper functions until we have a stable API to start the compiler.

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -190,35 +190,120 @@ where
     stable_mir::compiler_interface::run(&tables, || init(&tables, f))
 }
 
+/// Instantiate and run the compiler with the provided arguments and callback.
+///
+/// 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)
+/// # extern crate rustc_driver;
+/// # extern crate rustc_interface;
+/// # #[macro_use]
+/// # extern crate rustc_smir;
+/// # extern crate stable_mir;
+/// #
+/// # fn main() {
+/// #   use std::ops::ControlFlow;
+/// #   use stable_mir::CompilerError;
+///     fn analyze_code() -> ControlFlow<(), ()> {
+///         // Your code goes in here.
+/// #       ControlFlow::Continue(())
+///     }
+/// #   let args = vec!["--verbose".to_string()];
+///     let result = run!(args, analyze_code);
+/// #   assert_eq!(result, Err(CompilerError::Skipped))
+/// # }
+/// ```
+/// 2. A closure expression:
+/// ```ignore(needs-extern-crate)
+/// # extern crate rustc_driver;
+/// # extern crate rustc_interface;
+/// # #[macro_use]
+/// # extern crate rustc_smir;
+/// # extern crate stable_mir;
+/// #
+/// # fn main() {
+/// #   use std::ops::ControlFlow;
+/// #   use stable_mir::CompilerError;
+///     fn analyze_code(extra_args: Vec<String>) -> ControlFlow<(), ()> {
+/// #       let _ = extra_args;
+///         // Your code goes in here.
+/// #       ControlFlow::Continue(())
+///     }
+/// #   let args = vec!["--verbose".to_string()];
+/// #   let extra_args = vec![];
+///     let result = run!(args, || analyze_code(extra_args));
+/// #   assert_eq!(result, Err(CompilerError::Skipped))
+/// # }
+/// ```
 #[macro_export]
 macro_rules! run {
-    ($args:expr, $callback:expr) => {
-        run!($args, tcx, $callback)
+    ($args:expr, $callback_fn:ident) => {
+        run_driver!($args, || $callback_fn())
     };
-    ($args:expr, $tcx:ident, $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 = ()>
+        pub struct StableMir<B = (), C = (), F = fn($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C>>
         where
             B: Send,
             C: Send,
+            F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
         {
             args: Vec<String>,
-            callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>,
+            callback: Option<F>,
             result: Option<ControlFlow<B, C>>,
         }
 
-        impl<B, C> StableMir<B, C>
+        impl<B, C, F> StableMir<B, C, F>
         where
             B: Send,
             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: fn(TyCtxt<'_>) -> ControlFlow<B, C>) -> Self {
-                StableMir { args, callback, result: None }
+            pub fn new(args: Vec<String>, callback: F) -> Self {
+                StableMir { args, callback: Some(callback), result: None }
             }
 
             /// Runs the compiler against given target and tests it with `test_function`
@@ -238,10 +323,11 @@ macro_rules! run {
             }
         }
 
-        impl<B, C> Callbacks for StableMir<B, C>
+        impl<B, C, F> Callbacks for StableMir<B, C, F>
         where
             B: Send,
             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`)
@@ -251,20 +337,24 @@ macro_rules! run {
                 queries: &'tcx Queries<'tcx>,
             ) -> Compilation {
                 queries.global_ctxt().unwrap().enter(|tcx| {
-                    rustc_internal::run(tcx, || {
-                        self.result = Some((self.callback)(tcx));
-                    })
-                    .unwrap();
-                    if self.result.as_ref().is_some_and(|val| val.is_continue()) {
-                        Compilation::Continue
+                    if let Some(callback) = self.callback.take() {
+                        rustc_internal::run(tcx, || {
+                            self.result = Some(callback($(optional!($with_tcx tcx))?));
+                        })
+                        .unwrap();
+                        if self.result.as_ref().is_some_and(|val| val.is_continue()) {
+                            Compilation::Continue
+                        } else {
+                            Compilation::Stop
+                        }
                     } else {
-                        Compilation::Stop
+                        Compilation::Continue
                     }
                 })
             }
         }
 
-        StableMir::new($args, |$tcx| $callback).run()
+        StableMir::new($args, $callback).run()
     }};
 }