[sourcekitd] Fix range shifting vs semantic update timing issues

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

@@ -195,11 +195,14 @@ class EditTest : public ::testing::Test {
   void doubleOpenWithDelay(std::chrono::microseconds delay, bool close);
 
   void setupThreeAnnotations(const char *DocName, TestConsumer &Consumer) {
+    // The following is engineered so that the references to `mem` are at
+    // offsets 60, 70, and 80 for convenience. They're on the same line so
+    // that tests do not accidentally depend on line separation.
     const char *Contents =
     "struct S {\n"
     "  var mem: Int = 0\n"
     "  func test() {\n"
-    "    _ = (self.mem, self.mem, self.mem)\n"
+    "          _ = mem;  _ = mem;  _ = mem\n"
     "  }\n"
     "}\n";
     const char *Args[] = { "-parse-as-library" };
@@ -385,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, {
@@ -416,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]",
   });
@@ -426,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]",
@@ -550,8 +569,6 @@ TEST_F(EditTest, AnnotationsRangeShiftingAfterEditInsertEnd) {
   ASSERT_FALSE(waitForDocUpdate()) << "timed out";
   close(DocName);
 }
-// rdar://65934938 Failing in CI with ASan
-#if defined(__has_feature) && !__has_feature(address_sanitizer)
 TEST_F(EditTest, AnnotationsRangeShiftingAfterEditReplaceEnd) {
   const char *DocName = "test.swift";
   TestConsumer Consumer;
@@ -570,7 +587,6 @@ TEST_F(EditTest, AnnotationsRangeShiftingAfterEditReplaceEnd) {
   ASSERT_FALSE(waitForDocUpdate()) << "timed out";
   close(DocName);
 }
-#endif
 TEST_F(EditTest, AnnotationsRangeShiftingAfterEditDeleteEnd) {
   const char *DocName = "test.swift";
   TestConsumer Consumer;