[Macros] Teach name lookup in macro expansions to look in expansion context

Teach ASTScope how to reason about macro expansions. When we create an
ASTScope for a source file that represents a macro expansion, its
parent scope node is the macro expansion itself. When performing
unqualified lookup (of any form), find the starting source file based
on the location, not on the source file provided---this ensures that
we start lookups within the macro expansion (for example).

The effect of this is to enable macro expansions to work in nested
contexts, where they refer to names in the scope in which the macro is
expanded.

This is decidedly unhiegenic, but it fits with our syntactic model.

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -128,6 +128,7 @@ class ASTScopeImpl : public ASTAllocated<ASTScopeImpl> {
   friend class GenericTypeOrExtensionWherePortion;
   friend class IterableTypeBodyPortion;
   friend class ScopeCreator;
+  friend class ASTSourceFileScope;
 
 #pragma mark - tree state
 protected:

lib/AST/ASTScope.cpp

@@ -41,6 +41,10 @@ using namespace ast_scope;
 void ASTScope::unqualifiedLookup(
     SourceFile *SF, SourceLoc loc,
     namelookup::AbstractASTScopeDeclConsumer &consumer) {
+  if (loc.isValid()) {
+    SF = SF->getParentModule()->getSourceFileContainingLocation(loc);
+  }
+
   if (auto *s = SF->getASTContext().Stats)
     ++s->getFrontendCounters().NumASTScopeLookups;
   ASTScopeImpl::unqualifiedLookup(SF, loc, consumer);

lib/AST/ASTScopeCreation.cpp

@@ -246,7 +246,14 @@ void ASTSourceFileScope::expandFunctionBody(AbstractFunctionDecl *AFD) {
 
 ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
                                        ScopeCreator *scopeCreator)
-    : SF(SF), scopeCreator(scopeCreator) {}
+    : SF(SF), scopeCreator(scopeCreator) {
+  if (auto enclosingSF = SF->getEnclosingSourceFile()) {
+    SourceLoc parentLoc = SF->macroExpansion.getStartLoc();
+    if (auto parentScope = findStartingScopeForLookup(enclosingSF, parentLoc)) {
+      parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
+    }
+  }
+}
 
 #pragma mark NodeAdder
 

lib/AST/ASTVerifier.cpp

@@ -3690,6 +3690,8 @@ class Verifier : public ASTWalker {
       if (Parent.isNull())
           return;
 
+      // An alternative enclosing scope, used for macro expansions.
+      SourceRange AltEnclosing;
       if (Parent.getAsModule()) {
         return;
       } else if (Decl *D = Parent.getAsDecl()) {
@@ -3720,15 +3722,22 @@ class Verifier : public ASTWalker {
 
         if (E->isImplicit())
           return;
-        
+
+        if (auto expansion = dyn_cast<MacroExpansionExpr>(E)) {
+          if (auto rewritten = expansion->getRewritten())
+            AltEnclosing = rewritten->getSourceRange();
+        }
+
         Enclosing = E->getSourceRange();
       } else if (TypeRepr *TyR = Parent.getAsTypeRepr()) {
         Enclosing = TyR->getSourceRange();
       } else {
         llvm_unreachable("impossible parent node");
       }
 
-      if (!Ctx.SourceMgr.rangeContains(Enclosing, Current)) {
+      if (!Ctx.SourceMgr.rangeContains(Enclosing, Current) &&
+          !(AltEnclosing.isValid() &&
+            Ctx.SourceMgr.rangeContains(AltEnclosing, Current))) {
         Out << "child source range not contained within its parent: ";
         printEntity();
         Out << "\n  parent range: ";

test/Macros/macros.swift

@@ -2,9 +2,10 @@
 // REQUIRES: OS=macosx
 // REQUIRES: asserts
 
-let a = 1, b = 1
-let s = #stringify(a + b)
-// CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
-// CHECK-NEXT: argument_list
-// CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
+func test(a: Int, b: Int) {
+  let s = #stringify(a + b)
 
+  // CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
+  // CHECK-NEXT: argument_list
+  // CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
+}