Add support to return value in StableMIR interface and not crash due to compilation error

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -4,9 +4,10 @@
 //! 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;
 use crate::{
     rustc_smir::Tables,
     stable_mir::{self, with},
@@ -189,27 +190,45 @@ pub(crate) fn opaque<T: Debug>(value: &T) -> Opaque {
     Opaque(format!("{value:?}"))
 }
 
-pub struct StableMir {
+pub struct StableMir<B = (), C = ()>
+where
+    B: Send,
+    C: Send,
+{
     args: Vec<String>,
-    callback: fn(TyCtxt<'_>),
+    callback: fn(TyCtxt<'_>) -> ControlFlow<B, C>,
+    result: Option<ControlFlow<B, C>>,
 }
 
-impl StableMir {
+impl<B, C> StableMir<B, C>
+where
+    B: Send,
+    C: Send,
+{
     /// Creates a new `StableMir` instance, with given test_function and arguments.
-    pub fn new(args: Vec<String>, callback: fn(TyCtxt<'_>)) -> Self {
-        StableMir { args, callback }
+    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) {
-        rustc_driver::catch_fatal_errors(|| {
-            RunCompiler::new(&self.args.clone(), self).run().unwrap();
-        })
-        .unwrap();
+    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, 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 Callbacks for StableMir {
+impl<B, C> Callbacks for StableMir<B, C>
+where
+    B: Send,
+    C: Send,
+{
     /// Called after analysis. Return value instructs the compiler whether to
     /// continue the compilation afterwards (defaults to `Compilation::Continue`)
     fn after_analysis<'tcx>(
@@ -219,9 +238,14 @@ impl Callbacks for StableMir {
         queries: &'tcx Queries<'tcx>,
     ) -> Compilation {
         queries.global_ctxt().unwrap().enter(|tcx| {
-            rustc_internal::run(tcx, || (self.callback)(tcx));
-        });
-        // No need to keep going.
-        Compilation::Stop
+            rustc_internal::run(tcx, || {
+                self.result = Some((self.callback)(tcx));
+            });
+            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

@@ -10,12 +10,13 @@
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
 use crate::stable_mir::ty::{FloatTy, GenericParamDef, IntTy, Movability, RigidTy, TyKind, UintTy};
-use crate::stable_mir::{self, Context};
+use crate::stable_mir::{self, CompilerError, Context};
 use rustc_hir as hir;
 use rustc_middle::mir::interpret::{alloc_range, AllocId};
 use rustc_middle::mir::{self, ConstantKind};
 use rustc_middle::ty::{self, Ty, TyCtxt, Variance};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
+use rustc_span::ErrorGuaranteed;
 use rustc_target::abi::FieldIdx;
 use tracing::debug;
 
@@ -1452,3 +1453,9 @@ impl<'tcx> Stable<'tcx> for rustc_span::Span {
         opaque(self)
     }
 }
+
+impl<T> From<ErrorGuaranteed> for CompilerError<T> {
+    fn from(_error: ErrorGuaranteed) -> Self {
+        CompilerError::CompilationFailed
+    }
+}

compiler/rustc_smir/src/stable_mir/mod.rs

@@ -56,6 +56,20 @@ pub type TraitDecls = Vec<TraitDef>;
 /// A list of impl trait decls.
 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<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.
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct 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.
@@ -136,7 +139,7 @@ fn main() {
         CRATE_NAME.to_string(),
         path.to_string(),
     ];
-    rustc_internal::StableMir::new(args, test_stable_mir).run();
+    rustc_internal::StableMir::new(args, test_stable_mir).run().unwrap();
 }
 
 fn generate_input(path: &str) -> std::io::Result<()> {