[SourceKit] If diagnostics are 'stale' for a particular snapshot then ignore them and only return the syntactic parser diagnostics (#10388)

This makes sure that diagnostics returned for a particular state of source buffer are consistent and accurate.
rdar://32769873

test/SourceKit/CursorInfo/rdar_18677108-2.swift.response

@@ -5,7 +5,7 @@
     key.filepath: rdar_18677108-2-a.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected ')' in expression list",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse,
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema,
     key.diagnostics: [
       {
         key.line: 5,
@@ -22,7 +22,7 @@
     key.filepath: rdar_18677108-2-a.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected '}' at end of brace statement",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse,
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema,
     key.diagnostics: [
       {
         key.line: 4,
@@ -39,7 +39,7 @@
     key.filepath: rdar_18677108-2-a.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected '}' in class",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse,
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema,
     key.diagnostics: [
       {
         key.line: 2,

test/SourceKit/Sema/sema_diag_after_edit_fixit.swift

@@ -4,7 +4,7 @@ let a = 0; let b = 0 }; unresolved
 // Test that offsets of diagnostic ranges and fixits get updated correctly after the edit request
 
 // RUN: %sourcekitd-test -req=open %s -- %s == -req=print-diags %s \
-// RUN:    == -req=edit -pos=2:1 -replace="_" -length=5 %s -print-raw-response \
+// RUN:    == -req=edit -pos=2:1 -replace="_" -length=5 %s == -req=print-diags %s \
 // RUN: | %FileCheck %s
 
 // CHECK:      key.line: 2,

test/SourceKit/Sema/sema_playground.swift.response

@@ -12,6 +12,6 @@
     key.filepath: sema_playground.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected numeric value following '$'",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema
   }
 ]

test/SourceKit/SyntaxMapData/diags.swift.response

@@ -6,7 +6,7 @@
     key.filepath: parse_error.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected '(' in argument list of function declaration",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse,
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema,
     key.fixits: [
       {
         key.offset: 26,
@@ -37,7 +37,7 @@
     key.filepath: parse_error.swift,
     key.severity: source.diagnostic.severity.error,
     key.description: "expected '(' in argument list of function declaration",
-    key.diagnostic_stage: source.diagnostic.stage.swift.parse,
+    key.diagnostic_stage: source.diagnostic.stage.swift.sema,
     key.fixits: [
       {
         key.offset: 29,

tools/SourceKit/include/SourceKit/Core/Context.h

@@ -35,7 +35,8 @@ class Context {
 public:
   Context(StringRef RuntimeLibPath,
           llvm::function_ref<
-              std::unique_ptr<LangSupport>(Context &)> LangSupportFactoryFn);
+              std::unique_ptr<LangSupport>(Context &)> LangSupportFactoryFn,
+          bool shouldDispatchNotificationsOnMain = true);
   ~Context();
 
   StringRef getRuntimeLibPath() const { return RuntimeLibPath; }

tools/SourceKit/include/SourceKit/Core/NotificationCenter.h

@@ -14,6 +14,7 @@
 #define LLVM_SOURCEKIT_CORE_NOTIFICATIONCENTER_H
 
 #include "SourceKit/Core/LLVM.h"
+#include "llvm/Support/Mutex.h"
 #include <functional>
 #include <vector>
 
@@ -23,9 +24,14 @@ typedef std::function<void(StringRef DocumentName)>
     DocumentUpdateNotificationReceiver;
 
 class NotificationCenter {
+  bool DispatchToMain;
   std::vector<DocumentUpdateNotificationReceiver> DocUpdReceivers;
+  mutable llvm::sys::Mutex Mtx;
 
 public:
+  explicit NotificationCenter(bool dispatchToMain);
+  ~NotificationCenter();
+
   void addDocumentUpdateNotificationReceiver(
       DocumentUpdateNotificationReceiver Receiver);
 

tools/SourceKit/lib/Core/Context.cpp

@@ -18,8 +18,9 @@ using namespace SourceKit;
 
 SourceKit::Context::Context(StringRef RuntimeLibPath,
     llvm::function_ref<std::unique_ptr<LangSupport>(Context &)>
-    LangSupportFactoryFn) : RuntimeLibPath(RuntimeLibPath),
+    LangSupportFactoryFn,
-    NotificationCtr(new NotificationCenter()) {
+    bool shouldDispatchNotificationsOnMain) : RuntimeLibPath(RuntimeLibPath),
+    NotificationCtr(new NotificationCenter(shouldDispatchNotificationsOnMain)) {
   // Should be called last after everything is initialized.
   SwiftLang = LangSupportFactoryFn(*this);
 }

tools/SourceKit/lib/Core/NotificationCenter.cpp

@@ -15,20 +15,33 @@
 
 using namespace SourceKit;
 
+NotificationCenter::NotificationCenter(bool dispatchToMain)
+  : DispatchToMain(dispatchToMain) {
+}
+NotificationCenter::~NotificationCenter() {}
+
 void NotificationCenter::addDocumentUpdateNotificationReceiver(
     DocumentUpdateNotificationReceiver Receiver) {
 
-  WorkQueue::dispatchOnMain([this, Receiver]{
+  llvm::sys::ScopedLock L(Mtx);
-    DocUpdReceivers.push_back(Receiver);
+  DocUpdReceivers.push_back(Receiver);
-  });
 }
 
 void NotificationCenter::postDocumentUpdateNotification(
     StringRef DocumentName) const {
-  
+
-  std::string DocName = DocumentName;
+  std::vector<DocumentUpdateNotificationReceiver> recvs;
-  WorkQueue::dispatchOnMain([this, DocName]{
+  {
-    for (auto &Fn : DocUpdReceivers)
+    llvm::sys::ScopedLock L(Mtx);
-      Fn(DocName);
+    recvs = DocUpdReceivers;
-  });
+  }  
+  std::string docName = DocumentName;
+  auto sendNote = [recvs, docName]{
+    for (auto &Fn : recvs)
+      Fn(docName);
+  };
+  if (DispatchToMain)
+    WorkQueue::dispatchOnMain(sendNote);
+  else
+    sendNote();
 }

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -546,7 +546,7 @@ public:
 
   void readSemanticInfo(ImmutableTextSnapshotRef NewSnapshot,
                         std::vector<SwiftSemanticToken> &Tokens,
-                        std::vector<DiagnosticEntryInfo> &Diags,
+                        Optional<std::vector<DiagnosticEntryInfo>> &Diags,
                         ArrayRef<DiagnosticEntryInfo> ParserDiags);
 
   void processLatestSnapshotAsync(EditableTextBufferRef EditableBuffer);
@@ -564,7 +564,7 @@ private:
   std::vector<SwiftSemanticToken> takeSemanticTokens(
       ImmutableTextSnapshotRef NewSnapshot);
 
-  std::vector<DiagnosticEntryInfo> getSemanticDiagnostics(
+  Optional<std::vector<DiagnosticEntryInfo>> getSemanticDiagnostics(
       ImmutableTextSnapshotRef NewSnapshot,
       ArrayRef<DiagnosticEntryInfo> ParserDiags);
 };
@@ -656,7 +656,7 @@ uint64_t SwiftDocumentSemanticInfo::getASTGeneration() const {
 void SwiftDocumentSemanticInfo::readSemanticInfo(
     ImmutableTextSnapshotRef NewSnapshot,
     std::vector<SwiftSemanticToken> &Tokens,
-    std::vector<DiagnosticEntryInfo> &Diags,
+    Optional<std::vector<DiagnosticEntryInfo>> &Diags,
     ArrayRef<DiagnosticEntryInfo> ParserDiags) {
 
   llvm::sys::ScopedLock L(Mtx);
@@ -711,155 +711,73 @@ SwiftDocumentSemanticInfo::takeSemanticTokens(
   return std::move(SemaToks);
 }
 
-static bool
+Optional<std::vector<DiagnosticEntryInfo>>
-adjustDiagnosticRanges(SmallVectorImpl<std::pair<unsigned, unsigned>> &Ranges,
-                       unsigned ByteOffset, unsigned RemoveLen, int Delta) {
-  for (auto &Range : Ranges) {
-    unsigned RangeBegin = Range.first;
-    unsigned RangeEnd = Range.first +  Range.second;
-    unsigned RemoveEnd = ByteOffset + RemoveLen;
-    // If it intersects with the remove range, ignore the whole diagnostic.
-    if (!(RangeEnd < ByteOffset || RangeBegin > RemoveEnd))
-      return true; // Ignore.
-    if (RangeBegin > RemoveEnd)
-      Range.first += Delta;
-  }
-  return false;
-}
-
-static bool
-adjustDiagnosticFixits(SmallVectorImpl<DiagnosticEntryInfo::Fixit> &Fixits,
-                       unsigned ByteOffset, unsigned RemoveLen, int Delta) {
-  for (auto &Fixit : Fixits) {
-    unsigned FixitBegin = Fixit.Offset;
-    unsigned FixitEnd = Fixit.Offset +  Fixit.Length;
-    unsigned RemoveEnd = ByteOffset + RemoveLen;
-    // If it intersects with the remove range, ignore the whole diagnostic.
-    if (!(FixitEnd < ByteOffset || FixitBegin > RemoveEnd))
-      return true; // Ignore.
-    if (FixitBegin > RemoveEnd)
-      Fixit.Offset += Delta;
-  }
-  return false;
-}
-
-static bool
-adjustDiagnosticBase(DiagnosticEntryInfoBase &Diag,
-                     unsigned ByteOffset, unsigned RemoveLen, int Delta) {
-  if (Diag.Offset >= ByteOffset && Diag.Offset < ByteOffset+RemoveLen)
-    return true; // Ignore.
-  bool Ignore = adjustDiagnosticRanges(Diag.Ranges, ByteOffset, RemoveLen, Delta);
-  if (Ignore)
-    return true;
-  Ignore = adjustDiagnosticFixits(Diag.Fixits, ByteOffset, RemoveLen, Delta);
-  if (Ignore)
-    return true;
-  if (Diag.Offset > ByteOffset)
-    Diag.Offset += Delta;
-  return false;
-}
-
-static bool
-adjustDiagnostic(DiagnosticEntryInfo &Diag, StringRef Filename,
-                 unsigned ByteOffset, unsigned RemoveLen, int Delta) {
-  for (auto &Note : Diag.Notes) {
-    if (Filename != Note.Filename)
-      continue;
-    bool Ignore = adjustDiagnosticBase(Note, ByteOffset, RemoveLen, Delta);
-    if (Ignore)
-      return true;
-  }
-  return adjustDiagnosticBase(Diag, ByteOffset, RemoveLen, Delta);
-}
-
-static std::vector<DiagnosticEntryInfo>
-adjustDiagnostics(std::vector<DiagnosticEntryInfo> Diags, StringRef Filename,
-                  unsigned ByteOffset, unsigned RemoveLen, int Delta) {
-  std::vector<DiagnosticEntryInfo> NewDiags;
-  NewDiags.reserve(Diags.size());
-
-  for (auto &Diag : Diags) {
-    bool Ignore = adjustDiagnostic(Diag, Filename, ByteOffset, RemoveLen, Delta);
-    if (!Ignore) {
-      NewDiags.push_back(std::move(Diag));
-    }
-  }
-
-  return NewDiags;
-}
-
-std::vector<DiagnosticEntryInfo>
 SwiftDocumentSemanticInfo::getSemanticDiagnostics(
     ImmutableTextSnapshotRef NewSnapshot,
     ArrayRef<DiagnosticEntryInfo> ParserDiags) {
 
-  llvm::sys::ScopedLock L(Mtx);
+  std::vector<DiagnosticEntryInfo> curSemaDiags;
+  {
+    llvm::sys::ScopedLock L(Mtx);
 
-  if (SemaDiags.empty())
+    if (!DiagSnapshot || DiagSnapshot->getStamp() != NewSnapshot->getStamp()) {
-    return SemaDiags;
+      // The semantic diagnostics are out-of-date, ignore them.
-
+      return llvm::None;
-  assert(DiagSnapshot && "If we have diagnostics, we must have snapshot!");
-
-  if (!DiagSnapshot->precedesOrSame(NewSnapshot)) {
-    // It may happen that other thread has already updated the diagnostics to
-    // the version *after* NewSnapshot. This can happen in at least two cases:
-    //   (a) two or more editor.open or editor.replacetext queries are being
-    //       processed concurrently (not valid, but possible call pattern)
-    //   (b) while editor.replacetext processing is running, a concurrent
-    //       thread executes getBuffer()/getBufferForSnapshot() on the same
-    //       Snapshot/EditableBuffer (thus creating a new ImmutableTextBuffer)
-    //       and updates DiagSnapshot/SemaDiags
-    assert(NewSnapshot->precedesOrSame(DiagSnapshot));
-
-    // Since we cannot "adjust back" diagnostics, we just return an empty set.
-    // FIXME: add handling of the case#b above
-    return {};
-  }
-
-  SmallVector<unsigned, 16> ParserDiagLines;
-  for (auto Diag : ParserDiags)
-    ParserDiagLines.push_back(Diag.Line);
-  std::sort(ParserDiagLines.begin(), ParserDiagLines.end());
-
-  auto hasParserDiagAtLine = [&](unsigned Line) {
-    return std::binary_search(ParserDiagLines.begin(), ParserDiagLines.end(),
-                              Line);
-  };
-
-  // Adjust the position of the diagnostics.
-  DiagSnapshot->foreachReplaceUntil(NewSnapshot,
-    [&](ReplaceImmutableTextUpdateRef Upd) -> bool {
-      if (SemaDiags.empty())
-        return false;
-
-      unsigned ByteOffset = Upd->getByteOffset();
-      unsigned RemoveLen = Upd->getLength();
-      unsigned InsertLen = Upd->getText().size();
-      int Delta = InsertLen - RemoveLen;
-      SemaDiags = adjustDiagnostics(std::move(SemaDiags), Filename,
-                                    ByteOffset, RemoveLen, Delta);
-      return true;
-    });
-
-  if (!SemaDiags.empty()) {
-    auto ImmBuf = NewSnapshot->getBuffer();
-    for (auto &Diag : SemaDiags) {
-      std::tie(Diag.Line, Diag.Column) = ImmBuf->getLineAndColumn(Diag.Offset);
     }
 
-    // If there is a parser diagnostic in a line, ignore diagnostics in the same
+    curSemaDiags = SemaDiags;
-    // line that we got from the semantic pass.
-    // Note that the semantic pass also includes parser diagnostics so this
-    // avoids duplicates.
-    SemaDiags.erase(std::remove_if(SemaDiags.begin(), SemaDiags.end(),
-                                   [&](const DiagnosticEntryInfo &Diag) -> bool {
-                                     return hasParserDiagAtLine(Diag.Line);
-                                   }),
-                    SemaDiags.end());
   }
 
-  DiagSnapshot = NewSnapshot;
+  // Diagnostics from the AST and diagnostics from the parser are based on the
-  return SemaDiags;
+  // same source text snapshot. But diagnostics from the AST may have excluded
+  // the parser diagnostics due to a fatal error, e.g. if the source has a
+  // 'so such module' error, which will suppress other diagnostics.
+  // We don't want to turn off the suppression to avoid a flood of diagnostics
+  // when a module import fails, but we also don't want to lose the parser
+  // diagnostics in such a case, so merge the parser diagnostics with the sema
+  // ones.
+
+  auto orderDiagnosticEntryInfos = [](const DiagnosticEntryInfo &LHS,
+                                      const DiagnosticEntryInfo &RHS) -> bool {
+    if (LHS.Filename != RHS.Filename)
+      return LHS.Filename < RHS.Filename;
+    if (LHS.Offset != RHS.Offset)
+      return LHS.Offset < RHS.Offset;
+    return LHS.Description < RHS.Description;
+  };
+
+  std::vector<DiagnosticEntryInfo> sortedParserDiags;
+  sortedParserDiags.reserve(ParserDiags.size());
+  sortedParserDiags.insert(sortedParserDiags.end(), ParserDiags.begin(),
+                           ParserDiags.end());
+  std::stable_sort(sortedParserDiags.begin(), sortedParserDiags.end(),
+                   orderDiagnosticEntryInfos);
+
+  std::vector<DiagnosticEntryInfo> finalDiags;
+  finalDiags.reserve(sortedParserDiags.size()+curSemaDiags.size());
+
+  // Add sema diagnostics unless it is an existing parser diagnostic.
+  // Note that we want to merge and eliminate diagnostics from the 'sema' set
+  // that also show up in the 'parser' set, but we don't want to remove
+  // duplicate diagnostics from within the same set (e.g. duplicates existing in
+  // the 'sema' set). We want to report the diagnostics as the compiler reported
+  // them, even if there's some duplicate one. This is why we don't just do a
+  // simple append/sort/keep-uniques step.
+  for (const auto &curDE : curSemaDiags) {
+    bool existsAsParserDiag = std::binary_search(sortedParserDiags.begin(),
+                                                 sortedParserDiags.end(),
+                                             curDE, orderDiagnosticEntryInfos);
+    if (!existsAsParserDiag) {
+      finalDiags.push_back(curDE);
+    }
+  }
+
+  finalDiags.insert(finalDiags.end(),
+                    sortedParserDiags.begin(), sortedParserDiags.end());
+  std::stable_sort(finalDiags.begin(), finalDiags.end(),
+                   orderDiagnosticEntryInfos);
+
+  return finalDiags;
 }
 
 void SwiftDocumentSemanticInfo::updateSemanticInfo(
@@ -1809,10 +1727,7 @@ void SwiftEditorDocument::readSemanticInfo(ImmutableTextSnapshotRef Snapshot,
   }
 
   std::vector<SwiftSemanticToken> SemaToks;
-  std::vector<DiagnosticEntryInfo> SemaDiags;
+  Optional<std::vector<DiagnosticEntryInfo>> SemaDiags;
-
-  // FIXME: Parser diagnostics should be filtered out of the semantic ones,
-  // Then just merge the semantic ones with the current parse ones.
   Impl.SemanticInfo->readSemanticInfo(Snapshot, SemaToks, SemaDiags,
                                       Impl.ParserDiagnostics);
 
@@ -1829,16 +1744,17 @@ void SwiftEditorDocument::readSemanticInfo(ImmutableTextSnapshotRef Snapshot,
   static UIdent SemaDiagStage("source.diagnostic.stage.swift.sema");
   static UIdent ParseDiagStage("source.diagnostic.stage.swift.parse");
 
-  if (!SemaDiags.empty() || !SemaToks.empty()) {
+  // If there's no value returned for diagnostics it means they are out-of-date
+  // (based on a different snapshot).
+  if (SemaDiags.hasValue()) {
     Consumer.setDiagnosticStage(SemaDiagStage);
+    for (auto &Diag : SemaDiags.getValue())
+      Consumer.handleDiagnostic(Diag, SemaDiagStage);
   } else {
     Consumer.setDiagnosticStage(ParseDiagStage);
+    for (auto &Diag : Impl.ParserDiagnostics)
+      Consumer.handleDiagnostic(Diag, ParseDiagStage);
   }
-
-  for (auto &Diag : Impl.ParserDiagnostics)
-    Consumer.handleDiagnostic(Diag, ParseDiagStage);
-  for (auto &Diag : SemaDiags)
-    Consumer.handleDiagnostic(Diag, SemaDiagStage);
 }
 
 void SwiftEditorDocument::removeCachedAST() {

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -1055,26 +1055,10 @@ sourcekitd_variant_t LatestSemaAnnotations = {{0,0,0}};
 sourcekitd_variant_t LatestSemaDiags = {{0,0,0}};
 
 static void getSemanticInfoImpl(sourcekitd_variant_t Info) {
-  sourcekitd_variant_t annotations =
+  LatestSemaAnnotations =
     sourcekitd_variant_dictionary_get_value(Info, KeyAnnotations);
-  sourcekitd_variant_t diagnosticStage =
+  LatestSemaDiags =
-    sourcekitd_variant_dictionary_get_value(Info, KeyDiagnosticStage);
+    sourcekitd_variant_dictionary_get_value(Info, KeyDiagnostics);
-
-  auto hasSemaInfo = [&]{
-    if (sourcekitd_variant_get_type(annotations) != SOURCEKITD_VARIANT_TYPE_NULL)
-      return true;
-    if (sourcekitd_variant_get_type(diagnosticStage) == SOURCEKITD_VARIANT_TYPE_UID) {
-      return sourcekitd_variant_uid_get_value(diagnosticStage) == SemaDiagnosticStage;
-    }
-    return false;
-  };
-
-  if (hasSemaInfo()) {
-    LatestSemaAnnotations =
-      sourcekitd_variant_dictionary_get_value(Info, KeyAnnotations);
-    LatestSemaDiags =
-      sourcekitd_variant_dictionary_get_value(Info, KeyDiagnostics);
-  }
 }
 
 static void getSemanticInfo(sourcekitd_variant_t Info, StringRef Filename) {

unittests/SourceKit/SwiftLang/CMakeLists.txt

@@ -2,6 +2,7 @@ if(NOT SWIFT_HOST_VARIANT MATCHES "${SWIFT_DARWIN_EMBEDDED_VARIANTS}")
 
   add_swift_unittest(SourceKitSwiftLangTests
     CursorInfoTest.cpp
+    EditingTest.cpp
     )
 
   target_link_libraries(SourceKitSwiftLangTests

unittests/SourceKit/SwiftLang/EditingTest.cpp

@@ -0,0 +1,224 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "SourceKit/Core/Context.h"
+#include "SourceKit/Core/LangSupport.h"
+#include "SourceKit/Core/NotificationCenter.h"
+#include "SourceKit/Support/Concurrency.h"
+#include "SourceKit/SwiftLang/Factory.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/TargetSelect.h"
+#include "gtest/gtest.h"
+#include <mutex>
+#include <condition_variable>
+
+using namespace SourceKit;
+using namespace llvm;
+
+static StringRef getRuntimeLibPath() {
+  return sys::path::parent_path(SWIFTLIB_DIR);
+}
+
+namespace {
+
+class DiagConsumer : public EditorConsumer {
+public:
+  UIdent DiagStage;
+  std::vector<DiagnosticEntryInfo> Diags;
+
+private:
+  void handleRequestError(const char *Description) override {
+    llvm_unreachable("unexpected error");
+  }
+
+  bool handleSyntaxMap(unsigned Offset, unsigned Length, UIdent Kind) override {
+    return false;
+  }
+
+  bool handleSemanticAnnotation(unsigned Offset, unsigned Length,
+                                UIdent Kind, bool isSystem) override {
+    return false;
+  }
+
+  bool beginDocumentSubStructure(unsigned Offset, unsigned Length,
+                                 UIdent Kind, UIdent AccessLevel,
+                                 UIdent SetterAccessLevel,
+                                 unsigned NameOffset,
+                                 unsigned NameLength,
+                                 unsigned BodyOffset,
+                                 unsigned BodyLength,
+                                 StringRef DisplayName,
+                                 StringRef TypeName,
+                                 StringRef RuntimeName,
+                                 StringRef SelectorName,
+                                 ArrayRef<StringRef> InheritedTypes,
+                                 ArrayRef<UIdent> Attrs) override {
+    return false;
+  }
+
+  bool endDocumentSubStructure() override { return false; }
+
+  bool handleDocumentSubStructureElement(UIdent Kind,
+                                         unsigned Offset,
+                                         unsigned Length) override {
+    return false;
+  }
+
+  bool recordAffectedRange(unsigned Offset, unsigned Length) override {
+    return false;
+  }
+  
+  bool recordAffectedLineRange(unsigned Line, unsigned Length) override {
+    return false;
+  }
+
+  bool recordFormattedText(StringRef Text) override { return false; }
+
+  bool setDiagnosticStage(UIdent diagStage) override {
+    DiagStage = diagStage;
+    return true;
+  }
+  bool handleDiagnostic(const DiagnosticEntryInfo &Info,
+                        UIdent DiagStage) override {
+    Diags.push_back(Info);
+    return true;
+  }
+
+  bool handleSourceText(StringRef Text) override { return false; }
+};
+
+struct DocUpdateMutexState {
+  std::mutex Mtx;
+  std::condition_variable CV;
+  bool HasUpdate = false;
+};
+
+class EditTest : public ::testing::Test {
+  SourceKit::Context *Ctx;
+  std::shared_ptr<DocUpdateMutexState> DocUpdState;
+
+public:
+  EditTest() {
+    // This is avoiding destroying \p SourceKit::Context because another
+    // thread may be active trying to use it to post notifications.
+    // FIXME: Use shared_ptr ownership to avoid such issues.
+    Ctx = new SourceKit::Context(getRuntimeLibPath(),
+                                 SourceKit::createSwiftLangSupport,
+                                 /*dispatchOnMain=*/false);
+    auto localDocUpdState = std::make_shared<DocUpdateMutexState>();
+    Ctx->getNotificationCenter().addDocumentUpdateNotificationReceiver(
+        [localDocUpdState](StringRef docName) {
+          std::unique_lock<std::mutex> lk(localDocUpdState->Mtx);
+          localDocUpdState->HasUpdate = true;
+          localDocUpdState->CV.notify_one();
+        });
+    DocUpdState = localDocUpdState;
+  }
+
+  LangSupport &getLang() { return Ctx->getSwiftLangSupport(); }
+
+  void SetUp() override {
+    llvm::InitializeAllTargets();
+    llvm::InitializeAllTargetMCs();
+    llvm::InitializeAllAsmPrinters();
+    llvm::InitializeAllAsmParsers();
+  }
+
+  void addNotificationReceiver(DocumentUpdateNotificationReceiver Receiver) {
+    Ctx->getNotificationCenter().addDocumentUpdateNotificationReceiver(Receiver);
+  }
+
+  bool waitForDocUpdate() {
+    std::chrono::seconds secondsToWait(10);
+    std::unique_lock<std::mutex> lk(DocUpdState->Mtx);
+    auto when = std::chrono::system_clock::now() + secondsToWait;
+    return !DocUpdState->CV.wait_until(
+        lk, when, [&]() { return DocUpdState->HasUpdate; });
+  }
+
+  void open(const char *DocName, StringRef Text, ArrayRef<const char *> CArgs,
+            EditorConsumer &Consumer) {
+    auto Args = makeArgs(DocName, CArgs);
+    auto Buf = MemoryBuffer::getMemBufferCopy(Text, DocName);
+    getLang().editorOpen(DocName, Buf.get(), /*EnableSyntaxMap=*/false, Consumer,
+                         Args);
+  }
+
+  void replaceText(StringRef DocName, unsigned Offset, unsigned Length,
+                   StringRef Text, EditorConsumer &Consumer) {
+    auto Buf = MemoryBuffer::getMemBufferCopy(Text, DocName);
+    getLang().editorReplaceText(DocName, Buf.get(), Offset, Length, Consumer);
+  }
+
+  unsigned findOffset(StringRef Val, StringRef Text) {
+    auto pos = Text.find(Val);
+    assert(pos != StringRef::npos);
+    return pos;
+  }
+
+  void reset(DiagConsumer &Consumer) {
+    Consumer.Diags.clear();
+    Consumer.DiagStage = UIdent();
+    std::unique_lock<std::mutex> lk(DocUpdState->Mtx);
+    DocUpdState->HasUpdate = false;
+  }
+
+private:
+  std::vector<const char *> makeArgs(const char *DocName,
+                                     ArrayRef<const char *> CArgs) {
+    std::vector<const char *> Args = CArgs;
+    Args.push_back(DocName);
+    return Args;
+  }
+};
+
+} // anonymous namespace
+
+static UIdent SemaDiagStage("source.diagnostic.stage.swift.sema");
+static UIdent ParseDiagStage("source.diagnostic.stage.swift.parse");
+
+TEST_F(EditTest, DiagsAfterEdit) {
+  const char *DocName = "/test.swift";
+  const char *Contents =
+    "func foo\n"
+    "let v = 0\n";
+  const char *Args[] = { "-parse-as-library" };
+
+  DiagConsumer Consumer;
+  open(DocName, Contents, Args, Consumer);
+  ASSERT_TRUE(Consumer.Diags.size() == 1);
+  EXPECT_STREQ("expected '(' in argument list of function declaration", Consumer.Diags[0].Description.c_str());
+
+  reset(Consumer);
+  replaceText(DocName, findOffset("func foo", Contents)+strlen("func foo"), 0, "(){}", Consumer);
+  EXPECT_TRUE(Consumer.Diags.empty());
+
+  bool expired = waitForDocUpdate();
+  ASSERT_FALSE(expired);
+
+  reset(Consumer);
+  replaceText(DocName, 0, 0, StringRef(), Consumer);
+  EXPECT_TRUE(Consumer.Diags.empty());
+
+  // If typecheck completed for the edit we'll get 'sema stage', otherwise
+  // we'll get 'parser stage' and there will be another doc-update.
+  if (Consumer.DiagStage == ParseDiagStage) {
+    bool expired = waitForDocUpdate();
+    ASSERT_FALSE(expired);
+
+    reset(Consumer);
+    replaceText(DocName, 0, 0, StringRef(), Consumer);
+    EXPECT_TRUE(Consumer.Diags.empty());
+  }
+  EXPECT_EQ(SemaDiagStage, Consumer.DiagStage);
+}