[llvm] Introduce CounterExpressionBuilder::replace(C, Map) (PR #112698)

[NAKAMURA Takumi via llvm-commits]

https://github.com/chapuni updated https://github.com/llvm/llvm-project/pull/112698

>From 618e63946923a460b2684738e636c77b1706322a Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic at gmail.com>
Date: Thu, 17 Oct 2024 01:22:54 +0900
Subject: [PATCH 1/3] Introduce CounterExpressionBuilder::replace(C, Map)

This return a counter for each term in the expression replaced by
ReplaceMap.

At the moment, this doesn't update the Map, so Map is marked as `const`.
---
 .../ProfileData/Coverage/CoverageMapping.h    |  6 ++++
 .../ProfileData/Coverage/CoverageMapping.cpp  | 32 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h b/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
index fa07b3a9e8b14f..d6528bd407ef34 100644
--- a/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
+++ b/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h
@@ -34,6 +34,7 @@
 #include <cassert>
 #include <cstdint>
 #include <iterator>
+#include <map>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -213,6 +214,11 @@ class CounterExpressionBuilder {
   /// Return a counter that represents the expression that subtracts RHS from
   /// LHS.
   Counter subtract(Counter LHS, Counter RHS, bool Simplify = true);
+
+  using ReplaceMap = std::map<Counter, Counter>;
+
+  /// Return a counter for each term in the expression replaced by ReplaceMap.
+  Counter replace(Counter C, const ReplaceMap &Map);
 };
 
 using LineColPair = std::pair<unsigned, unsigned>;
diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
index c713371da81e40..b50f025d261e13 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -135,6 +135,38 @@ Counter CounterExpressionBuilder::subtract(Counter LHS, Counter RHS,
   return Simplify ? simplify(Cnt) : Cnt;
 }
 
+Counter CounterExpressionBuilder::replace(Counter C, const ReplaceMap &Map) {
+  auto I = Map.find(C);
+
+  // Replace C with the Map even if C is Expression.
+  if (I != Map.end())
+    return I->second;
+
+  // Traverse only Expression.
+  if (!C.isExpression())
+    return C;
+
+  auto CE = Expressions[C.getExpressionID()];
+  auto NewLHS = replace(CE.LHS, Map);
+  auto NewRHS = replace(CE.RHS, Map);
+
+  // Reconstruct Expression with induced subexpressions.
+  switch (CE.Kind) {
+  case CounterExpression::Add:
+    C = add(NewLHS, NewRHS);
+    break;
+  case CounterExpression::Subtract:
+    C = subtract(NewLHS, NewRHS);
+    break;
+  }
+
+  // Reconfirm if the reconstructed expression would hit the Map.
+  if ((I = Map.find(C)) != Map.end())
+    return I->second;
+
+  return C;
+}
+
 void CounterMappingContext::dump(const Counter &C, raw_ostream &OS) const {
   switch (C.getKind()) {
   case Counter::Zero:

>From 209ea4cfdb4f93d3c8fc60dc9af29af39fd24758 Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic at gmail.com>
Date: Sun, 20 Oct 2024 11:58:16 +0900
Subject: [PATCH 2/3] Update comments

---
 llvm/lib/ProfileData/Coverage/CoverageMapping.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
index b50f025d261e13..6f44797be20e32 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -138,14 +138,14 @@ Counter CounterExpressionBuilder::subtract(Counter LHS, Counter RHS,
 Counter CounterExpressionBuilder::replace(Counter C, const ReplaceMap &Map) {
   auto I = Map.find(C);
 
-  // Replace C with the Map even if C is Expression.
+  // Replace C with the value found in Map even if C is Expression.
   if (I != Map.end())
     return I->second;
 
-  // Traverse only Expression.
   if (!C.isExpression())
     return C;
 
+  // Traverse both sides of Expression.
   auto CE = Expressions[C.getExpressionID()];
   auto NewLHS = replace(CE.LHS, Map);
   auto NewRHS = replace(CE.RHS, Map);

>From f0afd04dd86573f1e7d868cc6e1c677d1779aa5f Mon Sep 17 00:00:00 2001
From: NAKAMURA Takumi <geek4civic at gmail.com>
Date: Sun, 20 Oct 2024 12:00:23 +0900
Subject: [PATCH 3/3] Use initializer statements

---
 llvm/lib/ProfileData/Coverage/CoverageMapping.cpp | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
index 6f44797be20e32..7ff2e5ba69e19d 100644
--- a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
+++ b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -136,10 +136,8 @@ Counter CounterExpressionBuilder::subtract(Counter LHS, Counter RHS,
 }
 
 Counter CounterExpressionBuilder::replace(Counter C, const ReplaceMap &Map) {
-  auto I = Map.find(C);
-
   // Replace C with the value found in Map even if C is Expression.
-  if (I != Map.end())
+  if (auto I = Map.find(C); I != Map.end())
     return I->second;
 
   if (!C.isExpression())
@@ -161,7 +159,7 @@ Counter CounterExpressionBuilder::replace(Counter C, const ReplaceMap &Map) {
   }
 
   // Reconfirm if the reconstructed expression would hit the Map.
-  if ((I = Map.find(C)) != Map.end())
+  if (auto I = Map.find(C); I != Map.end())
     return I->second;
 
   return C;