Maintain a buffer ID -> source file(s) mapping in the source manager

Now that every source file has a buffer ID, introduce the reverse mapping
so clients can find the source file(s) in their module that reference
that buffer ID.

Author: DougGregor

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

@@ -3446,6 +3446,8 @@ SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
     (void)problem;
   }
 
+  M.getASTContext().SourceMgr.recordSourceFile(bufferID, this);
+
   if (Kind == SourceFileKind::MacroExpansion ||
       Kind == SourceFileKind::DefaultArgument)
     M.addAuxiliaryFile(*this);

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