Link functions on MachO

Summary: Add first bits for linking functions on MachO.

(cherry picked from FBD21991721)

bolt/src/BinaryContext.cpp

@@ -1628,8 +1628,10 @@ BinarySection &BinaryContext::registerOrUpdateSection(StringRef Name,
     const auto Flag = Section->isAllocatable();
     Section->update(Data, Size, Alignment, ELFType, ELFFlags);
     DEBUG(dbgs() << *Section << "\n");
-    assert(Flag == Section->isAllocatable() &&
-           "can't change section allocation status");
+    // FIXME: Fix section flags/attributes for MachO.
+    if (isELF())
+      assert(Flag == Section->isAllocatable() &&
+             "can't change section allocation status");
     return *Section;
   }
 

bolt/src/BinaryPassManager.cpp

@@ -88,7 +88,7 @@ NeverPrint("never-print",
   cl::ReallyHidden,
   cl::cat(BoltOptCategory));
 
-static cl::opt<bool>
+cl::opt<bool>
 PrintAfterBranchFixup("print-after-branch-fixup",
   cl::desc("print function after fixing local branches"),
   cl::Hidden,

bolt/src/MachORewriteInstance.cpp

@@ -14,7 +14,9 @@
 #include "BinaryEmitter.h"
 #include "BinaryFunction.h"
 #include "BinaryPassManager.h"
+#include "ExecutableFileMemoryManager.h"
 #include "Utils.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCAsmLayout.h"
@@ -29,11 +31,13 @@ extern cl::opt<unsigned> AlignText;
 extern cl::opt<bool> KeepTmp;
 extern cl::opt<bool> NeverPrint;
 extern cl::opt<std::string> OutputFilename;
+extern cl::opt<bool> PrintAfterBranchFixup;
 extern cl::opt<bool> PrintFinalized;
 extern cl::opt<bool> PrintReordered;
 extern cl::opt<bool> PrintSections;
 extern cl::opt<bool> PrintDisasm;
 extern cl::opt<bool> PrintCFG;
+extern cl::opt<unsigned> Verbosity;
 } // namespace opts
 
 namespace llvm {
@@ -219,16 +223,40 @@ void MachORewriteInstance::runOptimizationPasses() {
   BinaryFunctionPassManager Manager(*BC);
   Manager.registerPass(
       llvm::make_unique<ReorderBasicBlocks>(opts::PrintReordered));
+  Manager.registerPass(
+      llvm::make_unique<FixupBranches>(opts::PrintAfterBranchFixup));
   // This pass should always run last.*
   Manager.registerPass(
       llvm::make_unique<FinalizeFunctions>(opts::PrintFinalized));
+
   Manager.runPasses();
 }
 
+void MachORewriteInstance::mapCodeSections(orc::VModuleKey Key) {
+  for (auto &BFI : BC->getBinaryFunctions()) {
+    BinaryFunction &Function = BFI.second;
+    if (!Function.isSimple())
+      continue;
+    assert(Function.isEmitted() && "Simple function has not been emitted");
+    ErrorOr<BinarySection &> FuncSection = Function.getCodeSection();
+    assert(FuncSection && "cannot find section for function");
+
+    FuncSection->setOutputAddress(Function.getAddress());
+    DEBUG(dbgs() << "BOLT: mapping 0x"
+                 << Twine::utohexstr(FuncSection->getAllocAddress()) << " to 0x"
+                 << Twine::utohexstr(Function.getAddress()) << '\n');
+    OLT->mapSectionAddress(Key, FuncSection->getSectionID(),
+                           Function.getAddress());
+    Function.setImageAddress(FuncSection->getAllocAddress());
+    Function.setImageSize(FuncSection->getOutputSize());
+  }
+}
+
 void MachORewriteInstance::emitAndLink() {
   std::error_code EC;
-  std::unique_ptr<ToolOutputFile> TempOut = llvm::make_unique<ToolOutputFile>(
-      opts::OutputFilename + ".bolt.o", EC, sys::fs::F_None);
+  std::unique_ptr<::llvm::ToolOutputFile> TempOut =
+      llvm::make_unique<::llvm::ToolOutputFile>(
+          opts::OutputFilename + ".bolt.o", EC, sys::fs::F_None);
   check_error(EC, "cannot create output object file");
 
   if (opts::KeepTmp)
@@ -257,6 +285,80 @@ void MachORewriteInstance::emitAndLink() {
       object::ObjectFile::createObjectFile(ObjectMemBuffer->getMemBufferRef()),
       "error creating in-memory object");
   assert(Obj && "createObjectFile cannot return nullptr");
+
+  auto Resolver = orc::createLegacyLookupResolver(
+      [&](const std::string &Name) -> JITSymbol {
+        llvm::errs() << "looking for " << Name << "\n";
+        assert(!BC->EFMM->ObjectsLoaded &&
+               "Linking multiple objects is unsupported");
+        DEBUG(dbgs() << "BOLT: looking for " << Name << "\n");
+        if (auto *I = BC->getBinaryDataByName(Name)) {
+          const uint64_t Address = I->isMoved() && !I->isJumpTable()
+                                       ? I->getOutputAddress()
+                                       : I->getAddress();
+          DEBUG(dbgs() << "Resolved to address 0x" << Twine::utohexstr(Address)
+                       << "\n");
+          return JITSymbol(Address, JITSymbolFlags());
+        }
+        DEBUG(dbgs() << "Resolved to address 0x0\n");
+        return JITSymbol(nullptr);
+      },
+      [](Error Err) { cantFail(std::move(Err), "lookup failed"); });
+
+  Resolver->setAllowsZeroSymbols(true);
+
+  MCAsmLayout FinalLayout(
+      static_cast<MCObjectStreamer *>(Streamer.get())->getAssembler());
+
+  SSP.reset(new decltype(SSP)::element_type());
+  ES.reset(new decltype(ES)::element_type(*SSP));
+  BC->EFMM.reset(new ExecutableFileMemoryManager(*BC, /*AllowStubs*/ false));
+
+  const orc::VModuleKey K = ES->allocateVModule();
+  OLT.reset(new decltype(OLT)::element_type(
+      *ES,
+      [this, &Resolver](orc::VModuleKey Key) {
+        orc::RTDyldObjectLinkingLayer::Resources R;
+        R.MemMgr = BC->EFMM;
+        R.Resolver = Resolver;
+        return R;
+      },
+      [&](orc::VModuleKey Key, const object::ObjectFile &Obj,
+          const RuntimeDyld::LoadedObjectInfo &) {
+        assert(Key == K && "Linking multiple objects is unsupported");
+        mapCodeSections(Key);
+      },
+      [&](orc::VModuleKey Key) {
+        assert(Key == K && "Linking multiple objects is unsupported");
+      }));
+
+  OLT->setProcessAllSections(true);
+  cantFail(OLT->addObject(K, std::move(ObjectMemBuffer)));
+  cantFail(OLT->emitAndFinalize(K));
+}
+
+void MachORewriteInstance::rewriteFile() {
+  std::error_code EC;
+  Out = llvm::make_unique<ToolOutputFile>(
+      opts::OutputFilename, EC, sys::fs::F_None,
+      sys::fs::all_read | sys::fs::all_write | sys::fs::all_exe);
+  check_error(EC, "cannot create output executable file");
+  raw_fd_ostream &OS = Out->os();
+  OS << InputFile->getData();
+
+  for (auto &BFI : BC->getBinaryFunctions()) {
+    BinaryFunction &Function = BFI.second;
+    if (!Function.isSimple())
+      continue;
+    assert(Function.isEmitted() && "Simple function has not been emitted");
+    if (Function.getImageSize() > Function.getMaxSize())
+      continue;
+    if (opts::Verbosity >= 2)
+      outs() << "BOLT: rewriting function \"" << Function << "\"\n";
+    OS.pwrite(reinterpret_cast<char *>(Function.getImageAddress()),
+              Function.getImageSize(), Function.getFileOffset());
+  }
+  Out->keep();
 }
 
 void MachORewriteInstance::adjustCommandLineOptions() {
@@ -272,6 +374,7 @@ void MachORewriteInstance::run() {
   postProcessFunctions();
   runOptimizationPasses();
   emitAndLink();
+  rewriteFile();
 }
 
 MachORewriteInstance::~MachORewriteInstance() {}

bolt/src/MachORewriteInstance.h

@@ -15,10 +15,22 @@
 #define LLVM_TOOLS_LLVM_BOLT_MACHO_REWRITE_INSTANCE_H
 
 #include "NameResolver.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/Object/MachO.h"
 #include <memory>
 
 namespace llvm {
+
+class ToolOutputFile;
+
+namespace orc {
+
+class SymbolStringPool;
+class ExecutionSession;
+class RTDyldObjectLinkingLayer;
+
+} // namespace orc
+
 namespace bolt {
 
 class BinaryContext;
@@ -29,8 +41,16 @@ class MachORewriteInstance {
 
   NameResolver NR;
 
+  std::unique_ptr<orc::SymbolStringPool> SSP;
+  std::unique_ptr<orc::ExecutionSession> ES;
+  std::unique_ptr<orc::RTDyldObjectLinkingLayer> OLT;
+
+  std::unique_ptr<ToolOutputFile> Out;
+
   static StringRef getOrgSecPrefix() { return ".bolt.org"; }
 
+  void mapCodeSections(orc::VModuleKey Key);
+
   void adjustCommandLineOptions();
   void readSpecialSections();
   void discoverFileObjects();
@@ -38,6 +58,7 @@ class MachORewriteInstance {
   void postProcessFunctions();
   void runOptimizationPasses();
   void emitAndLink();
+  void rewriteFile();
 
 public:
   explicit MachORewriteInstance(object::MachOObjectFile *InputFile);