Revert "[MemProf] Streamline and avoid unnecessary context id duplication (#107918)" (#108652)

This reverts commit 524a028, but
manually so that follow on PR108086 /
ae5f1a7
is retained (NFC patch to convert tuple to a struct).

Author: teresajohnson

llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp

@@ -1479,23 +1479,13 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::updateStackNodes() {
     // of length, and within each length, lexicographically by stack id. The
     // latter is so that we can specially handle calls that have identical stack
     // id sequences (either due to cloning or artificially because of the MIB
-    // context pruning). Those with the same Ids are then sorted by function to
-    // facilitate efficiently mapping them to the same context node.
-    // Because the functions are pointers, to ensure a stable sort first assign
-    // each function pointer to its first index in the Calls array, and then use
-    // that to sort by.
-    DenseMap<const FuncTy *, unsigned> FuncToIndex;
-    for (const auto &[Idx, CallCtxInfo] : enumerate(Calls))
-      FuncToIndex.insert({CallCtxInfo.Func, Idx});
-    std::stable_sort(
-        Calls.begin(), Calls.end(),
-        [&FuncToIndex](const CallContextInfo &A, const CallContextInfo &B) {
-          return A.StackIds.size() > B.StackIds.size() ||
-                 (A.StackIds.size() == B.StackIds.size() &&
-                  (A.StackIds < B.StackIds ||
-                   (A.StackIds == B.StackIds &&
-                    FuncToIndex[A.Func] < FuncToIndex[B.Func])));
-        });
+    // context pruning).
+    std::stable_sort(Calls.begin(), Calls.end(),
+                     [](const CallContextInfo &A, const CallContextInfo &B) {
+                       return A.StackIds.size() > B.StackIds.size() ||
+                              (A.StackIds.size() == B.StackIds.size() &&
+                               A.StackIds < B.StackIds);
+                     });
 
     // Find the node for the last stack id, which should be the same
     // across all calls recorded for this id, and is the id for this
@@ -1513,35 +1503,18 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::updateStackNodes() {
     DenseSet<uint32_t> LastNodeContextIds = LastNode->getContextIds();
     assert(!LastNodeContextIds.empty());
 
-#ifndef NDEBUG
-    // Save the set of functions seen for a particular set of the same stack
-    // ids. This is used to ensure that they have been correctly sorted to be
-    // adjacent in the Calls list, since we rely on that to efficiently place
-    // all such matching calls onto the same context node.
-    DenseSet<const FuncTy *> MatchingIdsFuncSet;
-#endif
+    // Map from function to the first call from the below list (with matching
+    // stack ids) found in that function. Note that calls from different
+    // functions can have the same stack ids because this is the list of stack
+    // ids that had (possibly pruned) nodes after building the graph from the
+    // allocation MIBs.
+    DenseMap<const FuncTy *, CallInfo> FuncToCallMap;
 
     for (unsigned I = 0; I < Calls.size(); I++) {
       auto &[Call, Ids, Func, SavedContextIds] = Calls[I];
       assert(SavedContextIds.empty());
       assert(LastId == Ids.back());
 
-#ifndef NDEBUG
-      // If this call has a different set of ids than the last one, clear the
-      // set used to ensure they are sorted properly.
-      if (I > 0 && Ids != Calls[I - 1].StackIds)
-        MatchingIdsFuncSet.clear();
-      else
-        // If the prior call had the same stack ids this set would not be empty.
-        // Check if we already have a call that "matches" because it is located
-        // in the same function. If the Calls list was sorted properly we should
-        // not encounter this situation as all such entries should be adjacent
-        // and processed in bulk further below.
-        assert(!MatchingIdsFuncSet.contains(Func));
-
-      MatchingIdsFuncSet.insert(Func);
-#endif
-
       // First compute the context ids for this stack id sequence (the
       // intersection of the context ids of the corresponding nodes).
       // Start with the remaining saved ids for the last node.
@@ -1610,27 +1583,23 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::updateStackNodes() {
           continue;
       }
 
+      // If the prior call had the same stack ids this map would not be empty.
+      // Check if we already have a call that "matches" because it is located
+      // in the same function.
+      if (FuncToCallMap.contains(Func)) {
+        // Record the matching call found for this call, and skip it. We
+        // will subsequently combine it into the same node.
+        CallToMatchingCall[Call] = FuncToCallMap[Func];
+        continue;
+      }
+
       // Check if the next set of stack ids is the same (since the Calls vector
       // of tuples is sorted by the stack ids we can just look at the next one).
-      // If so, save them in the CallToMatchingCall map so that they get
-      // assigned to the same context node, and skip them.
       bool DuplicateContextIds = false;
-      for (unsigned J = I + 1; J < Calls.size(); J++) {
-        auto &CallCtxInfo = Calls[J];
+      if (I + 1 < Calls.size()) {
+        auto &CallCtxInfo = Calls[I + 1];
         auto &NextIds = CallCtxInfo.StackIds;
-        if (NextIds != Ids)
-          break;
-        auto *NextFunc = CallCtxInfo.Func;
-        if (NextFunc != Func) {
-          // We have another Call with the same ids but that cannot share this
-          // node, must duplicate ids for it.
-          DuplicateContextIds = true;
-          break;
-        }
-        auto &NextCall = CallCtxInfo.Call;
-        CallToMatchingCall[NextCall] = Call;
-        // Update I so that it gets incremented correctly to skip this call.
-        I = J;
+        DuplicateContextIds = Ids == NextIds;
       }
 
       // If we don't have duplicate context ids, then we can assign all the
@@ -1654,7 +1623,14 @@ void CallsiteContextGraph<DerivedCCG, FuncTy, CallTy>::updateStackNodes() {
         set_subtract(LastNodeContextIds, StackSequenceContextIds);
         if (LastNodeContextIds.empty())
           break;
-      }
+        // No longer possibly in a sequence of calls with duplicate stack ids,
+        // clear the map.
+        FuncToCallMap.clear();
+      } else
+        // Record the call with its function, so we can locate it the next time
+        // we find a call from this function when processing the calls with the
+        // same stack ids.
+        FuncToCallMap[Func] = Call;
     }
   }