Reimplement ModuleDecl::getSourceFileContainingLocation() using SourceManager

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

@@ -156,9 +156,7 @@ class SourceFile final : public FileUnit {
   std::optional<StableHasher> InterfaceHasher;
 
   /// The ID for the memory buffer containing this file's source.
-  ///
-  /// May be -1, to indicate no association with a buffer.
-  int BufferID;
+  unsigned BufferID;
 
   /// The parsing options for the file.
   ParsingOptions ParsingOpts;
@@ -370,7 +368,7 @@ class SourceFile final : public FileUnit {
   /// \c #sourceLocation(file:) declarations.
   llvm::StringMap<SourceFilePathInfo> getInfoForUsedFilePaths() const;
 
-  SourceFile(ModuleDecl &M, SourceFileKind K, std::optional<unsigned> bufferID,
+  SourceFile(ModuleDecl &M, SourceFileKind K, unsigned bufferID,
              ParsingOptions parsingOpts = {}, bool isPrimary = false);
 
   ~SourceFile();
@@ -590,9 +588,7 @@ class SourceFile final : public FileUnit {
 
   /// The buffer ID for the file that was imported, or None if there
   /// is no associated buffer.
-  std::optional<unsigned> getBufferID() const {
-    if (BufferID == -1)
-      return std::nullopt;
+  unsigned getBufferID() const {
     return BufferID;
   }
 

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);

include/swift/Frontend/Frontend.h

@@ -789,7 +789,7 @@ class CompilerInstance {
   /// Creates a new source file for the main module.
   SourceFile *createSourceFileForMainModule(ModuleDecl *mod,
                                             SourceFileKind FileKind,
-                                            std::optional<unsigned> BufferID,
+                                            unsigned BufferID,
                                             bool isMainBuffer = false) const;
 
   /// Creates all the files to be added to the main module, appending them to

include/swift/IDE/IDERequests.h

@@ -50,7 +50,7 @@ struct CursorInfoOwner {
     return !(lhs == rhs);
   }
   bool isValid() const {
-    return File && File->getBufferID() && Loc.isValid();
+    return File && Loc.isValid();
   }
 };
 
@@ -110,7 +110,7 @@ struct RangeInfoOwner {
   }
 
   bool isValid() const {
-    return File && File->getBufferID() && StartLoc.isValid() && EndLoc.isValid();
+    return File && StartLoc.isValid() && EndLoc.isValid();
   }
 };
 

include/swift/Migrator/ASTMigratorPass.h

@@ -41,7 +41,7 @@ class ASTMigratorPass {
   ASTMigratorPass(EditorAdapter &Editor, SourceFile *SF,
                   const MigratorOptions &Opts)
     : Editor(Editor), SF(SF), Opts(Opts), Filename(SF->getFilename()),
-      BufferID(SF->getBufferID().value()),
+      BufferID(SF->getBufferID()),
       SM(SF->getASTContext().SourceMgr), Diags(SF->getASTContext().Diags) {}
 };
 

lib/AST/ASTContext.cpp

@@ -6622,9 +6622,8 @@ void VarDecl::setOriginalWrappedProperty(VarDecl *originalProperty) {
 static bool isSourceLocInOrignalBuffer(const Decl *D, SourceLoc Loc) {
   assert(Loc.isValid());
   auto bufferID = D->getDeclContext()->getParentSourceFile()->getBufferID();
-  assert(bufferID.has_value() && "Source buffer ID must be set");
   auto &SM = D->getASTContext().SourceMgr;
-  return SM.getRangeForBuffer(*bufferID).contains(Loc);
+  return SM.getRangeForBuffer(bufferID).contains(Loc);
 }
 #endif
 

lib/AST/ASTScopeCreation.cpp

@@ -272,7 +272,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
 
     if (SF->Kind == SourceFileKind::DefaultArgument) {
       auto genInfo = *SF->getASTContext().SourceMgr.getGeneratedSourceInfo(
-          *SF->getBufferID());
+          SF->getBufferID());
       parentLoc = ASTNode::getFromOpaqueValue(genInfo.astNode).getStartLoc();
       if (auto parentScope =
               findStartingScopeForLookup(enclosingSF, parentLoc)) {

lib/AST/ASTScopePrinting.cpp

@@ -43,13 +43,8 @@ void ASTScopeImpl::dump() const { print(llvm::errs(), 0, false); }
 void ASTScopeImpl::dumpOneScopeMapLocation(
     std::pair<unsigned, unsigned> lineColumn) {
   auto bufferID = getSourceFile()->getBufferID();
-  if (!bufferID) {
-    llvm::errs() << "***No buffer, dumping all scopes***";
-    print(llvm::errs());
-    return;
-  }
   SourceLoc loc = getSourceManager().getLocForLineCol(
-      *bufferID, lineColumn.first, lineColumn.second);
+      bufferID, lineColumn.first, lineColumn.second);
   if (loc.isInvalid())
     return;
 

lib/AST/ASTScopeSourceRange.cpp

@@ -203,17 +203,9 @@ SourceRange GenericParamScope::getSourceRangeOfThisASTNode(
 
 SourceRange ASTSourceFileScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  if (auto bufferID = SF->getBufferID()) {
-    auto charRange = getSourceManager().getRangeForBuffer(*bufferID);
-    return SourceRange(charRange.getStart(), charRange.getEnd());
-  }
-
-  if (SF->getTopLevelItems().empty())
-    return SourceRange();
-
-  // Use the source ranges of the declarations in the file.
-  return SourceRange(SF->getTopLevelItems().front().getStartLoc(),
-                     SF->getTopLevelItems().back().getEndLoc());
+  auto bufferID = SF->getBufferID();
+  auto charRange = getSourceManager().getRangeForBuffer(bufferID);
+  return SourceRange(charRange.getStart(), charRange.getEnd());
 }
 
 SourceRange GenericTypeOrExtensionScope::getSourceRangeOfThisASTNode(

lib/AST/DiagnosticEngine.cpp

@@ -1759,9 +1759,7 @@ BufferIndirectlyCausingDiagnosticRAII::BufferIndirectlyCausingDiagnosticRAII(
     const SourceFile &SF)
     : Diags(SF.getASTContext().Diags) {
   auto id = SF.getBufferID();
-  if (!id)
-    return;
-  auto loc = SF.getASTContext().SourceMgr.getLocForBufferStart(*id);
+  auto loc = SF.getASTContext().SourceMgr.getLocForBufferStart(id);
   if (loc.isValid())
     Diags.setBufferIndirectlyCausingDiagnosticToInput(loc);
 }