1 files changed, 49 insertions, 5 deletions

diff --git a/llvm/lib/CodeGen/MachineModuleInfo.cpp b/llvm/lib/CodeGen/MachineModuleInfo.cpp
index 244e3fbc4e8..4b4c69319a9 100644
--- a/llvm/lib/CodeGen/MachineModuleInfo.cpp
+++ b/llvm/lib/CodeGen/MachineModuleInfo.cpp
@@ -13,6 +13,7 @@
 #include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionInitializer.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
@@ -186,15 +187,19 @@ void MMIAddrLabelMapCallbackPtr::allUsesReplacedWith(Value *V2) {
 
 //===----------------------------------------------------------------------===//
 
-MachineModuleInfo::MachineModuleInfo(const MCAsmInfo &MAI,
+MachineModuleInfo::MachineModuleInfo(const TargetMachine &TM,
+                                     const MCAsmInfo &MAI,
                                      const MCRegisterInfo &MRI,
-                                     const MCObjectFileInfo *MOFI)
-  : ImmutablePass(ID), Context(&MAI, &MRI, MOFI, nullptr, false) {
+                                     const MCObjectFileInfo *MOFI,
+                                     MachineFunctionInitializer *MFI)
+  : ImmutablePass(ID), TM(TM), Context(&MAI, &MRI, MOFI, nullptr, false),
+    MFInitializer(MFI) {
   initializeMachineModuleInfoPass(*PassRegistry::getPassRegistry());
 }
 
 MachineModuleInfo::MachineModuleInfo()
-  : ImmutablePass(ID), Context(nullptr, nullptr, nullptr) {
+  : ImmutablePass(ID), TM(*((TargetMachine*)nullptr)),
+    Context(nullptr, nullptr, nullptr) {
   llvm_unreachable("This MachineModuleInfo constructor should never be called, "
                    "MMI should always be explicitly constructed by "
                    "LLVMTargetMachine");
@@ -213,7 +218,7 @@ bool MachineModuleInfo::doInitialization(Module &M) {
   DbgInfoAvailable = UsesVAFloatArgument = UsesMorestackAddr = false;
   PersonalityTypeCache = EHPersonality::Unknown;
   AddrLabelSymbols = nullptr;
-  TheModule = nullptr;
+  TheModule = &M;
 
   return false;
 }
@@ -461,3 +466,42 @@ try_next:;
   FilterIds.push_back(0); // terminator
   return FilterID;
 }
+
+MachineFunction &MachineModuleInfo::getMachineFunction(const Function &F) {
+  // Shortcut for the common case where a sequence of MachineFunctionPasses
+  // all query for the same Function.
+  if (LastRequest == &F)
+    return *LastResult;
+
+  auto I = MachineFunctions.insert(
+      std::make_pair(&F, std::unique_ptr<MachineFunction>()));
+  MachineFunction *MF;
+  if (I.second) {
+    // No pre-existing machine function, create a new one.
+    unsigned FunctionNum = (unsigned)MachineFunctions.size() - 1;
+    MF = new MachineFunction(&F, TM, FunctionNum, *this);
+    // Update the set entry.
+    I.first->second.reset(MF);
+
+    if (MFInitializer)
+      if (MFInitializer->initializeMachineFunction(*MF))
+        report_fatal_error("Unable to initialize machine function");
+  } else {
+    MF = I.first->second.get();
+  }
+
+  LastRequest = &F;
+  LastResult = MF;
+  return *MF;
+}
+
+void MachineModuleInfo::deleteMachineFunction(MachineFunction &MF) {
+  if (LastResult == &MF) {
+    LastRequest = nullptr;
+    LastResult = nullptr;
+  }
+
+  auto I = MachineFunctions.find(MF.getFunction());
+  assert(I != MachineFunctions.end() && "MachineFunction is known");
+  I->second.reset(nullptr);
+}