[Coverage] Fix handling of abnormal exits through repeat blocks

We do not correctly update the counter expression for conditionals in
repeat-while blocks in the following two situations:

Situation 1:

repeat {     // Region 1
  if (C1) {  // Region 2
    break
  }
} while C2   // Should be "Region 1 - Region 2", not "Region 1"

Situation 2:

repeat {     // Region 1
  if (C1) {  // Region 2
    continue
  }
} while C2   // Should be "Region 1", not "Region 1 + Region 2"

Fix both of these problems and add thorough regression tests.

Closes Swift PR #1244, rdar://problem/24572268

Author: vedantk

lib/SILGen/SILGenProfiling.cpp

@@ -514,7 +514,8 @@ struct CoverageMapping : public ASTWalker {
 
     } else if (auto *CS = dyn_cast<ContinueStmt>(S)) {
       // Continues create extra backedges, add them to the appropriate counters.
-      addToCounter(CS->getTarget(), getCurrentCounter());
+      if (!isa<RepeatWhileStmt>(CS->getTarget()))
+        addToCounter(CS->getTarget(), getCurrentCounter());
       if (auto *WS = dyn_cast<WhileStmt>(CS->getTarget())) {
         if (auto *E = getConditionNode(WS->getCond()))
           addToCounter(E, getCurrentCounter());
@@ -525,8 +526,12 @@ struct CoverageMapping : public ASTWalker {
 
     } else if (auto *BS = dyn_cast<BreakStmt>(S)) {
       // When we break from a loop, we need to adjust the exit count.
-      if (!isa<SwitchStmt>(BS->getTarget()))
+      if (auto *RWS = dyn_cast<RepeatWhileStmt>(BS->getTarget())) {
+        auto CondCount = CounterExpr::Sub(getCounter(RWS), getCurrentCounter());
+        addToCounter(RWS->getCond(), createCounter(std::move(CondCount)));
+      } else if (!isa<SwitchStmt>(BS->getTarget())) {
         addToCounter(BS->getTarget(), getCurrentCounter());
+      }
       terminateRegion(S);
 
     } else if (auto *FS = dyn_cast<FallthroughStmt>(S)) {

test/SILGen/coverage_while.swift

@@ -54,4 +54,53 @@ func foo() -> Int32 {
   return x
 }
 
+// rdar://problem/24572268
+// CHECK-LABEL: sil_coverage_map {{.*}}// coverage_while.goo
+// CHECK: [[@LINE+1]]:12 -> {{[0-9]+}}:2 : [[BAR:[0-9]+]]
+func goo() {
+  var x : Int32 = 0
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+2]]:4 : [[RWS1:[0-9]+]]
+    x += 1
+  } while x < 10 // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:17 : [[RWS1]]
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+6]]:4 : [[RWS2:[0-9]+]]
+    x += 1
+    if (x % 2 == 0) { // CHECK-DAG: [[@LINE]]:21 -> [[@LINE+2]]:6 : [[CONT1:[0-9]+]]
+      continue
+    } // CHECK-DAG: [[@LINE]]:6 -> [[@LINE+2]]:4 : ([[RWS2]] - [[CONT1]])
+    x += 1
+  } while x < 20 // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:17 : [[RWS2]]
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+6]]:4 : [[RWS3:[0-9]+]]
+    x += 1
+    if (x % 2 == 0) { // CHECK-DAG: [[@LINE]]:21 -> [[@LINE+2]]:6 : [[BRK1:[0-9]+]]
+      break
+    } // CHECK-DAG: [[@LINE]]:6 -> [[@LINE+2]]:4 : ([[RWS3]] - [[BRK1]])
+    x += 1
+  } while x < 30 // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:17 : ([[RWS3]] - [[BRK1]])
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+10]]:4 : [[RWS4:[0-9]+]]
+    x += 1
+    if (x % 2 == 0) { // CHECK-DAG: [[@LINE]]:21 -> [[@LINE+2]]:6 : [[CONT2:[0-9]+]]
+      continue
+    } // CHECK-DAG: [[@LINE]]:6 -> [[@LINE+6]]:4 : ([[RWS4]] - [[CONT2]])
+    x += 1
+    if (x % 7 == 0) { // CHECK-DAG: [[@LINE]]:21 -> [[@LINE+2]]:6 : [[BRK2:[0-9]+]]
+      break
+    } // CHECK-DAG: [[@LINE]]:6 -> [[@LINE+2]]:4 : (([[RWS4]] - [[CONT2]]) - [[BRK2]])
+    x += 1
+  } while x < 40 // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:17 : ([[RWS4]] - [[BRK2]])
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+1]]:4 : [[RWS5:[0-9]+]]
+  } while false // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:16 : [[RWS5]]
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+4]]:4 : [[RWS6:[0-9]+]]
+    repeat { // CHECK-DAG: [[@LINE]]:12 -> [[@LINE+2]]:6 : [[RWS7:[0-9]+]]
+      return
+    } while false // CHECK-DAG: [[@LINE]]:13 -> [[@LINE]]:18 : [[RWS7]]
+  } while false // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:16 : [[RWS6]]
+}
+
 foo()
+goo()