[sourcekit] Do not dequeue AST consumers expecting newer snapshots

When adding to the AST consumer queue, keep track of the snapshots
expected and only run such AST consumers when a new enough AST is built.
Progress is ensured because we always run the AST consumer that
triggered the build.

This prevents cases where enqueuing a consumer during an AST build has
different behaviour than enqueuing it after the AST has finished.

rdar://40340631

unittests/SourceKit/SwiftLang/EditingTest.cpp

@@ -142,12 +142,15 @@ public:
     Ctx->getNotificationCenter().addDocumentUpdateNotificationReceiver(Receiver);
   }
 
-  bool waitForDocUpdate() {
+  bool waitForDocUpdate(bool reset = false) {
     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(
+    auto result = !DocUpdState->CV.wait_until(
         lk, when, [&]() { return DocUpdState->HasUpdate; });
+    if (reset)
+      DocUpdState->HasUpdate = false;
+    return result;
   }
 
   void open(const char *DocName, StringRef Text, ArrayRef<const char *> CArgs,
@@ -158,6 +161,10 @@ public:
                          Args);
   }
 
+  void close(const char *DocName) {
+    getLang().editorClose(DocName, /*removeCache=*/false);
+  }
+
   void replaceText(StringRef DocName, unsigned Offset, unsigned Length,
                    StringRef Text, EditorConsumer &Consumer) {
     auto Buf = MemoryBuffer::getMemBufferCopy(Text, DocName);
@@ -177,6 +184,8 @@ public:
     DocUpdState->HasUpdate = false;
   }
 
+  void doubleOpenWithDelay(useconds_t delay, bool close);
+
 private:
   std::vector<const char *> makeArgs(const char *DocName,
                                      ArrayRef<const char *> CArgs) {
@@ -226,3 +235,65 @@ TEST_F(EditTest, DiagsAfterEdit) {
   }
   EXPECT_EQ(SemaDiagStage, Consumer.DiagStage);
 }
+
+void EditTest::doubleOpenWithDelay(useconds_t delay, bool closeDoc) {
+  const char *DocName = "/test.swift";
+  const char *Contents =
+    "func foo() { _ = unknown_name }\n";
+  const char *Args[] = { "-parse-as-library" };
+
+  DiagConsumer Consumer;
+  open(DocName, Contents, Args, Consumer);
+  ASSERT_EQ(0u, Consumer.Diags.size());
+  // Open again without closing; this reinitializes the semantic info on the doc
+  if (delay)
+    usleep(delay);
+  if (closeDoc)
+    close(DocName);
+  reset(Consumer);
+  open(DocName, Contents, Args, Consumer);
+  ASSERT_EQ(0u, Consumer.Diags.size());
+
+  // Wait for the document update from the second time we open the document. We
+  // may or may not get a notification from the first time it was opened, but
+  // only the second time will there be any semantic information available to
+  // be queried, since the semantic info from the first open is unreachable.
+  for (int i = 0; i < 2; ++i) {
+    bool expired = waitForDocUpdate(/*reset=*/true);
+    ASSERT_FALSE(expired) << "no second notification";
+    replaceText(DocName, 0, 0, StringRef(), Consumer);
+    if (!Consumer.Diags.empty())
+      break;
+    ASSERT_EQ(0, i) << "no diagnostics after second notification";
+  }
+
+  ASSERT_EQ(1u, Consumer.Diags.size());
+  EXPECT_STREQ("use of unresolved identifier 'unknown_name'", Consumer.Diags[0].Description.c_str());
+}
+
+TEST_F(EditTest, DiagsAfterCloseAndReopen) {
+  // Attempt to open the same file twice in a row. This tests (subject to
+  // timing) cases where:
+  // * the 2nd open happens before the first AST starts building
+  // * the 2nd open happens after the first AST starts building
+  // * the 2nd open happens after the AST finishes
+
+  // The middle case in particular verifies the ASTManager is only calling the
+  // correct ASTConsumers.
+
+  doubleOpenWithDelay(0, true);
+  doubleOpenWithDelay(1000, true);   // 1 ms
+  doubleOpenWithDelay(10000, true);  // 10 ms
+  doubleOpenWithDelay(100000, true); // 100 ms
+}
+
+TEST_F(EditTest, DiagsAfterReopen) {
+  // See description of DiagsAfterCloseAndReopen, but in this case we don't
+  // close the original document, causing it to reinitialize instead of create
+  // a fresh document.
+
+  doubleOpenWithDelay(0, false);
+  doubleOpenWithDelay(1000, false);   // 1 ms
+  doubleOpenWithDelay(10000, false);  // 10 ms
+  doubleOpenWithDelay(100000, false); // 100 ms
+}