[BOLT][NFC] Report errors from RewriteInstance `discoverStorage` and `run`

Further improve error handling in BOLT by reporting `RewriteInstance` errors in
a library and fuzzer-friendly way instead of exiting.

Follow-up to D119658

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D120224

5 files changed, 48 insertions, 34 deletions

diff --git a/bolt/include/bolt/Rewrite/RewriteInstance.h b/bolt/include/bolt/Rewrite/RewriteInstance.h
index 6c8d91042e0a..3b9df2c76ae2 100644
--- a/bolt/include/bolt/Rewrite/RewriteInstance.h
+++ b/bolt/include/bolt/Rewrite/RewriteInstance.h
@@ -56,7 +56,7 @@ public:
   Error setProfile(StringRef Filename);
 
   /// Run all the necessary steps to read, optimize and rewrite the binary.
-  void run();
+  Error run();
 
   /// Diff this instance against another one. Non-const since we may run passes
   /// to fold identical functions.
@@ -218,7 +218,7 @@ private:
 
   /// Detect addresses and offsets available in the binary for allocating
   /// new sections.
-  void discoverStorage();
+  Error discoverStorage();
 
   /// Adjust function sizes and set proper maximum size values after the whole
   /// symbol table has been processed.
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 51607eab8c8a..69ba24853898 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -448,7 +448,7 @@ static bool shouldDisassemble(const BinaryFunction &BF) {
   return !BF.isIgnored();
 }
 
-void RewriteInstance::discoverStorage() {
+Error RewriteInstance::discoverStorage() {
   NamedRegionTimer T("discoverStorage", "discover storage", TimerGroupName,
                      TimerGroupDesc, opts::TimeRewrite);
 
@@ -465,8 +465,11 @@ void RewriteInstance::discoverStorage() {
 
   NextAvailableAddress = 0;
   uint64_t NextAvailableOffset = 0;
-  ELF64LE::PhdrRange PHs =
-      cantFail(Obj.program_headers(), "program_headers() failed");
+  Expected<ELF64LE::PhdrRange> PHsOrErr = Obj.program_headers();
+  if (Error E = PHsOrErr.takeError())
+    return E;
+
+  ELF64LE::PhdrRange PHs = PHsOrErr.get();
   for (const ELF64LE::Phdr &Phdr : PHs) {
     switch (Phdr.p_type) {
     case ELF::PT_LOAD:
@@ -490,12 +493,18 @@ void RewriteInstance::discoverStorage() {
   }
 
   for (const SectionRef &Section : InputFile->sections()) {
-    StringRef SectionName = cantFail(Section.getName());
+    Expected<StringRef> SectionNameOrErr = Section.getName();
+    if (Error E = SectionNameOrErr.takeError())
+      return E;
+    StringRef SectionName = SectionNameOrErr.get();
     if (SectionName == ".text") {
       BC->OldTextSectionAddress = Section.getAddress();
       BC->OldTextSectionSize = Section.getSize();
 
-      StringRef SectionContents = cantFail(Section.getContents());
+      Expected<StringRef> SectionContentsOrErr = Section.getContents();
+      if (Error E = SectionContentsOrErr.takeError())
+        return E;
+      StringRef SectionContents = SectionContentsOrErr.get();
       BC->OldTextSectionOffset =
           SectionContents.data() - InputFile->getData().data();
     }
@@ -503,15 +512,15 @@ void RewriteInstance::discoverStorage() {
     if (!opts::HeatmapMode &&
         !(opts::AggregateOnly && BAT->enabledFor(InputFile)) &&
         (SectionName.startswith(getOrgSecPrefix()) ||
-         SectionName == getBOLTTextSectionName())) {
-      errs() << "BOLT-ERROR: input file was processed by BOLT. "
-                "Cannot re-optimize.\n";
-      exit(1);
-    }
+         SectionName == getBOLTTextSectionName()))
+      return createStringError(
+          errc::function_not_supported,
+          "BOLT-ERROR: input file was processed by BOLT. Cannot re-optimize");
   }
 
-  assert(NextAvailableAddress && NextAvailableOffset &&
-         "no PT_LOAD pheader seen");
+  if (!NextAvailableAddress || !NextAvailableOffset)
+    return createStringError(errc::executable_format_error,
+                             "no PT_LOAD pheader seen");
 
   outs() << "BOLT-INFO: first alloc address is 0x"
          << Twine::utohexstr(BC->FirstAllocAddress) << '\n';
@@ -566,11 +575,12 @@ void RewriteInstance::discoverStorage() {
 
   // Tools such as objcopy can strip section contents but leave header
   // entries. Check that at least .text is mapped in the file.
-  if (!getFileOffsetForAddress(BC->OldTextSectionAddress)) {
-    errs() << "BOLT-ERROR: input binary is not a valid ELF executable as its "
-              "text section is not mapped to a valid segment\n";
-    exit(1);
-  }
+  if (!getFileOffsetForAddress(BC->OldTextSectionAddress))
+    return createStringError(errc::executable_format_error,
+                             "BOLT-ERROR: input binary is not a valid ELF "
+                             "executable as its text section is not "
+                             "mapped to a valid segment");
+  return Error::success();
 }
 
 void RewriteInstance::parseSDTNotes() {
@@ -744,11 +754,8 @@ void RewriteInstance::patchBuildID() {
   outs() << "BOLT-INFO: patched build-id (flipped last bit)\n";
 }
 
-void RewriteInstance::run() {
-  if (!BC) {
-    errs() << "BOLT-ERROR: failed to create a binary context\n";
-    return;
-  }
+Error RewriteInstance::run() {
+  assert(BC && "failed to create a binary context");
 
   outs() << "BOLT-INFO: Target architecture: "
          << Triple::getArchTypeName(
@@ -756,7 +763,8 @@ void RewriteInstance::run() {
          << "\n";
   outs() << "BOLT-INFO: BOLT version: " << BoltRevision << "\n";
 
-  discoverStorage();
+  if (Error E = discoverStorage())
+    return E;
   readSpecialSections();
   adjustCommandLineOptions();
   discoverFileObjects();
@@ -767,7 +775,7 @@ void RewriteInstance::run() {
   // aggregation job.
   if (opts::AggregateOnly && BAT->enabledFor(InputFile)) {
     processProfileData();
-    return;
+    return Error::success();
   }
 
   selectFunctionsToProcess();
@@ -785,7 +793,7 @@ void RewriteInstance::run() {
   postProcessFunctions();
 
   if (opts::DiffOnly)
-    return;
+    return Error::success();
 
   runOptimizationPasses();
 
@@ -795,14 +803,15 @@ void RewriteInstance::run() {
 
   if (opts::LinuxKernelMode) {
     errs() << "BOLT-WARNING: not writing the output file for Linux Kernel\n";
-    return;
+    return Error::success();
   } else if (opts::OutputFilename == "/dev/null") {
     outs() << "BOLT-INFO: skipping writing final binary to disk\n";
-    return;
+    return Error::success();
   }
 
   // Rewrite allocatable contents and copy non-allocatable parts with mods.
   rewriteFile();
+  return Error::success();
 }
 
 void RewriteInstance::discoverFileObjects() {
diff --git a/bolt/tools/driver/llvm-bolt.cpp b/bolt/tools/driver/llvm-bolt.cpp
index 2c8d1d28f840..24e10c70ae2a 100644
--- a/bolt/tools/driver/llvm-bolt.cpp
+++ b/bolt/tools/driver/llvm-bolt.cpp
@@ -241,7 +241,8 @@ int main(int argc, char **argv) {
         exit(1);
       }
 
-      RI.run();
+      if (Error E = RI.run())
+        report_error(opts::InputFilename, std::move(E));
     } else if (auto *O = dyn_cast<MachOObjectFile>(&Binary)) {
       auto MachORIOrErr =
           MachORewriteInstance::createMachORewriteInstance(O, ToolPath);
@@ -292,12 +293,14 @@ int main(int argc, char **argv) {
              << "\n";
       outs() << "BOLT-DIFF: *** Binary 1 fdata:     " << opts::InputDataFilename
              << "\n";
-      RI1.run();
+      if (Error E = RI1.run())
+        report_error(opts::InputFilename, std::move(E));
       outs() << "BOLT-DIFF: *** Analyzing binary 2: " << opts::InputFilename2
              << "\n";
       outs() << "BOLT-DIFF: *** Binary 2 fdata:     "
              << opts::InputDataFilename2 << "\n";
-      RI2.run();
+      if (Error E = RI2.run())
+        report_error(opts::InputFilename2, std::move(E));
       RI1.compare(RI2);
     } else {
       report_error(opts::InputFilename2, object_error::invalid_file_type);
diff --git a/bolt/tools/heatmap/heatmap.cpp b/bolt/tools/heatmap/heatmap.cpp
index 0ab6a4fa52b7..fa28e858568e 100644
--- a/bolt/tools/heatmap/heatmap.cpp
+++ b/bolt/tools/heatmap/heatmap.cpp
@@ -94,7 +94,8 @@ int main(int argc, char **argv) {
     if (Error E = RI.setProfile(opts::PerfData))
       report_error(opts::PerfData, std::move(E));
 
-    RI.run();
+    if (Error E = RI.run())
+      report_error(opts::InputFilename, std::move(E));
   } else {
     report_error(opts::InputFilename, object_error::invalid_file_type);
   }
diff --git a/bolt/tools/llvm-bolt-fuzzer/llvm-bolt-fuzzer.cpp b/bolt/tools/llvm-bolt-fuzzer/llvm-bolt-fuzzer.cpp
index ac129a0d9a62..4f2abfaa7a79 100644
--- a/bolt/tools/llvm-bolt-fuzzer/llvm-bolt-fuzzer.cpp
+++ b/bolt/tools/llvm-bolt-fuzzer/llvm-bolt-fuzzer.cpp
@@ -52,7 +52,8 @@ extern "C" int LLVMFuzzerTestOneInput(const char *Data, size_t Size) {
     return 0;
   }
   RewriteInstance &RI = *RIOrErr.get();
-  RI.run();
+  if (Error E = RI.run())
+    consumeError(std::move(E));
   return 0;
 }