1 files changed, 38 insertions, 24 deletions

diff --git a/llvm/lib/ProfileData/CoverageMapping.cpp b/llvm/lib/ProfileData/CoverageMapping.cpp
index 115256358ed..3a568720b08 100644
--- a/llvm/lib/ProfileData/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/CoverageMapping.cpp
@@ -27,10 +27,6 @@ using namespace coverage;
 
 #define DEBUG_TYPE "coverage-mapping"
 
-CounterExpressionBuilder::CounterExpressionBuilder(unsigned NumCounterValues) {
-  Terms.resize(NumCounterValues);
-}
-
 Counter CounterExpressionBuilder::get(const CounterExpression &E) {
   for (unsigned I = 0, S = Expressions.size(); I < S; ++I) {
     if (Expressions[I] == E)
@@ -40,50 +36,68 @@ Counter CounterExpressionBuilder::get(const CounterExpression &E) {
   return Counter::getExpression(Expressions.size() - 1);
 }
 
-void CounterExpressionBuilder::extractTerms(Counter C, int Sign) {
+void CounterExpressionBuilder::extractTerms(
+    Counter C, int Sign, SmallVectorImpl<std::pair<unsigned, int>> &Terms) {
   switch (C.getKind()) {
   case Counter::Zero:
     break;
   case Counter::CounterValueReference:
-    Terms[C.getCounterID()] += Sign;
+    Terms.push_back(std::make_pair(C.getCounterID(), Sign));
     break;
   case Counter::Expression:
     const auto &E = Expressions[C.getExpressionID()];
-    extractTerms(E.LHS, Sign);
-    extractTerms(E.RHS, E.Kind == CounterExpression::Subtract ? -Sign : Sign);
+    extractTerms(E.LHS, Sign, Terms);
+    extractTerms(E.RHS, E.Kind == CounterExpression::Subtract ? -Sign : Sign,
+                 Terms);
     break;
   }
 }
 
 Counter CounterExpressionBuilder::simplify(Counter ExpressionTree) {
   // Gather constant terms.
-  for (auto &I : Terms)
-    I = 0;
-  extractTerms(ExpressionTree);
+  llvm::SmallVector<std::pair<unsigned, int>, 32> Terms;
+  extractTerms(ExpressionTree, +1, Terms);
+
+  // Group the terms by counter ID.
+  std::sort(Terms.begin(), Terms.end(),
+            [](const std::pair<unsigned, int> &LHS,
+               const std::pair<unsigned, int> &RHS) {
+    return LHS.first < RHS.first;
+  });
+
+  // Combine terms by counter ID to eliminate counters that sum to zero.
+  auto Prev = Terms.begin();
+  for (auto I = Prev + 1, E = Terms.end(); I != E; ++I) {
+    if (I->first == Prev->first) {
+      Prev->second += I->second;
+      continue;
+    }
+    ++Prev;
+    *Prev = *I;
+  }
+  Terms.erase(++Prev, Terms.end());
 
   Counter C;
-  // Create additions.
-  // Note: the additions are created first
-  // to avoid creation of a tree like ((0 - X) + Y) instead of (Y - X).
-  for (unsigned I = 0, S = Terms.size(); I < S; ++I) {
-    if (Terms[I] <= 0)
+  // Create additions. We do this before subtractions to avoid constructs like
+  // ((0 - X) + Y), as opposed to (Y - X).
+  for (auto Term : Terms) {
+    if (Term.second <= 0)
       continue;
-    for (int J = 0; J < Terms[I]; ++J) {
+    for (int I = 0; I < Term.second; ++I)
       if (C.isZero())
-        C = Counter::getCounter(I);
+        C = Counter::getCounter(Term.first);
       else
         C = get(CounterExpression(CounterExpression::Add, C,
-                                  Counter::getCounter(I)));
-    }
+                                  Counter::getCounter(Term.first)));
   }
 
   // Create subtractions.
-  for (unsigned I = 0, S = Terms.size(); I < S; ++I) {
-    if (Terms[I] >= 0)
+  for (auto Term : Terms) {
+    if (Term.second >= 0)
       continue;
-    for (int J = 0; J < (-Terms[I]); ++J)
+    for (int I = 0; I < -Term.second; ++I)
       C = get(CounterExpression(CounterExpression::Subtract, C,
-                                Counter::getCounter(I)));
+                                Counter::getCounter(Term.first)));
   }
   return C;
 }