[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

@@ -1137,6 +1137,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);
 }
@@ -1192,11 +1205,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

@@ -884,6 +884,9 @@ struct CounterCoverageMappingBuilder
   /// The map of statements to count values.
   llvm::DenseMap<const Stmt *, CounterPair> &CounterMap;
 
+  CounterExpressionBuilder::SubstMap MapToExpand;
+  unsigned NextCounterNum;
+
   MCDC::State &MCDCState;
 
   /// A stack of currently live regions.
@@ -919,15 +922,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);
   }
 
@@ -940,15 +939,35 @@ 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};
   }
 
   /// Returns {TrueCnt,FalseCnt} for "implicit default".
   /// FalseCnt is considered as the False count on SwitchStmt.
   std::pair<Counter, Counter>
   getSwitchImplicitDefaultCounterPair(const Stmt *Cond, Counter ParentCount,
                                       Counter CaseCountSum) {
+    if (llvm::EnableSingleByteCoverage)
+      return {Counter::getZero(), // Folded
+              Counter::getCounter(CounterMap[Cond].second = NextCounterNum++)};
+
     // Simplify is skipped while building the counters above: it can get
     // really slow on top of switches with thousands of cases. Instead,
     // trigger simplification by adding zero to the last counter.
@@ -962,6 +981,11 @@ struct CounterCoverageMappingBuilder
     if (OutCount == ParentCount)
       return true;
 
+    // Try comaparison with pre-replaced expressions.
+    if (Builder.subst(Builder.subtract(OutCount, ParentCount), MapToExpand)
+            .isZero())
+      return true;
+
     return false;
   }
 
@@ -1175,12 +1199,14 @@ struct CounterCoverageMappingBuilder
   /// and add it to the function's SourceRegions.
   /// Returns Counter that corresponds to SC.
   Counter createSwitchCaseRegion(const SwitchCase *SC, Counter ParentCount) {
+    Counter TrueCnt = getRegionCounter(SC);
+    Counter FalseCnt = (llvm::EnableSingleByteCoverage
+                            ? Counter::getZero() // Folded
+                            : subtractCounters(ParentCount, TrueCnt));
     // Push region onto RegionStack but immediately pop it (which adds it to
     // the function's SourceRegions) because it doesn't apply to any other
     // source other than the SwitchCase.
-    Counter TrueCnt = getRegionCounter(SC);
-    popRegions(pushRegion(TrueCnt, getStart(SC), SC->getColonLoc(),
-                          subtractCounters(ParentCount, TrueCnt)));
+    popRegions(pushRegion(TrueCnt, getStart(SC), SC->getColonLoc(), FalseCnt));
     return TrueCnt;
   }
 
@@ -1447,7 +1473,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

@@ -637,6 +637,9 @@ 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.isBranch())
+      MaxCounterID =
+          std::max(MaxCounterID, Ctx.getMaxCounterID(Region.FalseCount));
   }
   return MaxCounterID;
 }