Merge pull request swiftlang#76551 from DougGregor/one-loc-to-file-lookup

Reimplement ModuleDecl::getSourceFileContainingLocation() using SourceManager

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.

include/swift/Basic/SourceManager.h

@@ -18,15 +18,34 @@
 #include "clang/Basic/FileManager.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/SourceMgr.h"
 #include <map>
 #include <optional>
 #include <vector>
 
+namespace swift {
+  class SourceFile;
+}
+
+namespace llvm {
+  template <> struct PointerLikeTypeTraits<swift::SourceFile *> {
+  public:
+    static inline swift::SourceFile *getFromVoidPointer(void *P) {
+      return (swift::SourceFile *)P;
+    }
+    static inline void *getAsVoidPointer(swift::SourceFile *S) {
+      return (void *)S;
+    }
+    enum { NumLowBitsAvailable = /*swift::DeclContextAlignInBits=*/ 3 };
+  };
+}
+
 namespace swift {
 
 class CustomAttr;
 class DeclContext;
+class SourceFile;
 
 /// Augments a buffer that was created specifically to hold generated source
 /// code with the reasons for it being generated.
@@ -123,6 +142,13 @@ class SourceManager {
   /// is an unfortunate hack needed to allow for correct re-lexing.
   llvm::DenseSet<SourceLoc> RegexLiteralStartLocs;
 
+  /// Mapping from each buffer ID to the source files that describe it
+  /// semantically.
+  llvm::DenseMap<
+    unsigned,
+    llvm::TinyPtrVector<SourceFile *>
+  > bufferIDToSourceFiles;
+
   std::map<const char *, VirtualFile> VirtualFiles;
   mutable std::pair<const char *, const VirtualFile*> CachedVFile = {nullptr, nullptr};
 
@@ -323,6 +349,13 @@ class SourceManager {
   /// Adds a memory buffer to the SourceManager, taking ownership of it.
   unsigned addNewSourceBuffer(std::unique_ptr<llvm::MemoryBuffer> Buffer);
 
+  /// Record the source file as having the given buffer ID.
+  void recordSourceFile(unsigned bufferID, SourceFile *sourceFile);
+
+  /// Retrieve the source files for the given buffer ID.
+  llvm::TinyPtrVector<SourceFile *>
+  getSourceFilesForBufferID(unsigned bufferID) const;
+
   /// Add a \c #sourceLocation-defined virtual file region of \p Length.
   void createVirtualFile(SourceLoc Loc, StringRef Name, int LineOffset,
                          unsigned Length);

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>
@@ -3446,9 +3338,7 @@ SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
     (void)problem;
   }
 
-  if (Kind == SourceFileKind::MacroExpansion ||
-      Kind == SourceFileKind::DefaultArgument)
-    M.addAuxiliaryFile(*this);
+  M.getASTContext().SourceMgr.recordSourceFile(bufferID, this);
 }
 
 SourceFile::ParsingOptions

lib/Basic/SourceLoc.cpp

@@ -215,6 +215,19 @@ SourceManager::getIDForBufferIdentifier(StringRef BufIdentifier) const {
   return It->second;
 }
 
+void SourceManager::recordSourceFile(unsigned bufferID, SourceFile *sourceFile){
+  bufferIDToSourceFiles[bufferID].push_back(sourceFile);
+}
+
+llvm::TinyPtrVector<SourceFile *>
+SourceManager::getSourceFilesForBufferID(unsigned bufferID) const {
+  auto found = bufferIDToSourceFiles.find(bufferID);
+  if (found == bufferIDToSourceFiles.end())
+    return { };
+
+  return found->second;
+}
+
 SourceManager::~SourceManager() {
   for (auto &generated : GeneratedSourceInfos) {
     free((void*)generated.second.onDiskBufferCopyFileName.data());

lib/ClangImporter/ImportDecl.cpp

@@ -8141,9 +8141,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());