diff --git a/llvm/include/llvm/IR/Function.h b/llvm/include/llvm/IR/Function.h index af62476727a2..ae1061e9f98d 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> &MDs) const; + + /// \brief Drop metadata not in the given list. + /// + /// Drop all metadata from \c this not included in \c KnownIDs. + void dropUnknownMetadata(ArrayRef 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 227dfef9e983..ced989ab11df 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 a13afdd9b35b..edffef315115 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 InstructionMetadata; + /// Collection of per-function metadata used in this context. + DenseMap 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 eb1bd310511c..23a17a52b2b8 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> &MDs) const { + MDs.clear(); + + if (!hasMetadata()) + return; + + getContext().pImpl->FunctionMetadata[this].getAll(MDs); +} + +void Function::dropUnknownMetadata(ArrayRef KnownIDs) { + if (!hasMetadata()) + return; + if (KnownIDs.empty()) { + clearMetadata(); + return; + } + + SmallSet KnownSet; + KnownSet.insert(KnownIDs.begin(), KnownIDs.end()); + + auto &Store = getContext().pImpl->FunctionMetadata[this]; + assert(!Store.empty()); + + Store.remove_if([&KnownSet](const std::pair &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 8b4c5db73e3f..336822111b81 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, 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()); +} + }