Ensure MDNode used as key in metadata linking map cannot be RAUWed

As suggested in review for r255909, add a way to ensure that temporary
MD used as keys in the MetadataToID map during ThinLTO importing are not
RAUWed.

Add support for marking an MDNode as not replaceable. Clear the new
CanReplace flag when adding a temporary MD node to the MetadataToID map
and clear it when destroying the map.

llvm-svn: 256648

3 files changed, 24 insertions, 0 deletions

diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 824a3716b83..c7606fd488a 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -3085,6 +3085,11 @@ void BitcodeReader::saveMetadataList(
         assert(MetadataToIDs[MD] == ID && "Inconsistent metadata value id");
         continue;
       }
+      if (N && N->isTemporary())
+        // Ensure that we assert if someone tries to RAUW this temporary
+        // metadata while it is the key of a map. The flag will be set back
+        // to true when the saved metadata list is destroyed.
+        N->setCanReplace(false);
       MetadataToIDs[MD] = ID;
     }
   }
diff --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index ab1ba5e2035..f9543a65858 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -190,6 +190,8 @@ void ReplaceableMetadataImpl::moveRef(void *Ref, void *New,
 void ReplaceableMetadataImpl::replaceAllUsesWith(Metadata *MD) {
   assert(!(MD && isa<MDNode>(MD) && cast<MDNode>(MD)->isTemporary()) &&
          "Expected non-temp node");
+  assert(CanReplace &&
+         "Attempted to replace Metadata marked for no replacement");
 
   if (UseMap.empty())
     return;
diff --git a/llvm/lib/Linker/IRMover.cpp b/llvm/lib/Linker/IRMover.cpp
index fa6e37517fc..2e09d78a690 100644
--- a/llvm/lib/Linker/IRMover.cpp
+++ b/llvm/lib/Linker/IRMover.cpp
@@ -524,6 +524,23 @@ public:
       ValueMapperFlags = ValueMapperFlags | RF_HaveUnmaterializedMetadata;
   }
 
+  ~IRLinker() {
+    // In the case where we are not linking metadata, we unset the CanReplace
+    // flag on all temporary metadata in the MetadataToIDs map to ensure
+    // none was replaced while being a map key. Now that we are destructing
+    // the map, set the flag back to true, so that it is replaceable during
+    // metadata linking.
+    if (!shouldLinkMetadata()) {
+      for (auto MDI : MetadataToIDs) {
+        Metadata *MD = const_cast<Metadata *>(MDI.first);
+        MDNode *Node = dyn_cast<MDNode>(MD);
+        assert((Node && Node->isTemporary()) &&
+               "Found non-temp metadata in map when not linking metadata");
+        Node->setCanReplace(true);
+      }
+    }
+  }
+
   bool run();
   Value *materializeDeclFor(Value *V, bool ForAlias);
   void materializeInitFor(GlobalValue *New, GlobalValue *Old, bool ForAlias);