[llvm] [IR] Reduce memory usage in metadata use tracking (NFC) (PR #221866)

[Kazu Hirata via llvm-commits]

================
@@ -258,82 +258,104 @@ bool MetadataTracking::isReplaceable(const Metadata &MD) {
 }
 
 SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
-  SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
-  for (auto Pair : UseMap) {
-    OwnerTy Owner = Pair.second.first;
+  SmallVector<Metadata *> MDUsers;
+  for (const auto &[Ref, Owner] : UseMap) {
     if (Owner.isNull())
       continue;
     if (!isa<Metadata *>(Owner))
       continue;
     Metadata *OwnerMD = cast<Metadata *>(Owner);
     if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
-      MDUsersWithID.push_back(&UseMap[Pair.first]);
+      MDUsers.push_back(OwnerMD);
   }
-  llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
-    return UserA->second < UserB->second;
-  });
-  SmallVector<Metadata *> MDUsers;
-  for (auto *UserWithID : MDUsersWithID)
-    MDUsers.push_back(cast<Metadata *>(UserWithID->first));
   return MDUsers;
 }
 
 SmallVector<DbgVariableRecord *>
 ReplaceableUses::getAllDbgVariableRecordUsers() {
-  SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
-  for (auto Pair : UseMap) {
-    OwnerTy Owner = Pair.second.first;
-    if (Owner.isNull())
-      continue;
-    if (!isa<DebugValueUser *>(Owner))
-      continue;
-    DVRUsersWithID.push_back(&UseMap[Pair.first]);
-  }
   // Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
   // users of MetadataAsValues are ordered by their UseList, i.e. reverse order
   // of when they were added: we need to replicate that here. The structure of
   // debug-info output depends on the ordering of intrinsics, thus we need
   // to keep them consistent for comparisons sake.
-  llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
-    return UserA->second > UserB->second;
-  });
   SmallVector<DbgVariableRecord *> DVRUsers;
-  for (auto UserWithID : DVRUsersWithID)
-    DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
+  for (const auto &[Ref, Owner] : reverse(UseMap)) {
+    if (Owner.isNull())
+      continue;
+    if (!isa<DebugValueUser *>(Owner))
+      continue;
+    DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
+  }
   return DVRUsers;
 }
 
-void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
-  bool WasInserted =
-      UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
-          .second;
-  (void)WasInserted;
-  assert(WasInserted && "Expected to add a reference");
+ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref,
+                                                    bool EraseFromIndex) {
+  if (!IndexMap) {
+    // Search backward for temporal locality; recently added references are
+    // often dropped or moved first.
+    auto It = find_if(reverse(UseMap),
+                      [Ref](const UseEntry &U) { return U.Ref == Ref; });
+    return It != UseMap.rend() ? &*It : nullptr;
+  }
 
-  ++NextIndex;
-  assert(NextIndex != 0 && "Unexpected overflow");
+  auto It = IndexMap->find(Ref);
+  if (It == IndexMap->end())
+    return nullptr;
+  UseEntry *Entry = &UseMap[It->second];
+  if (EraseFromIndex)
+    IndexMap->erase(It);
+  return Entry;
+}
+
+bool ReplaceableUses::hasRef(void *Ref) const {
+  return const_cast<ReplaceableUses *>(this)->findRef(Ref) != nullptr;
+}
+
+void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
+  assert(Ref && "Expected live reference");
+  assert(!hasRef(Ref) && "Reference already tracked");
+  unsigned NewIdx = UseMap.size();
+  UseMap.push_back({Ref, Owner});
+  if (IndexMap) {
+    (*IndexMap)[Ref] = NewIdx;
+  } else if (UseMap.size() > IndexThreshold) {
+    // Build the index map once UseMap grows past the threshold.
+    IndexMap = std::make_unique<IndexMapTy>();
+    for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
+      (*IndexMap)[UseMap[I].Ref] = I;
+  }
 }
 
 void ReplaceableUses::dropRef(void *Ref) {
-  bool WasErased = UseMap.erase(Ref);
-  (void)WasErased;
-  assert(WasErased && "Expected to drop a reference");
+  UseEntry *Entry = findRef(Ref, /*EraseFromIndex=*/true);
+  assert(Entry && "Expected to find Ref");
+  unsigned Idx = Entry - UseMap.begin();
+  UseMap.erase(Entry);
+  if (IndexMap) {
+    if (UseMap.size() <= IndexThreshold / 2) {
+      // Tear down the index map with hysteresis to avoid thrashing.
+      IndexMap.reset();
----------------
kazutakahirata wrote:

Addressed in the latest revision.

https://github.com/llvm/llvm-project/pull/221866