Add MachineModuleInfo constructor with external MCContext

[hgreving]

Summary

Adds a constructor to MachineModuleInfo and MachineModuleInfoWapperPass that
takes an external MCContext. If provided, the external context will be used
throughout codegen instead of MMI's default one.

This enables external drivers to take ownership of data put on the MMI's context
during codegen. The internal context is otherwise used and destroyed upon
finish.

[hgreving]

Copy constructor.

[majnemer]

LGTM

llvm/lib/CodeGen/MachineModuleInfo.cpp

172–173

Probably should have a comment saying *not* to clear out ExternalContext.

[jmolloy]

Looks reasonable to me. I'd have probably done something like:

MCContext& Context; // not owned, used for all accesses;
optional<MCContext> LocalContextStorage;

MachineModuleInfo() :

LocalContextStorage(), Context(LocalContextStorage) {

}

MachineModuleInfo(MCContext* Context) :

Context(*Context) {

}

But your way is fine too.

llvm/include/llvm/CodeGen/MachineModuleInfo.h

 class MachineModuleInfo {
   friend class MachineModuleInfoWrapperPass;
   friend class MachineModuleAnalysis;
 
   const LLVMTargetMachine &TM;
 
   /// This is the MCContext used for the entire code generator.
   MCContext Context;
+  // This is an external context, that if assigned, will be used instead of the
+  // internal context.
+  MCContext *ExternalContext = nullptr;
 
   /// This is the LLVM Module being worked on.
   const Module *TheModule;
 
   /// This is the object-file-format-specific implementation of
   /// MachineModuleInfoImpl, which lets targets accumulate whatever info they
   /// want.
   MachineModuleInfoImpl *ObjFileMMI;
class MachineModuleInfo {
   const Function *LastRequest = nullptr; ///< Used for shortcut/cache.
   MachineFunction *LastResult = nullptr; ///< Used for shortcut/cache.
 
   MachineModuleInfo &operator=(MachineModuleInfo &&MMII) = delete;
 
 public:
   explicit MachineModuleInfo(const LLVMTargetMachine *TM = nullptr);
 
+  explicit MachineModuleInfo(const LLVMTargetMachine *TM,
+                             MCContext *ExtContext);
+
   MachineModuleInfo(MachineModuleInfo &&MMII);
 
   ~MachineModuleInfo();
 
   void initialize();
   void finalize();
 
   const LLVMTargetMachine &getTarget() const { return TM; }
 
-  const MCContext &getContext() const { return Context; }
-  MCContext &getContext() { return Context; }
+  const MCContext &getContext() const {
+    return ExternalContext ? *ExternalContext : Context;
+  }
+  MCContext &getContext() {
+    return ExternalContext ? *ExternalContext : Context;
+  }
 
   const Module *getModule() const { return TheModule; }
 
   /// Returns the MachineFunction constructed for the IR function \p F.
   /// Creates a new MachineFunction if none exists yet.
   MachineFunction &getOrCreateMachineFunction(Function &F);
 
   /// \brief Returns the MachineFunction associated to IR function \p F if there

 
 class MachineModuleInfoWrapperPass : public ImmutablePass {
   MachineModuleInfo MMI;
 
 public:
   static char ID; // Pass identification, replacement for typeid
   explicit MachineModuleInfoWrapperPass(const LLVMTargetMachine *TM = nullptr);
 
+  explicit MachineModuleInfoWrapperPass(const LLVMTargetMachine *TM,
+                                        MCContext *ExtContext);
+
   // Initialization and Finalization
   bool doInitialization(Module &) override;
   bool doFinalization(Module &) override;
 
   MachineModuleInfo &getMMI() { return MMI; }
   const MachineModuleInfo &getMMI() const { return MMI; }
 };
 

llvm/lib/CodeGen/MachineModuleInfo.cpp

 }
 
 void MachineModuleInfo::finalize() {
   Personalities.clear();
 
   delete AddrLabelSymbols;
   AddrLabelSymbols = nullptr;
 
   Context.reset();
+  // We don't clear the ExternalContext.
 
   delete ObjFileMMI;
   ObjFileMMI = nullptr;
 }
 
 MachineModuleInfo::MachineModuleInfo(MachineModuleInfo &&MMI)
     : TM(std::move(MMI.TM)),
       Context(MMI.TM.getMCAsmInfo(), MMI.TM.getMCRegisterInfo(),
               MMI.TM.getObjFileLowering(), nullptr, nullptr, false),
       MachineFunctions(std::move(MMI.MachineFunctions)) {
   ObjFileMMI = MMI.ObjFileMMI;
   CurCallSite = MMI.CurCallSite;
   UsesMSVCFloatingPoint = MMI.UsesMSVCFloatingPoint;
   UsesMorestackAddr = MMI.UsesMorestackAddr;
   HasSplitStack = MMI.HasSplitStack;
   HasNosplitStack = MMI.HasNosplitStack;
   AddrLabelSymbols = MMI.AddrLabelSymbols;
+  ExternalContext = MMI.ExternalContext;
   TheModule = MMI.TheModule;
 }
 
 MachineModuleInfo::MachineModuleInfo(const LLVMTargetMachine *TM)
     : TM(*TM), Context(TM->getMCAsmInfo(), TM->getMCRegisterInfo(),
                        TM->getObjFileLowering(), nullptr, nullptr, false) {
   initialize();
 }
 
+MachineModuleInfo::MachineModuleInfo(const LLVMTargetMachine *TM,
+                                     MCContext *ExtContext)
+    : TM(*TM), Context(TM->getMCAsmInfo(), TM->getMCRegisterInfo(),
+                       TM->getObjFileLowering(), nullptr, nullptr, false),
+      ExternalContext(ExtContext) {
+  initialize();
+}
+
 MachineModuleInfo::~MachineModuleInfo() { finalize(); }
 
 //===- Address of Block Management ----------------------------------------===//
 
 ArrayRef<MCSymbol *>
 MachineModuleInfo::getAddrLabelSymbolToEmit(const BasicBlock *BB) {
   // Lazily create AddrLabelSymbols.
   if (!AddrLabelSymbols)
-    AddrLabelSymbols = new MMIAddrLabelMap(Context);
+    AddrLabelSymbols = new MMIAddrLabelMap(getContext());
  return AddrLabelSymbols->getAddrLabelSymbolToEmit(const_cast<BasicBlock*>(BB));
 }
 
 /// \name Exception Handling
 /// \{
 
 void MachineModuleInfo::addPersonality(const Function *Personality) {
   for (unsigned i = 0; i < Personalities.size(); ++i)

 }
 
 MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass(
     const LLVMTargetMachine *TM)
     : ImmutablePass(ID), MMI(TM) {
   initializeMachineModuleInfoWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
+MachineModuleInfoWrapperPass::MachineModuleInfoWrapperPass(
+    const LLVMTargetMachine *TM, MCContext *ExtContext)
+    : ImmutablePass(ID), MMI(TM, ExtContext) {
+  initializeMachineModuleInfoWrapperPassPass(*PassRegistry::getPassRegistry());
+}
+
 // Handle the Pass registration stuff necessary to use DataLayout's.
 INITIALIZE_PASS(MachineModuleInfoWrapperPass, "machinemoduleinfo",
                 "Machine Module Information", false, false)
 char MachineModuleInfoWrapperPass::ID = 0;
 
 bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {
   MMI.initialize();
   MMI.TheModule = &M;