IR: Allow multiple global metadata attachments with the same type.

This will be necessary to allow the global merge pass to attach
multiple debug info metadata nodes to global variables once we reverse
the edge from DIGlobalVariable to GlobalVariable.

Differential Revision: http://reviews.llvm.org/D20414

llvm-svn: 271358

12 files changed, 147 insertions, 85 deletions

diff --git a/llvm/include/llvm/IR/GlobalObject.h b/llvm/include/llvm/IR/GlobalObject.h
index 038fda9218c..8c494cf3bc9 100644
--- a/llvm/include/llvm/IR/GlobalObject.h
+++ b/llvm/include/llvm/IR/GlobalObject.h
@@ -74,14 +74,23 @@ public:
   /// Check if this has any metadata.
   bool hasMetadata() const { return hasMetadataHashEntry(); }
 
-  /// Get the current metadata attachment, if any.
+  /// Get the current metadata attachments for the given kind, if any.
   ///
-  /// Returns \c nullptr if such an attachment is missing.
+  /// These functions require that the function have at most a single attachment
+  /// of the given kind, and return \c nullptr if such an attachment is missing.
   /// @{
   MDNode *getMetadata(unsigned KindID) const;
   MDNode *getMetadata(StringRef Kind) const;
   /// @}
 
+  /// Appends all attachments with the given ID to \c MDs in insertion order.
+  /// If the global has no attachments with the given ID, or if ID is invalid,
+  /// leaves MDs unchanged.
+  /// @{
+  void getMetadata(unsigned KindID, SmallVectorImpl<MDNode *> &MDs) const;
+  void getMetadata(StringRef Kind, SmallVectorImpl<MDNode *> &MDs) const;
+  /// @}
+
   /// Set a particular kind of metadata attachment.
   ///
   /// Sets the given attachment to \c MD, erasing it if \c MD is \c nullptr or
@@ -91,14 +100,19 @@ public:
   void setMetadata(StringRef Kind, MDNode *MD);
   /// @}
 
-  /// Get all current metadata attachments.
+  /// Add a metadata attachment.
+  /// @{
+  void addMetadata(unsigned KindID, MDNode &MD);
+  void addMetadata(StringRef Kind, MDNode &MD);
+  /// @}
+
+  /// Appends all attachments for the global to \c MDs, sorting by attachment
+  /// ID. Attachments with the same ID appear in insertion order.
   void
   getAllMetadata(SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) const;
 
-  /// Drop metadata not in the given list.
-  ///
-  /// Drop all metadata from \c this not included in \c KnownIDs.
-  void dropUnknownMetadata(ArrayRef<unsigned> KnownIDs);
+  /// Erase all metadata attachments with the given kind.
+  void eraseMetadata(unsigned KindID);
 
   void copyAttributesFrom(const GlobalValue *Src) override;
 
@@ -108,7 +122,6 @@ public:
            V->getValueID() == Value::GlobalVariableVal;
   }
 
-protected:
   void clearMetadata();
 
 private:
diff --git a/llvm/lib/AsmParser/LLParser.cpp b/llvm/lib/AsmParser/LLParser.cpp
index 09d012d5375..8081ef23297 100644
--- a/llvm/lib/AsmParser/LLParser.cpp
+++ b/llvm/lib/AsmParser/LLParser.cpp
@@ -1719,7 +1719,7 @@ bool LLParser::ParseGlobalObjectMetadataAttachment(GlobalObject &GO) {
   if (ParseMetadataAttachment(MDK, N))
     return true;
 
-  GO.setMetadata(MDK, N);
+  GO.addMetadata(MDK, *N);
   return false;
 }
 
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 55c0d064938..287be9b683a 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -4167,7 +4167,7 @@ std::error_code BitcodeReader::parseGlobalObjectAttachment(
     MDNode *MD = MetadataList.getMDNodeFwdRefOrNull(Record[I + 1]);
     if (!MD)
       return error("Invalid metadata attachment");
-    GO.setMetadata(K->second, MD);
+    GO.addMetadata(K->second, *MD);
   }
   return std::error_code();
 }
diff --git a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
index 6051173eed6..98eb7faf224 100644
--- a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
+++ b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
@@ -345,6 +345,7 @@ ValueEnumerator::ValueEnumerator(const Module &M,
 
   SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
   for (const GlobalVariable &GV : M.globals()) {
+    MDs.clear();
     GV.getAllMetadata(MDs);
     for (const auto &I : MDs)
       EnumerateMetadata(&GV, I.second);
@@ -356,6 +357,7 @@ ValueEnumerator::ValueEnumerator(const Module &M,
       EnumerateType(A.getType());
 
     // Enumerate metadata attached to this function.
+    MDs.clear();
     F.getAllMetadata(MDs);
     for (const auto &I : MDs)
       EnumerateMetadata(&F, I.second);
diff --git a/llvm/lib/IR/LLVMContextImpl.h b/llvm/lib/IR/LLVMContextImpl.h
index 02af08ac4f9..3e0ba4a59ed 100644
--- a/llvm/lib/IR/LLVMContextImpl.h
+++ b/llvm/lib/IR/LLVMContextImpl.h
@@ -997,6 +997,33 @@ public:
   }
 };
 
+/// Multimap-like storage for metadata attachments for globals. This differs
+/// from MDAttachmentMap in that it allows multiple attachments per metadata
+/// kind.
+class MDGlobalAttachmentMap {
+  struct Attachment {
+    unsigned MDKind;
+    TrackingMDNodeRef Node;
+  };
+  SmallVector<Attachment, 1> Attachments;
+
+public:
+  bool empty() const { return Attachments.empty(); }
+
+  /// Appends all attachments with the given ID to \c Result in insertion order.
+  /// If the global has no attachments with the given ID, or if ID is invalid,
+  /// leaves Result unchanged.
+  void get(unsigned ID, SmallVectorImpl<MDNode *> &Result);
+
+  void insert(unsigned ID, MDNode &MD);
+  void erase(unsigned ID);
+
+  /// Appends all attachments for the global to \c Result, sorting by attachment
+  /// ID. Attachments with the same ID appear in insertion order. This function
+  /// does \em not clear \c Result.
+  void getAll(SmallVectorImpl<std::pair<unsigned, MDNode *>> &Result) const;
+};
+
 class LLVMContextImpl {
 public:
   /// OwnedModules - The set of modules instantiated in this context, and which
@@ -1108,7 +1135,7 @@ public:
   DenseMap<const Instruction *, MDAttachmentMap> InstructionMetadata;
 
   /// Collection of per-GlobalObject metadata used in this context.
-  DenseMap<const GlobalObject *, MDAttachmentMap> GlobalObjectMetadata;
+  DenseMap<const GlobalObject *, MDGlobalAttachmentMap> GlobalObjectMetadata;
 
   /// DiscriminatorTable - This table maps file:line locations to an
   /// integer representing the next DWARF path discriminator to assign to
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index a77565cf087..865c68a2ef1 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -1125,6 +1125,43 @@ void MDAttachmentMap::getAll(
     array_pod_sort(Result.begin(), Result.end());
 }
 
+void MDGlobalAttachmentMap::insert(unsigned ID, MDNode &MD) {
+  Attachments.push_back({ID, TrackingMDNodeRef(&MD)});
+}
+
+void MDGlobalAttachmentMap::get(unsigned ID,
+                                SmallVectorImpl<MDNode *> &Result) {
+  for (auto A : Attachments)
+    if (A.MDKind == ID)
+      Result.push_back(A.Node);
+}
+
+void MDGlobalAttachmentMap::erase(unsigned ID) {
+  auto Follower = Attachments.begin();
+  for (auto Leader = Attachments.begin(), E = Attachments.end(); Leader != E;
+       ++Leader) {
+    if (Leader->MDKind != ID) {
+      if (Follower != Leader)
+        *Follower = std::move(*Leader);
+      ++Follower;
+    }
+  }
+  Attachments.resize(Follower - Attachments.begin());
+}
+
+void MDGlobalAttachmentMap::getAll(
+    SmallVectorImpl<std::pair<unsigned, MDNode *>> &Result) const {
+  for (auto &A : Attachments)
+    Result.emplace_back(A.MDKind, A.Node);
+
+  // Sort the resulting array so it is stable with respect to metadata IDs. We
+  // need to preserve the original insertion order though.
+  std::stable_sort(
+      Result.begin(), Result.end(),
+      [](const std::pair<unsigned, MDNode *> &A,
+         const std::pair<unsigned, MDNode *> &B) { return A.first < B.first; });
+}
+
 void Instruction::setMetadata(StringRef Kind, MDNode *Node) {
   if (!Node && !hasMetadata())
     return;
@@ -1281,27 +1318,30 @@ void Instruction::clearMetadataHashEntries() {
   setHasMetadataHashEntry(false);
 }
 
-MDNode *GlobalObject::getMetadata(unsigned KindID) const {
-  if (!hasMetadata())
-    return nullptr;
-  return getContext().pImpl->GlobalObjectMetadata[this].lookup(KindID);
+void GlobalObject::getMetadata(unsigned KindID,
+                               SmallVectorImpl<MDNode *> &MDs) const {
+  if (hasMetadata())
+    getContext().pImpl->GlobalObjectMetadata[this].get(KindID, MDs);
 }
 
-MDNode *GlobalObject::getMetadata(StringRef Kind) const {
-  if (!hasMetadata())
-    return nullptr;
-  return getMetadata(getContext().getMDKindID(Kind));
+void GlobalObject::getMetadata(StringRef Kind,
+                               SmallVectorImpl<MDNode *> &MDs) const {
+  if (hasMetadata())
+    getMetadata(getContext().getMDKindID(Kind), MDs);
 }
 
-void GlobalObject::setMetadata(unsigned KindID, MDNode *MD) {
-  if (MD) {
-    if (!hasMetadata())
-      setHasMetadataHashEntry(true);
+void GlobalObject::addMetadata(unsigned KindID, MDNode &MD) {
+  if (!hasMetadata())
+    setHasMetadataHashEntry(true);
 
-    getContext().pImpl->GlobalObjectMetadata[this].set(KindID, *MD);
-    return;
-  }
+  getContext().pImpl->GlobalObjectMetadata[this].insert(KindID, MD);
+}
 
+void GlobalObject::addMetadata(StringRef Kind, MDNode &MD) {
+  addMetadata(getContext().getMDKindID(Kind), MD);
+}
+
+void GlobalObject::eraseMetadata(unsigned KindID) {
   // Nothing to unset.
   if (!hasMetadata())
     return;
@@ -1312,12 +1352,6 @@ void GlobalObject::setMetadata(unsigned KindID, MDNode *MD) {
     clearMetadata();
 }
 
-void GlobalObject::setMetadata(StringRef Kind, MDNode *MD) {
-  if (!MD && !hasMetadata())
-    return;
-  setMetadata(getContext().getMDKindID(Kind), MD);
-}
-
 void GlobalObject::getAllMetadata(
     SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) const {
   MDs.clear();
@@ -1328,33 +1362,35 @@ void GlobalObject::getAllMetadata(
   getContext().pImpl->GlobalObjectMetadata[this].getAll(MDs);
 }
 
-void GlobalObject::dropUnknownMetadata(ArrayRef<unsigned> KnownIDs) {
+void GlobalObject::clearMetadata() {
   if (!hasMetadata())
     return;
-  if (KnownIDs.empty()) {
-    clearMetadata();
-    return;
-  }
+  getContext().pImpl->GlobalObjectMetadata.erase(this);
+  setHasMetadataHashEntry(false);
+}
 
-  SmallSet<unsigned, 5> KnownSet;
-  KnownSet.insert(KnownIDs.begin(), KnownIDs.end());
 
-  auto &Store = getContext().pImpl->GlobalObjectMetadata[this];
-  assert(!Store.empty());
+void GlobalObject::setMetadata(unsigned KindID, MDNode *N) {
+  eraseMetadata(KindID);
+  if (N)
+    addMetadata(KindID, *N);
+}
 
-  Store.remove_if([&KnownSet](const std::pair<unsigned, TrackingMDNodeRef> &I) {
-    return !KnownSet.count(I.first);
-  });
+void GlobalObject::setMetadata(StringRef Kind, MDNode *N) {
+  setMetadata(getContext().getMDKindID(Kind), N);
+}
 
-  if (Store.empty())
-    clearMetadata();
+MDNode *GlobalObject::getMetadata(unsigned KindID) const {
+  SmallVector<MDNode *, 1> MDs;
+  getMetadata(KindID, MDs);
+  assert(MDs.size() <= 1 && "Expected at most one metadata attachment");
+  if (MDs.empty())
+    return nullptr;
+  return MDs[0];
 }
 
-void GlobalObject::clearMetadata() {
-  if (!hasMetadata())
-    return;
-  getContext().pImpl->GlobalObjectMetadata.erase(this);
-  setHasMetadataHashEntry(false);
+MDNode *GlobalObject::getMetadata(StringRef Kind) const {
+  return getMetadata(getContext().getMDKindID(Kind));
 }
 
 void Function::setSubprogram(DISubprogram *SP) {
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index c47bb2518e0..fbf73509403 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -1991,6 +1991,7 @@ void Verifier::visitFunction(const Function &F) {
              "blockaddress may not be used with the entry block!", Entry);
     }
 
+    unsigned NumDebugAttachments = 0;
     // Visit metadata attachments.
     for (const auto &I : MDs) {
       // Verify that the attachment is legal.
@@ -1998,6 +1999,9 @@ void Verifier::visitFunction(const Function &F) {
       default:
         break;
       case LLVMContext::MD_dbg:
+        ++NumDebugAttachments;
+        AssertDI(NumDebugAttachments == 1,
+                 "function must have a single !dbg attachment", &F, I.second);
         AssertDI(isa<DISubprogram>(I.second),
                  "function !dbg attachment must be a subprogram", &F, I.second);
         break;
diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 24f7d69293e..b64fdbd7b1a 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -122,11 +122,11 @@ void llvm::CloneFunctionInto(Function *NewFunc, const Function *OldFunc,
   SmallVector<std::pair<unsigned, MDNode *>, 1> MDs;
   OldFunc->getAllMetadata(MDs);
   for (auto MD : MDs)
-    NewFunc->setMetadata(
+    NewFunc->addMetadata(
         MD.first,
-        MapMetadata(MD.second, VMap,
-                    ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges,
-                    TypeMapper, Materializer));
+        *MapMetadata(MD.second, VMap,
+                     ModuleLevelChanges ? RF_None : RF_NoModuleLevelChanges,
+                     TypeMapper, Materializer));
 
   // Loop over all of the basic blocks in the function, cloning them as
   // appropriate.  Note that we save BE this way in order to handle cloning of
diff --git a/llvm/lib/Transforms/Utils/ValueMapper.cpp b/llvm/lib/Transforms/Utils/ValueMapper.cpp
index d82f513419f..7b95ac04b5f 100644
--- a/llvm/lib/Transforms/Utils/ValueMapper.cpp
+++ b/llvm/lib/Transforms/Utils/ValueMapper.cpp
@@ -950,8 +950,9 @@ void Mapper::remapFunction(Function &F) {
   // Remap the metadata attachments.
   SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
   F.getAllMetadata(MDs);
+  F.clearMetadata();
   for (const auto &I : MDs)
-    F.setMetadata(I.first, cast_or_null<MDNode>(mapMetadata(I.second)));
+    F.addMetadata(I.first, *cast<MDNode>(mapMetadata(I.second)));
 
   // Remap the argument types.
   if (TypeMapper)
diff --git a/llvm/test/Assembler/metadata.ll b/llvm/test/Assembler/metadata.ll
index 2fe26be47de..a4b9c8af41d 100644
--- a/llvm/test/Assembler/metadata.ll
+++ b/llvm/test/Assembler/metadata.ll
@@ -1,8 +1,8 @@
 ; RUN: llvm-as < %s | llvm-dis | llvm-as | llvm-dis | FileCheck %s
 ; RUN: verify-uselistorder %s
 
-; CHECK: @global = global i32 0, !foo [[M2:![0-9]+]], !baz [[M3:![0-9]+]]
-@global = global i32 0, !foo !2, !baz !3
+; CHECK: @global = global i32 0, !foo [[M2:![0-9]+]], !foo [[M3:![0-9]+]], !baz [[M3]]
+@global = global i32 0, !foo !2, !foo !3, !baz !3
 
 ; CHECK-LABEL: @test
 ; CHECK: ret void, !foo [[M0:![0-9]+]], !bar [[M1:![0-9]+]]
diff --git a/llvm/test/Verifier/metadata-function-dbg.ll b/llvm/test/Verifier/metadata-function-dbg.ll
index b839e8708d1..77f7de26c87 100644
--- a/llvm/test/Verifier/metadata-function-dbg.ll
+++ b/llvm/test/Verifier/metadata-function-dbg.ll
@@ -1,6 +1,11 @@
 ; RUN: not llvm-as %s -disable-output 2>&1 | FileCheck %s
 
-define void @foo() !dbg !4 !dbg !4 {
+define void @foo() !dbg !4 {
+  unreachable
+}
+
+; CHECK:      function must have a single !dbg attachment
+define void @foo2() !dbg !4 !dbg !4 {
   unreachable
 }
 
diff --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index d14743dad70..a36e131b4bc 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -2242,32 +2242,6 @@ TEST_F(FunctionAttachmentTest, getAll) {
   EXPECT_EQ(T2, MDs[3].second);
 }
 
-TEST_F(FunctionAttachmentTest, dropUnknownMetadata) {
-  Function *F = getFunction("foo");
-
-  MDTuple *T1 = getTuple();
-  MDTuple *T2 = getTuple();
-  MDTuple *P = getTuple();
-  DISubprogram *SP = getSubprogram();
-
-  F->setMetadata("other1", T1);
-  F->setMetadata(LLVMContext::MD_dbg, SP);
-  F->setMetadata("other2", T2);
-  F->setMetadata(LLVMContext::MD_prof, P);
-
-  unsigned Known[] = {Context.getMDKindID("other2"), LLVMContext::MD_prof};
-  F->dropUnknownMetadata(Known);
-
-  EXPECT_EQ(T2, F->getMetadata("other2"));
-  EXPECT_EQ(P, F->getMetadata(LLVMContext::MD_prof));
-  EXPECT_EQ(nullptr, F->getMetadata("other1"));
-  EXPECT_EQ(nullptr, F->getMetadata(LLVMContext::MD_dbg));
-
-  F->setMetadata("other2", nullptr);
-  F->setMetadata(LLVMContext::MD_prof, nullptr);
-  EXPECT_FALSE(F->hasMetadata());
-}
-
 TEST_F(FunctionAttachmentTest, Verifier) {
   Function *F = getFunction("foo");
   F->setMetadata("attach", getTuple());