IR: Add Function metadata attachments

Add IR support for `Metadata` attachments.  Assembly and bitcode support
will follow shortly, but for now we just have unit tests.  This is part
of PR23340.

llvm-svn: 235783

5 files changed, 230 insertions, 1 deletions

diff --git a/llvm/include/llvm/IR/Function.h b/llvm/include/llvm/IR/Function.h
index af62476727a..ae1061e9f98 100644
--- a/llvm/include/llvm/IR/Function.h
+++ b/llvm/include/llvm/IR/Function.h
@@ -80,7 +80,8 @@ private:
   /// Bits from GlobalObject::GlobalObjectSubclassData.
   enum {
     /// Whether this function is materializable.
-    IsMaterializableBit = 1 << 0
+    IsMaterializableBit = 1 << 0,
+    HasMetadataHashEntryBit = 1 << 1
   };
   void setGlobalObjectBit(unsigned Mask, bool Value) {
     setGlobalObjectSubClassData((~Mask & getGlobalObjectSubClassData()) |
@@ -521,12 +522,50 @@ public:
   /// setjmp or other function that gcc recognizes as "returning twice".
   bool callsFunctionThatReturnsTwice() const;
 
+  /// \brief Check if this has any metadata.
+  bool hasMetadata() const { return hasMetadataHashEntry(); }
+
+  /// \brief Get the current metadata attachment, if any.
+  ///
+  /// Returns \c nullptr if such an attachment is missing.
+  /// @{
+  MDNode *getMetadata(unsigned KindID) const;
+  MDNode *getMetadata(StringRef Kind) const;
+  /// @}
+
+  /// \brief Set a particular kind of metadata attachment.
+  ///
+  /// Sets the given attachment to \c MD, erasing it if \c MD is \c nullptr or
+  /// replacing it if it already exists.
+  /// @{
+  void setMetadata(unsigned KindID, MDNode *MD);
+  void setMetadata(StringRef Kind, MDNode *MD);
+  /// @}
+
+  /// \brief Get all current metadata attachments.
+  void
+  getAllMetadata(SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) const;
+
+  /// \brief Drop metadata not in the given list.
+  ///
+  /// Drop all metadata from \c this not included in \c KnownIDs.
+  void dropUnknownMetadata(ArrayRef<unsigned> KnownIDs);
+
 private:
   // Shadow Value::setValueSubclassData with a private forwarding method so that
   // subclasses cannot accidentally use it.
   void setValueSubclassData(unsigned short D) {
     Value::setValueSubclassData(D);
   }
+
+  bool hasMetadataHashEntry() const {
+    return getGlobalObjectSubClassData() & HasMetadataHashEntryBit;
+  }
+  void setHasMetadataHashEntry(bool HasEntry) {
+    setGlobalObjectBit(HasMetadataHashEntryBit, HasEntry);
+  }
+
+  void clearMetadata();
 };
 
 inline ValueSymbolTable *
diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
index 227dfef9e98..ced989ab11d 100644
--- a/llvm/lib/IR/Function.cpp
+++ b/llvm/lib/IR/Function.cpp
@@ -324,6 +324,9 @@ void Function::dropAllReferences() {
   // Prefix and prologue data are stored in a side table.
   setPrefixData(nullptr);
   setPrologueData(nullptr);
+
+  // Metadata is stored in a side-table.
+  clearMetadata();
 }
 
 void Function::addAttribute(unsigned i, Attribute::AttrKind attr) {
diff --git a/llvm/lib/IR/LLVMContextImpl.h b/llvm/lib/IR/LLVMContextImpl.h
index a13afdd9b35..edffef31511 100644
--- a/llvm/lib/IR/LLVMContextImpl.h
+++ b/llvm/lib/IR/LLVMContextImpl.h
@@ -992,6 +992,9 @@ public:
   /// Collection of per-instruction metadata used in this context.
   DenseMap<const Instruction *, MDAttachmentMap> InstructionMetadata;
 
+  /// Collection of per-function metadata used in this context.
+  DenseMap<const Function *, MDAttachmentMap> FunctionMetadata;
+
   /// DiscriminatorTable - This table maps file:line locations to an
   /// integer representing the next DWARF path discriminator to assign to
   /// instructions in different blocks at the same location.
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index eb1bd310511..23a17a52b2b 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -1160,3 +1160,79 @@ void Instruction::clearMetadataHashEntries() {
   getContext().pImpl->InstructionMetadata.erase(this);
   setHasMetadataHashEntry(false);
 }
+
+MDNode *Function::getMetadata(unsigned KindID) const {
+  if (!hasMetadata())
+    return nullptr;
+  return getContext().pImpl->FunctionMetadata[this].lookup(KindID);
+}
+
+MDNode *Function::getMetadata(StringRef Kind) const {
+  if (!hasMetadata())
+    return nullptr;
+  return getMetadata(getContext().getMDKindID(Kind));
+}
+
+void Function::setMetadata(unsigned KindID, MDNode *MD) {
+  if (MD) {
+    if (!hasMetadata())
+      setHasMetadataHashEntry(true);
+
+    getContext().pImpl->FunctionMetadata[this].set(KindID, *MD);
+    return;
+  }
+
+  // Nothing to unset.
+  if (!hasMetadata())
+    return;
+
+  auto &Store = getContext().pImpl->FunctionMetadata[this];
+  Store.erase(KindID);
+  if (Store.empty())
+    clearMetadata();
+}
+
+void Function::setMetadata(StringRef Kind, MDNode *MD) {
+  if (!MD && !hasMetadata())
+    return;
+  setMetadata(getContext().getMDKindID(Kind), MD);
+}
+
+void Function::getAllMetadata(
+    SmallVectorImpl<std::pair<unsigned, MDNode *>> &MDs) const {
+  MDs.clear();
+
+  if (!hasMetadata())
+    return;
+
+  getContext().pImpl->FunctionMetadata[this].getAll(MDs);
+}
+
+void Function::dropUnknownMetadata(ArrayRef<unsigned> KnownIDs) {
+  if (!hasMetadata())
+    return;
+  if (KnownIDs.empty()) {
+    clearMetadata();
+    return;
+  }
+
+  SmallSet<unsigned, 5> KnownSet;
+  KnownSet.insert(KnownIDs.begin(), KnownIDs.end());
+
+  auto &Store = getContext().pImpl->FunctionMetadata[this];
+  assert(!Store.empty());
+
+  Store.remove_if([&KnownSet](const std::pair<unsigned, TrackingMDNodeRef> &I) {
+    return !KnownSet.count(I.first);
+  });
+
+  if (Store.empty())
+    clearMetadata();
+}
+
+void Function::clearMetadata() {
+  if (!hasMetadata())
+    return;
+  getContext().pImpl->FunctionMetadata.erase(this);
+  setHasMetadataHashEntry(false);
+}
diff --git a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
index 8b4c5db73e3..336822111b8 100644
--- a/llvm/unittests/IR/MetadataTest.cpp
+++ b/llvm/unittests/IR/MetadataTest.cpp
@@ -2125,4 +2125,112 @@ TEST(NamedMDNodeTest, Search) {
   EXPECT_STREQ("!llvm.NMD1 = !{!0, !1}\n",
                oss.str().c_str());
 }
+
+typedef MetadataTest FunctionAttachmentTest;
+TEST_F(FunctionAttachmentTest, setMetadata) {
+  Function *F = getFunction("foo");
+  ASSERT_FALSE(F->hasMetadata());
+  EXPECT_EQ(nullptr, F->getMetadata(LLVMContext::MD_dbg));
+  EXPECT_EQ(nullptr, F->getMetadata("dbg"));
+  EXPECT_EQ(nullptr, F->getMetadata("other"));
+
+  MDSubprogram *SP1 = getSubprogram();
+  MDSubprogram *SP2 = getSubprogram();
+  ASSERT_NE(SP1, SP2);
+
+  F->setMetadata("dbg", SP1);
+  EXPECT_TRUE(F->hasMetadata());
+  EXPECT_EQ(SP1, F->getMetadata(LLVMContext::MD_dbg));
+  EXPECT_EQ(SP1, F->getMetadata("dbg"));
+  EXPECT_EQ(nullptr, F->getMetadata("other"));
+
+  F->setMetadata(LLVMContext::MD_dbg, SP2);
+  EXPECT_TRUE(F->hasMetadata());
+  EXPECT_EQ(SP2, F->getMetadata(LLVMContext::MD_dbg));
+  EXPECT_EQ(SP2, F->getMetadata("dbg"));
+  EXPECT_EQ(nullptr, F->getMetadata("other"));
+
+  F->setMetadata("dbg", nullptr);
+  EXPECT_FALSE(F->hasMetadata());
+  EXPECT_EQ(nullptr, F->getMetadata(LLVMContext::MD_dbg));
+  EXPECT_EQ(nullptr, F->getMetadata("dbg"));
+  EXPECT_EQ(nullptr, F->getMetadata("other"));
+
+  MDTuple *T1 = getTuple();
+  MDTuple *T2 = getTuple();
+  ASSERT_NE(T1, T2);
+
+  F->setMetadata("other1", T1);
+  F->setMetadata("other2", T2);
+  EXPECT_TRUE(F->hasMetadata());
+  EXPECT_EQ(T1, F->getMetadata("other1"));
+  EXPECT_EQ(T2, F->getMetadata("other2"));
+  EXPECT_EQ(nullptr, F->getMetadata("dbg"));
+
+  F->setMetadata("other1", T2);
+  F->setMetadata("other2", T1);
+  EXPECT_EQ(T2, F->getMetadata("other1"));
+  EXPECT_EQ(T1, F->getMetadata("other2"));
+
+  F->setMetadata("other1", nullptr);
+  F->setMetadata("other2", nullptr);
+  EXPECT_FALSE(F->hasMetadata());
+  EXPECT_EQ(nullptr, F->getMetadata("other1"));
+  EXPECT_EQ(nullptr, F->getMetadata("other2"));
+}
+
+TEST_F(FunctionAttachmentTest, getAll) {
+  Function *F = getFunction("foo");
+
+  MDTuple *T1 = getTuple();
+  MDTuple *T2 = getTuple();
+  MDTuple *P = getTuple();
+  MDSubprogram *SP = getSubprogram();
+
+  F->setMetadata("other1", T2);
+  F->setMetadata(LLVMContext::MD_dbg, SP);
+  F->setMetadata("other2", T1);
+  F->setMetadata(LLVMContext::MD_prof, P);
+  F->setMetadata("other2", T2);
+  F->setMetadata("other1", T1);
+
+  SmallVector<std::pair<unsigned, MDNode *>, 4> MDs;
+  F->getAllMetadata(MDs);
+  ASSERT_EQ(4u, MDs.size());
+  EXPECT_EQ(LLVMContext::MD_dbg, MDs[0].first);
+  EXPECT_EQ(LLVMContext::MD_prof, MDs[1].first);
+  EXPECT_EQ(Context.getMDKindID("other1"), MDs[2].first);
+  EXPECT_EQ(Context.getMDKindID("other2"), MDs[3].first);
+  EXPECT_EQ(SP, MDs[0].second);
+  EXPECT_EQ(P, MDs[1].second);
+  EXPECT_EQ(T1, MDs[2].second);
+  EXPECT_EQ(T2, MDs[3].second);
+}
+
+TEST_F(FunctionAttachmentTest, dropUnknownMetadata) {
+  Function *F = getFunction("foo");
+
+  MDTuple *T1 = getTuple();
+  MDTuple *T2 = getTuple();
+  MDTuple *P = getTuple();
+  MDSubprogram *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());
+}
+
 }