@@ -190,35 +190,120 @@ where
190190 stable_mir:: compiler_interface:: run ( & tables, || init ( & tables, f) )
191191}
192192
193+ /// Instantiate and run the compiler with the provided arguments and callback.
194+ ///
195+ /// The callback will be invoked after the compiler ran all its analyses, but before code generation.
196+ /// Note that this macro accepts two different formats for the callback:
197+ /// 1. An ident that resolves to a function that accepts no argument and returns `ControlFlow<B, C>`
198+ /// ```ignore(needs-extern-crate)
199+ /// # extern crate rustc_driver;
200+ /// # extern crate rustc_interface;
201+ /// # #[macro_use]
202+ /// # extern crate rustc_smir;
203+ /// # extern crate stable_mir;
204+ /// #
205+ /// # fn main() {
206+ /// # use std::ops::ControlFlow;
207+ /// # use stable_mir::CompilerError;
208+ /// fn analyze_code() -> ControlFlow<(), ()> {
209+ /// // Your code goes in here.
210+ /// # ControlFlow::Continue(())
211+ /// }
212+ /// # let args = vec!["--verbose".to_string()];
213+ /// let result = run!(args, analyze_code);
214+ /// # assert_eq!(result, Err(CompilerError::Skipped))
215+ /// # }
216+ /// ```
217+ /// 2. A closure expression:
218+ /// ```ignore(needs-extern-crate)
219+ /// # extern crate rustc_driver;
220+ /// # extern crate rustc_interface;
221+ /// # #[macro_use]
222+ /// # extern crate rustc_smir;
223+ /// # extern crate stable_mir;
224+ /// #
225+ /// # fn main() {
226+ /// # use std::ops::ControlFlow;
227+ /// # use stable_mir::CompilerError;
228+ /// fn analyze_code(extra_args: Vec<String>) -> ControlFlow<(), ()> {
229+ /// # let _ = extra_args;
230+ /// // Your code goes in here.
231+ /// # ControlFlow::Continue(())
232+ /// }
233+ /// # let args = vec!["--verbose".to_string()];
234+ /// # let extra_args = vec![];
235+ /// let result = run!(args, || analyze_code(extra_args));
236+ /// # assert_eq!(result, Err(CompilerError::Skipped))
237+ /// # }
238+ /// ```
193239#[ macro_export]
194240macro_rules! run {
241+ ( $args: expr, $callback_fn: ident) => {
242+ run_driver!( $args, || $callback_fn( ) )
243+ } ;
244+ ( $args: expr, $callback: expr) => {
245+ run_driver!( $args, $callback)
246+ } ;
247+ }
248+
249+ /// Instantiate and run the compiler with the provided arguments and callback.
250+ ///
251+ /// This is similar to `run` but it invokes the callback with the compiler's `TyCtxt`,
252+ /// which can be used to invoke internal APIs.
253+ #[ macro_export]
254+ macro_rules! run_with_tcx {
255+ ( $args: expr, $callback_fn: ident) => {
256+ run_driver!( $args, |tcx| $callback_fn( tcx) , with_tcx)
257+ } ;
195258 ( $args: expr, $callback: expr) => {
196- run!( $args, tcx, $callback)
259+ run_driver!( $args, $callback, with_tcx)
260+ } ;
261+ }
262+
263+ /// Optionally include an ident. This is needed due to macro hygiene.
264+ #[ macro_export]
265+ #[ doc( hidden) ]
266+ macro_rules! optional {
267+ ( with_tcx $ident: ident) => {
268+ $ident
197269 } ;
198- ( $args: expr, $tcx: ident, $callback: expr) => { {
270+ }
271+
272+ /// Prefer using [run!] and [run_with_tcx] instead.
273+ ///
274+ /// This macro implements the instantiation of a StableMIR driver, and it will invoke
275+ /// the given callback after the compiler analyses.
276+ ///
277+ /// The third argument determines whether the callback requires `tcx` as an argument.
278+ #[ macro_export]
279+ #[ doc( hidden) ]
280+ macro_rules! run_driver {
281+ ( $args: expr, $callback: expr $( , $with_tcx: ident) ?) => { {
199282 use rustc_driver:: { Callbacks , Compilation , RunCompiler } ;
200283 use rustc_interface:: { interface, Queries } ;
201284 use stable_mir:: CompilerError ;
202285 use std:: ops:: ControlFlow ;
203286
204- pub struct StableMir <B = ( ) , C = ( ) >
287+ pub struct StableMir <B = ( ) , C = ( ) , F = fn ( $ ( optional! ( $with_tcx TyCtxt ) ) ? ) -> ControlFlow < B , C > >
205288 where
206289 B : Send ,
207290 C : Send ,
291+ F : FnOnce ( $( optional!( $with_tcx TyCtxt ) ) ?) -> ControlFlow <B , C > + Send ,
208292 {
209293 args: Vec <String >,
210- callback: fn ( TyCtxt < ' _> ) -> ControlFlow < B , C >,
294+ callback: Option < F >,
211295 result: Option <ControlFlow <B , C >>,
212296 }
213297
214- impl <B , C > StableMir <B , C >
298+ impl <B , C , F > StableMir <B , C , F >
215299 where
216300 B : Send ,
217301 C : Send ,
302+ F : FnOnce ( $( optional!( $with_tcx TyCtxt ) ) ?) -> ControlFlow <B , C > + Send ,
218303 {
219304 /// Creates a new `StableMir` instance, with given test_function and arguments.
220- pub fn new( args: Vec <String >, callback: fn ( TyCtxt < ' _> ) -> ControlFlow < B , C > ) -> Self {
221- StableMir { args, callback, result: None }
305+ pub fn new( args: Vec <String >, callback: F ) -> Self {
306+ StableMir { args, callback: Some ( callback ) , result: None }
222307 }
223308
224309 /// Runs the compiler against given target and tests it with `test_function`
@@ -238,10 +323,11 @@ macro_rules! run {
238323 }
239324 }
240325
241- impl <B , C > Callbacks for StableMir <B , C >
326+ impl <B , C , F > Callbacks for StableMir <B , C , F >
242327 where
243328 B : Send ,
244329 C : Send ,
330+ F : FnOnce ( $( optional!( $with_tcx TyCtxt ) ) ?) -> ControlFlow <B , C > + Send ,
245331 {
246332 /// Called after analysis. Return value instructs the compiler whether to
247333/// continue the compilation afterwards (defaults to `Compilation::Continue`)
@@ -251,20 +337,24 @@ macro_rules! run {
251337 queries: & ' tcx Queries <' tcx>,
252338 ) -> Compilation {
253339 queries. global_ctxt( ) . unwrap( ) . enter( |tcx| {
254- rustc_internal:: run( tcx, || {
255- self . result = Some ( ( self . callback) ( tcx) ) ;
256- } )
257- . unwrap( ) ;
258- if self . result. as_ref( ) . is_some_and( |val| val. is_continue( ) ) {
259- Compilation :: Continue
340+ if let Some ( callback) = self . callback. take( ) {
341+ rustc_internal:: run( tcx, || {
342+ self . result = Some ( callback( $( optional!( $with_tcx tcx) ) ?) ) ;
343+ } )
344+ . unwrap( ) ;
345+ if self . result. as_ref( ) . is_some_and( |val| val. is_continue( ) ) {
346+ Compilation :: Continue
347+ } else {
348+ Compilation :: Stop
349+ }
260350 } else {
261- Compilation :: Stop
351+ Compilation :: Continue
262352 }
263353 } )
264354 }
265355 }
266356
267- StableMir :: new( $args, |$tcx| $callback) . run( )
357+ StableMir :: new( $args, $callback) . run( )
268358 } } ;
269359}
270360
0 commit comments