[Coverage] Emit gap areas in braces-optional statements (PR35387)

Emit a gap area starting after the r-paren location and ending at the
start of the body for the braces-optional statements (for, for-each,
while, etc). The count for the gap area equal to the body's count. This
extends the fix in r317758.

Fixes PR35387, rdar://35570345

Testing: stage2 coverage-enabled build of clang, check-clang
llvm-svn: 319373

1 files changed, 61 insertions, 9 deletions

diff --git a/clang/lib/CodeGen/CoverageMappingGen.cpp b/clang/lib/CodeGen/CoverageMappingGen.cpp
index 6e0099bc201..89a30dc7040 100644
--- a/clang/lib/CodeGen/CoverageMappingGen.cpp
+++ b/clang/lib/CodeGen/CoverageMappingGen.cpp
@@ -112,6 +112,9 @@ struct SpellingRegion {
     ColumnEnd = SM.getSpellingColumnNumber(LocEnd);
   }
 
+  SpellingRegion(SourceManager &SM, SourceMappingRegion &R)
+      : SpellingRegion(SM, R.getStartLoc(), R.getEndLoc()) {}
+
   /// Check if the start and end locations appear in source order, i.e
   /// top->bottom, left->right.
   bool isInSourceOrder() const {
@@ -583,6 +586,7 @@ struct CounterCoverageMappingBuilder
           MostRecentLocation = getIncludeOrExpansionLoc(EndLoc);
 
         assert(SM.isWrittenInSameFile(Region.getStartLoc(), EndLoc));
+        assert(SpellingRegion(SM, Region).isInSourceOrder());
         SourceRegions.push_back(Region);
 
         if (ParentOfDeferredRegion) {
@@ -718,9 +722,11 @@ struct CounterCoverageMappingBuilder
       SourceLocation Loc = MostRecentLocation;
       while (isNestedIn(Loc, ParentFile)) {
         SourceLocation FileStart = getStartOfFileOrMacro(Loc);
-        if (StartLocs.insert(FileStart).second)
+        if (StartLocs.insert(FileStart).second) {
           SourceRegions.emplace_back(*ParentCounter, FileStart,
                                      getEndOfFileOrMacro(Loc));
+          assert(SpellingRegion(SM, SourceRegions.back()).isInSourceOrder());
+        }
         Loc = getIncludeOrExpansionLoc(Loc);
       }
     }
@@ -753,12 +759,31 @@ struct CounterCoverageMappingBuilder
     LastTerminatedRegion = {EndLoc, RegionStack.size()};
   }
 
+  /// Find a valid gap range between \p AfterLoc and \p BeforeLoc.
+  Optional<SourceRange> findGapAreaBetween(SourceLocation AfterLoc,
+                                           SourceLocation BeforeLoc) {
+    // If the start and end locations of the gap are both within the same macro
+    // file, the range may not be in source order.
+    if (AfterLoc.isMacroID() || BeforeLoc.isMacroID())
+      return None;
+    if (!SM.isWrittenInSameFile(AfterLoc, BeforeLoc))
+      return None;
+    return {{AfterLoc, BeforeLoc}};
+  }
+
+  /// Find the source range after \p AfterStmt and before \p BeforeStmt.
+  Optional<SourceRange> findGapAreaBetween(const Stmt *AfterStmt,
+                                           const Stmt *BeforeStmt) {
+    return findGapAreaBetween(getPreciseTokenLocEnd(getEnd(AfterStmt)),
+                              getStart(BeforeStmt));
+  }
+
   /// Emit a gap region between \p StartLoc and \p EndLoc with the given count.
   void fillGapAreaWithCount(SourceLocation StartLoc, SourceLocation EndLoc,
                             Counter Count) {
-    if (StartLoc == EndLoc || StartLoc.isMacroID() || EndLoc.isMacroID() ||
-        !SM.isWrittenInSameFile(StartLoc, EndLoc))
+    if (StartLoc == EndLoc)
       return;
+    assert(SpellingRegion(SM, StartLoc, EndLoc).isInSourceOrder());
     handleFileExit(StartLoc);
     size_t Index = pushRegion(Count, StartLoc, EndLoc);
     getRegion().setGap(true);
@@ -914,6 +939,11 @@ struct CounterCoverageMappingBuilder
     propagateCounts(CondCount, S->getCond());
     adjustForOutOfOrderTraversal(getEnd(S));
 
+    // The body count applies to the area immediately after the increment.
+    auto Gap = findGapAreaBetween(S->getCond(), S->getBody());
+    if (Gap)
+      fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
+
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
     if (OutCount != ParentCount)
@@ -968,6 +998,12 @@ struct CounterCoverageMappingBuilder
       adjustForOutOfOrderTraversal(getEnd(S));
     }
 
+    // The body count applies to the area immediately after the increment.
+    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
+                                  getStart(S->getBody()));
+    if (Gap)
+      fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
+
     Counter OutCount =
         addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
     if (OutCount != ParentCount)
@@ -987,6 +1023,12 @@ struct CounterCoverageMappingBuilder
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
+    // The body count applies to the area immediately after the range.
+    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
+                                  getStart(S->getBody()));
+    if (Gap)
+      fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
+
     Counter LoopCount =
         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
     Counter OutCount =
@@ -1007,6 +1049,12 @@ struct CounterCoverageMappingBuilder
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
     BreakContinue BC = BreakContinueStack.pop_back_val();
 
+    // The body count applies to the area immediately after the collection.
+    auto Gap = findGapAreaBetween(getPreciseTokenLocEnd(S->getRParenLoc()),
+                                  getStart(S->getBody()));
+    if (Gap)
+      fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
+
     Counter LoopCount =
         addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
     Counter OutCount =
@@ -1099,8 +1147,9 @@ struct CounterCoverageMappingBuilder
     propagateCounts(ParentCount, S->getCond());
 
     // The 'then' count applies to the area immediately after the condition.
-    fillGapAreaWithCount(getPreciseTokenLocEnd(getEnd(S->getCond())),
-                         getStart(S->getThen()), ThenCount);
+    auto Gap = findGapAreaBetween(S->getCond(), S->getThen());
+    if (Gap)
+      fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ThenCount);
 
     extendRegion(S->getThen());
     Counter OutCount = propagateCounts(ThenCount, S->getThen());
@@ -1108,8 +1157,9 @@ struct CounterCoverageMappingBuilder
     Counter ElseCount = subtractCounters(ParentCount, ThenCount);
     if (const Stmt *Else = S->getElse()) {
       // The 'else' count applies to the area immediately after the 'then'.
-      fillGapAreaWithCount(getPreciseTokenLocEnd(getEnd(S->getThen())),
-                           getStart(Else), ElseCount);
+      Gap = findGapAreaBetween(S->getThen(), Else);
+      if (Gap)
+        fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), ElseCount);
       extendRegion(Else);
       OutCount = addCounters(OutCount, propagateCounts(ElseCount, Else));
     } else
@@ -1148,8 +1198,10 @@ struct CounterCoverageMappingBuilder
 
     if (!isa<BinaryConditionalOperator>(E)) {
       // The 'then' count applies to the area immediately after the condition.
-      fillGapAreaWithCount(E->getQuestionLoc(), getStart(E->getTrueExpr()),
-                           TrueCount);
+      auto Gap =
+          findGapAreaBetween(E->getQuestionLoc(), getStart(E->getTrueExpr()));
+      if (Gap)
+        fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), TrueCount);
 
       extendRegion(E->getTrueExpr());
       propagateCounts(TrueCount, E->getTrueExpr());