[Coverage] Complete top-level deferred regions before labels

The area immediately after a terminated region in the function top-level
should have the same count as the label it precedes.

This solves another problem with wrapped segments. Consider:

  1| a:
  2|   return 0;
  3| b:
  4|   return 1;

Without a gap area starting after the first return, the wrapped segment
from line 2 would make it look like line 3 is executed, when it's not.

rdar://35373009

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@317759 91177308-0d34-0410-b5e6-96231b3b80d8

lib/CodeGen/CoverageMappingGen.cpp

@@ -445,6 +445,9 @@ struct CounterCoverageMappingBuilder
   /// expressions cross file or macro boundaries.
   SourceLocation MostRecentLocation;
 
+  /// Location of the last terminated region.
+  Optional<std::pair<SourceLocation, size_t>> LastTerminatedRegion;
+
   /// \brief Return a counter for the subtraction of \c RHS from \c LHS
   Counter subtractCounters(Counter LHS, Counter RHS) {
     return Builder.subtract(LHS, RHS);
@@ -520,6 +523,27 @@ struct CounterCoverageMappingBuilder
     return Index;
   }
 
+  /// Complete a deferred region created after a terminated region at the
+  /// top-level.
+  void completeTopLevelDeferredRegion(Counter Count,
+                                      SourceLocation DeferredEndLoc) {
+    if (DeferredRegion || !LastTerminatedRegion)
+      return;
+
+    if (LastTerminatedRegion->second != RegionStack.size())
+      return;
+
+    SourceLocation Start = LastTerminatedRegion->first;
+    if (SM.getFileID(Start) != SM.getMainFileID())
+      return;
+
+    SourceMappingRegion DR = RegionStack.back();
+    DR.setStartLoc(Start);
+    DR.setDeferred(false);
+    DeferredRegion = DR;
+    completeDeferred(Count, DeferredEndLoc);
+  }
+
   /// \brief Pop regions from the stack into the function's list of regions.
   ///
   /// Adds all regions from \c ParentIndex to the top of the stack to the
@@ -576,6 +600,12 @@ struct CounterCoverageMappingBuilder
         ParentOfDeferredRegion = true;
       }
       RegionStack.pop_back();
+
+      // If the zero region pushed after the last terminated region no longer
+      // exists, clear its cached information.
+      if (LastTerminatedRegion &&
+          RegionStack.size() < LastTerminatedRegion->second)
+        LastTerminatedRegion = None;
     }
     assert(!ParentOfDeferredRegion && "Deferred region with no parent");
   }
@@ -712,10 +742,13 @@ struct CounterCoverageMappingBuilder
   void terminateRegion(const Stmt *S) {
     extendRegion(S);
     SourceMappingRegion &Region = getRegion();
+    SourceLocation EndLoc = getEnd(S);
     if (!Region.hasEndLoc())
-      Region.setEndLoc(getEnd(S));
+      Region.setEndLoc(EndLoc);
     pushRegion(Counter::getZero());
-    getRegion().setDeferred(true);
+    auto &ZeroRegion = getRegion();
+    ZeroRegion.setDeferred(true);
+    LastTerminatedRegion = {EndLoc, RegionStack.size()};
   }
 
   /// Emit a gap region between \p StartLoc and \p EndLoc with the given count.
@@ -826,10 +859,12 @@ struct CounterCoverageMappingBuilder
   void VisitGotoStmt(const GotoStmt *S) { terminateRegion(S); }
 
   void VisitLabelStmt(const LabelStmt *S) {
+    Counter LabelCount = getRegionCounter(S);
     SourceLocation Start = getStart(S);
+    completeTopLevelDeferredRegion(LabelCount, Start);
     // We can't extendRegion here or we risk overlapping with our new region.
     handleFileExit(Start);
-    pushRegion(getRegionCounter(S), Start);
+    pushRegion(LabelCount, Start);
     Visit(S->getSubStmt());
   }
 

test/CoverageMapping/label.cpp

@@ -28,8 +28,8 @@ void test1(int x) {          // CHECK-NEXT: File 0, [[@LINE]]:19 -> {{[0-9]+}}:2
   if(x == 0)                 // CHECK-NEXT: File 0, [[@LINE]]:6 -> [[@LINE]]:12 = #0
     goto a;                  // CHECK: File 0, [[@LINE]]:5 -> [[@LINE]]:11 = #1
                              // CHECK-NEXT: File 0, [[@LINE-1]]:11 -> [[@LINE+1]]:3 = (#0 - #1)
-  goto b;                    // CHECK-NEXT: File 0, [[@LINE]]:3 -> [[@LINE+5]]:2 = (#0 - #1)
-                             // CHECK-NEXT: File 0, [[@LINE-1]]:3 -> [[@LINE+4]]:2 = #3
+  goto b;                    // CHECK: Gap,File 0, [[@LINE]]:3 -> [[@LINE+5]]:2 = #3
+                             // CHECK-NEXT: Gap,File 0, [[@LINE-1]]:9 -> [[@LINE+1]]:1 = #2
 a:                           // CHECK-NEXT: File 0, [[@LINE]]:1 -> [[@LINE+3]]:2 = #2
 b:                           // CHECK-NEXT: File 0, [[@LINE]]:1 -> [[@LINE+2]]:2 = #3
   x = x + 1;
@@ -58,7 +58,7 @@ int main() {                 // CHECK-NEXT: File 0, [[@LINE]]:12 -> {{[0-9]+}}:2
       goto e;                // CHECK: File 0, [[@LINE]]:7 -> [[@LINE]]:13 = #3
                              // CHECK-NEXT: File 0, [[@LINE-1]]:13 -> [[@LINE+1]]:5 = (#2 - #3)
     goto c;                  // CHECK-NEXT: File 0, [[@LINE]]:5 -> [[@LINE+8]]:4 = (#2 - #3)
-
+                             // CHECK-NEXT: Gap,File 0, [[@LINE-1]]:11 -> [[@LINE+1]]:3 = #4
   b:                         // CHECK-NEXT: File 0, [[@LINE]]:3 -> [[@LINE+6]]:4 = #4
     j = 2;
   c:                         // CHECK-NEXT: File 0, [[@LINE]]:3 -> [[@LINE+4]]:4 = #5

test/CoverageMapping/switchmacro.c

@@ -13,7 +13,7 @@ int foo(int i) { // CHECK-NEXT: File 0, [[@LINE]]:16 -> {{[0-9]+}}:2 = #0
     // CHECK-NEXT: File 0, [[@LINE+1]]:8 -> {{[0-9]+}}:11 = (#2 - #3)
     FOO(1);
   case 0:        // CHECK-NEXT: File 0, [[@LINE]]:3 -> [[@LINE+1]]:13 = ((#2 + #4) - #3)
-    return 2;
+    return 2;    // CHECK-NEXT: Gap,File 0, [[@LINE]]:13 -> [[@LINE+6]]:3 = #5
 
   // CHECK-NEXT: Expansion,File 0, [[@LINE+2]]:3 -> [[@LINE+2]]:6 = 0
   // CHECK-NEXT: File 0, [[@LINE+1]]:6 -> {{[0-9]+}}:11 = 0