From 2db01be5844ded1139eb4fb5f1434d0090b6ba38 Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Wed, 30 Aug 2023 16:50:07 -0700
Subject: [PATCH] Adjust StableMIR interface to return and not crash

Invoking StableMir::run() on a crate that has any compilation error was
crashing the entire process. Instead, return a `CompilerError` so the
user knows compilation did not succeed.

I believe ICE will also be converted to `CompilerError`.

I'm also adding a return value to the callback, because I think it will
be handy for users (at least it was for my current task of implementing
a tool to validate stable-mir). However, if people disagree,
I can remove that.
---
 compiler/rustc_smir/src/rustc_internal/mod.rs | 35 +++++++++++++------
 compiler/rustc_smir/src/rustc_smir/mod.rs     |  9 ++++-
 compiler/rustc_smir/src/stable_mir/mod.rs     |  4 +++
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 2e219e17a5f..4496a58327e 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -7,6 +7,7 @@ use std::fmt::Debug;
 use std::ops::Index;
 
 use crate::rustc_internal;
+use crate::stable_mir::CompilerError;
 use crate::{
     rustc_smir::Tables,
     stable_mir::{self, with},
@@ -17,6 +18,7 @@ use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::EarlyErrorHandler;
 pub use rustc_span::def_id::{CrateNum, DefId};
+use rustc_span::ErrorGuaranteed;
 
 fn with_tables<R>(mut f: impl FnMut(&mut Tables<'_>) -> R) -> R {
     let mut ret = None;
@@ -189,27 +191,38 @@ pub(crate) fn opaque<T: Debug>(value: &T) -> Opaque {
     Opaque(format!("{value:?}"))
 }
 
-pub struct StableMir {
+pub struct StableMir<T: Send>
+where
+    T: Send,
+{
     args: Vec<String>,
-    callback: fn(TyCtxt<'_>),
+    callback: fn(TyCtxt<'_>) -> T,
+    result: Option<T>,
 }
 
-impl StableMir {
+impl<T> StableMir<T>
+where
+    T: 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<'_>) -> T) -> Self {
+        StableMir { args, callback, result: None }
     }
 
     /// Runs the compiler against given target and tests it with `test_function`
-    pub fn run(&mut self) {
+    pub fn run(mut self) -> Result<T, CompilerError> {
         rustc_driver::catch_fatal_errors(|| {
-            RunCompiler::new(&self.args.clone(), self).run().unwrap();
+            RunCompiler::new(&self.args.clone(), &mut self).run().unwrap();
         })
-        .unwrap();
+        .map_err(|e| <ErrorGuaranteed as Into<CompilerError>>::into(e))?;
+        Ok(self.result.unwrap())
     }
 }
 
-impl Callbacks for StableMir {
+impl<T> Callbacks for StableMir<T>
+where
+    T: Send,
+{
     /// Called after analysis. Return value instructs the compiler whether to
     /// continue the compilation afterwards (defaults to `Compilation::Continue`)
     fn after_analysis<'tcx>(
@@ -219,7 +232,9 @@ impl Callbacks for StableMir {
         queries: &'tcx Queries<'tcx>,
     ) -> Compilation {
         queries.global_ctxt().unwrap().enter(|tcx| {
-            rustc_internal::run(tcx, || (self.callback)(tcx));
+            rustc_internal::run(tcx, || {
+                self.result = Some((self.callback)(tcx));
+            });
         });
         // No need to keep going.
         Compilation::Stop
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index b4c150c591c..827b51e8a36 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/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 From<ErrorGuaranteed> for CompilerError {
+    fn from(_error: ErrorGuaranteed) -> Self {
+        CompilerError
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index bfc2c15a42e..281d5aafb4c 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -56,6 +56,10 @@ 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 struct CompilerError;
+
 /// Holds information about a crate.
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Crate {