[PGO] Implement correct calculation of region execution count for `br…

…eak label;` and `continue label;`.

gen/pgo.cpp

@@ -328,6 +328,16 @@ struct ComputeRegionCounts : public RecursiveVisitor {
   };
   llvm::SmallVector<BreakContinue, 8> BreakContinueStack;
 
+  struct LoopLabel {
+    // If a label is used as break/continue target, this struct stores the
+    // BreakContinue stack index at the label point
+    LabelStatement *label;
+    size_t stackindex;
+    LoopLabel(LabelStatement *_label, size_t index)
+        : label(_label), stackindex(index) {}
+  };
+  llvm::SmallVector<LoopLabel, 8> LoopLabels;
+
   ComputeRegionCounts(llvm::DenseMap<const RootObject *, uint64_t> &CountMap,
                       CodeGenPGO &PGO)
       : PGO(PGO), RecordNextStmtCount(false), CountMap(CountMap) {}
@@ -381,22 +391,55 @@ struct ComputeRegionCounts : public RecursiveVisitor {
     // Counter tracks the block following the label.
     uint64_t BlockCount = setCount(PGO.getRegionCount(S));
     CountMap[S] = BlockCount;
+
+    // For each label pointing to a loop, store the current index of
+    // BreakContinueStack. This is needed for `break label;` and `continue
+    // label;` statements in loops.
+    // Assume all labels point to loops. (TODO: find predicate to filter which labels to add)
+    LoopLabels.push_back(LoopLabel(S, BreakContinueStack.size()));
+
     recurse(S->statement);
   }
 
   void visit(BreakStatement *S) override {
     RecordStmtCount(S);
     assert(!BreakContinueStack.empty() && "break not in a loop or switch!");
-    // FIXME: implement D-style break statements that break to a label
-    BreakContinueStack.back().BreakCount += CurrentCount;
+
+    if (S->target) {
+      auto it = std::find_if(
+          LoopLabels.begin(), LoopLabels.end(),
+          [S](const LoopLabel &LL) { return LL.label == S->target; });
+      assert(it != LoopLabels.end() && "It is not possible to break to a label "
+                                       "that has not been visited yet");
+      auto LL = *it;
+      assert(LL.stackindex < BreakContinueStack.size());
+      BreakContinueStack[LL.stackindex].BreakCount += CurrentCount;
+    } else {
+      BreakContinueStack.back().BreakCount += CurrentCount;
+    }
+
     CurrentCount = 0;
     RecordNextStmtCount = true;
   }
 
   void visit(ContinueStatement *S) override {
     RecordStmtCount(S);
     assert(!BreakContinueStack.empty() && "continue stmt not in a loop!");
-    BreakContinueStack.back().ContinueCount += CurrentCount;
+
+    if (S->target) {
+      auto it = std::find_if(
+          LoopLabels.begin(), LoopLabels.end(),
+          [S](const LoopLabel &LL) { return LL.label == S->target; });
+      assert(it != LoopLabels.end() &&
+             "It is not possible to continue to a label "
+             "that has not been visited yet");
+      auto LL = *it;
+      assert(LL.stackindex < BreakContinueStack.size());
+      BreakContinueStack[LL.stackindex].ContinueCount += CurrentCount;
+    } else {
+      BreakContinueStack.back().ContinueCount += CurrentCount;
+    }
+
     CurrentCount = 0;
     RecordNextStmtCount = true;
   }

tests/PGO/break.d

@@ -0,0 +1,106 @@
+// Test calculation of execution counts with loops with break-to-label and continue-to-label.
+
+// RUN: %ldc -c -output-ll -fprofile-instr-generate -of=%t.ll %s && FileCheck %s --check-prefix=PROFGEN < %t.ll
+
+// RUN: %ldc -fprofile-instr-generate=%t.profraw -run %s  \
+// RUN:   &&  %profdata merge %t.profraw -o %t.profdata \
+// RUN:   &&  %ldc -boundscheck=off -c -output-ll -of=%t2.ll -fprofile-instr-use=%t.profdata %s \
+// RUN:   &&  FileCheck %s -check-prefix=PROFUSE < %t2.ll
+
+extern(C):  // simplify name mangling for simpler string matching
+
+// PROFGEN-DAG: @[[BREAK:__(llvm_profile_counters|profc)_testbreak]] = private global [8 x i64] zeroinitializer
+// PROFGEN-DAG: @[[CONT:__(llvm_profile_counters|profc)_testcontinue]] = private global [8 x i64] zeroinitializer
+
+// PROFGEN-LABEL: @testbreak({{.*}})
+// PROFUSE-LABEL: @testbreak({{.*}})
+// PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 0
+// PROFUSE-SAME: !prof ![[BREAK0:[0-9]+]]
+void testbreak(bool a) {
+
+  // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 1
+outer:
+  // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 2
+  // PROFUSE: br {{.*}} !prof ![[BREAK2:[0-9]+]]
+  foreach (i; 0..4) {
+    // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 3
+    // PROFUSE: br {{.*}} !prof ![[BREAK3:[0-9]+]]
+    foreach (j; 0..4) {
+      // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 4
+      // PROFUSE: br {{.*}} !prof ![[BREAK4:[0-9]+]]
+      if (i>0)
+        break outer;
+
+      // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 5
+      // PROFUSE: br {{.*}} !prof ![[BREAK5:[0-9]+]]
+      if (a) {}
+    }
+
+    // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 6
+    // PROFUSE: br {{.*}} !prof ![[BREAK6:[0-9]+]]
+    if (a) {}
+  }
+
+  // PROFGEN: store {{.*}} @[[BREAK]], i64 0, i64 7
+  // PROFUSE: br {{.*}} !prof ![[BREAK7:[0-9]+]]
+  if (a) {}
+}
+
+// PROFGEN-LABEL: @testcontinue({{.*}})
+// PROFUSE-LABEL: @testcontinue({{.*}})
+// PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 0
+// PROFUSE-SAME: !prof ![[CONT0:[0-9]+]]
+void testcontinue(bool a) {
+
+  // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 1
+outer:
+  // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 2
+  // PROFUSE: br {{.*}} !prof ![[CONT2:[0-9]+]]
+  foreach (i; 0..4) {
+    // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 3
+    // PROFUSE: br {{.*}} !prof ![[CONT3:[0-9]+]]
+    foreach (j; 0..4) {
+      // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 4
+      // PROFUSE: br {{.*}} !prof ![[CONT4:[0-9]+]]
+      if (i>0)
+        continue outer;
+
+      // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 5
+      // PROFUSE: br {{.*}} !prof ![[CONT5:[0-9]+]]
+      if (a) {}
+    }
+
+    // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 6
+    // PROFUSE: br {{.*}} !prof ![[CONT6:[0-9]+]]
+    if (a) {}
+  }
+
+  // PROFGEN: store {{.*}} @[[CONT]], i64 0, i64 7
+  // PROFUSE: br {{.*}} !prof ![[CONT7:[0-9]+]]
+  if (a) {}
+}
+
+// PROFGEN-LABEL: @_Dmain(
+// PROFUSE-LABEL: @_Dmain(
+extern(D):
+void main() {
+  testbreak(false);
+  testcontinue(false);
+}
+
+
+// PROFUSE-DAG: ![[BREAK0]] = !{!"function_entry_count", i64 1}
+// PROFUSE-DAG: ![[BREAK2]] = !{!"branch_weights", i32 3, i32 1}
+// PROFUSE-DAG: ![[BREAK3]] = !{!"branch_weights", i32 6, i32 2}
+// PROFUSE-DAG: ![[BREAK4]] = !{!"branch_weights", i32 2, i32 5}
+// PROFUSE-DAG: ![[BREAK5]] = !{!"branch_weights", i32 1, i32 5}
+// PROFUSE-DAG: ![[BREAK6]] = !{!"branch_weights", i32 1, i32 2}
+// PROFUSE-DAG: ![[BREAK7]] = !{!"branch_weights", i32 1, i32 2}
+
+// PROFUSE-DAG: ![[CONT0]] = !{!"function_entry_count", i64 1}
+// PROFUSE-DAG: ![[CONT2]] = !{!"branch_weights", i32 5, i32 2}
+// PROFUSE-DAG: ![[CONT3]] = !{!"branch_weights", i32 8, i32 2}
+// PROFUSE-DAG: ![[CONT4]] = !{!"branch_weights", i32 4, i32 5}
+// PROFUSE-DAG: ![[CONT5]] = !{!"branch_weights", i32 1, i32 5}
+// PROFUSE-DAG: ![[CONT6]] = !{!"branch_weights", i32 1, i32 2}
+// PROFUSE-DAG: ![[CONT7]] = !{!"branch_weights", i32 1, i32 2}