diff --git a/tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp b/tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp
index 60636d3852d2c..2588ba5064f25 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp
+++ b/tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp
@@ -786,8 +786,10 @@ SwiftDocumentSemanticInfo::takeSemanticTokens(
           });
 
       std::vector<SwiftSemanticToken>::iterator ReplaceEnd;
-      if (Upd->getLength() == 0) {
+      if (ReplaceBegin == SemaToks.end()) {
         ReplaceEnd = ReplaceBegin;
+      } else if (Upd->getLength() == 0) {
+        ReplaceEnd = ReplaceBegin + 1;
       } else {
         ReplaceEnd = std::upper_bound(ReplaceBegin, SemaToks.end(),
             Upd->getByteOffset() + Upd->getLength(),
diff --git a/unittests/SourceKit/SwiftLang/EditingTest.cpp b/unittests/SourceKit/SwiftLang/EditingTest.cpp
index fed1fed2fee94..5dddaf91d2610 100644
--- a/unittests/SourceKit/SwiftLang/EditingTest.cpp
+++ b/unittests/SourceKit/SwiftLang/EditingTest.cpp
@@ -34,10 +34,18 @@ static StringRef getRuntimeLibPath() {
 
 namespace {
 
-class DiagConsumer : public EditorConsumer {
+struct Token {
+  unsigned Offset;
+  unsigned Length;
+  UIdent Kind;
+  bool IsSystem;
+};
+
+class TestConsumer : public EditorConsumer {
 public:
   UIdent DiagStage;
   std::vector<DiagnosticEntryInfo> Diags;
+  std::vector<Token> Annotations;
 
 private:
   void handleRequestError(const char *Description) override {
@@ -50,7 +58,9 @@ class DiagConsumer : public EditorConsumer {
   }
 
   void handleSemanticAnnotation(unsigned Offset, unsigned Length, UIdent Kind,
-                                bool isSystem) override {}
+                                bool isSystem) override {
+    Annotations.push_back({Offset, Length, Kind, isSystem});
+  }
 
   bool documentStructureEnabled() override { return false; }
 
@@ -174,15 +184,31 @@ class EditTest : public ::testing::Test {
     return pos;
   }
 
-  void reset(DiagConsumer &Consumer) {
+  void reset(TestConsumer &Consumer) {
     Consumer.Diags.clear();
     Consumer.DiagStage = UIdent();
+    Consumer.Annotations.clear();
     std::unique_lock<std::mutex> lk(DocUpdState->Mtx);
     DocUpdState->HasUpdate = false;
   }
 
   void doubleOpenWithDelay(std::chrono::microseconds delay, bool close);
 
+  void setupThreeAnnotations(const char *DocName, TestConsumer &Consumer) {
+    const char *Contents =
+    "struct S {\n"
+    "  var mem: Int = 0\n"
+    "  func test() {\n"
+    "    _ = (self.mem, self.mem, self.mem)\n"
+    "  }\n"
+    "}\n";
+    const char *Args[] = { "-parse-as-library" };
+
+    open(DocName, Contents, Args, Consumer);
+    ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+    reset(Consumer);
+  }
+
 private:
   std::vector<const char *> makeArgs(const char *DocName,
                                      ArrayRef<const char *> CArgs) {
@@ -204,7 +230,7 @@ TEST_F(EditTest, DiagsAfterEdit) {
     "let v = 0\n";
   const char *Args[] = { "-parse-as-library" };
 
-  DiagConsumer Consumer;
+  TestConsumer Consumer;
   open(DocName, Contents, Args, Consumer);
   ASSERT_EQ(1u, Consumer.Diags.size());
   EXPECT_STREQ("expected '(' in argument list of function declaration", Consumer.Diags[0].Description.c_str());
@@ -240,7 +266,7 @@ void EditTest::doubleOpenWithDelay(std::chrono::microseconds delay,
     "func foo() { _ = unknown_name }\n";
   const char *Args[] = { "-parse-as-library" };
 
-  DiagConsumer Consumer;
+  TestConsumer Consumer;
   open(DocName, Contents, Args, Consumer);
   ASSERT_LE(Consumer.Diags.size(), 1u);
   if (Consumer.Diags.size() > 0) {
@@ -310,3 +336,326 @@ TEST_F(EditTest, DiagsAfterReopen) {
   doubleOpenWithDelay(std::chrono::milliseconds(10), false);
   doubleOpenWithDelay(std::chrono::milliseconds(100), false);
 }
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const Token &Tok) {
+  OS << "[off=" << Tok.Offset << " len=" << Tok.Length << " " << Tok.Kind.getName();
+  if (Tok.IsSystem)
+    OS << " system";
+  OS << "]";
+  return OS;
+}
+
+template<typename Collection>
+std::string dumpStrings(const Collection &Annotations) {
+  std::string tmp;
+  llvm::raw_string_ostream OS(tmp);
+  for (auto &tok : Annotations) {
+    OS << tok << "\n";
+  }
+  return OS.str();
+}
+
+void checkTokens(llvm::ArrayRef<Token> Annotations, llvm::ArrayRef<const char *> Expected) {
+  EXPECT_EQ(Annotations.size(), Expected.size()) <<
+      dumpStrings(Annotations) << "  vs\n" << dumpStrings(Expected);
+  if (Annotations.size() != Expected.size())
+    return;
+
+  for (unsigned i = 0, e = Annotations.size(); i != e; ++i) {
+    std::string tok;
+    {
+      llvm::raw_string_ostream OS(tok);
+      OS << Annotations[i];
+    }
+    EXPECT_EQ(tok, Expected[i]);
+  }
+}
+
+TEST_F(EditTest, AnnotationsAfterOpen) {
+  const char *DocName = "test.swift";
+  const char *Contents =
+    "struct S {\n"
+    "  var mem: Int = 0\n"
+    "  func test() {\n"
+    "    _ = self.mem\n"
+    "  }\n"
+    "}\n";
+  const char *Args[] = { "-parse-as-library" };
+
+  TestConsumer Consumer;
+  open(DocName, Contents, Args, Consumer);
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  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]",
+    "[off=59 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  reset(Consumer);
+  replaceText(DocName, 0, 0, "", Consumer);
+  // FIXME: we currently "take" the annotations instead of "get"ing them.
+  EXPECT_EQ(0u, Consumer.Annotations.size());
+
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsAfterEdit) {
+  const char *DocName = "test.swift";
+  const char *Contents =
+    "struct S {\n"
+    "  var mem: Int = 0\n"
+    "  func test() {\n"
+    "    _ = self.me\n"
+    "  }\n"
+    "}\n";
+  const char *Args[] = { "-parse-as-library" };
+
+  TestConsumer Consumer;
+  open(DocName, Contents, Args, Consumer);
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  reset(Consumer);
+  replaceText(DocName, 0, 0, "", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+  });
+
+  reset(Consumer);
+  replaceText(DocName, 61, 0, "m", Consumer);
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+
+  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]",
+    "[off=59 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditInsertStart) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 70, 0, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+    // Shifted by 1
+    "[off=81 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditReplaceStart) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 70, 1, "XYZ", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted 3-1 = 2
+    "[off=82 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditDeleteStart) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 70, 2, "", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted -2
+    "[off=78 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditDeleteMiddle) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 71, 2, "", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted -2
+    "[off=78 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditInsertMiddle) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 71, 0, "XY", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted 2
+    "[off=82 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditInsertEnd) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 73, 0, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted 1
+    "[off=81 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditReplaceEnd) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 72, 1, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    "[off=80 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditDeleteEnd) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 72, 1, "", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+
+    // Shifted -1
+    "[off=79 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditLast) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 80, 0, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    "[off=70 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=79 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditLast2) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 83, 0, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    "[off=70 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=80 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsRangeShiftingAfterEditTouchMultiple) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 63, 17, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    // Removed, touches edit.
+    // "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    // "[off=70 len=3 source.lang.swift.ref.var.instance]",
+    // "[off=80 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}
+
+TEST_F(EditTest, AnnotationsRangeShiftingAfterMultipleEdits) {
+  const char *DocName = "test.swift";
+  TestConsumer Consumer;
+  setupThreeAnnotations(DocName, Consumer);
+  replaceText(DocName, 83, 0, "X", Consumer);
+  checkTokens(Consumer.Annotations, {
+    "[off=22 len=3 source.lang.swift.ref.struct system]",
+    "[off=60 len=3 source.lang.swift.ref.var.instance]",
+    "[off=70 len=3 source.lang.swift.ref.var.instance]",
+    // Removed, touches edit.
+    // "[off=80 len=3 source.lang.swift.ref.var.instance]",
+  });
+
+  // Re-sync
+  ASSERT_FALSE(waitForDocUpdate()) << "timed out";
+  close(DocName);
+}