Add tests and use ControlFlow

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -4,7 +4,7 @@
 //! until stable MIR is complete.
 
 use std::fmt::Debug;
-use std::ops::Index;
+use std::ops::{ControlFlow, Index};
 
 use crate::rustc_internal;
 use crate::stable_mir::CompilerError;
@@ -190,52 +190,44 @@ pub(crate) fn opaque<T: Debug>(value: &T) -> Opaque {
     Opaque(format!("{value:?}"))
 }
 
-pub struct StableMir<T: Send>
+pub struct StableMir<B = (), C = ()>
 where
-    T: Send,
+    B: Send,
+    C: Send,
 {
     args: Vec<String>,
-    callback: fn(TyCtxt<'_>) -> T,
-    after_analysis: Compilation,
-    result: Option<T>,
+    callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>,
+    result: Option<ControlFlow<B, C>>,
 }
 
-impl<T> StableMir<T>
+impl<B, C> StableMir<B, C>
 where
-    T: Send,
+    B: Send,
+    C: Send,
 {
     /// 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, 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
+    pub fn new(args: Vec<String>, callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>) -> Self {
+        StableMir { args, callback, result: None }
     }
 
     /// Runs the compiler against given target and tests it with `test_function`
-    pub fn run(&mut self) -> Result<T, CompilerError> {
+    pub fn run(&mut self) -> Result<C, CompilerError<B>> {
         let compiler_result =
             rustc_driver::catch_fatal_errors(|| RunCompiler::new(&self.args.clone(), self).run());
-        match compiler_result {
-            Ok(Ok(())) => Ok(self.result.take().unwrap()),
-            Ok(Err(_)) => Err(CompilerError::CompilationFailed),
-            Err(_) => Err(CompilerError::ICE),
+        match (compiler_result, self.result.take()) {
+            (Ok(Ok(())), Some(ControlFlow::Continue(value))) => Ok(value),
+            (Ok(Ok(())), Some(ControlFlow::Break(value))) => Err(CompilerError::Interrupted(value)),
+            (Ok(Ok(_)), None) => Err(CompilerError::Skipped),
+            (Ok(Err(_)), _) => Err(CompilerError::CompilationFailed),
+            (Err(_), _) => Err(CompilerError::ICE),
         }
     }
 }
 
-impl<T> Callbacks for StableMir<T>
+impl<B, C> Callbacks for StableMir<B, C>
 where
-    T: Send,
+    B: Send,
+    C: Send,
 {
     /// Called after analysis. Return value instructs the compiler whether to
     /// continue the compilation afterwards (defaults to `Compilation::Continue`)
@@ -249,8 +241,11 @@ where
             rustc_internal::run(tcx, || {
                 self.result = Some((self.callback)(tcx));
             });
-        });
-        // Let users define if they want to stop compilation.
-        self.after_analysis
+            if self.result.as_ref().is_some_and(|val| val.is_continue()) {
+                Compilation::Continue
+            } else {
+                Compilation::Stop
+            }
+        })
     }
 }

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -1454,7 +1454,7 @@ impl<'tcx> Stable<'tcx> for rustc_span::Span {
     }
 }
 
-impl From<ErrorGuaranteed> for CompilerError {
+impl<T> From<ErrorGuaranteed> for CompilerError<T> {
     fn from(_error: ErrorGuaranteed) -> Self {
         CompilerError::CompilationFailed
     }

compiler/rustc_smir/src/stable_mir/mod.rs

@@ -58,11 +58,16 @@ pub type ImplTraitDecls = Vec<ImplDef>;
 
 /// An error type used to represent an error that has already been reported by the compiler.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum CompilerError {
+pub enum CompilerError<T> {
     /// Internal compiler error (I.e.: Compiler crashed).
     ICE,
     /// Compilation failed.
     CompilationFailed,
+    /// Compilation was interrupted.
+    Interrupted(T),
+    /// Compilation skipped. This happens when users invoke rustc to retrieve information such as
+    /// --version.
+    Skipped,
 }
 
 /// Holds information about a crate.

tests/ui-fulldeps/stable-mir/compilation-result.rs

@@ -0,0 +1,77 @@
+// run-pass
+// Test StableMIR behavior when different results are given
+
+// ignore-stage1
+// ignore-cross-compile
+// ignore-remote
+// edition: 2021
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+
+extern crate rustc_middle;
+extern crate rustc_smir;
+
+use rustc_middle::ty::TyCtxt;
+use rustc_smir::{rustc_internal, stable_mir};
+use std::io::Write;
+use std::ops::ControlFlow;
+
+/// This test will generate and analyze a dummy crate using the stable mir.
+/// For that, it will first write the dummy crate into a file.
+/// Then it will create a `StableMir` using custom arguments and then
+/// it will run the compiler.
+fn main() {
+    let path = "input_compilation_result_test.rs";
+    generate_input(&path).unwrap();
+    let args = vec!["rustc".to_string(), path.to_string()];
+    test_continue(args.clone());
+    test_break(args.clone());
+    test_failed(args.clone());
+    test_skipped(args);
+}
+
+fn test_continue(args: Vec<String>) {
+    let continue_fn = |_: TyCtxt| ControlFlow::Continue::<(), bool>(true);
+    let result = rustc_internal::StableMir::new(args, continue_fn).run();
+    assert_eq!(result, Ok(true));
+}
+
+fn test_break(args: Vec<String>) {
+    let continue_fn = |_: TyCtxt| ControlFlow::Break::<bool, i32>(false);
+    let result = rustc_internal::StableMir::new(args, continue_fn).run();
+    assert_eq!(result, Err(stable_mir::CompilerError::Interrupted(false)));
+}
+
+fn test_skipped(mut args: Vec<String>) {
+    args.push("--version".to_string());
+    let unreach_fn = |_: TyCtxt| -> ControlFlow<()> { unreachable!() };
+    let result = rustc_internal::StableMir::new(args, unreach_fn).run();
+    assert_eq!(result, Err(stable_mir::CompilerError::Skipped));
+}
+
+fn test_failed(mut args: Vec<String>) {
+    args.push("--cfg=broken".to_string());
+    let unreach_fn = |_: TyCtxt| -> ControlFlow<()> { unreachable!() };
+    let result = rustc_internal::StableMir::new(args, unreach_fn).run();
+    assert_eq!(result, Err(stable_mir::CompilerError::CompilationFailed));
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+    // This should trigger a compilation failure when enabled.
+    #[cfg(broken)]
+    mod broken_mod {{
+        fn call_invalid() {{
+            invalid_fn();
+        }}
+    }}
+
+    fn main() {{}}
+    "#
+    )?;
+    Ok(())
+}

tests/ui-fulldeps/stable-mir/crate-info.rs

@@ -18,11 +18,12 @@ use rustc_middle::ty::TyCtxt;
 use rustc_smir::{rustc_internal, stable_mir};
 use std::assert_matches::assert_matches;
 use std::io::Write;
+use std::ops::ControlFlow;
 
 const CRATE_NAME: &str = "input";
 
 /// This function uses the Stable MIR APIs to get information about the test crate.
-fn test_stable_mir(tcx: TyCtxt<'_>) {
+fn test_stable_mir(tcx: TyCtxt<'_>) -> ControlFlow<()> {
     // Get the local crate using stable_mir API.
     let local = stable_mir::local_crate();
     assert_eq!(&local.name, CRATE_NAME);
@@ -108,6 +109,8 @@ fn test_stable_mir(tcx: TyCtxt<'_>) {
         stable_mir::mir::Terminator::Assert { .. } => {}
         other => panic!("{other:?}"),
     }
+
+    ControlFlow::Continue(())
 }
 
 // Use internal API to find a function in a crate.