[Coverage] Fix handling of return statements within repeat-while loops

In 2a34b7c, we introduced an incorrect test case for return statements
within repeat-while loops. Consider this situation:

repeat {     // Region 1
  return
} while C1   // Should be "zero", not "Region 1"

The fix requires maintaining a stack of active repeat-while loops. This
is an accepted idiom (c.f the clang CoverageMapping implementation). To
see why a stack is needed, consider:

repeat {       // Region 1
  repeat {     // Region 2
    if (C1) {  // Region 3
      return
    }
  } while C2   // Should be "Region 2 - *Region 3*", not "Region 2"
} while C3     // Should be "Region 1 - *Region 3*", not "Region 1"

rdar://problem/24572268

Author: vedantk

lib/SILGen/SILGenProfiling.cpp

@@ -229,6 +229,9 @@ struct CoverageMapping : public ASTWalker {
   /// \brief A stack of currently live regions.
   std::vector<SourceMappingRegion> RegionStack;
 
+  /// \brief A stack of active repeat-while loops.
+  std::vector<RepeatWhileStmt *> RepeatWhileStack;
+
   CounterExpr *ExitCounter;
 
   /// \brief Return true if \c Node has an associated counter.
@@ -271,6 +274,16 @@ struct CoverageMapping : public ASTWalker {
       Counter = CounterExpr::Add(createCounter(std::move(Counter)), Expr);
   }
 
+  /// \brief Subtract \c Expr from \c Node's counter.
+  void subtractFromCounter(ASTNode Node, CounterExpr &Expr) {
+    CounterExpr &Counter = getCounter(Node);
+    assert(!Counter.isZero() && "Cannot create a negative counter");
+    if (const CounterExpr *ReferencedCounter = Counter.getReferencedNode())
+      Counter = CounterExpr::Sub(*ReferencedCounter, Expr);
+    else
+      Counter = CounterExpr::Sub(createCounter(std::move(Counter)), Expr);
+  }
+
   /// \brief Return the current region's counter.
   CounterExpr &getCurrentCounter() { return getRegion().getCounter(); }
 
@@ -457,6 +470,7 @@ struct CoverageMapping : public ASTWalker {
       assignCounter(RWS, CounterExpr::Zero());
       CounterExpr &BodyCounter = assignCounter(RWS->getBody());
       assignCounter(RWS->getCond(), CounterExpr::Ref(BodyCounter));
+      RepeatWhileStack.push_back(RWS);
 
     } else if (auto *FS = dyn_cast<ForStmt>(S)) {
       assignCounter(FS, CounterExpr::Zero());
@@ -505,6 +519,10 @@ struct CoverageMapping : public ASTWalker {
       if (auto *E = getConditionNode(WS->getCond()))
         addToCounter(E, getExitCounter());
 
+    } else if (auto *RWS = dyn_cast<RepeatWhileStmt>(S)) {
+      assert(RepeatWhileStack.back() == RWS && "Malformed repeat-while stack");
+      RepeatWhileStack.pop_back();
+
     } else if (auto *FS = dyn_cast<ForStmt>(S)) {
       // Both the condition and the increment are reached through the backedge.
       if (Expr *E = FS->getCond().getPtrOrNull())
@@ -527,8 +545,7 @@ 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 (auto *RWS = dyn_cast<RepeatWhileStmt>(BS->getTarget())) {
-        auto CondCount = CounterExpr::Sub(getCounter(RWS), getCurrentCounter());
-        addToCounter(RWS->getCond(), createCounter(std::move(CondCount)));
+        subtractFromCounter(RWS->getCond(), getCurrentCounter());
       } else if (!isa<SwitchStmt>(BS->getTarget())) {
         addToCounter(BS->getTarget(), getCurrentCounter());
       }
@@ -548,6 +565,10 @@ struct CoverageMapping : public ASTWalker {
       replaceCount(CounterExpr::Ref(getCounter(S)), getEndLoc(S));
 
     } else if (isa<ReturnStmt>(S) || isa<FailStmt>(S) || isa<ThrowStmt>(S)) {
+      // When we return, we may need to adjust some loop condition counts.
+      for (auto *RWS : RepeatWhileStack)
+        subtractFromCounter(RWS->getCond(), getCurrentCounter());
+
       terminateRegion(S);
     }
     return S;

test/SILGen/coverage_while.swift

@@ -97,8 +97,18 @@ func goo() {
   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]]
+    } while false // CHECK-DAG: [[@LINE]]:13 -> [[@LINE]]:18 : zero
+  } while false // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:16 : ([[RWS6]] - [[RWS7]])
+
+  repeat { // CHECK-DAG: [[@LINE]]:10 -> [[@LINE+6]]:4 : [[RWS8:[0-9]+]]
+    repeat { // CHECK-DAG: [[@LINE]]:12 -> [[@LINE+4]]:6 : [[RWS9:[0-9]+]]
+      if (true) { // CHECK-DAG: [[@LINE]]:17 -> [[@LINE+2]]:8 : [[RET1:[0-9]+]]
+        return
+      }
+    } while false // CHECK-DAG: [[@LINE]]:13 -> [[@LINE]]:18 : ([[RWS9]] - [[RET1]])
+  } while false // CHECK-DAG: [[@LINE]]:11 -> [[@LINE]]:16 : ([[RWS8]] - [[RET1]])
+
+  // TODO(vsk): need tests for fail and throw statements.
 }
 
 foo()