[ASTScope] Match parent scopes in lookupCatchNode for brace stmts

Rather than looking for a given BraceStmtScope child for a particular
catch node, check whether the given catch node is the parent of the
innermost BraceStmtScope we've found, looking through an intermediate
source file if needed. This ensures it works correctly when the
BraceStmtScope is in a macro expansion.

rdar://149036108

Author: hamishknight

lib/AST/ASTScopeLookup.cpp

@@ -703,42 +703,24 @@ lookupEnclosingABIAttributeScope(SourceFile *sourceFile, SourceLoc loc) {
   return nullptr;
 }
 
-static std::pair<CatchNode, const BraceStmtScope *>
-getCatchNodeBody(const ASTScopeImpl *scope, CatchNode node) {
-  const auto &children = scope->getChildren();
-  if (children.empty())
-    return { CatchNode(), nullptr };
-
-  auto stmt = dyn_cast<BraceStmtScope>(children[0]);
-  if (!stmt || stmt->getStmt()->empty())
-    return { CatchNode(), nullptr };
-
-  return std::make_pair(node, stmt);
-}
-
 /// Retrieve the catch node associated with this scope, if any.
-static std::pair<CatchNode, const BraceStmtScope *>
-getCatchNode(const ASTScopeImpl *scope) {
+static CatchNode getCatchNode(const ASTScopeImpl *scope) {
   // Closures introduce a catch scope for errors initiated in their body.
   if (auto closureParams = dyn_cast<ClosureParametersScope>(scope)) {
-    if (auto closure = dyn_cast<ClosureExpr>(closureParams->closureExpr)) {
-      return getCatchNodeBody(scope, const_cast<ClosureExpr *>(closure));
-    }
+    if (auto closure = dyn_cast<ClosureExpr>(closureParams->closureExpr))
+      return closure;
   }
 
   // Functions introduce a catch scope for errors initiated in their body.
-  if (auto function = dyn_cast<FunctionBodyScope>(scope)) {
-    return getCatchNodeBody(
-        scope, const_cast<AbstractFunctionDecl *>(function->decl));
-  }
+  if (auto function = dyn_cast<FunctionBodyScope>(scope))
+    return function->decl;
 
   // Do..catch blocks introduce a catch scope for errors initiated in the `do`
   // body.
-  if (auto doCatch = dyn_cast<DoCatchStmtScope>(scope)) {
-    return getCatchNodeBody(scope, const_cast<DoCatchStmt *>(doCatch->stmt));
-  }
+  if (auto doCatch = dyn_cast<DoCatchStmtScope>(scope))
+    return doCatch->stmt;
 
-  return { CatchNode(), nullptr };
+  return CatchNode();
 }
 
 /// Check whether the given location precedes the start of the catch location
@@ -767,15 +749,22 @@ CatchNode ASTScopeImpl::lookupCatchNode(ModuleDecl *module, SourceLoc loc) {
   ASTScopeAssert(innermost->getWasExpanded(),
                  "If looking in a scope, it must have been expanded.");
 
-  // Look for a body scope that's the
+  // Look for a body scope that's the direct descendent of a catch node.
   const BraceStmtScope *innerBodyScope = nullptr;
   for (auto scope = innermost; scope; scope = scope->getParent().getPtrOrNull()) {
     // If we are at a catch node and in the body of the region from which that
     // node catches thrown errors, we have our result.
-    auto caught = getCatchNode(scope);
-    if (caught.first && caught.second == innerBodyScope &&
-        !locationIsPriorToStartOfCatchScope(loc, caught.first)) {
-      return caught.first;
+    if (innerBodyScope && innerBodyScope->getParent() == scope) {
+      // For a macro expansion, we may have an intermediate source file scope,
+      // we can look through it.
+      auto catchScope = scope;
+      if (auto *sfScope = dyn_cast<ASTSourceFileScope>(catchScope)) {
+        if (auto parent = sfScope->getParent())
+          catchScope = parent.get();
+      }
+      auto caught = getCatchNode(catchScope);
+      if (caught && !locationIsPriorToStartOfCatchScope(loc, caught))
+          return caught;
     }
 
     // If this is a try scope for a try! or try?, it catches the error.

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2683,6 +2683,16 @@ public struct BodyMacroWithControlFlow: BodyMacro {
   }
 }
 
+struct ThrowCancellationMacro: BodyMacro {
+  static func expansion(
+    of node: AttributeSyntax,
+    providingBodyFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    ["throw CancellationError()"]
+  }
+}
+
 @_spi(ExperimentalLanguageFeature)
 public struct TracedPreambleMacro: PreambleMacro {
   public static func expansion(

test/Macros/macro_expand.swift

@@ -134,6 +134,19 @@ class HasStoredPropertyClassInvalid {
   #AddStoredProperty((Self.self, 0).1) // expected-note {{in expansion of macro 'AddStoredProperty' here}}
   // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX[[@LINE-2]]_2_33_{{.*}}AddStoredPropertyfMf_.swift:1:22: error: covariant 'Self' type cannot be referenced from a stored property initializer
 }
+
+@attached(body)
+public macro ThrowCancellation() = #externalMacro(module: "MacroDefinition", type: "ThrowCancellationMacro")
+
+// https://github.com/swiftlang/swift/issues/79039 - Make sure we diagnose the
+// error mismatch.
+@ThrowCancellation // expected-note {{in expansion of macro 'ThrowCancellation' on global function 'issue79039()' here}}
+func issue79039() throws(DecodingError)
+// CHECK-DIAGS: @__swiftmacro_9MacroUser10issue7903917ThrowCancellationfMb_.swift:2:11: error: thrown expression type 'CancellationError' cannot be converted to error type 'DecodingError'
+
+@ThrowCancellation // expected-note {{in expansion of macro 'ThrowCancellation' on global function 'issue79039_2()' here}}
+func issue79039_2() throws(DecodingError) {}
+// CHECK-DIAGS: @__swiftmacro_9MacroUser12issue79039_217ThrowCancellationfMb_.swift:2:11: error: thrown expression type 'CancellationError' cannot be converted to error type 'DecodingError'
 #endif
 
 @freestanding(declaration)