1 files changed, 76 insertions, 40 deletions

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) {