[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

@@ -468,6 +468,10 @@ class SourceMappingRegion {
   /// The region's ending location.
   std::optional<SourceLoc> EndLoc;
 
+  /// Whether the region is within a macro expansion. Such regions do not
+  /// get recorded, but are needed to track the counters within the expansion.
+  bool IsInMacroExpansion = false;
+
   SourceMappingRegion(Kind RegionKind, std::optional<CounterExpr> Counter,
                       std::optional<SourceLoc> StartLoc)
       : RegionKind(RegionKind), Counter(Counter), StartLoc(StartLoc) {
@@ -516,6 +520,14 @@ class SourceMappingRegion {
   SourceMappingRegion(SourceMappingRegion &&Region) = default;
   SourceMappingRegion &operator=(SourceMappingRegion &&RHS) = default;
 
+  bool isInMacroExpansion() const {
+    return IsInMacroExpansion;
+  }
+
+  void setIsInMacroExpansion() {
+    IsInMacroExpansion = true;
+  }
+
   /// Whether this region is for scoping only.
   bool isForScopingOnly() const { return RegionKind == Kind::ScopingOnly; }
 
@@ -837,6 +849,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); }
 
@@ -993,6 +1011,10 @@ struct CoverageMapping : public ASTWalker {
 
   /// Push a region onto the stack.
   void pushRegion(SourceMappingRegion Region) {
+    // Note on the region whether we're currently in a macro expansion.
+    if (isInMacroExpansion())
+      Region.setIsInMacroExpansion();
+
     LLVM_DEBUG({
       llvm::dbgs() << "Pushed region: ";
       Region.print(llvm::dbgs(), SM);
@@ -1027,6 +1049,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 (Region.isInMacroExpansion())
+      return;
+
     // Don't bother recording regions that are only present for scoping.
     if (Region.isForScopingOnly())
       return;
@@ -1141,16 +1168,29 @@ struct CoverageMapping : public ASTWalker {
     if (SourceRegions.empty())
       return nullptr;
 
-    using MappedRegion = SILCoverageMap::MappedRegion;
+    auto FileSourceRange = SM.getRangeForBuffer(*SF->getBufferID());
+    auto isLocInFile = [&](SourceLoc Loc) {
+      return FileSourceRange.contains(Loc) || FileSourceRange.getEnd() == Loc;
+    };
 
+    using MappedRegion = SILCoverageMap::MappedRegion;
     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 (!isLocInFile(Range.Start) || !isLocInFile(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 {
@@ -1522,10 +1562,14 @@ struct CoverageMapping : public ASTWalker {
 
     if (hasCounter(E)) {
       pushRegion(SourceMappingRegion::forNode(E, SM));
-    } else if (isa<OptionalTryExpr>(E)) {
+    } else if (isa<OptionalTryExpr>(E) || isa<MacroExpansionExpr>(E)) {
       // If we have a `try?`, that doesn't already have a counter, record it
       // as a scoping-only region. We need it to scope child error branches,
       // but don't need it in the resulting set of regions.
+      //
+      // If we have a macro expansion, also push a scoping-only region. We'll
+      // discard any regions recorded within the macro, but will adjust for any
+      // control flow that may have happened within the macro.
       assignCounter(E, getCurrentCounter());
       pushRegion(SourceMappingRegion::scopingOnly(E, SM));
     }
@@ -1560,6 +1604,10 @@ struct CoverageMapping : public ASTWalker {
       // Already visited the children.
       return Action::SkipChildren(TE);
     }
+
+    if (isa<MacroExpansionExpr>(E))
+      MacroDepth += 1;
+
     return shouldWalkIntoExpr(E, Parent, Constant);
   }
 
@@ -1576,6 +1624,11 @@ struct CoverageMapping : public ASTWalker {
           Lexer::getLocForEndOfToken(SM, E->getEndLoc()));
     }
 
+    if (isa<MacroExpansionExpr>(E)) {
+      assert(isInMacroExpansion());
+      MacroDepth -= 1;
+    }
+
     if (hasCounter(E))
       exitRegion(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,61 @@ 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 the counter '2' is for the initialization.
+// CHECK-LABEL: sil hidden [lazy_getter] [noinline] @$s15coverage_macros2S1V2z3SiSgvg : $@convention(method) (@inout S1) -> Optional<Int>
+// CHECK:       switch_enum {{%[0-9]+}} : $Optional<Optional<Int>>, case #Optional.some!enumelt: {{bb[0-9]}}, case #Optional.none!enumelt: [[INITBB:bb[0-9]]]
+// CHECK:       [[INITBB]]
+// CHECK-NEXT:  increment_profiler_counter 2
+
+// 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
+
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.z
+  var z: Int = #id(.random() ? 3 : 4) // CHECK-NEXT: [[@LINE]]:16 -> [[@LINE]]:38 : 0
+
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.z1
+  var z1: Int? = try? #id(throwingFn()) // CHECK-NEXT: [[@LINE]]:18 -> [[@LINE]]:40 : 0
+
+  // CHECK-LABEL: sil_coverage_map{{.*}}variable initialization expression of coverage_macros.S1.z2
+  var z2: Int? = #id(try? throwingFn()) // CHECK-NEXT: [[@LINE]]:18 -> [[@LINE]]:40 : 0
+
+  // The counter for the initialization region is '2', see the match of s15coverage_macros2S1V2z3SiSgvg above.
+  // CHECK-LABEL: sil_coverage_map{{.*}}coverage_macros.S1.z3.getter
+  lazy var z3: Int? = #id(try? throwingFn()) // CHECK-NEXT: [[@LINE]]:23 -> [[@LINE]]:45 : 2
 }
 
 @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 +84,49 @@ 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
+}
+
+// CHECK-LABEL: sil_coverage_map{{.*}}s15coverage_macros5test4Si_SityKF
+func test4() throws -> (Int, Int) {                // CHECK-NEXT: [[@LINE]]:35 -> [[@LINE+3]]:2  : 0
+  let x = try #id(#fnCall(throwingFn, throwingFn)) // CHECK-NEXT: [[@LINE]]:51 -> [[@LINE+1]]:11 : ((0 - 1) - 2)
+  return x                                         // CHECK-NEXT: }
+}
+
+// 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: }
+}