Keep track of source files created for macro expansions and such.

Introduce a new source file kind to describe source files for macro
expansions, and include the macro expression that they expand. This
establishes a "parent" relationship

Also track every kind of auxiliary source file---whether for macro
expansions or other reasons---that is introduced into a module, adding
an operation that allows us to find the source file that contains a
given source location.

Author: DougGregor

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/SourceFile.h

@@ -13,6 +13,7 @@
 #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"
@@ -195,6 +196,15 @@ 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) {
@@ -326,7 +336,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 +498,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;
@@ -545,6 +561,7 @@ class SourceFile final : public FileUnit {
     case SourceFileKind::Library:
     case SourceFileKind::Interface:
     case SourceFileKind::SIL:
+    case SourceFileKind::MacroExpansion:
       return false;
     }
     llvm_unreachable("bad SourceFileKind");

lib/AST/Module.cpp

@@ -514,6 +514,123 @@ void ModuleDecl::addFile(FileUnit &newFile) {
   clearLookupCache();
 }
 
+void ModuleDecl::addAuxiliaryFile(SourceFile &sourceFile) {
+  AuxiliaryFiles.push_back(&sourceFile);
+}
+
+namespace {
+  /// Compare the source location ranges for two files, as an ordering to
+  /// use for fast searches.
+  struct SourceFileRangeComparison {
+    SourceManager *sourceMgr;
+
+    bool operator()(SourceFile *lhs, SourceFile *rhs) const {
+      auto lhsRange = sourceMgr->getRangeForBuffer(*lhs->getBufferID());
+      auto rhsRange = sourceMgr->getRangeForBuffer(*rhs->getBufferID());
+
+      std::less<const char *> pointerCompare;
+      return pointerCompare(
+          (const char *)lhsRange.getStart().getOpaquePointerValue(),
+          (const char *)rhsRange.getStart().getOpaquePointerValue());
+    }
+
+    bool operator()(SourceFile *lhs, SourceLoc rhsLoc) const {
+      auto lhsRange = sourceMgr->getRangeForBuffer(*lhs->getBufferID());
+
+      std::less<const char *> pointerCompare;
+      return pointerCompare(
+          (const char *)lhsRange.getEnd().getOpaquePointerValue(),
+          (const char *)rhsLoc.getOpaquePointerValue());
+    }
+
+    bool operator()(SourceLoc lhsLoc, SourceFile *rhs) const {
+      auto rhsRange = sourceMgr->getRangeForBuffer(*rhs->getBufferID());
+
+      std::less<const char *> pointerCompare;
+      return pointerCompare(
+          (const char *)lhsLoc.getOpaquePointerValue(),
+          (const char *)rhsRange.getEnd().getOpaquePointerValue());
+    }
+  };
+}
+
+class swift::ModuleSourceFileLocationMap {
+public:
+  unsigned numFiles = 0;
+  unsigned numAuxiliaryFiles = 0;
+  std::vector<SourceFile *> allSourceFiles;
+};
+
+void ModuleDecl::updateSourceFileLocationMap() {
+  // Allocate a source file location map, if we don't have one already.
+  if (!sourceFileLocationMap) {
+    ASTContext &ctx = getASTContext();
+    sourceFileLocationMap = ctx.Allocate<ModuleSourceFileLocationMap>();
+    ctx.addCleanup([sourceFileLocationMap=sourceFileLocationMap]() {
+      sourceFileLocationMap->~ModuleSourceFileLocationMap();
+    });
+  }
+
+  // If we are up-to-date, there's nothing to do.
+  if (sourceFileLocationMap->numFiles == getFiles().size() &&
+      sourceFileLocationMap->numAuxiliaryFiles ==
+          AuxiliaryFiles.size())
+    return;
+
+  // Rebuild the range structure.
+  sourceFileLocationMap->allSourceFiles.clear();
+
+  // First, add all of the source files with a backing buffer.
+  for (auto *fileUnit : getFiles()) {
+    if (auto sourceFile = dyn_cast<SourceFile>(fileUnit)) {
+      if (sourceFile->getBufferID())
+        sourceFileLocationMap->allSourceFiles.push_back(sourceFile);
+    }
+  }
+
+  // Next, add all of the macro expansion files.
+  for (auto *sourceFile : AuxiliaryFiles)
+    sourceFileLocationMap->allSourceFiles.push_back(sourceFile);
+
+  // Finally, sort them all so we can do a binary search for lookup.
+  std::sort(sourceFileLocationMap->allSourceFiles.begin(),
+            sourceFileLocationMap->allSourceFiles.end(),
+            SourceFileRangeComparison{&getASTContext().SourceMgr});
+}
+
+SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
+  if (loc.isInvalid())
+    return nullptr;
+
+  SourceLoc adjustedLoc = loc;
+
+  // Check whether this location is in a "replaced" range, in which case
+  // we want to use the original source file.
+  auto &sourceMgr = getASTContext().SourceMgr;
+  for (const auto &pair : sourceMgr.getReplacedRanges()) {
+    if (sourceMgr.rangeContainsTokenLoc(pair.second, loc)) {
+      adjustedLoc = pair.first.Start;
+      break;
+    }
+  }
+  updateSourceFileLocationMap();
+
+  auto found = std::lower_bound(sourceFileLocationMap->allSourceFiles.begin(),
+                                sourceFileLocationMap->allSourceFiles.end(),
+                                adjustedLoc,
+                                SourceFileRangeComparison{&sourceMgr});
+  if (found == sourceFileLocationMap->allSourceFiles.end())
+    return nullptr;
+
+  auto foundSourceFile = *found;
+  auto foundRange = sourceMgr.getRangeForBuffer(*foundSourceFile->getBufferID());
+  if (!foundRange.contains(adjustedLoc))
+    return nullptr;
+
+
+  return foundSourceFile;
+}
+
 ArrayRef<SourceFile *>
 PrimarySourceFilesRequest::evaluate(Evaluator &evaluator,
                                     ModuleDecl *mod) const {
@@ -747,6 +864,14 @@ void SourceFile::lookupClassMembers(ImportPath::Access accessPath,
   cache.lookupClassMembers(accessPath, consumer);
 }
 
+SourceFile *SourceFile::getEnclosingSourceFile() const {
+  if (Kind != SourceFileKind::MacroExpansion)
+    return nullptr;
+
+  auto sourceLoc = macroExpansion.getStartLoc();
+  return getParentModule()->getSourceFileContainingLocation(sourceLoc);
+}
+
 void ModuleDecl::lookupClassMember(ImportPath::Access accessPath,
                                    DeclName name,
                                    SmallVectorImpl<ValueDecl*> &results) const {
@@ -3075,19 +3200,26 @@ ModuleDecl::computeFileIDMap(bool shouldDiagnose) const {
 
 SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
                        Optional<unsigned> bufferID,
-                       ParsingOptions parsingOpts, bool isPrimary)
+                       ParsingOptions parsingOpts, bool isPrimary,
+                       ASTNode macroExpansion)
     : FileUnit(FileUnitKind::Source, M), BufferID(bufferID ? *bufferID : -1),
-      ParsingOpts(parsingOpts), IsPrimary(isPrimary), Kind(K) {
+      ParsingOpts(parsingOpts), IsPrimary(isPrimary),
+      macroExpansion(macroExpansion), Kind(K) {
   M.getASTContext().addDestructorCleanup(*this);
 
   assert(!IsPrimary || M.isMainModule() &&
          "A primary cannot appear outside the main module");
+  assert(macroExpansion.isNull() == (K != SourceFileKind::MacroExpansion) &&
+         "Macro expansions always need an expansion node");
 
   if (isScriptMode()) {
     bool problem = M.registerEntryPointFile(this, SourceLoc(), None);
     assert(!problem && "multiple main files?");
     (void)problem;
   }
+
+  if (Kind == SourceFileKind::MacroExpansion)
+    M.addAuxiliaryFile(*this);
 }
 
 SourceFile::ParsingOptions

lib/ClangImporter/ImportDecl.cpp

@@ -7227,6 +7227,7 @@ SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
   auto sourceFile = new (SwiftContext) SourceFile(
       module, SourceFileKind::Library, None);
   ClangSwiftAttrSourceFiles.insert({&module, sourceFile});
+  module.addAuxiliaryFile(*sourceFile);
   return *sourceFile;
 }
 

lib/IDETool/CompileInstance.cpp

@@ -136,6 +136,7 @@ getModifiedFunctionDeclList(const SourceFile &SF, SourceManager &tmpSM,
   auto tmpBufferID = tmpSM.addNewSourceBuffer(std::move(*tmpBuffer));
   SourceFile *tmpSF = new (tmpCtx)
       SourceFile(*tmpM, SF.Kind, tmpBufferID, SF.getParsingOptions());
+  tmpM->addAuxiliaryFile(*tmpSF);
 
   // If the top-level code has been changed, we can't do anything.
   if (SF.getInterfaceHash() != tmpSF->getInterfaceHash())

lib/IDETool/CompletionInstance.cpp

@@ -252,6 +252,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
   ModuleDecl *tmpM = ModuleDecl::create(Identifier(), *tmpCtx);
   SourceFile *tmpSF = new (*tmpCtx)
       SourceFile(*tmpM, oldSF->Kind, tmpBufferID, oldSF->getParsingOptions());
+  tmpM->addAuxiliaryFile(*tmpSF);
 
   // FIXME: Since we don't setup module loaders on the temporary AST context,
   // 'canImport()' conditional compilation directive always fails. That causes

lib/Parse/ParseDecl.cpp

@@ -8751,6 +8751,7 @@ parseDeclDeinit(ParseDeclOptions Flags, DeclAttributes &Attributes) {
       break;
     case SourceFileKind::Library:
     case SourceFileKind::Main:
+    case SourceFileKind::MacroExpansion:
       if (Tok.is(tok::identifier)) {
         diagnose(Tok, diag::destructor_has_name).fixItRemove(Tok.getLoc());
         consumeToken();

lib/Parse/Parser.cpp

@@ -370,6 +370,7 @@ static LexerMode sourceFileKindToLexerMode(SourceFileKind kind) {
       return LexerMode::SIL;
     case swift::SourceFileKind::Library:
     case swift::SourceFileKind::Main:
+    case swift::SourceFileKind::MacroExpansion:
       return LexerMode::Swift;
   }
   llvm_unreachable("covered switch");
@@ -1250,6 +1251,7 @@ struct ParserUnit::Implementation {
 
     auto *M = ModuleDecl::create(Ctx.getIdentifier(ModuleName), Ctx);
     SF = new (Ctx) SourceFile(*M, SFKind, BufferID, parsingOpts);
+    M->addAuxiliaryFile(*SF);
   }
 
   ~Implementation() {

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -788,6 +788,7 @@ addDistributedActorCodableConformance(
 
         case SourceFileKind::Library:
         case SourceFileKind::Main:
+        case SourceFileKind::MacroExpansion:
         case SourceFileKind::SIL:
           break;
         }

lib/Sema/TypeCheckConcurrency.cpp

@@ -45,6 +45,7 @@ static bool shouldInferAttributeInContext(const DeclContext *dc) {
           return false;
 
         case SourceFileKind::Library:
+        case SourceFileKind::MacroExpansion:
         case SourceFileKind::Main:
         case SourceFileKind::SIL:
           return true;
@@ -4486,6 +4487,7 @@ ProtocolConformance *GetImplicitSendableRequest::evaluate(
           return nullptr;
 
         case SourceFileKind::Library:
+        case SourceFileKind::MacroExpansion:
         case SourceFileKind::Main:
         case SourceFileKind::SIL:
           break;