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

scopes.

ASTScope lookup can no longer assume that all scopes are within the same source
file. Scope trees can now contain scopes that are in macro expansion buffers, which
live in different source files with independent source ranges from the root of a
given scope tree. To handle this when searching for a scope containing a given source
location, ASTScope lookup needs to walk up the chain of macro expansion buffers to find
the lowest common buffer in which to compare source locations.

This fixes a number of issues with unqualified lookup into and from within macro-expanded
code.

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

@@ -3692,6 +3692,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

@@ -1242,8 +1242,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()