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

[NAKAMURA Takumi via llvm-commits]

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

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`.

>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] 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: