From 510b41969eb54865e9cdbfd4ea11a29d45faf790 Mon Sep 17 00:00:00 2001
From: Vedant Kumar <vsk@apple.com>
Date: Fri, 16 Feb 2018 07:59:43 +0000
Subject: [PATCH] [Coverage] Handle break/continue outside of loop bodies

Teach the coverage mapping logic to handle break or continue statements
within for loop increments.

Fixes llvm.org/PR36406.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@325319 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/CodeGen/CoverageMappingGen.cpp | 24 ++++++++++++++++--------
 test/CoverageMapping/break.c       | 11 +++++++++++
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/lib/CodeGen/CoverageMappingGen.cpp b/lib/CodeGen/CoverageMappingGen.cpp
index 5eb48c2b57..d4ff7dac6c 100644
--- a/lib/CodeGen/CoverageMappingGen.cpp
+++ b/lib/CodeGen/CoverageMappingGen.cpp
@@ -982,20 +982,28 @@ struct CounterCoverageMappingBuilder
     Counter ParentCount = getRegion().getCounter();
     Counter BodyCount = getRegionCounter(S);
 
+    // The loop increment may contain a break or continue.
+    if (S->getInc())
+      BreakContinueStack.emplace_back();
+
     // Handle the body first so that we can get the backedge count.
-    BreakContinueStack.push_back(BreakContinue());
+    BreakContinueStack.emplace_back();
     extendRegion(S->getBody());
     Counter BackedgeCount = propagateCounts(BodyCount, S->getBody());
-    BreakContinue BC = BreakContinueStack.pop_back_val();
+    BreakContinue BodyBC = BreakContinueStack.pop_back_val();
 
     // The increment is essentially part of the body but it needs to include
     // the count for all the continue statements.
-    if (const Stmt *Inc = S->getInc())
-      propagateCounts(addCounters(BackedgeCount, BC.ContinueCount), Inc);
+    BreakContinue IncrementBC;
+    if (const Stmt *Inc = S->getInc()) {
+      propagateCounts(addCounters(BackedgeCount, BodyBC.ContinueCount), Inc);
+      IncrementBC = BreakContinueStack.pop_back_val();
+    }
 
     // Go back to handle the condition.
-    Counter CondCount =
-        addCounters(ParentCount, BackedgeCount, BC.ContinueCount);
+    Counter CondCount = addCounters(
+        addCounters(ParentCount, BackedgeCount, BodyBC.ContinueCount),
+        IncrementBC.ContinueCount);
     if (const Expr *Cond = S->getCond()) {
       propagateCounts(CondCount, Cond);
       adjustForOutOfOrderTraversal(getEnd(S));
@@ -1007,8 +1015,8 @@ struct CounterCoverageMappingBuilder
     if (Gap)
       fillGapAreaWithCount(Gap->getBegin(), Gap->getEnd(), BodyCount);
 
-    Counter OutCount =
-        addCounters(BC.BreakCount, subtractCounters(CondCount, BodyCount));
+    Counter OutCount = addCounters(BodyBC.BreakCount, IncrementBC.BreakCount,
+                                   subtractCounters(CondCount, BodyCount));
     if (OutCount != ParentCount)
       pushRegion(OutCount);
   }
diff --git a/test/CoverageMapping/break.c b/test/CoverageMapping/break.c
index d42c1bd082..08461d7ed2 100644
--- a/test/CoverageMapping/break.c
+++ b/test/CoverageMapping/break.c
@@ -31,3 +31,14 @@ int main() {         // CHECK: File 0, [[@LINE]]:12 -> {{[0-9]+}}:2 = #0
     ++cnt;
   }
 }
+
+// CHECK-LABEL: break_continue_in_increment:
+// CHECK:  [[@LINE+6]]:11 -> [[@LINE+6]]:45 = #1
+// CHECK:  [[@LINE+5]]:18 -> [[@LINE+5]]:19 = #1
+// CHECK:  [[@LINE+4]]:21 -> [[@LINE+4]]:26 = #2
+// CHECK:  [[@LINE+3]]:33 -> [[@LINE+3]]:41 = (#1 - #2)
+// CHECK:  [[@LINE+3]]:5 -> [[@LINE+3]]:6 = #1
+void break_continue_in_increment(int x) {
+  for (;; ({ if (x) break; else continue; }))
+    ;
+}
-- 
2.40.0