Reimplement ModuleDecl::getSourceFileContainingLocation() using SourceManager

ModuleDecl kept track of all of the source files in the module so that it
could find the source file containing a given location, which relied on
a sorted array all of these source files. SourceManager has its own
similar data structure for a similar query mapping the locations to
buffer IDs.

Replace ModuleDecl's dats structure with a use of the SourceManager's version
with the mapping from buffer IDs to source files.

Author: DougGregor

include/swift/AST/Module.h

@@ -162,10 +162,6 @@ enum class ResilienceStrategy : unsigned {
 
 class OverlayFile;
 
-/// A mapping used to find the source file that contains a particular source
-/// location.
-class ModuleSourceFileLocationMap;
-
 /// A unit that allows grouping of modules by a package name.
 ///
 /// PackageUnit is treated as an enclosing scope of ModuleDecl. Unlike other
@@ -302,13 +298,6 @@ 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;
 
@@ -410,9 +399,6 @@ 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);
@@ -566,9 +552,6 @@ 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.

lib/AST/Module.cpp

@@ -764,101 +764,14 @@ 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;
-  SourceFile *lastSourceFile = nullptr;
-};
-
-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.
-  ArrayRef<FileUnit *> files = Files;
-  if (sourceFileLocationMap->numFiles == files.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 : files) {
-    if (auto sourceFile = dyn_cast<SourceFile>(fileUnit)) {
-      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});
-
-  sourceFileLocationMap->numFiles = files.size();
-  sourceFileLocationMap->numAuxiliaryFiles = AuxiliaryFiles.size();
-}
-
 SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   if (loc.isInvalid())
     return nullptr;
 
+  auto &sourceMgr = getASTContext().SourceMgr;
+
   // Check whether this location is in a "replaced" range, in which case
   // we want to use the original source file.
-  auto &sourceMgr = getASTContext().SourceMgr;
   SourceLoc adjustedLoc = loc;
   for (const auto &pair : sourceMgr.getReplacedRanges()) {
     if (sourceMgr.rangeContainsTokenLoc(pair.second, loc)) {
@@ -867,35 +780,14 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
     }
   }
 
-  // Before we do any extra work, check the last source file we found a result
-  // in to see if it contains this.
-  if (sourceFileLocationMap) {
-    if (auto lastSourceFile = sourceFileLocationMap->lastSourceFile) {
-      auto range = sourceMgr.getRangeForBuffer(lastSourceFile->getBufferID());
-      if (range.contains(adjustedLoc))
-        return lastSourceFile;
-    }
+  auto bufferID = sourceMgr.findBufferContainingLoc(adjustedLoc);
+  auto sourceFiles = sourceMgr.getSourceFilesForBufferID(bufferID);
+  for (auto sourceFile: sourceFiles) {
+    if (sourceFile->getParentModule() == this)
+      return sourceFile;
   }
 
-  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());
-  // Positions inside an empty file or at EOF should still be considered within
-  // this file.
-  if (!foundRange.contains(adjustedLoc) && adjustedLoc != foundRange.getEnd())
-    return nullptr;
-
-  // Update the last source file.
-  sourceFileLocationMap->lastSourceFile = foundSourceFile;
-  return foundSourceFile;
+  return nullptr;
 }
 
 std::pair<unsigned, SourceLoc>
@@ -3447,10 +3339,6 @@ SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
   }
 
   M.getASTContext().SourceMgr.recordSourceFile(bufferID, this);
-
-  if (Kind == SourceFileKind::MacroExpansion ||
-      Kind == SourceFileKind::DefaultArgument)
-    M.addAuxiliaryFile(*this);
 }
 
 SourceFile::ParsingOptions

lib/ClangImporter/ImportDecl.cpp

@@ -8072,9 +8072,6 @@ SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
       SourceFile(module, SourceFileKind::Library, bufferID);
   ClangSwiftAttrSourceFiles.insert({&module, sourceFile});
 
-  // Record this attribute in the module.
-  module.addAuxiliaryFile(*sourceFile);
-
   return *sourceFile;
 }
 

lib/IDETool/CompileInstance.cpp

@@ -140,7 +140,6 @@ 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/IDEInspectionInstance.cpp

@@ -254,7 +254,6 @@ bool IDEInspectionInstance::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/Parser.cpp

@@ -1144,7 +1144,6 @@ struct ParserUnit::Implementation {
 
     auto *M = ModuleDecl::create(Ctx.getIdentifier(ModuleName), Ctx);
     SF = new (Ctx) SourceFile(*M, SFKind, BufferID, parsingOpts);
-    M->addAuxiliaryFile(*SF);
   }
 
   ~Implementation() {

lib/Serialization/ScanningLoaders.cpp

@@ -226,7 +226,6 @@ SwiftModuleScanner::scanInterfaceFile(Twine moduleInterfacePath,
         SourceFile::ParsingOptions parsingOpts;
         auto sourceFile = new (Ctx) SourceFile(
             *moduleDecl, SourceFileKind::Interface, bufferID, parsingOpts);
-        moduleDecl->addAuxiliaryFile(*sourceFile);
         std::vector<StringRef> ArgsRefs(Args.begin(), Args.end());
         std::vector<StringRef> compiledCandidatesRefs(compiledCandidates.begin(),
                                                       compiledCandidates.end());