ThinLTOBitcodeWriter: Include globals associated with type metadata globals in the merged module.

[pcc]

Summary

Globals that are associated with globals with type metadata need to appear
in the merged module because they will reference the global's section directly.

[tejohnson]

lgtm

llvm/lib/Transforms/IPO/ThinLTOBitcodeWriter.cpp

if (ThinLinkOS)
       WriteBitcodeToFile(M, *ThinLinkOS, /*ShouldPreserveUseListOrder=*/false,
                          &Index);
 
     return;
   }
 
   promoteTypeIds(M, ModuleId);
 
-  // Returns whether a global has attached type metadata. Such globals may
-  // participate in CFI or whole-program devirtualization, so they need to
-  // appear in the merged module instead of the thin LTO module.
+  // Returns whether a global or its associated global has attached type
+  // metadata. The former may participate in CFI or whole-program
+  // devirtualization, so they need to appear in the merged module instead of
+  // the thin LTO module. Similarly, globals that are associated with globals
+  // with type metadata need to appear in the merged module because they will
+  // reference the global's section directly.
   auto HasTypeMetadata = [](const GlobalObject *GO) {
+    if (MDNode *MD = GO->getMetadata(LLVMContext::MD_associated))
+      if (auto *AssocVM = dyn_cast_or_null<ValueAsMetadata>(MD->getOperand(0)))
+        if (auto *AssocGO = dyn_cast<GlobalObject>(AssocVM->getValue()))
+          if (AssocGO->hasMetadata(LLVMContext::MD_type))
+            return true;
     return GO->hasMetadata(LLVMContext::MD_type);
   };
 
   // Collect the set of virtual functions that are eligible for virtual constant
   // propagation. Each eligible function must not access memory, must return
   // an integer of width <=64 bits, must take at least one argument, must not
   // use its first argument (assumed to be "this") and all arguments other than
   // the first one must be of <=64 bit integer type.

llvm/test/Transforms/ThinLTOBitcodeWriter/associated.ll

+; RUN: opt -thinlto-bc -thinlto-split-lto-unit -o %t %s
+; RUN: llvm-modextract -b -n 0 -o - %t | llvm-dis | FileCheck --check-prefix=M0 %s
+; RUN: llvm-modextract -b -n 1 -o - %t | llvm-dis | FileCheck --check-prefix=M1 %s
+
+; M0: @g = external constant
+; M0-NOT: @assoc
+; M1: @g = constant i8 1
+; M1: @assoc = private constant i8 2
+
+@g = constant i8 1, !type !0
+@assoc = private constant i8 2, !associated !1
+
+!0 = !{i32 0, !"typeid"}
+!1 = !{i8* @g}