Merge pull request #70580 from hborla/macro-original-source-range

[Macros] Always use `GeneratedSourceInfo::originalSourceRange` for the insertion range of a macro expansion.

Author: hborla

include/swift/AST/SourceFile.h

@@ -581,6 +581,11 @@ class SourceFile final : public FileUnit {
   /// the \c SourceFileKind is \c MacroExpansion.
   ASTNode getMacroExpansion() const;
 
+  /// For source files created to hold the source code for a macro
+  /// expansion, this is the original source range replaced by the macro
+  /// expansion.
+  SourceRange getMacroInsertionRange() const;
+
   /// For source files created to hold the source code created by expanding
   /// an attached macro, this is the custom attribute that describes the macro
   /// expansion.

lib/AST/ASTScopeCreation.cpp

@@ -271,7 +271,6 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
   if (auto enclosingSF = SF->getEnclosingSourceFile()) {
     SourceLoc parentLoc;
     auto macroRole = SF->getFulfilledMacroRole();
-    auto expansion = SF->getMacroExpansion();
 
     // Determine the parent source location based on the macro role.
     AbstractFunctionDecl *bodyForDecl = nullptr;
@@ -283,35 +282,20 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     case MacroRole::MemberAttribute:
     case MacroRole::Conformance:
     case MacroRole::Extension:
-      parentLoc = expansion.getStartLoc();
+    case MacroRole::Member:
+    case MacroRole::Peer:
+    case MacroRole::Preamble:
+      parentLoc = SF->getMacroInsertionRange().End;;
       break;
-    case MacroRole::Preamble: {
-      // Preamble macro roles start at the beginning of the macro body.
-      auto func = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
-      parentLoc = func->getMacroExpandedBody()->getStartLoc();
-      break;
-    }
-    case MacroRole::Body:
+    case MacroRole::Body: {
+      // Use the end location of the function decl itself as the parentLoc
+      // for the new function body scope. This is different from the end
+      // location of the original source range, which is after the end of the
+      // function decl.
+      auto expansion = SF->getMacroExpansion();
       parentLoc = expansion.getEndLoc();
       bodyForDecl = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
       break;
-    case MacroRole::Peer: {
-      ASTContext &ctx = SF->getASTContext();
-      SourceManager &sourceMgr = ctx.SourceMgr;
-      auto generatedSourceInfo =
-          *sourceMgr.getGeneratedSourceInfo(*SF->getBufferID());
-
-      ASTNode node = ASTNode::getFromOpaqueValue(generatedSourceInfo.astNode);
-      parentLoc = Lexer::getLocForEndOfToken(sourceMgr, node.getEndLoc());
-      break;
-    }
-    case MacroRole::Member: {
-      // For synthesized member macros, take the end loc of the
-      // enclosing declaration (before the closing brace), because
-      // the macro expansion is inside this scope.
-      auto *decl = expansion.getAsDeclContext()->getAsDecl();
-      parentLoc = decl->getEndLoc();
-      break;
     }
     }
 

lib/AST/ASTVerifier.cpp

@@ -3816,8 +3816,9 @@ class Verifier : public ASTWalker {
         if (!Ctx.SourceMgr.rangeContains(Enclosing, Current)) {
           auto *expansionBuffer =
               D->getModuleContext()->getSourceFileContainingLocation(Current.Start);
-          if (auto expansion = expansionBuffer->getMacroExpansion()) {
-            Current = expansion.getSourceRange();
+
+          if (auto expansionRange = expansionBuffer->getMacroInsertionRange()) {
+            Current = expansionRange;
           }
         }
 

lib/AST/Module.cpp

@@ -1181,6 +1181,16 @@ ASTNode SourceFile::getMacroExpansion() const {
   return ASTNode::getFromOpaqueValue(genInfo.astNode);
 }
 
+SourceRange SourceFile::getMacroInsertionRange() const {
+  if (Kind != SourceFileKind::MacroExpansion)
+    return SourceRange();
+
+  auto generatedInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  auto origRange = generatedInfo.originalSourceRange;
+  return {origRange.getStart(), origRange.getEnd()};
+}
+
 CustomAttr *SourceFile::getAttachedMacroAttribute() const {
   if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;