1 files changed, 15 insertions, 0 deletions
diff --git a/clang/lib/CodeGen/CodeGenPGO.cpp b/clang/lib/CodeGen/CodeGenPGO.cpp
index 349160ad352..cdd411cfa35 100644
--- a/clang/lib/CodeGen/CodeGenPGO.cpp
+++ b/clang/lib/CodeGen/CodeGenPGO.cpp
@@ -315,6 +315,10 @@ namespace {
       (*CounterMap)[S->getBody()] = NextCounter++;
       Visit(S->getBody());
     }
+    void VisitBlockDecl(const BlockDecl *S) {
+      (*CounterMap)[S->getBody()] = NextCounter++;
+      Visit(S->getBody());
+    }
     /// Assign a counter to track the block following a label.
     void VisitLabelStmt(const LabelStmt *S) {
       (*CounterMap)[S] = NextCounter++;
@@ -481,6 +485,13 @@ namespace {
       Visit(S->getBody());
     }
 
+    void VisitBlockDecl(const BlockDecl *S) {
+      RegionCounter Cnt(PGO, S->getBody());
+      Cnt.beginRegion();
+      (*CountMap)[S->getBody()] = PGO.getCurrentRegionCount();
+      Visit(S->getBody());
+    }
+
     void VisitReturnStmt(const ReturnStmt *S) {
       RecordStmtCount(S);
       if (S->getRetValue())
@@ -802,6 +813,8 @@ void CodeGenPGO::mapRegionCounters(const Decl *D) {
     Walker.VisitFunctionDecl(FD);
   else if (const ObjCMethodDecl *MD = dyn_cast_or_null<ObjCMethodDecl>(D))
     Walker.VisitObjCMethodDecl(MD);
+  else if (const BlockDecl *BD = dyn_cast_or_null<BlockDecl>(D))
+    Walker.VisitBlockDecl(BD);
   NumRegionCounters = Walker.NextCounter;
 }
 
@@ -812,6 +825,8 @@ void CodeGenPGO::computeRegionCounts(const Decl *D) {
     Walker.VisitFunctionDecl(FD);
   else if (const ObjCMethodDecl *MD = dyn_cast_or_null<ObjCMethodDecl>(D))
     Walker.VisitObjCMethodDecl(MD);
+  else if (const BlockDecl *BD = dyn_cast_or_null<BlockDecl>(D))
+    Walker.VisitBlockDecl(BD);
 }
 
 void CodeGenPGO::emitCounterVariables() {