Use StringMap in EditorDiagConsumer

As noted by the comment, storing a `StringRef` is
indeed brittle. Update to use a StringMap.

rdar://111589090

Author: hamishknight

include/swift/AST/DiagnosticConsumer.h

@@ -122,7 +122,9 @@ class DiagnosticConsumer {
   /// Invoked whenever the frontend emits a diagnostic.
   ///
   /// \param SM The source manager associated with the source locations in
-  /// this diagnostic.
+  /// this diagnostic. NOTE: Do not persist either the SourceManager, or the
+  /// buffer names from the SourceManager, since it may not outlive the
+  /// DiagnosticConsumer (this is the case when building module interfaces).
   ///
   /// \param Info Information describing the diagnostic.
   virtual void handleDiagnostic(SourceManager &SM,

lib/Frontend/SerializedDiagnosticConsumer.cpp

@@ -228,11 +228,12 @@ namespace serialized_diagnostics {
 unsigned SerializedDiagnosticConsumer::getEmitFile(
     SourceManager &SM, StringRef Filename, unsigned bufferID
 ) {
-  // NOTE: Using Filename.data() here relies on SourceMgr using
-  // const char* as buffer identifiers.  This is fast, but may
-  // be brittle.  We can always switch over to using a StringMap.
-  // Note that the logic in EditorDiagConsumer::getBufferInfo
-  // will also need changing.
+  // FIXME: Using Filename.data() here is wrong, since the provided
+  // SourceManager may not live as long as this consumer (which is
+  // the case if it's a diagnostic produced from building a module
+  // interface). We ought to switch over to using a StringMap once
+  // buffer names are unique (currently not the case for
+  // pretty-printed decl buffers).
   unsigned &existingEntry = State->Files[Filename.data()];
   if (existingEntry)
     return existingEntry;

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -91,10 +91,6 @@ BufferInfoSharedPtr
 EditorDiagConsumer::getBufferInfo(StringRef FileName,
                                   llvm::Optional<unsigned> BufferID,
                                   swift::SourceManager &SM) {
-  // NOTE: Using StringRef as a key here relies on SourceMgr using const char*
-  // as buffer identifiers. This is fast, but may be brittle.  We can always
-  // switch over to using a StringMap. Note that the logic in
-  // SerializedDiagnosticConsumer::getEmitFile will also need changing.
   auto Result = BufferInfos.find(FileName);
   if (Result != BufferInfos.end())
     return Result->second;

tools/SourceKit/lib/SwiftLang/SwiftEditorDiagConsumer.h

@@ -16,7 +16,7 @@
 #include "SourceKit/Core/LangSupport.h"
 #include "swift/AST/DiagnosticConsumer.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/StringMap.h"
 
 namespace SourceKit {
 
@@ -27,7 +27,7 @@ class EditorDiagConsumer : public swift::DiagnosticConsumer {
   llvm::DenseMap<unsigned, DiagnosticsTy> BufferDiagnostics;
   DiagnosticsTy InvalidLocDiagnostics;
 
-  llvm::MapVector<StringRef, BufferInfoSharedPtr> BufferInfos;
+  llvm::StringMap<BufferInfoSharedPtr> BufferInfos;
 
   int LastDiagBufferID = -1;
   unsigned LastDiagIndex = 0;