Update aosp/master LLVM for rebase to r235153

Change-Id: I9bf53792f9fc30570e81a8d80d296c681d005ea7
(cherry picked from commit 0c7f116bb6950ef819323d855415b2f2b0aad987)

6 files changed, 287 insertions, 104 deletions

diff --git a/lib/ExecutionEngine/Orc/Android.mk b/lib/ExecutionEngine/Orc/Android.mk
index 61c1daf..f28f359 100644
--- a/lib/ExecutionEngine/Orc/Android.mk
+++ b/lib/ExecutionEngine/Orc/Android.mk
@@ -6,6 +6,7 @@ include $(CLEAR_VARS)
 
 LOCAL_SRC_FILES := \
   CloneSubModule.cpp \
+  ExecutionUtils.cpp \
   IndirectionUtils.cpp \
   OrcMCJITReplacement.cpp \
   OrcTargetSupport.cpp
diff --git a/lib/ExecutionEngine/Orc/CMakeLists.txt b/lib/ExecutionEngine/Orc/CMakeLists.txt
index b0a8445..b38b459 100644
--- a/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_llvm_library(LLVMOrcJIT
   CloneSubModule.cpp
+  ExecutionUtils.cpp
   IndirectionUtils.cpp
   OrcMCJITReplacement.cpp
   OrcTargetSupport.cpp
diff --git a/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
new file mode 100644
index 0000000..b7220db
--- /dev/null
+++ b/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -0,0 +1,102 @@
+//===---- ExecutionUtils.cpp - Utilities for executing functions in Orc ---===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Module.h"
+
+namespace llvm {
+namespace orc {
+
+CtorDtorIterator::CtorDtorIterator(const GlobalVariable *GV, bool End)
+  : InitList(
+      GV ? dyn_cast_or_null<ConstantArray>(GV->getInitializer()) : nullptr),
+    I((InitList && End) ? InitList->getNumOperands() : 0) {
+}
+
+bool CtorDtorIterator::operator==(const CtorDtorIterator &Other) const {
+  assert(InitList == Other.InitList && "Incomparable iterators.");
+  return I == Other.I;
+}
+
+bool CtorDtorIterator::operator!=(const CtorDtorIterator &Other) const {
+  return !(*this == Other);
+}
+
+CtorDtorIterator& CtorDtorIterator::operator++() {
+  ++I;
+  return *this;
+}
+
+CtorDtorIterator CtorDtorIterator::operator++(int) {
+  CtorDtorIterator Temp = *this;
+  ++I;
+  return Temp;
+}
+
+CtorDtorIterator::Element CtorDtorIterator::operator*() const {
+  ConstantStruct *CS = dyn_cast<ConstantStruct>(InitList->getOperand(I));
+  assert(CS && "Unrecognized type in llvm.global_ctors/llvm.global_dtors");
+
+  Constant *FuncC = CS->getOperand(1);
+  Function *Func = nullptr;
+
+  // Extract function pointer, pulling off any casts.
+  while (FuncC) {
+    if (Function *F = dyn_cast_or_null<Function>(FuncC)) {
+      Func = F;
+      break;
+    } else if (ConstantExpr *CE = dyn_cast_or_null<ConstantExpr>(FuncC)) {
+      if (CE->isCast())
+        FuncC = dyn_cast_or_null<ConstantExpr>(CE->getOperand(0));
+      else
+        break;
+    } else {
+      // This isn't anything we recognize. Bail out with Func left set to null.
+      break;
+    }
+  }
+
+  ConstantInt *Priority = dyn_cast<ConstantInt>(CS->getOperand(0));
+  Value *Data = CS->getOperand(2);
+  return Element(Priority->getZExtValue(), Func, Data);
+}
+
+iterator_range<CtorDtorIterator> getConstructors(const Module &M) {
+  const GlobalVariable *CtorsList = M.getNamedGlobal("llvm.global_ctors");
+  return make_range(CtorDtorIterator(CtorsList, false),
+                    CtorDtorIterator(CtorsList, true));
+}
+
+iterator_range<CtorDtorIterator> getDestructors(const Module &M) {
+  const GlobalVariable *DtorsList = M.getNamedGlobal("llvm.global_dtors");
+  return make_range(CtorDtorIterator(DtorsList, false),
+                    CtorDtorIterator(DtorsList, true));
+}
+
+void LocalCXXRuntimeOverrides::runDestructors() {
+  auto& CXXDestructorDataPairs = DSOHandleOverride;
+  for (auto &P : CXXDestructorDataPairs)
+    P.first(P.second);
+  CXXDestructorDataPairs.clear();
+}
+
+int LocalCXXRuntimeOverrides::CXAAtExitOverride(DestructorPtr Destructor,
+                                                void *Arg, void *DSOHandle) {
+  auto& CXXDestructorDataPairs =
+    *reinterpret_cast<CXXDestructorDataPairList*>(DSOHandle);
+  CXXDestructorDataPairs.push_back(std::make_pair(Destructor, Arg));
+  return 0;
+}
+
+} // End namespace orc.
+} // End namespace llvm.
diff --git a/lib/ExecutionEngine/Orc/IndirectionUtils.cpp b/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
index 8cf490f..ebeedef 100644
--- a/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
+++ b/lib/ExecutionEngine/Orc/IndirectionUtils.cpp
@@ -14,17 +14,25 @@
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/IRBuilder.h"
 #include <set>
+#include <sstream>
 
 namespace llvm {
 namespace orc {
 
-GlobalVariable* createImplPointer(Function &F, const Twine &Name,
-                                  Constant *Initializer) {
-  assert(F.getParent() && "Function isn't in a module.");
+Constant* createIRTypedAddress(FunctionType &FT, TargetAddress Addr) {
+  Constant *AddrIntVal =
+    ConstantInt::get(Type::getInt64Ty(FT.getContext()), Addr);
+  Constant *AddrPtrVal =
+    ConstantExpr::getCast(Instruction::IntToPtr, AddrIntVal,
+                          PointerType::get(&FT, 0));
+  return AddrPtrVal;
+}
+
+GlobalVariable* createImplPointer(PointerType &PT, Module &M,
+                                  const Twine &Name, Constant *Initializer) {
   if (!Initializer)
-    Initializer = Constant::getNullValue(F.getType());
-  Module &M = *F.getParent();
-  return new GlobalVariable(M, F.getType(), false, GlobalValue::ExternalLinkage,
+    Initializer = Constant::getNullValue(&PT);
+  return new GlobalVariable(M, &PT, false, GlobalValue::ExternalLinkage,
                             Initializer, Name, nullptr,
                             GlobalValue::NotThreadLocal, 0, true);
 }
@@ -44,8 +52,41 @@ void makeStub(Function &F, GlobalVariable &ImplPointer) {
   Builder.CreateRet(Call);
 }
 
+// Utility class for renaming global values and functions during partitioning.
+class GlobalRenamer {
+public:
+
+  static bool needsRenaming(const Value &New) {
+    if (!New.hasName() || New.getName().startswith("\01L"))
+      return true;
+    return false;
+  }
+
+  const std::string& getRename(const Value &Orig) {
+    // See if we have a name for this global.
+    {
+      auto I = Names.find(&Orig);
+      if (I != Names.end())
+        return I->second;
+    }
+
+    // Nope. Create a new one.
+    // FIXME: Use a more robust uniquing scheme. (This may blow up if the user
+    //        writes a "__orc_anon[[:digit:]]* method).
+    unsigned ID = Names.size();
+    std::ostringstream NameStream;
+    NameStream << "__orc_anon" << ID++;
+    auto I = Names.insert(std::make_pair(&Orig, NameStream.str()));
+    return I.first->second;
+  }
+private:
+  DenseMap<const Value*, std::string> Names;
+};
+
 void partition(Module &M, const ModulePartitionMap &PMap) {
 
+  GlobalRenamer Renamer;
+
   for (auto &KVPair : PMap) {
 
     auto ExtractGlobalVars =
@@ -54,20 +95,26 @@ void partition(Module &M, const ModulePartitionMap &PMap) {
         if (KVPair.second.count(&Orig)) {
           copyGVInitializer(New, Orig, VMap);
         }
-        if (New.getLinkage() == GlobalValue::PrivateLinkage) {
+        if (New.hasLocalLinkage()) {
+          if (Renamer.needsRenaming(New))
+            New.setName(Renamer.getRename(Orig));
           New.setLinkage(GlobalValue::ExternalLinkage);
           New.setVisibility(GlobalValue::HiddenVisibility);
         }
+        assert(!Renamer.needsRenaming(New) && "Invalid global name.");
       };
 
     auto ExtractFunctions =
       [&](Function &New, const Function &Orig, ValueToValueMapTy &VMap) {
         if (KVPair.second.count(&Orig))
           copyFunctionBody(New, Orig, VMap);
-        if (New.getLinkage() == GlobalValue::InternalLinkage) {
+        if (New.hasLocalLinkage()) {
+          if (Renamer.needsRenaming(New))
+            New.setName(Renamer.getRename(Orig));
           New.setLinkage(GlobalValue::ExternalLinkage);
           New.setVisibility(GlobalValue::HiddenVisibility);
         }
+        assert(!Renamer.needsRenaming(New) && "Invalid function name.");
       };
 
     CloneSubModule(*KVPair.first, M, ExtractGlobalVars, ExtractFunctions,
diff --git a/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h b/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
index 00e39bb..4023344 100644
--- a/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
+++ b/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
@@ -26,15 +26,21 @@ namespace orc {
 
 class OrcMCJITReplacement : public ExecutionEngine {
 
-  class ForwardingRTDyldMM : public RTDyldMemoryManager {
+  // OrcMCJITReplacement needs to do a little extra book-keeping to ensure that
+  // Orc's automatic finalization doesn't kick in earlier than MCJIT clients are
+  // expecting - see finalizeMemory.
+  class MCJITReplacementMemMgr : public MCJITMemoryManager {
   public:
-    ForwardingRTDyldMM(OrcMCJITReplacement &M) : M(M) {}
+    MCJITReplacementMemMgr(OrcMCJITReplacement &M,
+                           std::shared_ptr<MCJITMemoryManager> ClientMM)
+      : M(M), ClientMM(std::move(ClientMM)) {}
 
     uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
                                  unsigned SectionID,
                                  StringRef SectionName) override {
       uint8_t *Addr =
-          M.MM->allocateCodeSection(Size, Alignment, SectionID, SectionName);
+          ClientMM->allocateCodeSection(Size, Alignment, SectionID,
+                                        SectionName);
       M.SectionsAllocatedSinceLastLoad.insert(Addr);
       return Addr;
     }
@@ -42,43 +48,35 @@ class OrcMCJITReplacement : public ExecutionEngine {
     uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
                                  unsigned SectionID, StringRef SectionName,
                                  bool IsReadOnly) override {
-      uint8_t *Addr = M.MM->allocateDataSection(Size, Alignment, SectionID,
-                                                SectionName, IsReadOnly);
+      uint8_t *Addr = ClientMM->allocateDataSection(Size, Alignment, SectionID,
+                                                    SectionName, IsReadOnly);
       M.SectionsAllocatedSinceLastLoad.insert(Addr);
       return Addr;
     }
 
     void reserveAllocationSpace(uintptr_t CodeSize, uintptr_t DataSizeRO,
                                 uintptr_t DataSizeRW) override {
-      return M.MM->reserveAllocationSpace(CodeSize, DataSizeRO, DataSizeRW);
+      return ClientMM->reserveAllocationSpace(CodeSize, DataSizeRO,
+                                                DataSizeRW);
     }
 
     bool needsToReserveAllocationSpace() override {
-      return M.MM->needsToReserveAllocationSpace();
+      return ClientMM->needsToReserveAllocationSpace();
     }
 
     void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                           size_t Size) override {
-      return M.MM->registerEHFrames(Addr, LoadAddr, Size);
+      return ClientMM->registerEHFrames(Addr, LoadAddr, Size);
     }
 
     void deregisterEHFrames(uint8_t *Addr, uint64_t LoadAddr,
                             size_t Size) override {
-      return M.MM->deregisterEHFrames(Addr, LoadAddr, Size);
-    }
-
-    uint64_t getSymbolAddress(const std::string &Name) override {
-      return M.getSymbolAddressWithoutMangling(Name);
-    }
-
-    void *getPointerToNamedFunction(const std::string &Name,
-                                    bool AbortOnFailure = true) override {
-      return M.MM->getPointerToNamedFunction(Name, AbortOnFailure);
+      return ClientMM->deregisterEHFrames(Addr, LoadAddr, Size);
     }
 
     void notifyObjectLoaded(ExecutionEngine *EE,
                             const object::ObjectFile &O) override {
-      return M.MM->notifyObjectLoaded(EE, O);
+      return ClientMM->notifyObjectLoaded(EE, O);
     }
 
     bool finalizeMemory(std::string *ErrMsg = nullptr) override {
@@ -96,21 +94,41 @@ class OrcMCJITReplacement : public ExecutionEngine {
       // get more than one set of objects loaded but not yet finalized is if
       // they were loaded during relocation of another set.
       if (M.UnfinalizedSections.size() == 1)
-        return M.MM->finalizeMemory(ErrMsg);
+        return ClientMM->finalizeMemory(ErrMsg);
       return false;
     }
 
   private:
     OrcMCJITReplacement &M;
+    std::shared_ptr<MCJITMemoryManager> ClientMM;
+  };
+
+  class LinkingResolver : public RuntimeDyld::SymbolResolver {
+  public:
+    LinkingResolver(OrcMCJITReplacement &M) : M(M) {}
+
+    RuntimeDyld::SymbolInfo findSymbol(const std::string &Name) override {
+      return M.findMangledSymbol(Name);
+    }
+
+    RuntimeDyld::SymbolInfo
+    findSymbolInLogicalDylib(const std::string &Name) override {
+      return M.ClientResolver->findSymbolInLogicalDylib(Name);
+    }
+
+  private:
+    OrcMCJITReplacement &M;
   };
 
 private:
 
   static ExecutionEngine *
   createOrcMCJITReplacement(std::string *ErrorMsg,
-                            std::unique_ptr<RTDyldMemoryManager> OrcJMM,
+                            std::shared_ptr<MCJITMemoryManager> MemMgr,
+                            std::shared_ptr<RuntimeDyld::SymbolResolver> Resolver,
                             std::unique_ptr<TargetMachine> TM) {
-    return new OrcMCJITReplacement(std::move(OrcJMM), std::move(TM));
+    return new OrcMCJITReplacement(std::move(MemMgr), std::move(Resolver),
+                                   std::move(TM));
   }
 
 public:
@@ -118,12 +136,15 @@ public:
     OrcMCJITReplacementCtor = createOrcMCJITReplacement;
   }
 
-  OrcMCJITReplacement(std::unique_ptr<RTDyldMemoryManager> MM,
-                      std::unique_ptr<TargetMachine> TM)
-      : TM(std::move(TM)), MM(std::move(MM)), Mang(this->TM->getDataLayout()),
+  OrcMCJITReplacement(
+                    std::shared_ptr<MCJITMemoryManager> MemMgr,
+                    std::shared_ptr<RuntimeDyld::SymbolResolver> ClientResolver,
+                    std::unique_ptr<TargetMachine> TM)
+      : TM(std::move(TM)), MemMgr(*this, std::move(MemMgr)),
+        Resolver(*this), ClientResolver(std::move(ClientResolver)),
+        Mang(this->TM->getDataLayout()),
         NotifyObjectLoaded(*this), NotifyFinalized(*this),
-        ObjectLayer(ObjectLayerT::CreateRTDyldMMFtor(), NotifyObjectLoaded,
-                    NotifyFinalized),
+        ObjectLayer(NotifyObjectLoaded, NotifyFinalized),
         CompileLayer(ObjectLayer, SimpleCompiler(*this->TM)),
         LazyEmitLayer(CompileLayer) {
     setDataLayout(this->TM->getDataLayout());
@@ -139,15 +160,13 @@ public:
     Modules.push_back(std::move(M));
     std::vector<Module *> Ms;
     Ms.push_back(&*Modules.back());
-    LazyEmitLayer.addModuleSet(std::move(Ms),
-                               llvm::make_unique<ForwardingRTDyldMM>(*this));
+    LazyEmitLayer.addModuleSet(std::move(Ms), &MemMgr, &Resolver);
   }
 
   void addObjectFile(std::unique_ptr<object::ObjectFile> O) override {
     std::vector<std::unique_ptr<object::ObjectFile>> Objs;
     Objs.push_back(std::move(O));
-    ObjectLayer.addObjectSet(std::move(Objs),
-                             llvm::make_unique<ForwardingRTDyldMM>(*this));
+    ObjectLayer.addObjectSet(std::move(Objs), &MemMgr, &Resolver);
   }
 
   void addObjectFile(object::OwningBinary<object::ObjectFile> O) override {
@@ -157,8 +176,7 @@ public:
     std::vector<std::unique_ptr<object::ObjectFile>> Objs;
     Objs.push_back(std::move(Obj));
     auto H =
-      ObjectLayer.addObjectSet(std::move(Objs),
-                               llvm::make_unique<ForwardingRTDyldMM>(*this));
+      ObjectLayer.addObjectSet(std::move(Objs), &MemMgr, &Resolver);
 
     std::vector<std::unique_ptr<MemoryBuffer>> Bufs;
     Bufs.push_back(std::move(Buf));
@@ -170,7 +188,11 @@ public:
   }
 
   uint64_t getSymbolAddress(StringRef Name) {
-    return getSymbolAddressWithoutMangling(Mangle(Name));
+    return findSymbol(Name).getAddress();
+  }
+
+  RuntimeDyld::SymbolInfo findSymbol(StringRef Name) {
+    return findMangledSymbol(Mangle(Name));
   }
 
   void finalizeObject() override {
@@ -214,18 +236,19 @@ public:
   }
 
 private:
-  uint64_t getSymbolAddressWithoutMangling(StringRef Name) {
-    if (uint64_t Addr = LazyEmitLayer.findSymbol(Name, false).getAddress())
-      return Addr;
-    if (uint64_t Addr = MM->getSymbolAddress(Name))
-      return Addr;
-    if (uint64_t Addr = scanArchives(Name))
-      return Addr;
 
-    return 0;
+  RuntimeDyld::SymbolInfo findMangledSymbol(StringRef Name) {
+    if (auto Sym = LazyEmitLayer.findSymbol(Name, false))
+      return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
+    if (auto Sym = ClientResolver->findSymbol(Name))
+      return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
+    if (auto Sym = scanArchives(Name))
+      return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
+
+    return nullptr;
   }
 
-  uint64_t scanArchives(StringRef Name) {
+  JITSymbol scanArchives(StringRef Name) {
     for (object::OwningBinary<object::Archive> &OB : Archives) {
       object::Archive *A = OB.getBinary();
       // Look for our symbols in each Archive
@@ -241,14 +264,13 @@ private:
           std::vector<std::unique_ptr<object::ObjectFile>> ObjSet;
           ObjSet.push_back(std::unique_ptr<object::ObjectFile>(
               static_cast<object::ObjectFile *>(ChildBin.release())));
-          ObjectLayer.addObjectSet(
-              std::move(ObjSet), llvm::make_unique<ForwardingRTDyldMM>(*this));
-          if (uint64_t Addr = ObjectLayer.findSymbol(Name, true).getAddress())
-            return Addr;
+          ObjectLayer.addObjectSet(std::move(ObjSet), &MemMgr, &Resolver);
+          if (auto Sym = ObjectLayer.findSymbol(Name, true))
+            return Sym;
         }
       }
     }
-    return 0;
+    return nullptr;
   }
 
   class NotifyObjectLoadedT {
@@ -267,7 +289,7 @@ private:
       assert(Objects.size() == Infos.size() &&
              "Incorrect number of Infos for Objects.");
       for (unsigned I = 0; I < Objects.size(); ++I)
-        M.MM->notifyObjectLoaded(&M, *Objects[I]);
+        M.MemMgr.notifyObjectLoaded(&M, *Objects[I]);
     };
 
   private:
@@ -299,7 +321,9 @@ private:
   typedef LazyEmittingLayer<CompileLayerT> LazyEmitLayerT;
 
   std::unique_ptr<TargetMachine> TM;
-  std::unique_ptr<RTDyldMemoryManager> MM;
+  MCJITReplacementMemMgr MemMgr;
+  LinkingResolver Resolver;
+  std::shared_ptr<RuntimeDyld::SymbolResolver> ClientResolver;
   Mangler Mang;
 
   NotifyObjectLoadedT NotifyObjectLoaded;
diff --git a/lib/ExecutionEngine/Orc/OrcTargetSupport.cpp b/lib/ExecutionEngine/Orc/OrcTargetSupport.cpp
index 6fe5301..fc56e67 100644
--- a/lib/ExecutionEngine/Orc/OrcTargetSupport.cpp
+++ b/lib/ExecutionEngine/Orc/OrcTargetSupport.cpp
@@ -6,39 +6,6 @@ using namespace llvm::orc;
 
 namespace {
 
-std::array<const char *, 12> X86GPRsToSave = {{
-    "rbp", "rbx", "r12", "r13", "r14", "r15", // Callee saved.
-    "rdi", "rsi", "rdx", "rcx", "r8", "r9",   // Int args.
-}};
-
-std::array<const char *, 8> X86XMMsToSave = {{
-    "xmm0", "xmm1", "xmm2", "xmm3", "xmm4", "xmm5", "xmm6", "xmm7" // FP args
-}};
-
-template <typename OStream> unsigned saveX86Regs(OStream &OS) {
-  for (const auto &GPR : X86GPRsToSave)
-    OS << "  pushq   %" << GPR << "\n";
-
-  OS << "  subq    $" << (16 * X86XMMsToSave.size()) << ", %rsp\n";
-
-  for (unsigned i = 0; i < X86XMMsToSave.size(); ++i)
-    OS << "  movdqu  %" << X86XMMsToSave[i] << ", "
-       << (16 * (X86XMMsToSave.size() - i - 1)) << "(%rsp)\n";
-
-  return (8 * X86GPRsToSave.size()) + (16 * X86XMMsToSave.size());
-}
-
-template <typename OStream> void restoreX86Regs(OStream &OS) {
-  for (unsigned i = 0; i < X86XMMsToSave.size(); ++i)
-    OS << "  movdqu  " << (16 * i) << "(%rsp), %"
-       << X86XMMsToSave[(X86XMMsToSave.size() - i - 1)] << "\n";
-  OS << "  addq    $" << (16 * X86XMMsToSave.size()) << ", %rsp\n";
-
-  for (unsigned i = 0; i < X86GPRsToSave.size(); ++i)
-    OS << "  popq    %" << X86GPRsToSave[X86GPRsToSave.size() - i - 1] << "\n";
-}
-
-template <typename TargetT>
 uint64_t executeCompileCallback(JITCompileCallbackManagerBase *JCBM,
                                 TargetAddress CallbackID) {
   return JCBM->executeCompileCallback(CallbackID);
@@ -53,14 +20,28 @@ const char* OrcX86_64::ResolverBlockName = "orc_resolver_block";
 
 void OrcX86_64::insertResolverBlock(
     Module &M, JITCompileCallbackManagerBase &JCBM) {
+
+  // Trampoline code-sequence length, used to get trampoline address from return
+  // address.
   const unsigned X86_64_TrampolineLength = 6;
-  auto CallbackPtr = executeCompileCallback<OrcX86_64>;
+
+  // List of x86-64 GPRs to save. Note - RBP saved separately below.
+  std::array<const char *, 14> GPRs = {{
+      "rax", "rbx", "rcx", "rdx",
+      "rsi", "rdi", "r8", "r9",
+      "r10", "r11", "r12", "r13",
+      "r14", "r15"
+    }};
+
+  // Address of the executeCompileCallback function.
   uint64_t CallbackAddr =
-      static_cast<uint64_t>(reinterpret_cast<uintptr_t>(CallbackPtr));
+      static_cast<uint64_t>(
+        reinterpret_cast<uintptr_t>(executeCompileCallback));
 
   std::ostringstream AsmStream;
   Triple TT(M.getTargetTriple());
 
+  // Switch to text section.
   if (TT.getOS() == Triple::Darwin)
     AsmStream << ".section __TEXT,__text,regular,pure_instructions\n"
               << ".align 4, 0x90\n";
@@ -68,24 +49,51 @@ void OrcX86_64::insertResolverBlock(
     AsmStream << ".text\n"
               << ".align 16, 0x90\n";
 
+  // Bake in a pointer to the callback manager immediately before the
+  // start of the resolver function.
   AsmStream << "jit_callback_manager_addr:\n"
-            << "  .quad " << &JCBM << "\n"
-            << ResolverBlockName << ":\n";
-
-  uint64_t ReturnAddrOffset = saveX86Regs(AsmStream);
-
-  // Compute index, load object address, and call JIT.
-  AsmStream << "  leaq    jit_callback_manager_addr(%rip), %rdi\n"
+            << "  .quad " << &JCBM << "\n";
+
+  // Start the resolver function.
+  AsmStream << ResolverBlockName << ":\n"
+            << "  pushq   %rbp\n"
+            << "  movq    %rsp, %rbp\n";
+
+  // Store the GPRs.
+  for (const auto &GPR : GPRs)
+    AsmStream << "  pushq   %" << GPR << "\n";
+
+  // Store floating-point state with FXSAVE.
+  // Note: We need to keep the stack 16-byte aligned, so if we've emitted an odd
+  //       number of 64-bit pushes so far (GPRs.size() plus 1 for RBP) then add
+  //       an extra 64 bits of padding to the FXSave area.
+  unsigned Padding = (GPRs.size() + 1) % 2 ? 8 : 0;
+  unsigned FXSaveSize = 512 + Padding;
+  AsmStream << "  subq    $" << FXSaveSize << ", %rsp\n"
+            << "  fxsave  (%rsp)\n"
+
+  // Load callback manager address, compute trampoline address, call JIT.
+            << "  lea     jit_callback_manager_addr(%rip), %rdi\n"
             << "  movq    (%rdi), %rdi\n"
-            << "  movq    " << ReturnAddrOffset << "(%rsp), %rsi\n"
+            << "  movq    0x8(%rbp), %rsi\n"
             << "  subq    $" << X86_64_TrampolineLength << ", %rsi\n"
             << "  movabsq $" << CallbackAddr << ", %rax\n"
             << "  callq   *%rax\n"
-            << "  movq    %rax, " << ReturnAddrOffset << "(%rsp)\n";
 
-  restoreX86Regs(AsmStream);
+  // Replace the return to the trampoline with the return address of the
+  // compiled function body.
+            << "  movq    %rax, 0x8(%rbp)\n"
+
+  // Restore the floating point state.
+            << "  fxrstor (%rsp)\n"
+            << "  addq    $" << FXSaveSize << ", %rsp\n";
+
+  for (const auto &GPR : make_range(GPRs.rbegin(), GPRs.rend()))
+    AsmStream << "  popq    %" << GPR << "\n";
 
-  AsmStream << "  retq\n";
+  // Restore original RBP and return to compiled function body.
+  AsmStream << "  popq    %rbp\n"
+            << "  retq\n";
 
   M.appendModuleInlineAsm(AsmStream.str());
 }