Merge pull request #61071 from hamishknight/walk-that-way

Author: hamishknight

lib/SIL/IR/SILProfiler.cpp

@@ -187,10 +187,25 @@ visitFunctionDecl(ASTWalker &Walker, AbstractFunctionDecl *AFD, F Func) {
   return ASTWalker::Action::SkipChildren();
 }
 
-/// Whether to skip visitation of an expression. Children of skipped exprs
-/// should still be visited.
-static bool skipExpr(Expr *E) {
-  return !E->getStartLoc().isValid() || !E->getEndLoc().isValid();
+/// Whether to skip visitation of an expression. If the expression should be
+/// skipped, a walker action is returned that determines whether or not the
+/// children should also be skipped.
+static Optional<ASTWalker::PreWalkResult<Expr *>>
+shouldSkipExpr(Expr *E, ASTWalker::ParentTy Parent) {
+  using Action = ASTWalker::Action;
+  using Result = ASTWalker::PreWalkResult<Expr *>;
+
+  // Profiling for closures should be handled separately. Do not visit
+  // closure expressions twice.
+  if (isa<AbstractClosureExpr>(E) && !Parent.isNull())
+    return Result(Action::SkipChildren(E));
+
+  // Expressions with no location should be skipped, but we still want to visit
+  // their children.
+  if (E->getStartLoc().isInvalid() || E->getEndLoc().isInvalid())
+    return Result(Action::Continue(E));
+
+  return None;
 }
 
 /// Whether the children of an unmapped decl should still be walked.
@@ -262,14 +277,16 @@ struct MapRegionCounters : public ASTWalker {
     return Action::Continue(S);
   }
 
-  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-    if (skipExpr(E))
-      return Action::Continue(E);
+  PreWalkAction walkToParameterListPre(ParameterList *PL) override {
+    // We don't walk into parameter lists. Default arguments should be visited
+    // directly.
+    // FIXME: We don't yet profile default argument generators at all.
+    return Action::SkipChildren();
+  }
 
-    // Profiling for closures should be handled separately. Do not visit
-    // closure expressions twice.
-    if (isa<AbstractClosureExpr>(E) && !Parent.isNull())
-      return Action::SkipChildren(E);
+  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+    if (auto SkipAction = shouldSkipExpr(E, Parent))
+      return *SkipAction;
 
     // If AST visitation begins with an expression, the counter map must be
     // empty. Set up a counter for the root.
@@ -648,14 +665,16 @@ struct PGOMapping : public ASTWalker {
     return Action::Continue(S);
   }
 
-  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-    if (skipExpr(E))
-      return Action::Continue(E);
+  PreWalkAction walkToParameterListPre(ParameterList *PL) override {
+    // We don't walk into parameter lists. Default arguments should be visited
+    // directly.
+    // FIXME: We don't yet profile default argument generators at all.
+    return Action::SkipChildren();
+  }
 
-    // Profiling for closures should be handled separately. Do not visit
-    // closure expressions twice.
-    if (isa<AbstractClosureExpr>(E) && !Parent.isNull())
-      return Action::SkipChildren(E);
+  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+    if (auto SkipAction = shouldSkipExpr(E, Parent))
+      return *SkipAction;
 
     unsigned parent = getParentCounter();
 
@@ -1146,14 +1165,21 @@ struct CoverageMapping : public ASTWalker {
     return Action::Continue(S);
   }
 
-  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
-    if (skipExpr(E))
-      return Action::Continue(E);
+  PreWalkAction walkToParameterListPre(ParameterList *PL) override {
+    // We don't walk into parameter lists. Default arguments should be visited
+    // directly.
+    // FIXME: We don't yet generate coverage for default argument generators at
+    // all. This is inconsistent with property initializers, which are
+    // effectively default values too. Seems like coverage doesn't offer much
+    // benefit in these cases, as they're unlikely to have side effects, and
+    // the values can be exercized explicitly, but we should probably at least
+    // have a consistent behavior for both no matter what we choose here.
+    return Action::SkipChildren();
+  }
 
-    // Profiling for closures should be handled separately. Do not visit
-    // closure expressions twice.
-    if (isa<AbstractClosureExpr>(E) && !Parent.isNull())
-      return Action::SkipChildren(E);
+  PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
+    if (auto SkipAction = shouldSkipExpr(E, Parent))
+      return *SkipAction;
 
     // If we're in an 'incomplete' region, update it to include this node. This
     // ensures we only create the region if needed.
@@ -1166,28 +1192,28 @@ struct CoverageMapping : public ASTWalker {
       assert(RegionStack.empty() &&
              "Mapped a region before visiting the root?");
       assignCounter(E);
-      pushRegion(E);
-    }
-
-    // If there isn't an active region, we may be visiting a default
-    // initializer for a function argument.
-    if (!RegionStack.empty()) {
-      if (auto *IE = dyn_cast<IfExpr>(E)) {
-        CounterExpr &ThenCounter = assignCounter(IE->getThenExpr());
-        assignCounter(IE->getElseExpr(),
-                      CounterExpr::Sub(getCurrentCounter(), ThenCounter));
-      }
     }
 
     if (isa<LazyInitializerExpr>(E))
       assignCounter(E);
 
-    if (hasCounter(E) && !Parent.isNull())
+    if (hasCounter(E))
       pushRegion(E);
+
+    assert(!RegionStack.empty() && "Must be within a region");
+
+    if (auto *IE = dyn_cast<IfExpr>(E)) {
+      CounterExpr &ThenCounter = assignCounter(IE->getThenExpr());
+      assignCounter(IE->getElseExpr(),
+                    CounterExpr::Sub(getCurrentCounter(), ThenCounter));
+    }
     return Action::Continue(E);
   }
 
   PostWalkResult<Expr *> walkToExprPost(Expr *E) override {
+    if (shouldSkipExpr(E, Parent))
+      return Action::Continue(E);
+
     if (hasCounter(E))
       popRegions(E);
 

test/Profiler/coverage_pc_macro.swift

@@ -0,0 +1,21 @@
+// RUN: %target-swift-frontend -Xllvm -sil-full-demangle -pc-macro -profile-generate -profile-coverage-mapping -emit-sorted-sil -emit-sil -module-name coverage_pc_macro %s | %FileCheck %s
+// RUN: %target-swift-frontend -pc-macro -profile-generate -profile-coverage-mapping -emit-ir %s
+
+// Make sure the PCMacro transform doesn't cause coverage to crash.
+
+public func __builtin_pc_before(_ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) {
+  print("[\(moduleId):\(fileId)] [\(sl):\(sc)-\(el):\(ec)] pc before")
+}
+
+public func __builtin_pc_after(_ sl : Int, _ el : Int, _ sc : Int, _ ec: Int, _ moduleId : Int, _ fileId : Int) {
+  print("[\(moduleId):\(fileId)] [\(sl):\(sc)-\(el):\(ec)] pc after")
+}
+
+// CHECK-LABEL: sil_coverage_map {{.*}} "$s17coverage_pc_macro3fooySiSbF"
+func foo(_ x: Bool) -> Int {
+  if x {
+    return 5
+  } else {
+    return 7
+  }
+}