[Coverage] Avoid recording regions from macro expansions

Ignore any regions recorded while inside a macro
expansion, but account for any control flow that
may have happened such that the exit count is
correctly adjusted. This allows e.g throwing function
calls to behave correctly within the expansion.

rdar://129081384

Author: hamishknight

lib/SIL/IR/SILProfiler.cpp

@@ -837,6 +837,12 @@ struct CoverageMapping : public ASTWalker {
 
   Stmt *ImplicitTopLevelBody = nullptr;
 
+  /// The number of parent MacroExpansionExprs.
+  unsigned MacroDepth = 0;
+
+  /// Whether the current walk is within a macro expansion.
+  bool isInMacroExpansion() const { return MacroDepth > 0; }
+
   /// Return true if \c Ref has an associated counter.
   bool hasCounter(ProfileCounterRef Ref) { return CounterExprs.count(Ref); }
 
@@ -1027,6 +1033,11 @@ struct CoverageMapping : public ASTWalker {
       llvm::dbgs() << "\n";
     });
 
+    // Don't record regions in macro expansions, they don't have source
+    // locations that can be meaningfully mapped to source code.
+    if (isInMacroExpansion())
+      return;
+
     // Don't bother recording regions that are only present for scoping.
     if (Region.isForScopingOnly())
       return;
@@ -1142,15 +1153,26 @@ struct CoverageMapping : public ASTWalker {
       return nullptr;
 
     using MappedRegion = SILCoverageMap::MappedRegion;
+    auto FileSourceRange = SM.getRangeForBuffer(*SF->getBufferID());
 
     std::vector<MappedRegion> Regions;
     SourceRange OuterRange;
     for (const auto &Region : SourceRegions) {
       assert(Region.hasStartLoc() && "invalid region");
       assert(Region.hasEndLoc() && "incomplete region");
 
-      // Build up the outer range from the union of all coverage regions.
       SourceRange Range(Region.getStartLoc(), Region.getEndLoc());
+
+      // Make sure we haven't ended up with any source locations outside the
+      // SourceFile (e.g for generated code such as macros), asserting in an
+      // asserts build, dropping in a non-asserts build.
+      if (!FileSourceRange.contains(Range.Start) ||
+          !FileSourceRange.contains(Range.End)) {
+        assert(false && "range outside of file");
+        continue;
+      }
+
+      // Build up the outer range from the union of all coverage regions.
       if (!OuterRange) {
         OuterRange = Range;
       } else {
@@ -1517,6 +1539,16 @@ struct CoverageMapping : public ASTWalker {
       assignKnownCounter(E);
     }
 
+    // If we have a macro expansion, push a scoping-only region, and note that
+    // we're in a macro expansion. We'll discard any regions recorded within the
+    // macro, but will adjust for any control flow that may have happened within
+    // the macro.
+    if (isa<MacroExpansionExpr>(E)) {
+      assignCounter(E, getCurrentCounter());
+      pushRegion(SourceMappingRegion::scopingOnly(E, SM));
+      MacroDepth += 1;
+    }
+
     if (isa<LazyInitializerExpr>(E))
       assignKnownCounter(E);
 
@@ -1579,6 +1611,11 @@ struct CoverageMapping : public ASTWalker {
     if (hasCounter(E))
       exitRegion(E);
 
+    if (isa<MacroExpansionExpr>(E)) {
+      assert(isInMacroExpansion());
+      MacroDepth -= 1;
+    }
+
     return Action::Continue(E);
   }
 };

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -212,6 +212,25 @@ public class NestedDeclInExprMacro: ExpressionMacro {
   }
 }
 
+public class NullaryFunctionCallMacro: ExpressionMacro {
+  public static func expansion(
+    of macro: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) -> ExprSyntax {
+    let calls = macro.arguments.compactMap(\.expression).map { "\($0)()" }
+    return "(\(raw: calls.joined(separator: ", ")))"
+  }
+}
+
+public class TupleMacro: ExpressionMacro {
+  public static func expansion(
+    of macro: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) -> ExprSyntax {
+    return "(\(raw: macro.arguments.map { "\($0)" }.joined()))"
+  }
+}
+
 enum CustomError: Error, CustomStringConvertible {
   case message(String)
 

test/Profiler/coverage_macros.swift

@@ -19,20 +19,41 @@ macro addMembers() = #externalMacro(module: "MacroDefinition", type: "AddMembers
 @freestanding(expression)
 macro nestedDeclInExpr() -> () -> Void = #externalMacro(module: "MacroDefinition", type: "NestedDeclInExprMacro")
 
+@freestanding(expression)
+macro fnCall<T>(_: () throws -> T) -> T = #externalMacro(module: "MacroDefinition", type: "NullaryFunctionCallMacro")
+
+@freestanding(expression)
+macro fnCall<T>(_: () throws -> T, _: () throws -> T) -> (T, T) = #externalMacro(module: "MacroDefinition", type: "NullaryFunctionCallMacro")
+
+@freestanding(expression)
+macro id<T>(_: T) -> T = #externalMacro(module: "MacroDefinition", type: "TupleMacro")
+
+// This needs to be matched up here due to the sorting of the SIL; just make
+// sure we emit the counter increments for the error branches.
+// CHECK-LABEL: sil hidden @$s15coverage_macros5test2Si_SityKF
+// CHECK:       increment_profiler_counter 0,{{.*}}s15coverage_macros5test2Si_SityKF
+// CHECK:       {{bb[0-9]+}}({{%[0-9]+}} : $any Error):
+// CHECK-NEXT:    increment_profiler_counter 1,{{.*}}s15coverage_macros5test2Si_SityKF
+// CHECK:       {{bb[0-9]+}}({{%[0-9]+}} : $any Error):
+// CHECK-NEXT:    increment_profiler_counter 2,{{.*}}s15coverage_macros5test2Si_SityKF
+
 // Note we use implicit-check-not, so this test ensures we don't emit
 // coverage maps for the macro expansions.
 
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros10throwingFnSiyKF
+func throwingFn() throws -> Int { 0 }
+
 struct S1 {
-  // CHECK: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.x
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.x
   var x: Int = 0
 
-  // CHECK: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.y
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.y
   var y: Int = 0
 }
 
 @addMembers
 struct S2 {
-  // CHECK: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S2.(_storage
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S2.(_storage
   private var _storage = S1()
 
   @accessViaStorage
@@ -43,7 +64,43 @@ struct S2 {
   var y: Int = 17
 }
 
-// CHECK: sil_coverage_map{{.*}}s15coverage_macros3fooyyF
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros3fooyyF
 func foo() {
   _ = #nestedDeclInExpr()
 }
+
+// For cases where control flow happens in the macro expansion, we drop
+// all the regions that occur within the expansion, but account for any
+// control flow changes when exiting the expansion.
+//
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros5test1yyKF
+func test1() throws {         // CHECK-NEXT: [[@LINE]]:21 -> [[@LINE+2]]:2 : 0
+  _ = try #fnCall(throwingFn) // CHECK-NEXT: [[@LINE]]:30 -> [[@LINE+1]]:2 : (0 - 1)
+}                             // CHECK-NEXT: }
+
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros5test2Si_SityKF
+func test2() throws -> (Int, Int) {            // CHECK-NEXT: [[@LINE]]:35 -> [[@LINE+3]]:2  : 0
+  let x = try #fnCall(throwingFn, throwingFn)  // CHECK-NEXT: [[@LINE]]:46 -> [[@LINE+1]]:11 : ((0 - 1) - 2)
+  return x                                     // CHECK-NEXT: }
+}
+
+// In this case the control flow is entirely contained within the macro, with
+// the same exit count, so no change.
+//
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros5test3SiyF
+func test3() -> Int {            // CHECK-NEXT: [[@LINE]]:21 -> [[@LINE+3]]:2 : 0
+  let x = #id(.random() ? 3 : 4) // CHECK-NEXT: }
+  return x
+}
+
+// 0: entry, 1: first else, 2: first error branch, 3: second else,
+// 4: second error branch, 5: third else, 6: third error branch,
+// 7: fourth error branch
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros5test5Si_S2ityKF
+func test5() throws -> (Int, Int, Int) {        // CHECK-NEXT: [[@LINE]]:40 -> [[@LINE+6]]:2  : 0
+  let x = #id(.random() ? try throwingFn() : 4) // CHECK-NEXT: [[@LINE]]:48 -> [[@LINE+4]]:19 : (0 - 2)
+  let y = #id(.random() ? 5 : try throwingFn()) // CHECK-NEXT: [[@LINE]]:48 -> [[@LINE+3]]:19 : ((0 - 2) - 4)
+  let z = #id(.random() ? try throwingFn()
+                        : try throwingFn())     // CHECK-NEXT: [[@LINE]]:44 -> [[@LINE+1]]:19 : ((((0 - 2) - 4) - 6) - 7)
+  return (x, y, z)                              // CHECK-NEXT: }
+}