Use StringMap in EditorDiagConsumer

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,

include/swift/Basic/Edit.h

@@ -13,20 +13,34 @@
 #ifndef SWIFT_BASIC_EDIT_H
 #define SWIFT_BASIC_EDIT_H
 
+#include "llvm/ADT/ArrayRef.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/SourceLoc.h"
 
 namespace swift {
   class SourceManager;
   class CharSourceRange;
 
-  struct SingleEdit {
-    SourceManager &SM;
-    CharSourceRange Range;
-    std::string Text;
+  /// A set of edits made to a source file.
+  class SourceEdits final {
+  public:
+    struct Edit {
+      std::string Path;
+      std::string Text;
+      unsigned Offset;
+      unsigned Length;
+    };
+
+  private:
+    std::vector<Edit> Edits;
+
+  public:
+    ArrayRef<Edit> getEdits() const { return Edits; }
+
+    void addEdit(SourceManager &SM, CharSourceRange Range, StringRef Text);
   };
 
-  void writeEditsInJson(ArrayRef<SingleEdit> Edits, llvm::raw_ostream &OS);
+  void writeEditsInJson(const SourceEdits &Edits, llvm::raw_ostream &OS);
 }
 
 #endif // SWIFT_BASIC_EDIT_H

lib/Basic/Edit.cpp

@@ -14,28 +14,37 @@
 #include "swift/Basic/Edit.h"
 #include "swift/Basic/SourceManager.h"
 
+using namespace swift;
+
+void SourceEdits::addEdit(SourceManager &SM, CharSourceRange Range, 
+                          StringRef Text) {
+  SourceLoc Loc = Range.getStart();
+  unsigned BufID = SM.findBufferContainingLoc(Loc);
+  unsigned Offset = SM.getLocOffsetInBuffer(Loc, BufID);
+  unsigned Length = Range.getByteLength();
+  StringRef Path = SM.getIdentifierForBuffer(BufID);
+
+  // NOTE: We cannot store SourceManager here since this logic is used by a
+  // DiagnosticConsumer where the SourceManager may not remain valid. This is
+  // the case when e.g build swift interfaces, we create a fresh
+  // CompilerInstance for a limited scope, but diagnostics are passed outside of
+  // it.
+  Edits.push_back({Path.str(), Text.str(), Offset, Length});
+}
+
 void swift::
-writeEditsInJson(ArrayRef<SingleEdit> AllEdits, llvm::raw_ostream &OS) {
+writeEditsInJson(const SourceEdits &AllEdits, llvm::raw_ostream &OS) {
   OS << "[\n";
-  for (auto &Edit : AllEdits) {
-    SourceManager &SM = Edit.SM;
-    CharSourceRange Range = Edit.Range;
-    StringRef Text = Edit.Text;
-    SourceLoc Loc = Range.getStart();
-    unsigned BufID = SM.findBufferContainingLoc(Loc);
-    unsigned Offset = SM.getLocOffsetInBuffer(Loc, BufID);
-    unsigned Length = Range.getByteLength();
-    StringRef Path(SM.getIdentifierForBuffer(BufID));
-
+  for (auto &Edit : AllEdits.getEdits()) {
     OS << " {\n";
     OS << "  \"file\": \"";
-    OS.write_escaped(Path) << "\",\n";
-    OS << "  \"offset\": " << Offset << ",\n";
-    if (Length != 0)
-      OS << "  \"remove\": " << Length << ",\n";
-    if (!Text.empty()) {
+    OS.write_escaped(Edit.Path) << "\",\n";
+    OS << "  \"offset\": " << Edit.Offset << ",\n";
+    if (Edit.Length != 0)
+      OS << "  \"remove\": " << Edit.Length << ",\n";
+    if (!Edit.Text.empty()) {
       OS << "  \"text\": \"";
-      OS.write_escaped(Text) << "\",\n";
+      OS.write_escaped(Edit.Text) << "\",\n";
     }
     OS << " },\n";
   }

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;

lib/FrontendTool/FrontendTool.cpp

@@ -216,7 +216,7 @@ class JSONFixitWriter
   std::string FixitsOutputPath;
   std::unique_ptr<llvm::raw_ostream> OSPtr;
   bool FixitAll;
-  std::vector<SingleEdit> AllEdits;
+  SourceEdits AllEdits;
 
 public:
   JSONFixitWriter(std::string fixitsOutputPath,
@@ -229,9 +229,8 @@ class JSONFixitWriter
                         const DiagnosticInfo &Info) override {
     if (!(FixitAll || shouldTakeFixit(Info)))
       return;
-    for (const auto &Fix : Info.FixIts) {
-      AllEdits.push_back({SM, Fix.getRange(), Fix.getText().str()});
-    }
+    for (const auto &Fix : Info.FixIts)
+      AllEdits.addEdit(SM, Fix.getRange(), Fix.getText());
   }
 
   bool finishProcessing() override {
@@ -251,7 +250,7 @@ class JSONFixitWriter
       return true;
     }
 
-    swift::writeEditsInJson(llvm::makeArrayRef(AllEdits), *OS);
+    swift::writeEditsInJson(AllEdits, *OS);
     return false;
   }
 };

lib/IDE/Utils.cpp

@@ -722,14 +722,14 @@ void swift::ide::SourceEditConsumer::acceptMacroExpansionBuffer(
 
 struct swift::ide::SourceEditJsonConsumer::Implementation {
   llvm::raw_ostream &OS;
-  std::vector<SingleEdit> AllEdits;
+  SourceEdits AllEdits;
   Implementation(llvm::raw_ostream &OS) : OS(OS) {}
   ~Implementation() {
     writeEditsInJson(AllEdits, OS);
   }
   void accept(SourceManager &SM, CharSourceRange Range,
               llvm::StringRef Text) {
-    AllEdits.push_back({SM, Range, Text.str()});
+    AllEdits.addEdit(SM, Range, Text);
   }
 };
 

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;