Introduce CounterExpressionBuilder::subst(C, Map)

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>;

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: