Merge pull request swiftlang#63232 from hborla/member-macro-ast-scope

[ASTScope] Teach ASTScope lookup to find source locations inside macro-expanded scopes.

Author: hborla

lib/AST/ASTScopeLookup.cpp

@@ -30,6 +30,7 @@
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/Parse/Lexer.h"
 #include "llvm/Support/Compiler.h"
 
 using namespace swift;
@@ -90,14 +91,78 @@ static SourceLoc translateLocForReplacedRange(SourceManager &sourceMgr,
 NullablePtr<ASTScopeImpl>
 ASTScopeImpl::findChildContaining(SourceLoc loc,
                                   SourceManager &sourceMgr) const {
+  auto *moduleDecl = this->getSourceFile()->getParentModule();
+  auto *locSourceFile = moduleDecl->getSourceFileContainingLocation(loc);
+
   // Use binary search to find the child that contains this location.
   auto *const *child = llvm::lower_bound(
       getChildren(), loc,
-      [&sourceMgr](const ASTScopeImpl *scope, SourceLoc loc) {
+      [&](const ASTScopeImpl *scope, SourceLoc loc) {
         auto rangeOfScope = scope->getCharSourceRangeOfScope(sourceMgr);
         ASTScopeAssert(!sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(),
                                                    rangeOfScope.getStart()),
                        "Source range is backwards");
+
+        // If the scope source range and the loc are in two different source
+        // files, one or both of them are in a macro expansion buffer.
+
+        // Note that `scope->getSourceFile()` returns the root of the source tree,
+        // not the source file containing the location of the ASTScope.
+        auto scopeStart = scope->getSourceRangeOfThisASTNode().Start;
+        auto *scopeSourceFile = moduleDecl->getSourceFileContainingLocation(scopeStart);
+
+        if (scopeSourceFile != locSourceFile) {
+          // To compare a source location that is possibly inside a macro expansion
+          // with a source range that is also possibly in a macro expansion (not
+          // necessarily the same one as before) we need to find the LCA in the
+          // source file tree of macro expansions, and compare the original source
+          // ranges within that common ancestor. We can't walk all the way up to the
+          // source file containing the parent scope we're searching the children of,
+          // because two independent (possibly nested) macro expansions can have the
+          // same original source range in that file; freestanding and peer macros
+          // mean that we can have arbitrarily nested macro expansions that all add
+          // declarations to the same scope, that all originate from a single macro
+          // invocation in the original source file.
+
+          // A map from enclosing source files to original source ranges of the macro
+          // expansions within that file, recording the chain of macro expansions for
+          // the given scope.
+          llvm::SmallDenseMap<const SourceFile *, SourceRange> scopeExpansions;
+
+          // Walk up the chain of macro expansion buffers for the scope, recording the
+          // original source range of the macro expansion along the way using generated
+          // source info.
+          auto *scopeExpansion = scopeSourceFile;
+          scopeExpansions[scopeExpansion] = scope->getSourceRangeOfThisASTNode();
+          while (auto *ancestor = scopeExpansion->getEnclosingSourceFile()) {
+            auto generatedInfo =
+                sourceMgr.getGeneratedSourceInfo(*scopeExpansion->getBufferID());
+            scopeExpansions[ancestor] = generatedInfo->originalSourceRange;
+            scopeExpansion = ancestor;
+          }
+
+          // Walk up the chain of macro expansion buffers for the source loc we're
+          // searching for to find the LCA using `scopeExpansions`.
+          auto *potentialLCA = locSourceFile;
+          auto expansionLoc = loc;
+          while (potentialLCA) {
+            auto scopeExpansion = scopeExpansions.find(potentialLCA);
+            if (scopeExpansion != scopeExpansions.end()) {
+              // Take the original expansion range within the LCA of the loc and
+              // the scope to compare.
+              rangeOfScope =
+                  Lexer::getCharSourceRangeFromSourceRange(sourceMgr, scopeExpansion->second);
+              loc = expansionLoc;
+              break;
+            }
+
+            auto generatedInfo =
+                sourceMgr.getGeneratedSourceInfo(*potentialLCA->getBufferID());
+            expansionLoc = generatedInfo->originalSourceRange.Start;
+            potentialLCA = potentialLCA->getEnclosingSourceFile();
+          }
+        }
+
         loc = translateLocForReplacedRange(sourceMgr, rangeOfScope, loc);
         return (rangeOfScope.getEnd() == loc ||
                 sourceMgr.isBeforeInBuffer(rangeOfScope.getEnd(), loc));

lib/AST/ASTVerifier.cpp

@@ -3736,6 +3736,18 @@ class Verifier : public ASTWalker {
         return;
       } else if (Decl *D = Parent.getAsDecl()) {
         Enclosing = D->getSourceRange();
+
+        // If the current source range is in a macro expansion buffer, its enclosing
+        // context can be in the source file where the macro expansion originated. In
+        // this case, grab the source range of the original ASTNode that was expanded.
+        if (!Ctx.SourceMgr.rangeContains(Enclosing, Current)) {
+          auto *expansionBuffer =
+              D->getModuleContext()->getSourceFileContainingLocation(Current.Start);
+          if (auto expansion = expansionBuffer->getMacroExpansion()) {
+            Current = expansion.getSourceRange();
+          }
+        }
+
         if (D->isImplicit())
           return;
         // FIXME: This is not working well for decl parents.

lib/AST/DeclContext.cpp

@@ -984,6 +984,11 @@ void IterableDeclContext::addMemberSilently(Decl *member, Decl *hint,
     if (getASTContext().SourceMgr.isBeforeInBuffer(prevEnd, nextStart))
       return;
 
+    // Synthesized member macros can add new members in a macro expansion buffer.
+    auto *memberSourceFile = member->getInnermostDeclContext()->getParentSourceFile();
+    if (memberSourceFile->getFulfilledMacroRole() == MacroRole::SynthesizedMembers)
+      return;
+
     llvm::errs() << "Source ranges out of order in addMember():\n";
     prev->dump(llvm::errs());
     next->dump(llvm::errs());

lib/Sema/TypeCheckMacros.cpp

@@ -1230,8 +1230,10 @@ bool swift::expandSynthesizedMembers(CustomAttr *attr, MacroDecl *macro,
   bool synthesizedMembers = false;
   auto topLevelDecls = macroSourceFile->getTopLevelDecls();
   for (auto member : topLevelDecls) {
+    // Note that synthesized members are not considered implicit. They have
+    // proper source ranges that should be validated, and ASTScope does not
+    // expand implicit scopes to the parent scope tree.
     member->setDeclContext(decl->getInnermostDeclContext());
-    member->setImplicit();
 
     if (auto *nominal = dyn_cast<NominalTypeDecl>(decl)) {
       nominal->addMember(member);

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -360,12 +360,15 @@ public struct AddMembers: MemberDeclarationMacro {
 
     let storageVariable: VariableDeclSyntax =
       """
-      var storage = Storage()
+      private var storage = Storage()
       """
 
     let instanceMethod: FunctionDeclSyntax =
       """
-      func method() { print("synthesized method") }
+      func getStorage() -> Storage {
+        print("synthesized method")
+        return storage
+      }
       """
 
     return [

test/Macros/macro_expand_synthesized_members.swift

@@ -14,13 +14,12 @@
 @addMembers
 struct S {
   func useSynthesized() {
-    print(type(of: storage))
-    method()
+    print(type(of: getStorage()))
   }
 }
 
 let s = S()
 
-// CHECK: Storage
 // CHECK: synthesized method
+// CHECK: Storage
 s.useSynthesized()