From ba0fc204633044e6bdaa1416288d9548d74f9d9b Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson@google.com>
Date: Wed, 30 Dec 2015 19:32:24 +0000
Subject: [PATCH] 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.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@256648 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/llvm/IR/Metadata.h           | 13 ++++++++++++-
 lib/Bitcode/Reader/BitcodeReader.cpp |  5 +++++
 lib/IR/Metadata.cpp                  |  2 ++
 lib/Linker/IRMover.cpp               | 17 +++++++++++++++++
 4 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/include/llvm/IR/Metadata.h b/include/llvm/IR/Metadata.h
index 2ea591383f8..0715229a875 100644
--- a/include/llvm/IR/Metadata.h
+++ b/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
diff --git a/lib/Bitcode/Reader/BitcodeReader.cpp b/lib/Bitcode/Reader/BitcodeReader.cpp
index 824a3716b83..c7606fd488a 100644
--- a/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/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/lib/IR/Metadata.cpp b/lib/IR/Metadata.cpp
index ab1ba5e2035..f9543a65858 100644
--- a/lib/IR/Metadata.cpp
+++ b/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/lib/Linker/IRMover.cpp b/lib/Linker/IRMover.cpp
index fa6e37517fc..2e09d78a690 100644
--- a/lib/Linker/IRMover.cpp
+++ b/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);
-- 
2.34.1