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

llvm/include/llvm/IR/Metadata.h

@@ -283,14 +283,20 @@ private:
   LLVMContext &Context;
   uint64_t NextIndex;
   SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
+  /// Flag that can be set to false if this metadata should not be
+  /// RAUW'ed, e.g. if it is used as the key of a map.
+  bool CanReplace;
 
 public:
   ReplaceableMetadataImpl(LLVMContext &Context)
-      : Context(Context), NextIndex(0) {}
+      : Context(Context), NextIndex(0), CanReplace(true) {}
   ~ReplaceableMetadataImpl() {
     assert(UseMap.empty() && "Cannot destroy in-use replaceable metadata");
   }
 
+  /// Set the CanReplace flag to the given value.
+  void setCanReplace(bool Replaceable) { CanReplace = Replaceable; }
+
   LLVMContext &getContext() const { return Context; }
 
   /// \brief Replace all uses of this with MD.
@@ -901,6 +907,11 @@ public:
     Context.getReplaceableUses()->replaceAllUsesWith(MD);
   }
 
+  /// Set the CanReplace flag to the given value.
+  void setCanReplace(bool Replaceable) {
+    Context.getReplaceableUses()->setCanReplace(Replaceable);
+  }
+
   /// \brief Resolve cycles.
   ///
   /// Once all forward declarations have been resolved, force cycles to be

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;
     }
   }

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;

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