[NFC]Do not use recursion for CounterMappingContext::evaluate (#66961)

This causes stack overflows for real-world coverage reports. Ran $ build/bin/llvm-lit -a llvm/test/tools/llvm-cov locally and passed. Co-authored-by: Qing Shen <[email protected]>
llvm · Sep 22, 2023 · 487c784 · 487c784
1 parent 8466eb7
commit 487c784
Showing 1 changed file with 54 additions and 20 deletions.
diff --git a/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp b/llvm/lib/ProfileData/Coverage/CoverageMapping.cpp
@@ -167,27 +167,61 @@ void CounterMappingContext::dump(const Counter &C, raw_ostream &OS) const {
 }
 
 Expected<int64_t> CounterMappingContext::evaluate(const Counter &C) const {
-  switch (C.getKind()) {
-  case Counter::Zero:
-    return 0;
-  case Counter::CounterValueReference:
-    if (C.getCounterID() >= CounterValues.size())
-      return errorCodeToError(errc::argument_out_of_domain);
-    return CounterValues[C.getCounterID()];
-  case Counter::Expression: {
-    if (C.getExpressionID() >= Expressions.size())
-      return errorCodeToError(errc::argument_out_of_domain);
-    const auto &E = Expressions[C.getExpressionID()];
-    Expected<int64_t> LHS = evaluate(E.LHS);
-    if (!LHS)
-      return LHS;
-    Expected<int64_t> RHS = evaluate(E.RHS);
-    if (!RHS)
-      return RHS;
-    return E.Kind == CounterExpression::Subtract ? *LHS - *RHS : *LHS + *RHS;
-  }
+  struct StackElem {
+    Counter ICounter;
+    int64_t LHS = 0;
+    enum {
+      KNeverVisited = 0,
+      KVisitedOnce = 1,
+      KVisitedTwice = 2,
+    } VisitCount = KNeverVisited;
+  };
+
+  std::stack<StackElem> CounterStack;
+  CounterStack.push({C});
+
+  int64_t LastPoppedValue;
+
+  while (!CounterStack.empty()) {
+    StackElem &Current = CounterStack.top();
+
+    switch (Current.ICounter.getKind()) {
+    case Counter::Zero:
+      LastPoppedValue = 0;
+      CounterStack.pop();
+      break;
+    case Counter::CounterValueReference:
+      if (Current.ICounter.getCounterID() >= CounterValues.size())
+        return errorCodeToError(errc::argument_out_of_domain);
+      LastPoppedValue = CounterValues[Current.ICounter.getCounterID()];
+      CounterStack.pop();
+      break;
+    case Counter::Expression: {
+      if (Current.ICounter.getExpressionID() >= Expressions.size())
+        return errorCodeToError(errc::argument_out_of_domain);
+      const auto &E = Expressions[Current.ICounter.getExpressionID()];
+      if (Current.VisitCount == StackElem::KNeverVisited) {
+        CounterStack.push(StackElem{E.LHS});
+        Current.VisitCount = StackElem::KVisitedOnce;
+      } else if (Current.VisitCount == StackElem::KVisitedOnce) {
+        Current.LHS = LastPoppedValue;
+        CounterStack.push(StackElem{E.RHS});
+        Current.VisitCount = StackElem::KVisitedTwice;
+      } else {
+        int64_t LHS = Current.LHS;
+        int64_t RHS = LastPoppedValue;
+        LastPoppedValue =
+            E.Kind == CounterExpression::Subtract ? LHS - RHS : LHS + RHS;
+        CounterStack.pop();
+      }
+      break;
+    }
+    default:
+      llvm_unreachable("Unhandled CounterKind");
+    }
   }
-  llvm_unreachable("Unhandled CounterKind");
+
+  return LastPoppedValue;
 }
 
 unsigned CounterMappingContext::getMaxCounterID(const Counter &C) const {