1 files changed, 31 insertions, 21 deletions

diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index 9016af7..a2007a2 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -437,10 +437,14 @@ class IndexBitcodeWriter : public BitcodeWriterBase {
   /// index and a value id generated by this class to use in references.
   std::map<GlobalValue::GUID, unsigned> GUIDToValueIdMap;
 
-  // The sorted stack id indices actually used in the summary entries being
-  // written, which will be a subset of those in the full index in the case of
-  // distributed indexes.
-  std::vector<unsigned> StackIdIndices;
+  // The stack ids used by this index, which will be a subset of those in
+  // the full index in the case of distributed indexes.
+  std::vector<uint64_t> StackIds;
+
+  // Keep a map of the stack id indices used by records being written for this
+  // index to the index of the corresponding stack id in the above StackIds
+  // vector. Ensures we write each referenced stack id once.
+  DenseMap<unsigned, unsigned> StackIdIndicesToIndex;
 
   /// Tracks the last value id recorded in the GUIDToValueMap.
   unsigned GlobalValueId = 0;
@@ -464,6 +468,19 @@ public:
       : BitcodeWriterBase(Stream, StrtabBuilder), Index(Index),
         DecSummaries(DecSummaries),
         ModuleToSummariesForIndex(ModuleToSummariesForIndex) {
+
+    // See if the StackIdIndex was already added to the StackId map and
+    // vector. If not, record it.
+    auto RecordStackIdReference = [&](unsigned StackIdIndex) {
+      // If the StackIdIndex is not yet in the map, the below insert ensures
+      // that it will point to the new StackIds vector entry we push to just
+      // below.
+      auto Inserted =
+          StackIdIndicesToIndex.insert({StackIdIndex, StackIds.size()});
+      if (Inserted.second)
+        StackIds.push_back(Index.getStackIdAtIndex(StackIdIndex));
+    };
+
     // Assign unique value ids to all summaries to be written, for use
     // in writing out the call graph edges. Save the mapping from GUID
     // to the new global value id to use when writing those edges, which
@@ -494,17 +511,13 @@ public:
           continue;
         }
         for (auto Idx : CI.StackIdIndices)
-          StackIdIndices.push_back(Idx);
+          RecordStackIdReference(Idx);
       }
       for (auto &AI : FS->allocs())
         for (auto &MIB : AI.MIBs)
           for (auto Idx : MIB.StackIdIndices)
-            StackIdIndices.push_back(Idx);
+            RecordStackIdReference(Idx);
     });
-    llvm::sort(StackIdIndices);
-    StackIdIndices.erase(
-        std::unique(StackIdIndices.begin(), StackIdIndices.end()),
-        StackIdIndices.end());
   }
 
   /// The below iterator returns the GUID and associated summary.
@@ -4492,18 +4505,15 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
                       ArrayRef<uint64_t>{GVI.second, GVI.first});
   }
 
-  if (!StackIdIndices.empty()) {
+  // Write the stack ids used by this index, which will be a subset of those in
+  // the full index in the case of distributed indexes.
+  if (!StackIds.empty()) {
     auto StackIdAbbv = std::make_shared<BitCodeAbbrev>();
     StackIdAbbv->Add(BitCodeAbbrevOp(bitc::FS_STACK_IDS));
     // numids x stackid
     StackIdAbbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::Array));
     StackIdAbbv->Add(BitCodeAbbrevOp(BitCodeAbbrevOp::VBR, 8));
     unsigned StackIdAbbvId = Stream.EmitAbbrev(std::move(StackIdAbbv));
-    // Write the stack ids used by this index, which will be a subset of those in
-    // the full index in the case of distributed indexes.
-    std::vector<uint64_t> StackIds;
-    for (auto &I : StackIdIndices)
-      StackIds.push_back(Index.getStackIdAtIndex(I));
     Stream.EmitRecord(bitc::FS_STACK_IDS, StackIds, StackIdAbbvId);
   }
 
@@ -4669,11 +4679,11 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
           return *ValueID;
         },
         /*GetStackIndex*/ [&](unsigned I) {
-          // Get the corresponding index into the list of StackIdIndices
-          // actually being written for this combined index (which may be a
-          // subset in the case of distributed indexes).
-          auto Lower = llvm::lower_bound(StackIdIndices, I);
-          return std::distance(StackIdIndices.begin(), Lower);
+          // Get the corresponding index into the list of StackIds actually
+          // being written for this combined index (which may be a subset in
+          // the case of distributed indexes).
+          assert(StackIdIndicesToIndex.contains(I));
+          return StackIdIndicesToIndex[I];
         });
 
     NameVals.push_back(*ValueId);