BitcodeReader: Avoid referencing unresolved nodes from distinct ones

Each reference to an unresolved MDNode is expensive, since the RAUW
support in MDNode uses a separate allocation and side map.  Since
a distinct MDNode doesn't require its operands on creation (unlike
uniuqed nodes, there's no need to check for structural equivalence),
use nullptr for any of its unresolved operands.  Besides reducing the
burden on MDNode maps, this can avoid allocating temporary MDNodes in
the first place.

We need some way to track operands.  Invent DistinctMDOperandPlaceholder
for this purpose, which is a Metadata subclass that holds an ID and
points at its single user.  DistinctMDOperandPlaceholder::replaceUseWith
is just like RAUW, but its name highlights that there is only ever
exactly one use.

There is no support for moving (or, obviously, copying) these.  Move
support would be possible but expensive; leaving it unimplemented
prevents user error.  In the BitcodeReader I originally considered
allocating on a BumpPtrAllocator and keeping a vector of pointers to
them, and then I realized that std::deque implements exactly this.

A couple of obvious follow-ups:

  - Change ValueEnumerator to emit distinct nodes first to take more
    advantage of this optimization.  (How convenient... I think I might
    have a couple of patches for this.)

  - Change DIBuilder and its consumers (like CGDebugInfo in clang) to
    use something like this when constructing debug info in the first
    place.

llvm-svn: 267270

llvm/include/llvm/IR/Metadata.def

@@ -77,6 +77,7 @@ HANDLE_METADATA_LEAF(MDString)
 HANDLE_METADATA_BRANCH(ValueAsMetadata)
 HANDLE_METADATA_LEAF(ConstantAsMetadata)
 HANDLE_METADATA_LEAF(LocalAsMetadata)
+HANDLE_METADATA_LEAF(DistinctMDOperandPlaceholder)
 HANDLE_MDNODE_BRANCH(MDNode)
 HANDLE_MDNODE_LEAF_UNIQUABLE(MDTuple)
 HANDLE_SPECIALIZED_MDNODE_LEAF_UNIQUABLE(DILocation)

llvm/include/llvm/IR/Metadata.h

@@ -1193,6 +1193,52 @@ public:
   typedef MDTupleTypedArrayWrapper<CLASS> CLASS##Array;
 #include "llvm/IR/Metadata.def"
 
+/// Placeholder metadata for operands of distinct MDNodes.
+///
+/// This is a lightweight placeholder for an operand of a distinct node.  It's
+/// purpose is to help track forward references when creating a distinct node.
+/// This allows distinct nodes involved in a cycle to be constructed before
+/// their operands without requiring a heavyweight temporary node with
+/// full-blown RAUW support.
+///
+/// Each placeholder supports only a single MDNode user.  Clients should pass
+/// an ID, retrieved via \a getID(), to indicate the "real" operand that this
+/// should be replaced with.
+///
+/// While it would be possible to implement move operators, they would be
+/// fairly expensive.  Leave them unimplemented to discourage their use
+/// (clients can use std::deque, std::list, BumpPtrAllocator, etc.).
+class DistinctMDOperandPlaceholder : public Metadata {
+  friend class MetadataTracking;
+
+  Metadata **Use = nullptr;
+
+  DistinctMDOperandPlaceholder() = delete;
+  DistinctMDOperandPlaceholder(DistinctMDOperandPlaceholder &&) = delete;
+  DistinctMDOperandPlaceholder(const DistinctMDOperandPlaceholder &) = delete;
+
+public:
+  explicit DistinctMDOperandPlaceholder(unsigned ID)
+      : Metadata(DistinctMDOperandPlaceholderKind, Distinct) {
+    SubclassData32 = ID;
+  }
+
+  ~DistinctMDOperandPlaceholder() {
+    if (Use)
+      *Use = nullptr;
+  }
+
+  unsigned getID() const { return SubclassData32; }
+
+  /// Replace the use of this with MD.
+  void replaceUseWith(Metadata *MD) {
+    if (!Use)
+      return;
+    *Use = MD;
+    Use = nullptr;
+  }
+};
+
 //===----------------------------------------------------------------------===//
 /// \brief A tuple of MDNodes.
 ///

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -133,13 +133,26 @@ public:
     return MetadataPtrs[i];
   }
 
+  Metadata *lookup(unsigned I) const {
+    return I < MetadataPtrs.size() ? MetadataPtrs[I] : nullptr;
+  }
+
   void shrinkTo(unsigned N) {
     assert(N <= size() && "Invalid shrinkTo request!");
     assert(!AnyFwdRefs && "Unexpected forward refs");
     MetadataPtrs.resize(N);
   }
 
+  /// Return the given metadata, creating a replaceable forward reference if
+  /// necessary.
   Metadata *getMetadataFwdRef(unsigned Idx);
+
+  /// Return the the given metadata only if it is fully resolved.
+  ///
+  /// Gives the same result as \a lookup(), unless \a MDNode::isResolved()
+  /// would give \c false.
+  Metadata *getMetadataIfResolved(unsigned Idx);
+
   MDNode *getMDNodeFwdRefOrNull(unsigned Idx);
   void assignValue(Metadata *MD, unsigned Idx);
   void tryToResolveCycles();
@@ -1090,6 +1103,14 @@ Metadata *BitcodeReaderMetadataList::getMetadataFwdRef(unsigned Idx) {
   return MD;
 }
 
+Metadata *BitcodeReaderMetadataList::getMetadataIfResolved(unsigned Idx) {
+  Metadata *MD = lookup(Idx);
+  if (auto *N = dyn_cast_or_null<MDNode>(MD))
+    if (!N->isResolved())
+      return nullptr;
+  return MD;
+}
+
 MDNode *BitcodeReaderMetadataList::getMDNodeFwdRefOrNull(unsigned Idx) {
   return dyn_cast_or_null<MDNode>(getMetadataFwdRef(Idx));
 }
@@ -1933,6 +1954,31 @@ std::error_code BitcodeReader::parseMetadataStrings(ArrayRef<uint64_t> Record,
   return std::error_code();
 }
 
+namespace {
+class PlaceholderQueue {
+  // Placeholders would thrash around when moved, so store in a std::deque
+  // instead of some sort of vector.
+  std::deque<DistinctMDOperandPlaceholder> PHs;
+
+public:
+  DistinctMDOperandPlaceholder &getPlaceholderOp(unsigned ID);
+  void flush(BitcodeReaderMetadataList &MetadataList);
+};
+} // end namespace
+
+DistinctMDOperandPlaceholder &PlaceholderQueue::getPlaceholderOp(unsigned ID) {
+  PHs.emplace_back(ID);
+  return PHs.back();
+}
+
+void PlaceholderQueue::flush(BitcodeReaderMetadataList &MetadataList) {
+  while (!PHs.empty()) {
+    PHs.front().replaceUseWith(
+        MetadataList.getMetadataFwdRef(PHs.front().getID()));
+    PHs.pop_front();
+  }
+}
+
 /// Parse a METADATA_BLOCK. If ModuleLevel is true then we are parsing
 /// module level metadata.
 std::error_code BitcodeReader::parseMetadata(bool ModuleLevel) {
@@ -1948,12 +1994,19 @@ std::error_code BitcodeReader::parseMetadata(bool ModuleLevel) {
   std::vector<std::pair<DICompileUnit *, Metadata *>> CUSubprograms;
   SmallVector<uint64_t, 64> Record;
 
-  auto getMD = [&](unsigned ID) -> Metadata * {
-    return MetadataList.getMetadataFwdRef(ID);
+  PlaceholderQueue Placeholders;
+  bool IsDistinct;
+  auto getMD = [&](unsigned ID, bool AllowPlaceholders = true) -> Metadata * {
+    if (!IsDistinct || !AllowPlaceholders)
+      return MetadataList.getMetadataFwdRef(ID);
+    if (auto *MD = MetadataList.getMetadataIfResolved(ID))
+      return MD;
+    return &Placeholders.getPlaceholderOp(ID);
   };
-  auto getMDOrNull = [&](unsigned ID) -> Metadata *{
+  auto getMDOrNull = [&](unsigned ID,
+                         bool AllowPlaceholders = true) -> Metadata * {
     if (ID)
-      return getMD(ID - 1);
+      return getMD(ID - 1, AllowPlaceholders);
     return nullptr;
   };
   auto getMDString = [&](unsigned ID) -> MDString *{
@@ -1982,6 +2035,7 @@ std::error_code BitcodeReader::parseMetadata(bool ModuleLevel) {
               SP->replaceOperandWith(7, CU_SP.first);
 
       MetadataList.tryToResolveCycles();
+      Placeholders.flush(MetadataList);
       return std::error_code();
     case BitstreamEntry::Record:
       // The interesting case.
@@ -1992,7 +2046,7 @@ std::error_code BitcodeReader::parseMetadata(bool ModuleLevel) {
     Record.clear();
     StringRef Blob;
     unsigned Code = Stream.readRecord(Entry.ID, Record, &Blob);
-    bool IsDistinct = false;
+    IsDistinct = false;
     switch (Code) {
     default:  // Default behavior: ignore.
       break;
@@ -2275,7 +2329,8 @@ std::error_code BitcodeReader::parseMetadata(bool ModuleLevel) {
       MetadataList.assignValue(CU, NextMetadataNo++);
 
       // Move the Upgrade the list of subprograms.
-      if (Metadata *SPs = getMDOrNull(Record[11]))
+      if (Metadata *SPs =
+              getMDOrNull(Record[11], /* AllowPlaceholders = */ false))
         CUSubprograms.push_back({CU, SPs});
       break;
     }

llvm/lib/IR/Metadata.cpp

@@ -126,6 +126,12 @@ bool MetadataTracking::track(void *Ref, Metadata &MD, OwnerTy Owner) {
     R->addRef(Ref, Owner);
     return true;
   }
+  if (auto *PH = dyn_cast<DistinctMDOperandPlaceholder>(&MD)) {
+    assert(!PH->Use && "Placeholders can only be used once");
+    assert(!Owner && "Unexpected callback to owner");
+    PH->Use = static_cast<Metadata **>(Ref);
+    return true;
+  }
   return false;
 }
 
@@ -133,6 +139,8 @@ void MetadataTracking::untrack(void *Ref, Metadata &MD) {
   assert(Ref && "Expected live reference");
   if (auto *R = ReplaceableMetadataImpl::getIfExists(MD))
     R->dropRef(Ref);
+  else if (auto *PH = dyn_cast<DistinctMDOperandPlaceholder>(&MD))
+    PH->Use = nullptr;
 }
 
 bool MetadataTracking::retrack(void *Ref, Metadata &MD, void *New) {
@@ -143,6 +151,8 @@ bool MetadataTracking::retrack(void *Ref, Metadata &MD, void *New) {
     R->moveRef(Ref, New, MD);
     return true;
   }
+  assert(!isa<DistinctMDOperandPlaceholder>(MD) &&
+         "Unexpected move of an MDOperand");
   assert(!isReplaceable(MD) &&
          "Expected un-replaceable metadata, since we didn't move a reference");
   return false;

llvm/unittests/IR/MetadataTest.cpp

@@ -2309,4 +2309,79 @@ TEST_F(FunctionAttachmentTest, SubprogramAttachment) {
   EXPECT_EQ(SP, F->getMetadata(LLVMContext::MD_dbg));
 }
 
+typedef MetadataTest DistinctMDOperandPlaceholderTest;
+TEST_F(DistinctMDOperandPlaceholderTest, getID) {
+  EXPECT_EQ(7u, DistinctMDOperandPlaceholder(7).getID());
+}
+
+TEST_F(DistinctMDOperandPlaceholderTest, replaceUseWith) {
+  // Set up some placeholders.
+  DistinctMDOperandPlaceholder PH0(7);
+  DistinctMDOperandPlaceholder PH1(3);
+  DistinctMDOperandPlaceholder PH2(0);
+  Metadata *Ops[] = {&PH0, &PH1, &PH2};
+  auto *D = MDTuple::getDistinct(Context, Ops);
+  ASSERT_EQ(&PH0, D->getOperand(0));
+  ASSERT_EQ(&PH1, D->getOperand(1));
+  ASSERT_EQ(&PH2, D->getOperand(2));
+
+  // Replace them.
+  auto *N0 = MDTuple::get(Context, None);
+  auto *N1 = MDTuple::get(Context, N0);
+  PH0.replaceUseWith(N0);
+  PH1.replaceUseWith(N1);
+  PH2.replaceUseWith(nullptr);
+  EXPECT_EQ(N0, D->getOperand(0));
+  EXPECT_EQ(N1, D->getOperand(1));
+  EXPECT_EQ(nullptr, D->getOperand(2));
+}
+
+TEST_F(DistinctMDOperandPlaceholderTest, replaceUseWithNoUser) {
+  // There is no user, but we can still call replace.
+  DistinctMDOperandPlaceholder(7).replaceUseWith(MDTuple::get(Context, None));
+}
+
+#ifdef GTEST_HAS_DEATH_TEST
+TEST_F(DistinctMDOperandPlaceholderTest, MetadataAsValue) {
+  // This shouldn't crash.
+  DistinctMDOperandPlaceholder PH(7);
+  EXPECT_DEATH(MetadataAsValue::get(Context, &PH),
+               "Unexpected callback to owner");
+}
+
+TEST_F(DistinctMDOperandPlaceholderTest, UniquedMDNode) {
+  // This shouldn't crash.
+  DistinctMDOperandPlaceholder PH(7);
+  EXPECT_DEATH(MDTuple::get(Context, &PH), "Unexpected callback to owner");
+}
+
+TEST_F(DistinctMDOperandPlaceholderTest, SecondDistinctMDNode) {
+  // This shouldn't crash.
+  DistinctMDOperandPlaceholder PH(7);
+  MDTuple::getDistinct(Context, &PH);
+  EXPECT_DEATH(MDTuple::getDistinct(Context, &PH),
+               "Placeholders can only be used once");
+}
+
+TEST_F(DistinctMDOperandPlaceholderTest, TrackingMDRefAndDistinctMDNode) {
+  // TrackingMDRef doesn't install an owner callback, so it can't be detected
+  // as an invalid use.  However, using a placeholder in a TrackingMDRef *and*
+  // a distinct node isn't possible and we should assert.
+  //
+  // (There's no positive test for using TrackingMDRef because it's not a
+  // useful thing to do.)
+  {
+    DistinctMDOperandPlaceholder PH(7);
+    MDTuple::getDistinct(Context, &PH);
+    EXPECT_DEATH(TrackingMDRef Ref(&PH), "Placeholders can only be used once");
+  }
+  {
+    DistinctMDOperandPlaceholder PH(7);
+    TrackingMDRef Ref(&PH);
+    EXPECT_DEATH(MDTuple::getDistinct(Context, &PH),
+                 "Placeholders can only be used once");
+  }
+}
+#endif
+
 } // end namespace