1 files changed, 7 insertions, 8 deletions
diff --git a/llvm/lib/Transforms/Utils/ValueMapper.cpp b/llvm/lib/Transforms/Utils/ValueMapper.cpp
index 9882f81..5e50536 100644
--- a/llvm/lib/Transforms/Utils/ValueMapper.cpp
+++ b/llvm/lib/Transforms/Utils/ValueMapper.cpp
@@ -120,12 +120,12 @@ class Mapper {
   SmallVector<WorklistEntry, 4> Worklist;
   SmallVector<DelayedBasicBlock, 1> DelayedBBs;
   SmallVector<Constant *, 16> AppendingInits;
-  const MetadataSetTy *IdentityMD;
+  const MetadataPredicate *IdentityMD;
 
 public:
   Mapper(ValueToValueMapTy &VM, RemapFlags Flags,
          ValueMapTypeRemapper *TypeMapper, ValueMaterializer *Materializer,
-         const MetadataSetTy *IdentityMD)
+         const MetadataPredicate *IdentityMD)
       : Flags(Flags), TypeMapper(TypeMapper),
         MCs(1, MappingContext(VM, Materializer)), IdentityMD(IdentityMD) {}
 
@@ -904,11 +904,10 @@ std::optional<Metadata *> Mapper::mapSimpleMetadata(const Metadata *MD) {
     return wrapConstantAsMetadata(*CMD, mapValue(CMD->getValue()));
   }
 
-  // Map metadata from IdentityMD on first use. We need to add these nodes to
-  // the mapping as otherwise metadata nodes numbering gets messed up. This is
-  // still economical because the amount of data in IdentityMD may be a lot
-  // larger than what will actually get used.
-  if (IdentityMD && IdentityMD->contains(MD))
+  // Map metadata matching IdentityMD predicate on first use. We need to add
+  // these nodes to the mapping as otherwise metadata nodes numbering gets
+  // messed up.
+  if (IdentityMD && (*IdentityMD)(MD))
     return getVM().MD()[MD] = TrackingMDRef(const_cast<Metadata *>(MD));
 
   assert(isa<MDNode>(MD) && "Expected a metadata node");
@@ -1211,7 +1210,7 @@ public:
 ValueMapper::ValueMapper(ValueToValueMapTy &VM, RemapFlags Flags,
                          ValueMapTypeRemapper *TypeMapper,
                          ValueMaterializer *Materializer,
-                         const MetadataSetTy *IdentityMD)
+                         const MetadataPredicate *IdentityMD)
     : pImpl(new Mapper(VM, Flags, TypeMapper, Materializer, IdentityMD)) {}
 
 ValueMapper::~ValueMapper() { delete getAsMapper(pImpl); }