[CodeGen] Fix MachineModuleInfo's move constructor to be more safe with MCContext ownership.

llvm/include/llvm/CodeGen/MachineModuleInfo.h

@@ -86,7 +86,7 @@ class MachineModuleInfo {
   const LLVMTargetMachine &TM;
 
   /// This is the MCContext used for the entire code generator.
-  MCContext Context;
+  std::unique_ptr<MCContext> Context;
   // This is an external context, that if assigned, will be used instead of the
   // internal context.
   MCContext *ExternalContext = nullptr;
@@ -124,10 +124,10 @@ class MachineModuleInfo {
   const LLVMTargetMachine &getTarget() const { return TM; }
 
   const MCContext &getContext() const {
-    return ExternalContext ? *ExternalContext : Context;
+    return ExternalContext ? *ExternalContext : *Context;
   }
   MCContext &getContext() {
-    return ExternalContext ? *ExternalContext : Context;
+    return ExternalContext ? *ExternalContext : *Context;
   }
 
   const Module *getModule() const { return TheModule; }

llvm/lib/CodeGen/MachineModuleInfo.cpp

@@ -30,39 +30,39 @@ void MachineModuleInfo::initialize() {
 }
 
 void MachineModuleInfo::finalize() {
-  Context.reset();
+  if (Context)
+    Context->reset();
   // We don't clear the ExternalContext.
 
   delete ObjFileMMI;
   ObjFileMMI = nullptr;
 }
 
 MachineModuleInfo::MachineModuleInfo(MachineModuleInfo &&MMI)
-    : TM(std::move(MMI.TM)),
-      Context(TM.getTargetTriple(), TM.getMCAsmInfo(), TM.getMCRegisterInfo(),
-              TM.getMCSubtargetInfo(), nullptr, &TM.Options.MCOptions, false),
+    : TM(std::move(MMI.TM)), Context(std::move(MMI.Context)),
       MachineFunctions(std::move(MMI.MachineFunctions)) {
-  Context.setObjectFileInfo(TM.getObjFileLowering());
   ObjFileMMI = MMI.ObjFileMMI;
   ExternalContext = MMI.ExternalContext;
   TheModule = MMI.TheModule;
 }
 
 MachineModuleInfo::MachineModuleInfo(const LLVMTargetMachine *TM)
-    : TM(*TM), Context(TM->getTargetTriple(), TM->getMCAsmInfo(),
-                       TM->getMCRegisterInfo(), TM->getMCSubtargetInfo(),
-                       nullptr, &TM->Options.MCOptions, false) {
-  Context.setObjectFileInfo(TM->getObjFileLowering());
+    : TM(*TM),
+      Context(std::make_unique<MCContext>(
+          TM->getTargetTriple(), TM->getMCAsmInfo(), TM->getMCRegisterInfo(),
+          TM->getMCSubtargetInfo(), nullptr, &TM->Options.MCOptions, false)) {
+  Context->setObjectFileInfo(TM->getObjFileLowering());
   initialize();
 }
 
 MachineModuleInfo::MachineModuleInfo(const LLVMTargetMachine *TM,
                                      MCContext *ExtContext)
-    : TM(*TM), Context(TM->getTargetTriple(), TM->getMCAsmInfo(),
-                       TM->getMCRegisterInfo(), TM->getMCSubtargetInfo(),
-                       nullptr, &TM->Options.MCOptions, false),
+    : TM(*TM),
+      Context(std::make_unique<MCContext>(
+          TM->getTargetTriple(), TM->getMCAsmInfo(), TM->getMCRegisterInfo(),
+          TM->getMCSubtargetInfo(), nullptr, &TM->Options.MCOptions, false)),
       ExternalContext(ExtContext) {
-  Context.setObjectFileInfo(TM->getObjFileLowering());
+  Context->setObjectFileInfo(TM->getObjFileLowering());
   initialize();
 }
 

llvm/unittests/CodeGen/CMakeLists.txt

@@ -35,6 +35,7 @@ add_llvm_unittest(CodeGenTests
   MachineDomTreeUpdaterTest.cpp
   MachineInstrBundleIteratorTest.cpp
   MachineInstrTest.cpp
+  MachineModuleInfoTest.cpp
   MachineOperandTest.cpp
   RegAllocScoreTest.cpp
   PassManagerTest.cpp

llvm/unittests/CodeGen/MachineModuleInfoTest.cpp

@@ -0,0 +1,79 @@
+//===- llvm/unittest/CodeGen/PassManager.cpp - PassManager tests ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// Test MachineModuleInfo.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/TargetParser/Host.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+std::unique_ptr<Module> parseIR(LLVMContext &Context, const char *IR) {
+  SMDiagnostic Err;
+  return parseAssemblyString(IR, Err, Context);
+}
+
+class MachineModuleInfoTest : public ::testing::Test {
+protected:
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+  std::unique_ptr<TargetMachine> TM;
+
+public:
+  MachineModuleInfoTest()
+      : M(parseIR(Context, "define void @f() {\n"
+                           "entry:\n"
+                           "  call void @g()\n"
+                           "  ret void\n"
+                           "}\n"
+                           "define void @g() {\n"
+                           "  ret void\n"
+                           "}\n")) {
+    // MachineModuleInfo needs a TargetMachine instance.
+    llvm::InitializeAllTargets();
+
+    std::string TripleName = Triple::normalize(sys::getDefaultTargetTriple());
+    std::string Error;
+    const Target *TheTarget = TargetRegistry::lookupTarget(TripleName, Error);
+    if (!TheTarget)
+      return;
+
+    TargetOptions Options;
+    TM.reset(TheTarget->createTargetMachine(TripleName, "", "", Options,
+                                            std::nullopt));
+  }
+};
+
+TEST_F(MachineModuleInfoTest, MachineModuleInfoMoveConstructorMovesMCContext) {
+  if (!TM)
+    GTEST_SKIP();
+
+  LLVMTargetMachine *LLVMTM = static_cast<LLVMTargetMachine *>(TM.get());
+  M->setDataLayout(TM->createDataLayout());
+
+  MachineModuleInfo MMI(LLVMTM);
+
+  MCContext *OriginalCtx = &MMI.getContext();
+
+  MachineModuleInfo MovedMMI(std::move(MMI));
+  MCContext *MovedCtx = &MovedMMI.getContext();
+
+  EXPECT_EQ(OriginalCtx, MovedCtx);
+}
+
+} // end namespace