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

[Amir]

Summary

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

Follow-up to D119658

[Amir]

Turn assert into error

[rafauler]

LGTM. One nit below.

bolt/lib/Rewrite/RewriteInstance.cpp

758

[Amir]

Address nit

bolt/include/bolt/Rewrite/RewriteInstance.h

public:
   createRewriteInstance(llvm::object::ELFObjectFileBase *File, const int Argc,
                         const char *const *Argv, StringRef ToolPath);
   ~RewriteInstance();
 
   /// Assign profile from \p Filename to this instance.
   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.
   void compare(RewriteInstance &RI2);
 
   /// Return binary context.
   const BinaryContext &getBinaryContext() const { return *BC; }
 
private:
   /// corresponding to \p OldAddress address in the original binary.
   uint64_t getNewFunctionOrDataAddress(uint64_t OldAddress);
 
   /// Return value for the symbol \p Name in the output.
   uint64_t getNewValueForSymbol(const StringRef Name);
 
   /// 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.
   void adjustFunctionBoundaries();
 
   /// Make .eh_frame section relocatable.
   void relocateEHFrameSection();
 

bolt/lib/Rewrite/RewriteInstance.cpp

if (BF.isPseudo())
     return false;
 
   if (opts::processAllFunctions())
     return true;
 
   return !BF.isIgnored();
 }
 
-void RewriteInstance::discoverStorage() {
+Error RewriteInstance::discoverStorage() {
   NamedRegionTimer T("discoverStorage", "discover storage", TimerGroupName,
                      TimerGroupDesc, opts::TimeRewrite);
 
   // Stubs are harmful because RuntimeDyld may try to increase the size of
   // sections accounting for stubs when we need those sections to match the
   // same size seen in the input binary, in case this section is a copy
   // of the original one seen in the binary.
   BC->EFMM.reset(new ExecutableFileMemoryManager(*BC, /*AllowStubs*/ false));
 
   auto ELF64LEFile = dyn_cast<ELF64LEObjectFile>(InputFile);
   const ELFFile<ELF64LE> &Obj = ELF64LEFile->getELFFile();
 
   BC->StartFunctionAddress = Obj.getHeader().e_entry;
 
   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:
       BC->FirstAllocAddress = std::min(BC->FirstAllocAddress,
                                        static_cast<uint64_t>(Phdr.p_vaddr));
       NextAvailableAddress = std::max(NextAvailableAddress,
                                       Phdr.p_vaddr + Phdr.p_memsz);
       NextAvailableOffset = std::max(NextAvailableOffset,
                                      Phdr.p_offset + Phdr.p_filesz);
 
       BC->SegmentMapInfo[Phdr.p_vaddr] = SegmentInfo{Phdr.p_vaddr,
                                                      Phdr.p_memsz,
                                                      Phdr.p_offset,
                                                      Phdr.p_filesz,
                                                      Phdr.p_align};
       break;
     case ELF::PT_INTERP:
       BC->HasInterpHeader = true;
       break;
     }
   }
 
   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();
     }
 
     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 &&
+  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';
 
   FirstNonAllocatableOffset = NextAvailableOffset;
 
   NextAvailableAddress = alignTo(NextAvailableAddress, BC->PageAlign);
   NextAvailableOffset = alignTo(NextAvailableOffset, BC->PageAlign);
Error RewriteInstance::discoverStorage() {
   NextAvailableOffset = alignTo(NextAvailableOffset, 64);
 
   NewTextSegmentAddress = NextAvailableAddress;
   NewTextSegmentOffset = NextAvailableOffset;
   BC->LayoutStartAddress = NextAvailableAddress;
 
   // 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() {
   if (!SDTSection)
     return;
 
   StringRef Buf = SDTSection->getContents();
   DataExtractor DE = DataExtractor(Buf, BC->AsmInfo->isLittleEndian(),
void RewriteInstance::patchBuildID() {
 
   char LastIDByte = BuildID[BuildID.size() - 1];
   LastIDByte ^= 1;
   OS.pwrite(&LastIDByte, 1, FileOffset + IDOffset + BuildID.size() - 1);
 
   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(
                 (llvm::Triple::ArchType)InputFile->getArch())
          << "\n";
   outs() << "BOLT-INFO: BOLT version: " << BoltRevision << "\n";
 
-  discoverStorage();
+  if (Error E = discoverStorage())
+    return E;
   readSpecialSections();
   adjustCommandLineOptions();
   discoverFileObjects();
 
   preprocessProfileData();
 
   // Skip disassembling if we have a translation table and we are running an
   // aggregation job.
   if (opts::AggregateOnly && BAT->enabledFor(InputFile)) {
     processProfileData();
-    return;
+    return Error::success();
   }
 
   selectFunctionsToProcess();
 
   readDebugInfo();
 
   disassembleFunctions();
 
   processProfileDataPreCFG();
 
   buildFunctionsCFG();
 
   processProfileData();
 
   postProcessFunctions();
 
   if (opts::DiffOnly)
-    return;
+    return Error::success();
 
   runOptimizationPasses();
 
   emitAndLink();
 
   updateMetadata();
 
   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() {
   NamedRegionTimer T("discoverFileObjects", "discover file objects",
                      TimerGroupName, TimerGroupDesc, opts::TimeRewrite);
   FileSymRefs.clear();
   BC->getBinaryFunctions().clear();
   BC->clearBinaryData();

bolt/tools/driver/llvm-bolt.cpp

if (auto *e = dyn_cast<ELFObjectFileBase>(&Binary)) {
         if (Error E = RI.setProfile(opts::InputDataFilename))
           report_error(opts::InputDataFilename, std::move(E));
       }
       if (opts::AggregateOnly && opts::PerfData.empty()) {
         errs() << ToolName << ": missing required -perfdata option.\n";
         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);
       if (Error E = MachORIOrErr.takeError())
         report_error(opts::InputFilename, std::move(E));
       MachORewriteInstance &MachORI = *MachORIOrErr.get();
 
       if (!opts::InputDataFilename.empty())
if (auto *ELFObj2 = dyn_cast<ELFObjectFileBase>(&Binary2)) {
         report_error(opts::InputFilename2, std::move(E));
       RewriteInstance &RI2 = *RI2OrErr.get();
       if (Error E = RI2.setProfile(opts::InputDataFilename2))
         report_error(opts::InputDataFilename2, std::move(E));
       outs() << "BOLT-DIFF: *** Analyzing binary 1: " << opts::InputFilename
              << "\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);
     }
   } else {
     report_error(opts::InputFilename, object_error::invalid_file_type);
   }
 
   return EXIT_SUCCESS;
 }

bolt/tools/heatmap/heatmap.cpp

auto RIOrErr =
         RewriteInstance::createRewriteInstance(e, argc, argv, ToolPath);
     if (Error E = RIOrErr.takeError())
       report_error("RewriteInstance", std::move(E));
 
     RewriteInstance &RI = *RIOrErr.get();
     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);
   }
 
   return EXIT_SUCCESS;
 }

bolt/tools/llvm-bolt-fuzzer/llvm-bolt-fuzzer.cpp

extern "C" int LLVMFuzzerTestOneInput(const char *Data, size_t Size) {
   // Fuzz RewriteInstance.
   auto RIOrErr =
       RewriteInstance::createRewriteInstance(E, argc, argv, "llvm-bolt");
   if (Error E = RIOrErr.takeError()) {
     consumeError(std::move(E));
     return 0;
   }
   RewriteInstance &RI = *RIOrErr.get();
-  RI.run();
+  if (Error E = RI.run())
+    consumeError(std::move(E));
   return 0;
 }
 
 extern "C" LLVM_ATTRIBUTE_USED int LLVMFuzzerInitialize(int *argc,
                                                         char ***argv) {
   llvm::InitializeAllTargetInfos();
   llvm::InitializeAllTargetMCs();
   llvm::InitializeAllAsmParsers();
   llvm::InitializeAllDisassemblers();
 
   llvm::InitializeAllTargets();
   llvm::InitializeAllAsmPrinters();
 
   return 0;
 }