[Coverage] Consolidate visitation logic for functions and nominal types (swiftlang#16012)

This should address an ASAN failure which arose due to the PGOMapping
ASTWalker not being updated in sync with the other profiling-related
walkers.

rdar://39534066

Author: vedantk

lib/SIL/SILProfiler.cpp

@@ -157,6 +157,33 @@ std::pair<bool, Expr *> visitClosureExpr(ASTWalker &Walker,
   return {true, CE};
 }
 
+/// Special logic for handling function visitation.
+///
+/// To avoid creating duplicate mappings, a function decl is only profiled if
+/// it hasn't been reached via recursive walk, or if it's a constructor for a
+/// nominal type (these are profiled in a group).
+///
+/// Apply \p Func is the function can be visited.
+template <typename F>
+bool visitFunctionDecl(ASTWalker &Walker, AbstractFunctionDecl *AFD, F Func) {
+  bool continueWalk = Walker.Parent.isNull() || isa<ConstructorDecl>(AFD);
+  if (continueWalk)
+    Func();
+  return continueWalk;
+}
+
+/// Special logic for handling nominal type visitation.
+///
+/// Apply \p Func if the nominal type can be visited (i.e it has not been
+/// reached via recursive walk).
+template <typename F>
+bool visitNominalTypeDecl(ASTWalker &Walker, NominalTypeDecl *NTD, F Func) {
+  bool continueWalk = Walker.Parent.isNull();
+  if (continueWalk)
+    Func();
+  return continueWalk;
+}
+
 /// An ASTWalker that maps ASTNodes to profiling counters.
 struct MapRegionCounters : public ASTWalker {
   /// The next counter value to assign.
@@ -176,18 +203,13 @@ struct MapRegionCounters : public ASTWalker {
       return false;
 
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      // Don't map a nested function unless it's a constructor.
-      bool continueWalk = Parent.isNull() || isa<ConstructorDecl>(AFD);
-      if (continueWalk)
-        CounterMap[AFD->getBody()] = NextCounter++;
-      return continueWalk;
+      return visitFunctionDecl(
+          *this, AFD, [&] { CounterMap[AFD->getBody()] = NextCounter++; });
     } else if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
       CounterMap[TLCD->getBody()] = NextCounter++;
-    } else if (isa<NominalTypeDecl>(D)) {
-      bool continueWalk = Parent.isNull();
-      if (continueWalk)
-        WithinNominalType = true;
-      return continueWalk;
+    } else if (auto *NTD = dyn_cast<NominalTypeDecl>(D)) {
+      return visitNominalTypeDecl(*this, NTD,
+                                  [&] { WithinNominalType = true; });
     }
     return true;
   }
@@ -411,17 +433,22 @@ struct PGOMapping : public ASTWalker {
     if (isUnmapped(D))
       return false;
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      auto node = AFD->getBody();
-      CounterMap[node] = NextCounter++;
-      auto count = loadExecutionCount(node);
-      LoadedCounterMap[node] = count;
+      return visitFunctionDecl(*this, AFD, [&] {
+        auto node = AFD->getBody();
+        CounterMap[node] = NextCounter++;
+        auto count = loadExecutionCount(node);
+        LoadedCounterMap[node] = count;
+      });
     }
     if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
       auto node = TLCD->getBody();
       CounterMap[node] = NextCounter++;
       auto count = loadExecutionCount(node);
       LoadedCounterMap[node] = count;
     }
+    if (auto *NTD = dyn_cast<NominalTypeDecl>(D)) {
+      return visitNominalTypeDecl(*this, NTD, [&] {});
+    }
     return true;
   }
 
@@ -768,26 +795,21 @@ struct CoverageMapping : public ASTWalker {
       return false;
 
     if (auto *AFD = dyn_cast<AbstractFunctionDecl>(D)) {
-      // Don't map a nested function unless it's a constructor.
-      bool continueWalk = Parent.isNull() || isa<ConstructorDecl>(AFD);
-      if (continueWalk) {
+      return visitFunctionDecl(*this, AFD, [&] {
         CounterExpr &funcCounter = assignCounter(AFD->getBody());
 
         if (isa<ConstructorDecl>(AFD))
           addToCounter(ParentNominalType, funcCounter);
-      }
-      return continueWalk;
+      });
     } else if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
       assignCounter(TLCD->getBody());
       ImplicitTopLevelBody = TLCD->getBody();
     } else if (auto *NTD = dyn_cast<NominalTypeDecl>(D)) {
-      bool continueWalk = Parent.isNull();
-      if (continueWalk) {
+      return visitNominalTypeDecl(*this, NTD, [&] {
         ParentNominalType = NTD;
         assignCounter(NTD, CounterExpr::Zero());
         pushRegion(NTD);
-      }
-      return continueWalk;
+      });
     }
     return true;
   }