2 files changed, 79 insertions, 1 deletions

diff --git a/llvm/include/llvm/Support/Error.h b/llvm/include/llvm/Support/Error.h
index 33a3bc9404e..2e1699e8a33 100644
--- a/llvm/include/llvm/Support/Error.h
+++ b/llvm/include/llvm/Support/Error.h
@@ -519,7 +519,8 @@ inline void handleAllErrors(Error E) {
 /// (allowing clean deallocation of resources, etc.), while reporting error
 /// information to the user.
 template <typename... HandlerTs>
-void logAllUnhandledErrors(Error E, raw_ostream &OS, std::string ErrorBanner) {
+void logAllUnhandledErrors(Error E, raw_ostream &OS,
+                           const std::string &ErrorBanner) {
   if (!E)
     return;
   OS << ErrorBanner;
@@ -742,6 +743,55 @@ inline std::error_code errorToErrorCode(Error Err) {
   return EC;
 }
 
+/// Helper for check-and-exit error handling.
+///
+/// For tool use only. NOT FOR USE IN LIBRARY CODE.
+///
+class ExitOnError {
+public:
+
+  /// Create an error on exit helper.
+  ExitOnError(std::string Banner = "", int DefaultErrorExitCode = 1)
+    : Banner(Banner),
+      GetExitCode([=](const Error&) { return DefaultErrorExitCode; }) {}
+
+  /// Set the banner string for any errors caught by operator().
+  void setBanner(std::string Banner) {
+    this->Banner = Banner;
+  }
+
+  /// Set the exit-code mapper function.
+  void setExitCodeMapper(std::function<int(const Error&)> GetExitCode) {
+    this->GetExitCode = GetExitCode;
+  }
+
+  /// Check Err. If it's in a failure state log the error(s) and exit.
+  void operator()(Error Err) const {
+    checkError(std::move(Err));
+  }
+
+  /// Check E. If it's in a success state return the contained value. If it's
+  /// in a failure state log the error(s) and exit.
+  template <typename T>
+  T operator()(Expected<T> E) const {
+    checkError(E.takeError());
+    return std::move(*E);
+  }
+
+private:
+
+  void checkError(Error Err) const {
+    if (Err) {
+      int ExitCode = GetExitCode(Err);
+      logAllUnhandledErrors(std::move(Err), errs(), Banner);
+      exit(ExitCode);
+    }
+  }
+
+  std::string Banner;
+  std::function<int(const Error&)> GetExitCode;
+};
+
 } // namespace llvm
 
 #endif // LLVM_SUPPORT_ERROR_H
diff --git a/llvm/unittests/Support/ErrorTest.cpp b/llvm/unittests/Support/ErrorTest.cpp
index d5af2cab967..343a2d240c1 100644
--- a/llvm/unittests/Support/ErrorTest.cpp
+++ b/llvm/unittests/Support/ErrorTest.cpp
@@ -38,6 +38,7 @@ namespace {
 //   - consume_error to consume a "safe" error without any output.
 //   - handleAllUnhandledErrors to assert that all errors are handled.
 //   - logAllUnhandledErrors to log errors to a stream.
+//   - ExitOnError tests.
 //
 // Expected tests:
 //   - Expected<T> with T.
@@ -50,6 +51,7 @@ namespace {
 //   - std::error_code to Error (ECError) in failure mode.
 //   - Error to std::error_code in success mode.
 //   - Error (ECError) to std::error_code in failure mode.
+//
 
 // Custom error class with a default base class and some random 'info' attached.
 class CustomError : public ErrorInfo<CustomError> {
@@ -376,6 +378,32 @@ TEST(Error, CheckErrorUtilities) {
     EXPECT_EQ(ErrorInfo, 7)
         << "Failed to handle Error returned from handleErrors.";
   }
+
+  // Test ExitOnError
+  {
+    ExitOnError ExitOnErr;
+    ExitOnErr.setBanner("Error in tool:");
+    ExitOnErr.setExitCodeMapper(
+      [](const Error &E) {
+        if (E.isA<CustomSubError>())
+          return 2;
+        return 1;
+      });
+
+    // Make sure we don't bail on success.
+    ExitOnErr(Error::success());
+    EXPECT_EQ(ExitOnErr(Expected<int>(7)), 7)
+      << "exitOnError returned an invalid value for Expected";
+
+    // Exit tests.
+    EXPECT_EXIT(ExitOnErr(make_error<CustomError>(7)),
+                ::testing::ExitedWithCode(1), "Error in tool:")
+      << "exitOnError returned an unexpected error result";
+
+    EXPECT_EXIT(ExitOnErr(Expected<int>(make_error<CustomSubError>(0, 0))),
+                ::testing::ExitedWithCode(2), "Error in tool:")
+      << "exitOnError returned an unexpected error result";
+  }
 }
 
 // Test Expected behavior.