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

docs/Generics/generics.tex

@@ -1351,7 +1351,7 @@ \subsection*{Module System}
 \apiref{SourceFile}{class}
 Represents a parsed source file from disk. Inherits from \texttt{FileUnit}.
 \begin{itemize}
-\item \texttt{getTopLevelDecls()} returns an array of all top-level declarations in this source file.
+\item \texttt{getTopLevelItems()} returns an array of all top-level items in this source file.
 \item \texttt{isPrimary()} returns \texttt{true} if this is a primary file, \texttt{false} if this is a secondary file.
 \item \texttt{isScriptMode()} answers if this is the main file of a module.
 \item \texttt{getScope()} returns the root of the scope tree for unqualified lookup.

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:

include/swift/AST/Module.h

@@ -106,7 +106,8 @@ enum class SourceFileKind {
   Library,  ///< A normal .swift file.
   Main,     ///< A .swift file that can have top-level code.
   SIL,      ///< Came from a .sil file.
-  Interface ///< Came from a .swiftinterface file, representing another module.
+  Interface, ///< Came from a .swiftinterface file, representing another module.
+  MacroExpansion, ///< Came from a macro expansion.
 };
 
 /// Contains information about where a particular path is used in
@@ -158,6 +159,10 @@ enum class ResilienceStrategy : unsigned {
 
 class OverlayFile;
 
+/// A mapping used to find the source file that contains a particular source
+/// location.
+class ModuleSourceFileLocationMap;
+
 /// The minimum unit of compilation.
 ///
 /// A module is made up of several file-units, which are all part of the same
@@ -229,6 +234,13 @@ class ModuleDecl
 
   SmallVector<FileUnit *, 2> Files;
 
+  /// Mapping used to find the source file associated with a given source
+  /// location.
+  ModuleSourceFileLocationMap *sourceFileLocationMap = nullptr;
+
+  /// The set of auxiliary source files build as part of this module.
+  SmallVector<SourceFile *, 2> AuxiliaryFiles;
+
   llvm::SmallDenseMap<Identifier, SmallVector<OverlayFile *, 1>>
     declaredCrossImports;
 
@@ -328,6 +340,13 @@ class ModuleDecl
   /// SynthesizedFileUnit instead.
   void addFile(FileUnit &newFile);
 
+  /// Add an auxiliary source file, introduced as part of the translation.
+  void addAuxiliaryFile(SourceFile &sourceFile);
+
+  /// Produces the source file that contains the given source location, or
+  /// \c nullptr if the source location isn't in this module.
+  SourceFile *getSourceFileContainingLocation(SourceLoc loc);
+
   /// Creates a map from \c #filePath strings to corresponding \c #fileID
   /// strings, diagnosing any conflicts.
   ///
@@ -414,6 +433,9 @@ class ModuleDecl
   /// present overlays as if they were part of their underlying module.
   std::pair<ModuleDecl *, Identifier> getDeclaringModuleAndBystander();
 
+  /// Update the source-file location map to make it current.
+  void updateSourceFileLocationMap();
+
 public:
   ///  If this is a traditional (non-cross-import) overlay, get its underlying
   ///  module if one exists.

include/swift/AST/ParseRequests.h

@@ -25,6 +25,8 @@
 
 namespace swift {
 
+struct ASTNode;
+
 /// Report that a request of the given kind is being evaluated, so it
 /// can be recorded by the stats reporter.
 template<typename Request>
@@ -85,13 +87,13 @@ class ParseAbstractFunctionBodyRequest
 };
 
 struct SourceFileParsingResult {
-  ArrayRef<Decl *> TopLevelDecls;
+  ArrayRef<ASTNode> TopLevelItems;
   Optional<ArrayRef<Token>> CollectedTokens;
   Optional<StableHasher> InterfaceHasher;
   Optional<syntax::SourceFileSyntax> SyntaxRoot;
 };
 
-/// Parse the top-level decls of a SourceFile.
+/// Parse the top-level items of a SourceFile.
 class ParseSourceFileRequest
     : public SimpleRequest<
           ParseSourceFileRequest, SourceFileParsingResult(SourceFile *),
@@ -116,6 +118,25 @@ class ParseSourceFileRequest
   readDependencySource(const evaluator::DependencyRecorder &) const;
 };
 
+/// Parse the top-level items of a SourceFile.
+class ParseTopLevelDeclsRequest
+    : public SimpleRequest<
+          ParseTopLevelDeclsRequest, ArrayRef<Decl *>(SourceFile *),
+          RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  ArrayRef<Decl *> evaluate(Evaluator &evaluator, SourceFile *SF) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 void simple_display(llvm::raw_ostream &out,
                     const CodeCompletionCallbacksFactory *factory);
 

include/swift/AST/ParseTypeIDZone.def

@@ -25,3 +25,6 @@ SWIFT_REQUEST(Parse, ParseAbstractFunctionBodyRequest,
 SWIFT_REQUEST(Parse, ParseSourceFileRequest,
               SourceFileParsingResult(SourceFile *), SeparatelyCached,
               NoLocationInfo)
+SWIFT_REQUEST(Parse, ParseTopLevelDeclsRequest,
+              ArrayRef<Decl *>(SourceFile *), Cached,
+              NoLocationInfo)

include/swift/AST/SourceFile.h

@@ -13,13 +13,15 @@
 #ifndef SWIFT_AST_SOURCEFILE_H
 #define SWIFT_AST_SOURCEFILE_H
 
+#include "swift/AST/ASTNode.h"
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/Import.h"
 #include "swift/AST/SynthesizedFileUnit.h"
 #include "swift/Basic/Debug.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/STLExtras.h"
 
 namespace swift {
 
@@ -151,11 +153,11 @@ class SourceFile final : public FileUnit {
   /// been validated.
   llvm::SetVector<ValueDecl *> UnvalidatedDeclsWithOpaqueReturnTypes;
 
-  /// The list of top-level declarations in the source file. This is \c None if
+  /// The list of top-level items in the source file. This is \c None if
   /// they have not yet been parsed.
   /// FIXME: Once addTopLevelDecl/prependTopLevelDecl
   /// have been removed, this can become an optional ArrayRef.
-  Optional<std::vector<Decl *>> Decls;
+  Optional<std::vector<ASTNode>> Items;
 
   /// The list of hoisted declarations. See Decl::isHoisted().
   /// This is only used by lldb.
@@ -195,42 +197,48 @@ class SourceFile final : public FileUnit {
   friend ASTContext;
 
 public:
+  /// For source files created to hold the source code created by expanding
+  /// a macro, this is the AST node that describes the macro expansion.
+  ///
+  /// The source location of this AST node is the place in the source that
+  /// triggered the creation of the macro expansion whose resulting source
+  /// code is in this source file. This field is only valid when
+  /// the \c SourceFileKind is \c MacroExpansion.
+  const ASTNode macroExpansion;
+
   /// Appends the given declaration to the end of the top-level decls list. Do
   /// not add any additional uses of this function.
-  void addTopLevelDecl(Decl *d) {
-    // Force decl parsing if we haven't already.
-    (void)getTopLevelDecls();
-    Decls->push_back(d);
-  }
+  void addTopLevelDecl(Decl *d);
 
   /// Prepends a declaration to the top-level decls list.
   ///
   /// FIXME: This entrypoint exists to support LLDB. Calls to this function are
   /// always a mistake, and additional uses should not be added.
   ///
   /// See rdar://58355191
-  void prependTopLevelDecl(Decl *d) {
-    // Force decl parsing if we haven't already.
-    (void)getTopLevelDecls();
-    Decls->insert(Decls->begin(), d);
-  }
+  void prependTopLevelDecl(Decl *d);
 
   /// Add a hoisted declaration. See Decl::isHoisted().
   void addHoistedDecl(Decl *d);
 
+  /// Retrieves an immutable view of the list of top-level items in this file.
+  ArrayRef<ASTNode> getTopLevelItems() const;
+
   /// Retrieves an immutable view of the list of top-level decls in this file.
+  ///
+  /// NOTE: Please use getTopLevelItems() instead.
   ArrayRef<Decl *> getTopLevelDecls() const;
 
   /// Retrieves an immutable view of the list of hoisted decls in this file.
   /// See Decl::isHoisted().
   ArrayRef<Decl *> getHoistedDecls() const;
 
-  /// Retrieves an immutable view of the top-level decls if they have already
+  /// Retrieves an immutable view of the top-level items if they have already
   /// been parsed, or \c None if they haven't. Should only be used for dumping.
-  Optional<ArrayRef<Decl *>> getCachedTopLevelDecls() const {
-    if (!Decls)
+  Optional<ArrayRef<ASTNode>> getCachedTopLevelItems() const {
+    if (!Items)
       return None;
-    return llvm::makeArrayRef(*Decls);
+    return llvm::makeArrayRef(*Items);
   }
 
   /// Retrieve the parsing options for the file.
@@ -326,7 +334,8 @@ class SourceFile final : public FileUnit {
   llvm::StringMap<SourceFilePathInfo> getInfoForUsedFilePaths() const;
 
   SourceFile(ModuleDecl &M, SourceFileKind K, Optional<unsigned> bufferID,
-             ParsingOptions parsingOpts = {}, bool isPrimary = false);
+             ParsingOptions parsingOpts = {}, bool isPrimary = false,
+             ASTNode macroExpansion = ASTNode());
 
   ~SourceFile();
 
@@ -487,6 +496,11 @@ class SourceFile final : public FileUnit {
     return BufferID;
   }
 
+  /// When this source file is enclosed within another source file, for example
+  /// because it describes a macro expansion, return the source file it was
+  /// enclosed in.
+  SourceFile *getEnclosingSourceFile() const;
+
   /// If this buffer corresponds to a file on disk, returns the path.
   /// Otherwise, return an empty string.
   StringRef getFilename() const;
@@ -506,7 +520,7 @@ class SourceFile final : public FileUnit {
     // FIXME: Ideally the parser state should be an output of
     // ParseSourceFileRequest, but the evaluator doesn't currently support
     // move-only outputs for cached requests.
-    (void)getTopLevelDecls();
+    (void)getTopLevelItems();
 
     auto *state = DelayedParserState.get();
     assert(state && "Didn't set any delayed parser state!");
@@ -545,6 +559,7 @@ class SourceFile final : public FileUnit {
     case SourceFileKind::Library:
     case SourceFileKind::Interface:
     case SourceFileKind::SIL:
+    case SourceFileKind::MacroExpansion:
       return false;
     }
     llvm_unreachable("bad SourceFileKind");

include/swift/Parse/Parser.h

@@ -79,6 +79,8 @@ namespace swift {
     InactiveConditionalBlock,
     /// The body of the active clause of an #if/#else/#endif block
     ActiveConditionalBlock,
+    /// The top-level of a macro expansion "file".
+    MacroExpansion,
   };
 
 /// The receiver will be fed with consumed tokens while parsing. The main purpose
@@ -969,8 +971,10 @@ class Parser {
   /// Returns true if the parser is at the start of a SIL decl.
   bool isStartOfSILDecl();
 
-  /// Parse the top-level Swift decls into the provided vector.
-  void parseTopLevel(SmallVectorImpl<Decl *> &decls);
+  /// Parse the top-level Swift items into the provided vector.
+  ///
+  /// Each item will be a declaration, statement, or expression.
+  void parseTopLevelItems(SmallVectorImpl<ASTNode> &items);
 
   /// Parse the top-level SIL decls into the SIL module.
   /// \returns \c true if there was a parsing error.

lib/AST/ASTDumper.cpp

@@ -804,13 +804,21 @@ namespace {
       PrintWithColorRAII(OS, ASTNodeColor) << "source_file ";
       PrintWithColorRAII(OS, LocationColor) << '\"' << SF.getFilename() << '\"';
 
-      if (auto decls = SF.getCachedTopLevelDecls()) {
-        for (Decl *D : *decls) {
-          if (D->isImplicit())
+      if (auto items = SF.getCachedTopLevelItems()) {
+        for (auto item : *items) {
+          if (item.isImplicit())
             continue;
 
           OS << '\n';
-          printRec(D);
+
+          if (auto decl = item.dyn_cast<Decl *>()) {
+            printRec(decl);
+          } else if (auto stmt = item.dyn_cast<Stmt *>()) {
+            stmt->dump(OS, &SF.getASTContext(), Indent + 2);
+          } else {
+            auto expr = item.get<Expr *>();
+            expr->dump(OS, Indent + 2);
+          }
         }
       }
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
@@ -1470,7 +1478,7 @@ void SourceFile::dump(llvm::raw_ostream &OS, bool parseIfNeeded) const {
   // parsing request as by default the dumping logic tries not to kick any
   // requests.
   if (parseIfNeeded)
-    (void)getTopLevelDecls();
+    (void)getTopLevelItems();
 
   PrintDecl(OS).visitSourceFile(*this);
   llvm::errs() << '\n';

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
 
@@ -692,9 +699,9 @@ ASTSourceFileScope::expandAScopeThatCreatesANewInsertionPoint(
   SourceLoc endLoc = getSourceRangeOfThisASTNode().End;
 
   ASTScopeImpl *insertionPoint = this;
-  for (auto *d : SF->getTopLevelDecls()) {
+  for (auto node : SF->getTopLevelItems()) {
     insertionPoint = scopeCreator.addToScopeTreeAndReturnInsertionPoint(
-      ASTNode(d), insertionPoint, endLoc);
+      node, insertionPoint, endLoc);
   }
 
   return {insertionPoint, "Next time decls are added they go here."};

lib/AST/ASTScopeSourceRange.cpp

@@ -181,12 +181,12 @@ SourceRange ASTSourceFileScope::getSourceRangeOfThisASTNode(
     return SourceRange(charRange.getStart(), charRange.getEnd());
   }
 
-  if (SF->getTopLevelDecls().empty())
+  if (SF->getTopLevelItems().empty())
     return SourceRange();
 
   // Use the source ranges of the declarations in the file.
-  return SourceRange(SF->getTopLevelDecls().front()->getStartLoc(),
-                     SF->getTopLevelDecls().back()->getEndLoc());
+  return SourceRange(SF->getTopLevelItems().front().getStartLoc(),
+                     SF->getTopLevelItems().back().getEndLoc());
 }
 
 SourceRange GenericTypeOrExtensionScope::getSourceRangeOfThisASTNode(