SMIR: Allow users to pick if compilation continues

Currently we stop compilation, but some users might want to keep going.
This is needed for us to test against rustc tests. Other tools, such as
Kani, also implements parts of their logic as a backend so it is
important for compilation to continue.

1 files changed, 20 insertions, 8 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 1c5a0924f4a..26127c5eb85 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -196,6 +196,7 @@ where
 {
     args: Vec<String>,
     callback: fn(TyCtxt<'_>) -> T,
+    after_analysis: Compilation,
     result: Option<T>,
 }
 
@@ -205,16 +206,27 @@ where
 {
     /// Creates a new `StableMir` instance, with given test_function and arguments.
     pub fn new(args: Vec<String>, callback: fn(TyCtxt<'_>) -> T) -> Self {
-        StableMir { args, callback, result: None }
+        StableMir { args, callback, result: None, after_analysis: Compilation::Stop }
+    }
+
+    /// Configure object to stop compilation after callback is called.
+    pub fn stop_compilation(&mut self) -> &mut Self {
+        self.after_analysis = Compilation::Stop;
+        self
+    }
+
+    /// Configure object to continue compilation after callback is called.
+    pub fn continue_compilation(&mut self) -> &mut Self {
+        self.after_analysis = Compilation::Continue;
+        self
     }
 
     /// Runs the compiler against given target and tests it with `test_function`
-    pub fn run(mut self) -> Result<T, CompilerError> {
-        let compiler_result = rustc_driver::catch_fatal_errors(|| {
-            RunCompiler::new(&self.args.clone(), &mut self).run()
-        });
+    pub fn run(&mut self) -> Result<T, CompilerError> {
+        let compiler_result =
+            rustc_driver::catch_fatal_errors(|| RunCompiler::new(&self.args.clone(), self).run());
         match compiler_result {
-            Ok(Ok(())) => Ok(self.result.unwrap()),
+            Ok(Ok(())) => Ok(self.result.take().unwrap()),
             Ok(Err(_)) => Err(CompilerError::CompilationFailed),
             Err(_) => Err(CompilerError::ICE),
         }
@@ -238,7 +250,7 @@ where
                 self.result = Some((self.callback)(tcx));
             });
         });
-        // No need to keep going.
-        Compilation::Stop
+        // Let users define if they want to stop compilation.
+        self.after_analysis
     }
 }