[Coverage] Make additional counters available for BranchRegion. NFC.

clang/lib/CodeGen/CodeGenFunction.h

@@ -1632,11 +1632,17 @@ class CodeGenFunction : public CodeGenTypeCache {
   /// Increment the profiler's counter for the given statement by \p StepV.
   /// If \p StepV is null, the default increment is 1.
   void incrementProfileCounter(const Stmt *S, llvm::Value *StepV = nullptr) {
+    incrementProfileCounter(false, S, false, StepV);
+  }
+
+  void incrementProfileCounter(bool UseSkipPath, const Stmt *S,
+                               bool UseBoth = false,
+                               llvm::Value *StepV = nullptr) {
     if (CGM.getCodeGenOpts().hasProfileClangInstr() &&
         !CurFn->hasFnAttribute(llvm::Attribute::NoProfile) &&
         !CurFn->hasFnAttribute(llvm::Attribute::SkipProfile)) {
       auto AL = ApplyDebugLocation::CreateArtificial(*this);
-      PGO.emitCounterSetOrIncrement(Builder, S, StepV);
+      PGO.emitCounterSetOrIncrement(Builder, S, UseSkipPath, UseBoth, StepV);
     }
     PGO.setCurrentStmt(S);
   }

clang/lib/CodeGen/CodeGenPGO.cpp

@@ -1138,6 +1138,19 @@ void CodeGenPGO::emitCounterRegionMapping(const Decl *D) {
   if (CoverageMapping.empty())
     return;
 
+  // Scan max(FalseCnt) and update NumRegionCounters.
+  unsigned MaxNumCounters = NumRegionCounters;
+  for (const auto [_, V] : *RegionCounterMap) {
+    auto HasCounters = V.getIsCounterPair();
+    assert((!HasCounters.first ||
+            MaxNumCounters > (V.first & CounterPair::Mask)) &&
+           "TrueCnt should not be reassigned");
+    if (HasCounters.second)
+      MaxNumCounters =
+          std::max(MaxNumCounters, (V.second & CounterPair::Mask) + 1);
+  }
+  NumRegionCounters = MaxNumCounters;
+
   CGM.getCoverageMapping()->addFunctionMappingRecord(
       FuncNameVar, FuncName, FunctionHash, CoverageMapping);
 }
@@ -1193,11 +1206,26 @@ std::pair<bool, bool> CodeGenPGO::getIsCounterPair(const Stmt *S) const {
 }
 
 void CodeGenPGO::emitCounterSetOrIncrement(CGBuilderTy &Builder, const Stmt *S,
+                                           bool UseSkipPath, bool UseBoth,
                                            llvm::Value *StepV) {
-  if (!RegionCounterMap || !Builder.GetInsertBlock())
+  if (!RegionCounterMap)
     return;
 
-  unsigned Counter = (*RegionCounterMap)[S].first;
+  unsigned Counter;
+  auto &TheMap = (*RegionCounterMap)[S];
+  auto IsCounter = TheMap.getIsCounterPair();
+  if (!UseSkipPath) {
+    if (!IsCounter.first)
+      return;
+    Counter = (TheMap.first & CounterPair::Mask);
+  } else {
+    if (!IsCounter.second)
+      return;
+    Counter = (TheMap.second & CounterPair::Mask);
+  }
+
+  if (!Builder.GetInsertBlock())
+    return;
 
   // Make sure that pointer to global is passed in with zero addrspace
   // This is relevant during GPU profiling

clang/lib/CodeGen/CodeGenPGO.h

@@ -112,6 +112,7 @@ class CodeGenPGO {
 public:
   std::pair<bool, bool> getIsCounterPair(const Stmt *S) const;
   void emitCounterSetOrIncrement(CGBuilderTy &Builder, const Stmt *S,
+                                 bool UseFalsePath, bool UseBoth,
                                  llvm::Value *StepV);
   void emitMCDCTestVectorBitmapUpdate(CGBuilderTy &Builder, const Expr *S,
                                       Address MCDCCondBitmapAddr,

clang/lib/CodeGen/CoverageMappingGen.cpp

@@ -887,6 +887,9 @@ struct CounterCoverageMappingBuilder
   /// The map of statements to count values.
   llvm::DenseMap<const Stmt *, CounterPair> &CounterMap;
 
+  CounterExpressionBuilder::ReplaceMap MapToExpand;
+  unsigned NextCounterNum;
+
   MCDC::State &MCDCState;
 
   /// A stack of currently live regions.
@@ -922,15 +925,11 @@ struct CounterCoverageMappingBuilder
 
   /// Return a counter for the sum of \c LHS and \c RHS.
   Counter addCounters(Counter LHS, Counter RHS, bool Simplify = true) {
-    assert(!llvm::EnableSingleByteCoverage &&
-           "cannot add counters when single byte coverage mode is enabled");
     return Builder.add(LHS, RHS, Simplify);
   }
 
   Counter addCounters(Counter C1, Counter C2, Counter C3,
                       bool Simplify = true) {
-    assert(!llvm::EnableSingleByteCoverage &&
-           "cannot add counters when single byte coverage mode is enabled");
     return addCounters(addCounters(C1, C2, Simplify), C3, Simplify);
   }
 
@@ -943,19 +942,43 @@ struct CounterCoverageMappingBuilder
 
   std::pair<Counter, Counter> getBranchCounterPair(const Stmt *S,
                                                    Counter ParentCnt) {
-    Counter ExecCnt = getRegionCounter(S);
-    return {ExecCnt, Builder.subtract(ParentCnt, ExecCnt)};
+    auto &TheMap = CounterMap[S];
+    auto ExecCnt = Counter::getCounter(TheMap.first);
+    auto SkipExpr = Builder.subtract(ParentCnt, ExecCnt);
+
+    if (!llvm::EnableSingleByteCoverage || !SkipExpr.isExpression()) {
+      assert(
+          !TheMap.getIsCounterPair().second &&
+          "SkipCnt shouldn't be allocated but refer to an existing counter.");
+      return {ExecCnt, SkipExpr};
+    }
+
+    // Assign second if second is not assigned yet.
+    if (!TheMap.getIsCounterPair().second)
+      TheMap.second = NextCounterNum++;
+
+    Counter SkipCnt = Counter::getCounter(TheMap.second);
+    MapToExpand[SkipCnt] = SkipExpr;
+    return {ExecCnt, SkipCnt};
   }
 
   Counter getSwitchImplicitDefaultCounter(const Stmt *Cond, Counter ParentCount,
                                           Counter CaseCountSum) {
-    return Builder.subtract(ParentCount, CaseCountSum);
+    return (
+        llvm::EnableSingleByteCoverage
+            ? Counter::getCounter(CounterMap[Cond].second = NextCounterNum++)
+            : Builder.subtract(ParentCount, CaseCountSum));
   }
 
   bool IsCounterEqual(Counter OutCount, Counter ParentCount) {
     if (OutCount == ParentCount)
       return true;
 
+    // Try comaparison with pre-replaced expressions.
+    if (Builder.replace(Builder.subtract(OutCount, ParentCount), MapToExpand)
+            .isZero())
+      return true;
+
     return false;
   }
 
@@ -1438,7 +1461,8 @@ struct CounterCoverageMappingBuilder
       llvm::DenseMap<const Stmt *, CounterPair> &CounterMap,
       MCDC::State &MCDCState, SourceManager &SM, const LangOptions &LangOpts)
       : CoverageMappingBuilder(CVM, SM, LangOpts), CounterMap(CounterMap),
-        MCDCState(MCDCState), MCDCBuilder(CVM.getCodeGenModule(), MCDCState) {}
+        NextCounterNum(CounterMap.size()), MCDCState(MCDCState),
+        MCDCBuilder(CVM.getCodeGenModule(), MCDCState) {}
 
   /// Write the mapping data to the output stream
   void write(llvm::raw_ostream &OS) {

llvm/lib/ProfileData/Coverage/CoverageMapping.cpp

@@ -638,6 +638,10 @@ static unsigned getMaxCounterID(const CounterMappingContext &Ctx,
   unsigned MaxCounterID = 0;
   for (const auto &Region : Record.MappingRegions) {
     MaxCounterID = std::max(MaxCounterID, Ctx.getMaxCounterID(Region.Count));
+    if (Region.Kind == CounterMappingRegion::BranchRegion ||
+        Region.Kind == CounterMappingRegion::MCDCBranchRegion)
+      MaxCounterID =
+          std::max(MaxCounterID, Ctx.getMaxCounterID(Region.FalseCount));
   }
   return MaxCounterID;
 }