From cfd38c42a1d524fa6fb879190f3a48c6d9412770 Mon Sep 17 00:00:00 2001
From: Vedant Kumar <vsk@apple.com>
Date: Tue, 18 Aug 2020 14:04:31 -0700
Subject: [PATCH] [Profiler] Assign counter to IfStmt condition

Push a region to cover any expressions within an IfStmt condition.

rdar://67280997
---
 lib/SIL/IR/SILProfiler.cpp              | 2 ++
 test/Profiler/coverage_exceptions.swift | 1 +
 test/Profiler/coverage_if.swift         | 1 +
 test/Profiler/coverage_smoke.swift      | 5 +----
 test/Profiler/coverage_toplevel.swift   | 3 ++-
 5 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/lib/SIL/IR/SILProfiler.cpp b/lib/SIL/IR/SILProfiler.cpp
index 91ba2f397d5..a5026ffa377 100644
--- a/lib/SIL/IR/SILProfiler.cpp
+++ b/lib/SIL/IR/SILProfiler.cpp
@@ -928,6 +928,8 @@ public:
         pushRegion(BS);
 
     } else if (auto *IS = dyn_cast<IfStmt>(S)) {
+      if (auto *Cond = getConditionNode(IS->getCond()))
+        assignCounter(Cond, CounterExpr::Ref(getCurrentCounter()));
       assignCounter(IS, CounterExpr::Zero());
       CounterExpr &ThenCounter = assignCounter(IS->getThenStmt());
       if (IS->getElseStmt())
diff --git a/test/Profiler/coverage_exceptions.swift b/test/Profiler/coverage_exceptions.swift
index 1762cbb1bec..aa1fd867ca0 100644
--- a/test/Profiler/coverage_exceptions.swift
+++ b/test/Profiler/coverage_exceptions.swift
@@ -70,6 +70,7 @@ func goo(_ b: Bool) -> Int { // CHECK-NEXT: [[@LINE]]:28 {{.*}} : 0
   do {                       // CHECK-NEXT: [[@LINE]]:6 -> [[@LINE+2]]:4 : 0
     throw SomeErr.Err1
   } catch {                  // CHECK-NEXT: [[@LINE]]:11 {{.*}} : 1
+                             // CHECK-NEXT: [[@LINE+1]]:8 {{.*}} : 1
     if b {                   // CHECK-NEXT: [[@LINE]]:10 {{.*}} : 2
       return 1
     }                        // CHECK-NEXT: [[@LINE]]:6 {{.*}} : (1 - 2)
diff --git a/test/Profiler/coverage_if.swift b/test/Profiler/coverage_if.swift
index 5dfb738ad45..a19813f4312 100644
--- a/test/Profiler/coverage_if.swift
+++ b/test/Profiler/coverage_if.swift
@@ -3,6 +3,7 @@
 
 // CHECK-LABEL: sil_coverage_map {{.*}}// coverage_if.foo
 func foo(x : Bool) { // CHECK: [[@LINE]]:20 -> {{[0-9]+}}:2 : 0
+           // CHECK: [[@LINE+1]]:6 -> [[@LINE+1]]:9 : 0
   if (x) { // CHECK: [[@LINE]]:10 -> [[@LINE+1]]:4 : 1
   }
 
diff --git a/test/Profiler/coverage_smoke.swift b/test/Profiler/coverage_smoke.swift
index 09df48a1678..4edf0177839 100644
--- a/test/Profiler/coverage_smoke.swift
+++ b/test/Profiler/coverage_smoke.swift
@@ -13,7 +13,6 @@
 // RUN: %llvm-profdata show %t/default.profdata -function=main | %FileCheck %s --check-prefix=CHECK-MAIN
 // RUN: %llvm-cov show %t/main -instr-profile=%t/default.profdata | %FileCheck %s --check-prefix=CHECK-COV
 // RUN: %llvm-cov report %t/main -instr-profile=%t/default.profdata -show-functions %s | %FileCheck %s --check-prefix=CHECK-REPORT
-// RUN: rm -rf %t
 
 // REQUIRES: profile_runtime
 // REQUIRES: executable_test
@@ -164,9 +163,7 @@ func catchError2(_ b: Bool) -> Int {
   do {
     throw CustomError.Err // CHECK-COV: {{ *}}[[@LINE]]|{{ *}}2
   } catch {
-    // reviews.llvm.org/D85036 regressed coverage reporting for the
-    // following line (rdar://67280997).
-    if b {                // CHECK-COV: {{ *}}[[@LINE]]|{{ *}}1
+    if b {                // CHECK-COV: {{ *}}[[@LINE]]|{{ *}}2
       return 1            // CHECK-COV: {{ *}}[[@LINE]]|{{ *}}1
     }
   }
diff --git a/test/Profiler/coverage_toplevel.swift b/test/Profiler/coverage_toplevel.swift
index a7795c6d1a7..e285e03f242 100644
--- a/test/Profiler/coverage_toplevel.swift
+++ b/test/Profiler/coverage_toplevel.swift
@@ -21,7 +21,8 @@ while (i < 10) {
 // CHECK-NEXT:  [[@LINE+1]]:21 -> [[@LINE+1]]:22 : (0 - 1)
 var i2 = true ? 1 : 0;
 
-// CHECK: sil_coverage_map{{.*}}__tlcd_line:[[@LINE+4]]:1
+// CHECK: sil_coverage_map{{.*}}__tlcd_line:[[@LINE+5]]:1
+// CHECK-NEXT:  [[@LINE+4]]:4 -> [[@LINE+4]]:10 : 0
 // CHECK-NEXT:  [[@LINE+3]]:11 -> [[@LINE+5]]:2 : 1
 // CHECK-NEXT:  [[@LINE+2]]:1 -> [[@LINE+4]]:2 : 0
 // CHECK-NEXT:  [[@LINE+3]]:2 -> [[@LINE+3]]:2 : 0