[Profiler] Fix coverage mapping for switches

lib/SIL/IR/SILProfiler.cpp

@@ -259,8 +259,6 @@ struct MapRegionCounters : public ASTWalker {
       mapRegion(RWS->getBody());
     } else if (auto *FES = dyn_cast<ForEachStmt>(S)) {
       mapRegion(FES->getBody());
-    } else if (auto *SS = dyn_cast<SwitchStmt>(S)) {
-      mapRegion(SS);
     } else if (auto *CS = dyn_cast<CaseStmt>(S)) {
       mapRegion(getProfilerStmtForCase(CS));
     }
@@ -647,8 +645,6 @@ struct PGOMapping : public ASTWalker {
     } else if (auto *FES = dyn_cast<ForEachStmt>(S)) {
       setKnownExecutionCount(FES->getBody());
       setExecutionCount(FES, parentCount);
-    } else if (auto *SS = dyn_cast<SwitchStmt>(S)) {
-      setKnownExecutionCount(SS);
     } else if (auto *CS = dyn_cast<CaseStmt>(S)) {
       setKnownExecutionCount(getProfilerStmtForCase(CS));
     }
@@ -695,6 +691,9 @@ struct PGOMapping : public ASTWalker {
   }
 };
 
+/// Produce coverage mapping information for a function. This involves taking
+/// the counters computed by MapRegionCounters, and annotating the source with
+/// regions that are defined in terms of those counters.
 struct CoverageMapping : public ASTWalker {
 private:
   const SourceManager &SM;
@@ -820,17 +819,7 @@ struct CoverageMapping : public ASTWalker {
     if (ControlFlowAdjust)
       Count = &createCounter(CounterExpr::Sub(*Count, *ControlFlowAdjust));
 
-    if (Count->isSemanticallyZero()) {
-      // If the counter is semantically zero, form an 'incomplete' region with
-      // no starting location. This prevents forming unreachable regions unless
-      // there is a following statement or expression to extend the region.
-      RegionStack.emplace_back(ASTNode(), *Count, None, None);
-    } else {
-      // Otherwise, we have a non-zero counter, so form a new region starting
-      // at the end of the previous scope. This ensures the user covers both
-      // branches of a condition.
-      RegionStack.emplace_back(ASTNode(), *Count, getEndLoc(Scope), None);
-    }
+    replaceCount(Count, getEndLoc(Scope));
   }
 
   /// Push a region covering \c Node onto the stack.
@@ -844,10 +833,24 @@ struct CoverageMapping : public ASTWalker {
     });
   }
 
-  /// Replace the current region's count by pushing an incomplete region.
-  void replaceCount(CounterExpr &&Expr, Optional<SourceLoc> Start = None) {
-    CounterExpr &Counter = createCounter(std::move(Expr));
-    RegionStack.emplace_back(ASTNode(), Counter, Start, None);
+  /// Replace the current region at \p Start with a new counter. If \p Start is
+  /// \c None, or the counter is semantically zero, an 'incomplete' region is
+  /// formed, which is not recorded unless followed by additional AST nodes.
+  void replaceCount(CounterExpr *Counter, Optional<SourceLoc> Start) {
+    // If the counter is semantically zero, form an 'incomplete' region with
+    // no starting location. This prevents forming unreachable regions unless
+    // there is a following statement or expression to extend the region.
+    if (Start && Counter->isSemanticallyZero())
+      Start = None;
+
+    RegionStack.emplace_back(ASTNode(), *Counter, Start, None);
+  }
+
+  /// Replace the current region at \p Start with a new counter. If \p Start is
+  /// \c None, or the counter is semantically zero, an 'incomplete' region is
+  /// formed, which is not recorded unless followed by additional AST nodes.
+  void replaceCount(CounterExpr &&Expr, Optional<SourceLoc> Start) {
+    replaceCount(&createCounter(std::move(Expr)), Start);
   }
 
   /// Get the location for the end of the last token in \c Node.
@@ -909,7 +912,7 @@ struct CoverageMapping : public ASTWalker {
     SourceMappingRegion &Region = getRegion();
     if (!Region.hasEndLoc())
       Region.setEndLoc(getEndLoc(S));
-    replaceCount(CounterExpr::Zero());
+    replaceCount(CounterExpr::Zero(), /*Start*/ None);
   }
 
   Expr *getConditionNode(StmtCondition SC) {
@@ -976,6 +979,8 @@ struct CoverageMapping : public ASTWalker {
     if (S->isImplicit() && S != ImplicitTopLevelBody)
       return {true, S};
 
+    // If we're in an 'incomplete' region, update it to include this node. This
+    // ensures we only create the region if needed.
     if (!RegionStack.empty())
       extendRegion(S);
 
@@ -986,7 +991,13 @@ struct CoverageMapping : public ASTWalker {
     } else if (auto *IS = dyn_cast<IfStmt>(S)) {
       if (auto *Cond = getConditionNode(IS->getCond()))
         assignCounter(Cond, CounterExpr::Ref(getCurrentCounter()));
+
+      // The counter for the if statement itself tracks the number of jumps to
+      // it by break statements.
       assignCounter(IS, CounterExpr::Zero());
+
+      // We emit a counter for the then block, and define the else block in
+      // terms of it.
       CounterExpr &ThenCounter = assignCounter(IS->getThenStmt());
       if (IS->getElseStmt())
         assignCounter(IS->getElseStmt(),
@@ -996,23 +1007,38 @@ struct CoverageMapping : public ASTWalker {
       assignCounter(GS->getBody());
 
     } else if (auto *WS = dyn_cast<WhileStmt>(S)) {
+      // The counter for the while statement itself tracks the number of jumps
+      // to it by break and continue statements.
       assignCounter(WS, CounterExpr::Zero());
+
       if (auto *E = getConditionNode(WS->getCond()))
         assignCounter(E, CounterExpr::Ref(getCurrentCounter()));
       assignCounter(WS->getBody());
 
     } else if (auto *RWS = dyn_cast<RepeatWhileStmt>(S)) {
+      // The counter for the while statement itself tracks the number of jumps
+      // to it by break and continue statements.
       assignCounter(RWS, CounterExpr::Zero());
+
       CounterExpr &BodyCounter = assignCounter(RWS->getBody());
       assignCounter(RWS->getCond(), CounterExpr::Ref(BodyCounter));
       RepeatWhileStack.push_back(RWS);
 
     } else if (auto *FES = dyn_cast<ForEachStmt>(S)) {
+      // The counter for the for statement itself tracks the number of jumps
+      // to it by break and continue statements.
       assignCounter(FES, CounterExpr::Zero());
       assignCounter(FES->getBody());
 
     } else if (auto *SS = dyn_cast<SwitchStmt>(S)) {
-      assignCounter(SS);
+      // The counter for the switch statement itself tracks the number of jumps
+      // to it by break statements, including the implicit breaks at the end of
+      // cases.
+      assignCounter(SS, CounterExpr::Zero());
+
+      assignCounter(SS->getSubjectExpr(),
+                    CounterExpr::Ref(getCurrentCounter()));
+
       // Assign counters for cases so they're available for fallthrough.
       for (CaseStmt *Case : SS->getCases())
         assignCounter(Case);
@@ -1021,7 +1047,10 @@ struct CoverageMapping : public ASTWalker {
       if (caseStmt->getParentKind() == CaseParentKind::Switch)
         pushRegion(S);
     } else if (auto *DS = dyn_cast<DoStmt>(S)) {
+      // The counter for the do statement itself tracks the number of jumps
+      // to it by break statements.
       assignCounter(DS, CounterExpr::Zero());
+
       assignCounter(DS->getBody(), CounterExpr::Ref(getCurrentCounter()));
 
     } else if (auto *DCS = dyn_cast<DoCatchStmt>(S)) {
@@ -1075,7 +1104,8 @@ struct CoverageMapping : public ASTWalker {
       Stmt *BreakTarget = BS->getTarget();
       if (auto *RWS = dyn_cast<RepeatWhileStmt>(BreakTarget)) {
         subtractFromCounter(RWS->getCond(), getCurrentCounter());
-      } else if (!isa<SwitchStmt>(BreakTarget)) {
+      } else {
+        // Update the exit counter for the target.
         addToCounter(BS->getTarget(), getCurrentCounter());
       }
 
@@ -1093,9 +1123,12 @@ struct CoverageMapping : public ASTWalker {
       replaceCount(CounterExpr::Ref(getCounter(S)), getEndLoc(S));
 
     } else if (auto caseStmt = dyn_cast<CaseStmt>(S)) {
-      if (caseStmt->getParentKind() == CaseParentKind::Switch)
+      if (caseStmt->getParentKind() == CaseParentKind::Switch) {
+        // The end of a case block is an implicit break, update the exit
+        // counter to reflect this.
+        addToCounter(caseStmt->getParentStmt(), getCurrentCounter());
         popRegions(S);
-
+      }
     } else if (auto *DCS = dyn_cast<DoCatchStmt>(S)) {
       assert(DoCatchStack.back() == DCS && "Malformed do-catch stack");
       DoCatchStack.pop_back();
@@ -1125,6 +1158,8 @@ struct CoverageMapping : public ASTWalker {
     if (isa<AbstractClosureExpr>(E) && !Parent.isNull())
       return {false, E};
 
+    // If we're in an 'incomplete' region, update it to include this node. This
+    // ensures we only create the region if needed.
     if (!RegionStack.empty())
       extendRegion(E);
 

lib/SILGen/SILGenPattern.cpp

@@ -2814,7 +2814,6 @@ void SILGenFunction::emitSwitchStmt(SwitchStmt *S) {
   emission.emitAddressOnlyAllocations();
 
   SILBasicBlock *contBB = createBasicBlock();
-  emitProfilerIncrement(S);
   JumpDest contDest(contBB, Cleanups.getCleanupsDepth(), CleanupLocation(S));
 
   LexicalScope switchScope(*this, CleanupLocation(S));

test/Profiler/coverage_switch.swift

@@ -1,19 +1,61 @@
 // RUN: %target-swift-frontend -Xllvm -sil-full-demangle -profile-generate -profile-coverage-mapping -emit-sorted-sil -emit-sil -module-name coverage_switch %s | %FileCheck %s
 // RUN: %target-swift-frontend -profile-generate -profile-coverage-mapping -emit-ir %s
 
+// CHECK-LABEL: sil hidden @$s15coverage_switch2f1yys5Int32VF : $@convention(thin) (Int32) -> ()
+// CHECK:       string_literal
+// CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 0
+// CHECK-NEXT:  int_instrprof_increment
+
+// CHECK:       integer_literal $Builtin.Int32, 1
+// CHECK:       cmp_eq_Int32
+// CHECK:       cond_br {{%[0-9]+}}, [[CASE1:bb[0-9]+]], [[NOTCASE1:bb[0-9]+]]
+
+// CHECK:       [[NOTCASE1]]
+// CHECK:       integer_literal $Builtin.Int32, 2
+// CHECK:       cmp_eq_Int32
+// CHECK:       cond_br {{%[0-9]+}}, [[CASE2:bb[0-9]+]], [[NOTCASE2:bb[0-9]+]]
+
+// CHECK:       [[NOTCASE2]]
+// CHECK-NEXT:  string_literal
+// CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 3
+// CHECK-NEXT:  int_instrprof_increment
+// CHECK-NEXT:  br [[DEFAULT:bb[0-9]+]]
+
+// CHECK:       [[CASE2]]
+// CHECK-NEXT:  string_literal
+// CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 2
+// CHECK-NEXT:  int_instrprof_increment
+// CHECK-NEXT:  br [[DEFAULT]]
+
+// CHECK:       [[DEFAULT]]
+// CHECK:       function_ref @$s15coverage_switch2f1yys5Int32VF : $@convention(thin) (Int32) -> ()
+
+// CHECK:       [[CASE1]]
+// CHECK-NEXT:  string_literal
+// CHECK-NEXT:  integer_literal $Builtin.Int64, 0
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 4
+// CHECK-NEXT:  integer_literal $Builtin.Int32, 1
+// CHECK-NEXT:  int_instrprof_increment
+
 // CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f1
-func f1(_ x : Int32) {
-  switch (x) {
-  case 1: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:10 : 2
+func f1(_ x : Int32) { // CHECK-NEXT: [[@LINE]]:22 -> [[@LINE+11]]:2 : 0
+  switch (x) {         // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:13 : 0
+  case 1:              // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:10 : 1
     break
-  case 2: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:16 : 3
+  case 2:              // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:16 : 2
     fallthrough
-  default: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:14 : (3 + 4)
+  default:             // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:14 : (2 + 3)
     f1(x - 1)
-  } // CHECK: [[@LINE]]:4 -> [[@LINE+3]]:2 : 1
+  } // CHECK-NEXT: [[@LINE]]:4 -> [[@LINE+3]]:2 : ((1 + 2) + 3)
 
   var y = x
-}
+} // CHECK-NEXT: }
 
 enum Algebraic {
   case Type1(Int32, Int32)
@@ -25,41 +67,96 @@ enum Algebraic {
 func nop() {}
 
 // CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f2
-func f2(_ x : Algebraic) -> Int32 {
-  switch(x) {
-  case let .Type1(y, z): // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:10 : 2
+func f2(_ x : Algebraic) -> Int32 { // CHECK-NEXT: [[@LINE]]:35 -> [[@LINE+15]]:2 : 0
+  switch(x) {                       // CHECK-NEXT: [[@LINE]]:9 -> [[@LINE]]:12 : 0
+  case let .Type1(y, z):            // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:10 : 1
     nop()
-  case .Type2(let b): // CHECK: [[@LINE]]:3 -> [[@LINE+2]]:16 : 3
+  case .Type2(let b):               // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+2]]:16 : 2
     nop()
     fallthrough
-  case .Type3: // CHECK: [[@LINE]]:3 -> [[@LINE+3]]:6 : (3 + 4)
-    if (false) { // CHECK: [[@LINE]]:16 -> [[@LINE+2]]:6 : 5
-      fallthrough
-    }
-  case .Type4: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:10 : (5 + 6)
+  case .Type3:                      // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+3]]:6 : (2 + 3)
+    if (false) {                    // CHECK-NEXT: [[@LINE]]:8 -> [[@LINE]]:15 : (2 + 3)
+      fallthrough                   // CHECK-NEXT: [[@LINE-1]]:16 -> [[@LINE+1]]:6 : 4
+    }                               // CHECK-NEXT: [[@LINE]]:6 -> [[@LINE]]:6 : ((2 + 3) - 4)
+  case .Type4:                      // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:10 : (4 + 5)
     break
-  } // CHECK: [[@LINE]]:4 -> [[@LINE+1]]:11 : 1
+  } // CHECK-NEXT: [[@LINE]]:4 -> [[@LINE+1]]:11 : (((1 + 2) + 3) + 5)
   return 0
-}
+} // CHECK-NEXT: }
 
 public enum Simple {
   case First, Second
 }
 
 // CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f3
-func f3(_ x : Simple) -> Int32 {
-  switch (x) {
-  case .First: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:13 : 2
+func f3(_ x : Simple) -> Int32 { // CHECK-NEXT: [[@LINE]]:32 -> [[@LINE+9]]:2 : 0
+  switch (x) {                   // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:13 : 0
+  case .First:                   // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:13 : 1
     return 1
-  case .Second: // CHECK: [[@LINE]]:3 -> [[@LINE+1]]:10 : 3
+  case .Second:                  // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:10 : 2
     break
-  } // CHECK: [[@LINE]]:4 -> [[@LINE+2]]:11 : 1
+  } // CHECK: [[@LINE]]:4 -> [[@LINE+2]]:11 : 2
 
   return 0
-}
+} // CHECK-NEXT: }
 
 f1(3)
 f2(Algebraic.Type1(1, 1))
 f2(Algebraic.Type2(false))
 f2(Algebraic.Type3)
 f3(Simple.Second)
+
+// CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f4
+func f4(_ x: Int) throws -> Int { // CHECK-NEXT: [[@LINE]]:33 -> [[@LINE+21]]:2 : 0
+  y: do {                         // CHECK-NEXT: [[@LINE]]:9 -> [[@LINE+18]]:4 : 0
+    switch x {                    // CHECK-NEXT: [[@LINE]]:12 -> [[@LINE]]:13 : 0
+    case 1, 2, 3:                 // CHECK-NEXT: [[@LINE]]:5 -> [[@LINE+4]]:18 : 1
+      if .random() {              // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:19 : 1
+        return 5                  // CHECK-NEXT: [[@LINE-1]]:20 -> [[@LINE+1]]:8 : 2
+      }                           // CHECK-NEXT: [[@LINE]]:8 -> [[@LINE+1]]:18 : (1 - 2)
+      fallthrough
+    case 4:                       // CHECK-NEXT: [[@LINE]]:5 -> [[@LINE+4]]:8 : ((1 + 3) - 2)
+      if .random() {              // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:19 : ((1 + 3) - 2)
+        struct E : Error {}       // CHECK-NEXT: [[@LINE-1]]:20 -> [[@LINE+2]]:8 : 4
+        throw E()
+      }                           // CHECK-NEXT: [[@LINE]]:8 -> [[@LINE]]:8 : (((1 + 3) - 2) - 4)
+    default:                      // CHECK-NEXT: [[@LINE]]:5 -> [[@LINE+3]]:8 : 5
+      if .random() {              // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:19 : 5
+        break y                   // CHECK-NEXT: [[@LINE-1]]:20 -> [[@LINE+1]]:8 : 6
+      }                           // CHECK-NEXT: [[@LINE]]:8 -> [[@LINE]]:8 : (5 - 6)
+    }
+    f1(0)                         // CHECK-NEXT: [[@LINE-1]]:6 -> [[@LINE+1]]:4 : (((((1 + 3) + 5) - 2) - 4) - 6)
+  }
+  return 1                        // CHECK-NEXT: [[@LINE-1]]:4 -> [[@LINE]]:11 : ((((1 + 3) + 5) - 2) - 4)
+} // CHECK-NEXT: }
+
+// CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f5
+func f5(_ x: Simple) -> Int { // CHECK-NEXT: [[@LINE]]:29 -> [[@LINE+7]]:2 : 0
+  switch x {                  // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:11 : 0
+  case .First:                // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:13 : 1
+    return 0
+  case .Second:               // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:13 : 2
+    return 1
+  }
+} // CHECK-NEXT: }
+
+// CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f6
+func f6(_ x: Simple) -> Int { // CHECK-NEXT: [[@LINE]]:29 -> [[@LINE+5]]:2 : 0
+  switch x {                  // CHECK-NEXT: [[@LINE]]:10 -> [[@LINE]]:11 : 0
+  case .First:  return 0      // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE]]:25 : 1
+  case .Second: return 1      // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE]]:25 : 2
+  }
+} // CHECK-NEXT: }
+
+// CHECK-LABEL: sil_coverage_map {{.*}}// coverage_switch.f7
+func f7() -> Int { // CHECK-NEXT: [[@LINE]]:18 -> [[@LINE+10]]:2 : 0
+  switch .random() ? Simple.First : .Second {
+    // CHECK-NEXT: [[@LINE-1]]:10 -> [[@LINE-1]]:44 : 0
+    // CHECK-NEXT: [[@LINE-2]]:22 -> [[@LINE-2]]:34 : 1
+    // CHECK-NEXT: [[@LINE-3]]:37 -> [[@LINE-3]]:44 : (0 - 1)
+  case .First: // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:13 : 2
+    return 0
+  case .Second: // CHECK-NEXT: [[@LINE]]:3 -> [[@LINE+1]]:13 : 3
+    return 1
+  }
+} // CHECK-NEXT: }