[sourcekitd] Fix range shifting "race" with a fast semantic update

If a semantic update finishes fast enough, the token snapshot may be
identical to the edit snapshot, but because of getBufferForSnapshot
consolidating edits into a new buffer, we were not detecting that case
properly, and it could cause an assertion failure (or potentially
incorrect range shifting in a release build). This would have reproduced
very rarely in practice, but I can reproduce it by putting `sleep(2)`
calls right before we read the semantic info in open and edit requests.

Incidentally, fix sourcekit-test and unit tests for the (rare) case
where an open or edit already has updated semantic info.

Author: benlangmuir

tools/SourceKit/lib/Support/ImmutableTextBuffer.cpp

@@ -107,6 +107,19 @@ bool ImmutableTextSnapshot::foreachReplaceUntil(
 
   assert(EndSnapshot);
   ImmutableTextUpdateRef Upd = DiffEnd;
+
+  // Since `getBufferForSnapshot` may produce a snapshot with edits consolidated
+  // into a buffer, we need to check the stamp explicitly for "equal" snapshots.
+  // Otherwise, the following code, which should be a no-op, will behave
+  // incorrectly:
+  //
+  //     auto snap1 = editableBuffer->getSnapshot();
+  //     auto buffer = editableBuffer->getBuffer();
+  //     auto snap2 = editableBuffer->getSnapshot();
+  //     snap2->forEachReplaceUntil(snap1, ...); // should be no-op
+  if (getStamp() == EndSnapshot->getStamp())
+    return true;
+
   while (Upd != EndSnapshot->DiffEnd) {
     Upd = Upd->getNext();
     if (!Upd) {

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

@@ -1350,9 +1350,24 @@ static void getSemanticInfoImpl(sourcekitd_variant_t Info) {
     sourcekitd_variant_dictionary_get_value(Info, KeyDiagnostics);
 }
 
-static void getSemanticInfo(sourcekitd_variant_t Info, StringRef Filename) {
-  getSemanticInfoImpl(Info);
+static void getSemanticInfoImplAfterDocUpdate(sourcekitd_variant_t EditOrOpen,
+                                              sourcekitd_variant_t DocUpdate) {
+  if (sourcekitd_variant_dictionary_get_uid(EditOrOpen, KeyDiagnosticStage) ==
+      SemaDiagnosticStage) {
+    // FIXME: currently we only return annotations once, so if the original edit
+    // or open request was slow enough, it may "take" the annotations. If that
+    // is fixed, we can skip checking the diagnostic stage and always use the
+    // DocUpdate variant.
+    assert(sourcekitd_variant_get_type(sourcekitd_variant_dictionary_get_value(
+               DocUpdate, KeyAnnotations)) == SOURCEKITD_VARIANT_TYPE_NULL);
+
+    getSemanticInfoImpl(EditOrOpen);
+  } else {
+    getSemanticInfoImpl(DocUpdate);
+  }
+}
 
+static void getSemanticInfo(sourcekitd_variant_t Info, StringRef Filename) {
   // Wait for the notification that semantic info is available.
   // But only for 1 min.
   bool expired = semaSemaphore.wait(60 * 1000);
@@ -1364,7 +1379,9 @@ static void getSemanticInfo(sourcekitd_variant_t Info, StringRef Filename) {
     llvm::report_fatal_error(
       llvm::Twine("Got notification for different doc name: ") + semaName);
   }
-  getSemanticInfoImpl(sourcekitd_response_get_value(semaResponse));
+
+  getSemanticInfoImplAfterDocUpdate(
+      Info, sourcekitd_response_get_value(semaResponse));
 }
 
 static int printAnnotations() {

unittests/SourceKit/SwiftLang/EditingTest.cpp

@@ -388,8 +388,13 @@ TEST_F(EditTest, AnnotationsAfterOpen) {
   TestConsumer Consumer;
   open(DocName, Contents, Args, Consumer);
   ASSERT_FALSE(waitForDocUpdate()) << "timed out";
-  reset(Consumer);
-  replaceText(DocName, 0, 0, "", Consumer);
+  if (Consumer.DiagStage == SemaDiagStage) {
+    // AST already built, annotations are on Consumer.
+  } else {
+    // Re-query.
+    reset(Consumer);
+    replaceText(DocName, 0, 0, "", Consumer);
+  }
 
   ASSERT_EQ(0u, Consumer.Diags.size());
   checkTokens(Consumer.Annotations, {
@@ -419,8 +424,13 @@ TEST_F(EditTest, AnnotationsAfterEdit) {
   TestConsumer Consumer;
   open(DocName, Contents, Args, Consumer);
   ASSERT_FALSE(waitForDocUpdate()) << "timed out";
-  reset(Consumer);
-  replaceText(DocName, 0, 0, "", Consumer);
+  if (Consumer.DiagStage == SemaDiagStage) {
+    // AST already built, annotations are on Consumer.
+  } else {
+    // Re-query.
+    reset(Consumer);
+    replaceText(DocName, 0, 0, "", Consumer);
+  }
   checkTokens(Consumer.Annotations, {
     "[off=22 len=3 source.lang.swift.ref.struct system]",
   });
@@ -429,8 +439,14 @@ TEST_F(EditTest, AnnotationsAfterEdit) {
   replaceText(DocName, 61, 0, "m", Consumer);
   ASSERT_FALSE(waitForDocUpdate()) << "timed out";
 
-  reset(Consumer);
-  replaceText(DocName, 0, 0, "", Consumer);
+  if (Consumer.DiagStage == SemaDiagStage) {
+    // AST already built, annotations are on Consumer.
+  } else {
+    // Re-query.
+    reset(Consumer);
+    replaceText(DocName, 0, 0, "", Consumer);
+  }
+
   ASSERT_EQ(0u, Consumer.Diags.size());
   checkTokens(Consumer.Annotations, {
     "[off=22 len=3 source.lang.swift.ref.struct system]",