[clangd] Duplicate lines of semantic highlightings sent removed.

[jvikstrom]

Summary

Added a function for diffing highlightings and removing duplicate lines. It also returns empty lines if the IDE should remove previous highlighting tokens on the line. Integrated with state into the highlighting generation flow. Only works correctly as long as there are no multiline tokens.

[hokein]

Thanks for bringing this up and implementing it.

I haven't looked into the details of the patch, just some high-level comments:

• the scope of the patch seems a bit unclear to me, what's the problem we are trying to solve?
• the patch looks like calculating the diff/delta highlightings (pre highlightings vs. new highlightings) from a server-only perspective, which seems incorrect. At least we should make sure that LSP client and server share the same understanding of source diff change, otherwise client may render the delta information in a wrong way. To do that, I guess we may start from the DidChangeTextDocument notification (where the client sends source changes to server);
• I'm not sure that we should start implement this at the moment -- I don't see that we will get much performance gain, we save some traffic cost between LSP client and server, but at the same time we are spending more CPU resource to compute highlighting diff in server side. Maybe there are some other clever ways (like traversing the decls in source diff change);

We probably need more investigations, like trying the current implementation on some real-project files in a real client (e.g. theia) to figure out and understand whether the latency is a real issue.

[sammccall]

clang-tools-extra/clangd/SemanticHighlighting.cpp

264

I think this is just Lhs.R.start < Lhs.R.end

(if it's not, we should add the operator<)

is it enforced somewhere that you don't have two highlights at the same spot?

286

holding the lock while doing all the diffing seems dubious

You could reasonably put the cache as a map<file, highlights> in ClangdServer, then you don't have to deal with not having the cache for the current file.

You'd need to lock the map itself, but note that no races on updating individual entries are possible because onMainAST is called synchronously from the (per-file) worker thread.

303

Whatever you do about storage, please pull out the diff(old, new) function so you can unit test it.

303

(llvm uses unsigned for indices. It's a terrible idea, but it's used fairly consistently...)

305

I can believe this is correct, but it's hard to verify as it's a bit monolithic and... index-arithmetic-y.

could you factor out something like:

using TokRange = ArrayRef<HighlightingToken>;

// The highlights for the current line.
TokRange OldLine = {PrevHighlightings.data(), 0}, NewLine = {Highlightings.data(), 0};
// iterate over lines until we run out of data
for (unsigned Line = 0; OldLine.end() < PrevHighlightings.end() || NewLine.end() < PrevHighlightings.end(); ++Line) {
  // Get the current line highlights
  OldLine = getLineRange(PrevHighlightings, Line, OldLine);
  NewLine = getLineRange(Highlightings, Line, NewLine);
  if (OldLine != NewLine)
    emitLine(NewLine, Line);
}

// get the highlights for Line, which follows PrevLine
TokRange getLineRange(TokRange AllTokens, unsigned Line, TokRange PrevLine) {
  return makeArrayRef(PrevLine.end(), AllTokens.end()).take_while( Tok => Tok.R.start.line == Line);
}

[jvikstrom]

Made diffing function shorter, added multiple previous highlighting entries.

[jvikstrom]

Removed unused code snippets.

[jvikstrom]

clang-tools-extra/clangd/SemanticHighlighting.cpp

264

It's not enforced anywhere, it depends on how the HighlightingTokenCollector is implemented, it should not currently take any multiples of tokens at the same spot.

Even if there are tokens at the same location it should still satisfy Compare though?

286

I did not know onMainAST was called synchronously per file.
I feel like after a while a lot of memory is going to be consumed by keeping ASTs for files that were closed long ago in the cache (say when you've opened 1000 files in a session, all those files will have one Highlightings entry here)
Could solve that with LRU (just keep a set of pair<long long, std::string> (unix time of adding/editing to cache and filename)... would actually have to be a struct, would need to make the equal operator check the filename only)

But maybe that isn't a big concern/a concern at all?

[hokein]

based on the offline discussion, now I understand the patch better (thanks).

some comments on the interface.

clang-tools-extra/clangd/ClangdLSPServer.h

60

how about creating a new structure LineHighlightingTokens to represent the line-based tokens?

clang-tools-extra/clangd/ClangdServer.cpp

99

If we make it as a pointer, the bool SemanticHighlighting is not needed (just follow the what the other field FIndex does).

clang-tools-extra/clangd/SemanticHighlighting.cpp

327

we could get the canonical path from onMainAST callback (the first parameter PathRef Path).

clang-tools-extra/clangd/SemanticHighlighting.h

73

nit: llvm::StringMap.

I think this map is storing all file highlightings, maybe call it FileToHighlightings?

87

this method does two things,

1. calculate the diff between old and new
2. replace the old highlighting with new highlightings

but the name just reflects 1).

I think we just narrow it to 2) only:

// calculate the new highlightings, and return the old one.
std::vector<HighlightingToken> updateFile(PathRef Path, std::vector<HighlightingToken> NewHighlightings);

// to get diff, in `onMainAST`
diffHighlightings(differ.updateFile(Path, NewHighlightings), NewHighlightings);

95

this utility method doesn't use any internal member of the class, we could pull it out or make it static.

[sammccall]

As discussed offline, it's a little scary that the highlight state for diffing lives in a separate map from TUScheduler's workers, but it's hard to fix without introducing lots of abstractions or layering violations.

It does seem like a good idea to move the diffing up to ClangdLSPServer, as it's fairly bound to protocol details. It needs to be reset in didClose (and maybe didOpen too, for safety - I think there's a race when a document is closed, then tokens are delivered...).

I'd tend to put the map/mutex directly in ClangdLSPServer rather than encapsulating it in the Differ object (again, because it's a protocol/lifetime detail rather than diffing logic itself) but up to you.

[jvikstrom]

Moved highlighting state to LSP layer. Removed class holding state. Addressed comments.

[hokein]

the code looks clearer now!

clang-tools-extra/clangd/ClangdLSPServer.cpp

617

nit: I would create a new {} block merely for the hightlightings.

1121

this seems unsafe, we get a reference of the map value, we might access it without the mutex being guarded.

std::vector<HighlightingToken> Old;
{
    std::lock_guard<std::mutex> Lock(HighlightingsMutex);
    Old = std::move(FileToHighlightings[File]);
}

1136

FileToPrevHighlightings[File] = std::move(Highlightings);

clang-tools-extra/clangd/ClangdLSPServer.h

142

nit: I'd drop Prev.

clang-tools-extra/clangd/SemanticHighlighting.cpp

263

nit: I'm not sure I understand the comment.

267

nit: this implies that OldLine and AllTokens must ref to the same container, could you refine the OldLine as a start Iterator?

289

IIUC, we are initializing an empty ArrayRef, if so, how about using NewLine(Highlightings.data(), /*length*/0)? NewLine(Highlightings.data(), Highlightings.data()) is a bit weird, it took me a while to understand the purpose.

291

nit: split it into a new statement.

293

we don't change Current and Prev below, how about re-using Highlightings and PrevHighlightings?

clang-tools-extra/clangd/SemanticHighlighting.h

72

I believe Assumes the highlightings are sorted by the tokens' ranges. is a requirement for this function?

If so, mark it more explicitly in the comment, like

///
/// REQUIRED: OldHighlightings and NewHighlightings are sorted.

75

nit: I'd use name the parameters symmetrically, NewHighlightings and OldHighlightings.

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

31–32

we don't need ParsedAST now I think.

308

this is a complicated structure, could you add some comments describing the test strategy here?

[jvikstrom]

Address comments.

clang-tools-extra/clangd/ClangdLSPServer.cpp

1121

Aren't the parsing callbacks thread safe in the same file? I think @sammccall said so above.

Although actually, the map might reallocate and move around the memory if things are inserted/deleted invalidating the reference maybe (I have no idea of how StringMap memory allocation works).

So I guess it doesn't hurt to be safe and copy it.

clang-tools-extra/clangd/SemanticHighlighting.cpp

289

I couldn't do that because the ArrayRef(const T *data, size_t length) and ArrayRef(const T *begin, const T *end) were ambiguous. Added a cast to cast 0 to a size_t which solved it.

[hokein]

clang-tools-extra/clangd/SemanticHighlighting.cpp

289

you could also do it like NewLine(Highlightings.data(), /*length*/0UL);, which saves you a static_cast<>.

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

326

so the empty lines are stored separately, which is not easy to figure out from the testing code snippet. I think we should make the empty lines more obvious in the code snippet (maybe use an Empty annotations?) , and explicitly verify the empty lines.

What do you think refining the test like below, just annotate the diff tokens? I find it is easy to spot the diff tokens.

struct Testcase {
   Code Before;
   Code After;
};

std::vector<TestCase> cases = {
   { 
       R"cpp(
             int a;
             int b; 
             int c;
      )cpp",
      R"cpp(
            int a;
            $Empty[[]] // int b
            int $Variable[[C]];
      )cpp"
  }
}

oldHighlightings = getSemanticHighlightings(OldAST);
newHighlightings = getSemanticHighlightings(NewAST);
diffHighlightings = diff...;
// verify the diffHighlightings has the expected empty lines ("Empty" annotations).
// verify the diffHighlightings has the expected highlightings (the regular annotations);

[jvikstrom]

Rewrote test to be more readable.

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

326

Moved everything into the checkDiffedHighlights as well.

[hokein]

clang-tools-extra/clangd/SemanticHighlighting.cpp

309

nit: remove the redundant ;.

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

34

could we split it into three functions?

• getExpectedHighlightings
• getActualHighlightings
• getExpectedEmptyLines

92

nit: OrgCode => OldCode, use llvm::StringRef

120

I think the tokens are also sorted, as we pass two sorted lists of highlightings to diffHighlightings.

310

nit: the code indentation is strange

311

The annotations in the OldCode are not used -- you only use the actual highlightings in the checkDiffedHighlights.

[jvikstrom]

Address comments.

[hokein]

clang-tools-extra/clangd/SemanticHighlighting.cpp

143

This means if we won't emit most of symbols if the AST is ill-formed? I'm not sure whether we should include it in this patch, it seems out of scope. could we leave it out in this patch?

clang-tools-extra/clangd/SemanticHighlighting.h

69

The comment seems a bit out of date (as the code is updated during review), and it should mention the diff strategy of this function.

// Return a line-by-line diff between two highlightings.
//  - if the tokens on a line are the same in both hightlightings, we omit this line
//  - if a line exists in NewHighligtings but not in OldHighligtings, we emit the tokens on this line
//  - if a line not exists in NewHighligtings but in OldHighligtings, we emit an empty line (to tell client not highlighting this line)

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

49

nit: let's sort the tokens here to make the result deterministic.

88

nit: you could use for-range loop to simplify the code.

for (const auto& EmptyRange : EmptyRanges)
   EmptyLines.push_back(EmptyRange.start.line);

96

nit: since OldActualTokens and NewActualTokens are used only once in diffHighlightings, I'd inline them there (instead of creating two new variables).

109

if we do sort in getExpectedTokens, we don't need this here.

309

nit: Use an explicit struct, then you wouldn't need a comment.

struct {
   llvm::StringRef OldCode;
   llvm::StringRef NewCode;
} TestCases[] = {
   ...
}

310

could we add one more case?

// Before
int a;
int b;
int c;

// After
int a;
int $Variable[[new_b]];
int c;

362

nit: const auto&

[jvikstrom]

Address comments.

[hokein]

Thanks! Looks good from my side.

@ilya-biryukov will be nice if you could take a second look on the patch. We plan to land it before the release cut today.

clang-tools-extra/clangd/ClangdLSPServer.cpp

1124

use std::move() to save a copy here. we'd drop the old highlighting anyway (replace it with new highlightings).

[hokein]

one missing point -- could you please update the patch description?

[ilya-biryukov]

A few suggestions from me, I haven't looked into the diffing algorithm and the tests yet.

clang-tools-extra/clangd/ClangdLSPServer.cpp

466

Should can't we handle this on didClose instead?

1124

NIT: maybe std::move() into Old?

If we had exceptions, that would be complicated, but we don't have them

1135

To follow the same pattern as diagnostics, could we store before calling publish...?

clang-tools-extra/clangd/SemanticHighlighting.cpp

277

Do you have any ideas on how we can fix this?

I would simply split those tokens into multiple tokens of the same kind at newlines boundaries, but maybe there are other ways to handle this.

In any case, could you put a suggested way to fix this (in case someone will want to pick it up, they'll have a starting point)

clang-tools-extra/clangd/SemanticHighlighting.h

79

Are we locked into the line-based diff implementation?
It works nicely while editing inside the same line, but seem to do a bad job on a common case of inserting/removing lines.

Does the protocol have a way to communicate this cleanly?

[jvikstrom]

clang-tools-extra/clangd/ClangdLSPServer.cpp

466

We are removing in didClose but the problem is that there is a race condition (I think).

If someone does some edits and closes the document right after, the highlightings for the final edit might finish being generated after the FileToHighlightings have earsed the highlightings for the file. So next time when opening the file we will have those final highlightings that were generated for the last edit in the map.
I don't really know how we could handle this in didClose without having another map and with an open/closed bit and checking that every time we generate new highlightings. But we'd still have to set the bit to be open every didOpen so I don't really see the benefit.

However I've head that ASTWorked is synchronous for a single file, is that the case for the didClose call as well? Because in that case there is no race condition.

clang-tools-extra/clangd/SemanticHighlighting.cpp

277

Yeah, I would split the tokens into multiple lines as well. Or enforce that a token is single line and handle it in addToken in the collector. (i.e. change the HighlightingToken struct)

It's a bit annoying to solve because we'd have to get the line lengths of every line that the multiline length token covers when doing that.

clang-tools-extra/clangd/SemanticHighlighting.h

79

We are looked into some kind of line based diff implementation as the protocol sends token line by line.
There should be away of solving the inserting/removing lines, but I'll have a look at that in a follow up.
Theia seems to do a good job of moving tokens to where they should be automatically when inserting a new line though. I want to be able to see how vscode handles it first as well though before.

[ilya-biryukov]

clang-tools-extra/clangd/SemanticHighlighting.cpp

275

NIT: maybe rename to New and Old, the suffix of the name could be easily inferred from the variable type (clangd has hover/go-to-definition to find the type quickly, the code is rather small so its should always be visible on the screen too).

277

NIT: add assertions that highlightings are sorted.

287

NIT: maybe mention diff in the name somehow. I was confused at first, as I thought it's all highlightings grouped by line, not the diff.

293

NIT: maybe just inline NewEnd and OldEnd to save a few lines?

295

Could we skip directly to the next interesting line?
The simplest way to do this I could come up with is this:

auto NextLine = [&]() {
  int NextNew = NewLine.end() != NewHighlightings.end() ? NewLine.end()->start.line : numeric_limits<int>::max();
  int NextOld = ...;
  return std::min(NextNew, NextOld);
}

for (int Line = 0; ...; Line = NextLine()) {
}

[jvikstrom]

Added additional note to FIXME on how to solve and also move(ing) Old.

[ilya-biryukov]

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

307

Can we test this in a more direct manner by specifying:

1. annotated input for old highlightings,
2. annotated input for new highlightings,
3. the expected diff?

The resulting tests don't have to be real C++ then, allowing to express what we're testing in a more direct manner.

{/*Old*/ "$Variable[[a]]", /*New*/ "$Class[[a]]", /*Diff*/ "{{/*line */0, "$Class[[a]]"}}

It also seems that the contents of the lines could even be checked automatically (they should be equal to the corresponding line from /*New*/, right?), that leaves us with even simpler inputs:

{/*Old*/ "$Variable[[a]]", /*New*/ "$Class[[a]]", /*DiffLines*/ "{0}}

391

NIT: remove braces

[ilya-biryukov]

clang-tools-extra/clangd/ClangdLSPServer.cpp

466

You are correct, there is actually a race condition. We worked hard to eliminate it for diagnostics, but highlightings are going through a different code path in ASTWorker, not guarded by DiagsMu and ReportDiagnostics.

And, unfortunately, I don't think this guard here prevents it in all cases. In particular, there is still a possibility (albeit very low, I guess) that the old highlightings you are trying to remove here are still being computed. If they are reported after this erase() runs, we will end up with inconsistent highlightings.

Ideally, we would guard the diagnostics callback with the same mutex, but given our current layering it seems like a hard problem, will need to think what's the simplest way to fix this.

clang-tools-extra/clangd/SemanticHighlighting.h

79

It seems there no compact way to send interesting diffs then, although the clients can do a reasonably good job of moving the older version of highlightings on changes until the server sends them a new version.

The diff-based approach only seems to be helping with changes on a single line.

Which is ok, thanks for the explanation.

[jvikstrom]

Address comments and fixed bug where we'd send diffs outside of the file.

[jvikstrom]

Remove braces.

[jvikstrom]

Fixed test case that was actually (kinda) broken even though it didn't fail.

(The last two edits in semantic-highlighting test. They removed the = 2 part and did not insert an extra ; which produced invalid code. It was still fine because the highlightings that were supposed to be generated were still generated. This just makes the code not be invalid)

[jvikstrom]

Removed one lock from onSemanticHighlighting.

[jvikstrom]

clang-tools-extra/clangd/SemanticHighlighting.cpp

293

Wouldn't that violate code style: https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop ?
Or maybe that doesn't matter? (I have no preference either way, just picked this way as I could remeber reading from the styleguide)

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

307

That's a great idea on how to write these tests.

[hokein]

clang-tools-extra/clangd/ClangdLSPServer.cpp

466

The race condition of highlighting sounds bad (especially when a user opens a large file and closes it immediately, then the highlighting is finished and we emit it to the client). No need to fix it in this patch, just add a FIXME.

Can we use the same mechanism for Diagnostic to guard the highlighting here? or use the DiagsMu and ReportDiagnostics to guard the callback.onMainAST() as well (which is probably not a good idea)...

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

307

hmm, I have a different opinion here (I'd prefer the previous one), let's chat.

[ilya-biryukov]

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

311

@hokein rightfully pointed out that mentioning all changed lines makes the tests unreadable.
An alternative idea we all came up with is to force people to put ^ on each of the changed lines inside the NewCode, i.e.

{/*Before*/ R"(
  $Var[[a]]
  $Func[[b]]
"),
 /*After*/ R"(
  $Var[[a]]
 ^$Func[[b]]
)"} // and no list of lines is needed!

Could we do that here?

One interesting case that we can't test this way to removing lines from the end of the file. But for that particular case, could we just write a separate test case?

[ilya-biryukov]

The fix for a race condition on remove has landed in rL366577, this revision would need a small update after it.

[jvikstrom]

Made tests more readable. Updated to work with rL366577.

[jvikstrom]

In D64475#1593481, @ilya-biryukov wrote:

The fix for a race condition on remove has landed in rL366577, this revision would need a small update after it.

Fixed to work with that patch.
Should the diffing be moved somewhere else that is not inside the publish function though? That would require moving the state outside the LSP server though and handle the onClose/onOpen events somewhere else though.
Maybe it's ok to keep the diffing in the publish lock?

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

311

We don't want to send diffs that are beyond the end of the file. There is a testcase for that as well (the count of newlines in the new code is sent to the differ as the number of lines in the file).

[ilya-biryukov]

In D64475#1602195, @jvikstrom wrote:

In D64475#1593481, @ilya-biryukov wrote:

The fix for a race condition on remove has landed in rL366577, this revision would need a small update after it.

Fixed to work with that patch.
Should the diffing be moved somewhere else that is not inside the publish function though? That would require moving the state outside the LSP server though and handle the onClose/onOpen events somewhere else though.
Maybe it's ok to keep the diffing in the publish lock?

Yeah, let's keep it in publish for now. Diffing in there does not change complexity of the operations.
If it ever turns out that the constant factor is too high, we can figure out ways to fight this.

[ilya-biryukov]

Mostly final NITs from me, but there is an important one about removing the erase call on didOpen.

clang-tools-extra/clangd/ClangdLSPServer.cpp

467

Now that the patch landed, this is obsolete. Could you remove?

clang-tools-extra/clangd/ClangdServer.cpp

80

NIT: use unsigned instead of unsigned int

clang-tools-extra/clangd/ClangdServer.h

56

NIT: could you document NLines

clang-tools-extra/clangd/SemanticHighlighting.cpp

294

Could we try to find a better name for this? Having Line and NextLine() which represent line numbers and NewLine which represents highlightings, creates some confusion.

310

Line != intmax is redundant (NLines is <= intmax by definition).
Maybe remove it altogether to simplify the loop condition?

clang-tools-extra/clangd/SemanticHighlighting.h

80

Could you document what NLines is? Specifically, whether it's the number of lines for New or Old.

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

93

NIT: add documentation on how this should be used
Most importantly, the fact that we need to put ^ on all changed lines should be mentioned.

99

NIT: use llvm::DenseMap

[hokein]

a few more comments from my side.

clang-tools-extra/clangd/SemanticHighlighting.cpp

309

nit: I believe theia is crashing because of this? could we file a bug to theia? I think it would be nice for client to be robust on this case.

clang-tools-extra/clangd/SemanticHighlighting.h

70

nit: add , before this

73

could you refine this sentence, I can't parse it.

74

and here as well.

80

could we use a more describe name for NLines? maybe MaxLines

[jvikstrom]

Address comments.

[jvikstrom]

clang-tools-extra/clangd/SemanticHighlighting.cpp

294

I renamed the Line and NextLine() instead. Could also rename NewLine to something like NewLineHighlightings but I feel like it almost becomes to verbose at that point.

309

I seem to be misremembering, when I first started testing in theia I think I was crashing the client (due to a broken base64 encoding function, will take a look and see what was actually happening, can't quite remember)
It actually seems like theia just ignores any out of bounds highlightings.

So this could be removed, it will just sometimes return highlightings that are outside of the file and hopefully any future clients would handle that as well.

[hokein]

clang-tools-extra/clangd/ClangdServer.cpp

81

from the comment of the getLineNumber, this is not cheap to call this method. We may use SM.getBufferData(MainFileID).count('\n') to count the line numbers.

// This requires building and caching a table of line offsets for the
/// MemoryBuffer, so this is not cheap: use only when about to emit a
/// diagnostic.

clang-tools-extra/clangd/ClangdServer.h

55

nit: drop the needed ..., this is an interface, don't mention the details of subclass (diffs are subclass implementation details).

Maybe name it "NumLines"?

[jvikstrom]

Stopped using expensive getLineNumber function.

[ilya-biryukov]

clang-tools-extra/clangd/ClangdLSPServer.cpp

1125

NIT: avoid copying (from Highlightings into the map) under a lock, make the copy outside the lock instead

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

408

Could you also add a separate test that checks diffs when the number of lines is getting smaller (i.e. taking NLines into account)

[ilya-biryukov]

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

408

I would expect this to be better handled outside checkDiffedHighlights to avoid over-generalizing this function. But up to you.

[jvikstrom]

Copy outside lock.

[jvikstrom]

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

408

That's what this test does:

{
                    R"(
        $Class[[A]]
        $Variable[[A]]
        $Class[[A]]
        $Variable[[A]]
      )",
                    R"(
        $Class[[A]]
        $Variable[[A]]
      )"},

But do you mean I should do a testcase like:

{
                    R"(
        $Class[[A]]
        $Variable[[A]]
        $Class[[A]]
        $Variable[[A]]
      )",
                    R"(
        $Class[[A]]
        $Variable[[A]]
        $Class[[A]]
        $Variable[[A]]
      )"},

And just set NLines = 2 instead when doing the diffing?

[hokein]

clang-tools-extra/clangd/ClangdLSPServer.cpp

1125

I think we can even avoid copy, since we only use the Highlightings for calculating the diff.

Maybe just

{
    std::lock_guard<std::mutex> Lock(HighlightingsMutex);
    Old = std::move(FileToHighlightings[File]);
}
// calculate the diff.
{  
     std::lock_guard<std::mutex> Lock(HighlightingsMutex);
     FileToHighlightings[File] = std::move(Highlightings);
}

[jvikstrom]

clang-tools-extra/clangd/ClangdLSPServer.cpp

1125

I think I changed this to only lock once (and copy instead) at the same time me and @ilya-biryukov were talking about the race condition (?)

[ilya-biryukov]

LGTM from my side

clang-tools-extra/clangd/ClangdLSPServer.cpp

1125

It means there's a window where nobody can access the highlightings for this file. Which is probably fine, we really do use this only for computing the diff and nothing else.

But doing the copy shouldn't matter here either, so leaving as is also looks ok from my side.

If you decide to avoid an extra copy, please add comments to FileToHighlightings (e.g. only used to compute diffs, must not be used outside onHighlightingsReady and didRemove).
It is declared really close to FixItsMap and looks very similar, but the latter is used in completely different ways; that could lead to confusion.

Don't remember exact details about race conditions, but we should be good now that it's called inside Publish()

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

408

Ah, the test actually needs the set of changed lines to be exactly the same as marked lines.
In that case, all is good, I've just missed this detail.

[hokein]

clang-tools-extra/clangd/ClangdLSPServer.cpp

1125

I'm fine with the current state, avoiding the copy cost would make the code here a bit mess...

clang-tools-extra/clangd/ClangdLSPServer.h

public:
   ///
   /// \return Whether we shut down cleanly with a 'shutdown' -> 'exit' sequence.
   bool run();
 
 private:
   // Implement DiagnosticsConsumer.
   void onDiagnosticsReady(PathRef File, std::vector<Diag> Diagnostics) override;
   void onFileUpdated(PathRef File, const TUStatus &Status) override;
-  void
-  onHighlightingsReady(PathRef File,
-                       std::vector<HighlightingToken> Highlightings) override;
+  void onHighlightingsReady(PathRef File,
+                            std::vector<HighlightingToken> Highlightings,
+                            int NLines) override;
 
   // LSP methods. Notifications have signature void(const Params&).
   // Calls have signature void(const Params&, Callback<Response>).
   void onInitialize(const InitializeParams &, Callback<llvm::json::Value>);
   void onShutdown(const ShutdownParams &, Callback<std::nullptr_t>);
   void onSync(const NoParams &, Callback<std::nullptr_t>);
   void onDocumentDidOpen(const DidOpenTextDocumentParams &);
   void onDocumentDidChange(const DidChangeTextDocumentParams &);
private:
   /// Language Server client.
   bool ShutdownRequestReceived = false;
 
   std::mutex FixItsMutex;
   typedef std::map<clangd::Diagnostic, std::vector<Fix>, LSPDiagnosticCompare>
       DiagnosticToReplacementMap;
   /// Caches FixIts per file and diagnostics
   llvm::StringMap<DiagnosticToReplacementMap> FixItsMap;
+  std::mutex HighlightingsMutex;
+  llvm::StringMap<std::vector<HighlightingToken>> FileToHighlightings;
 
   // Most code should not deal with Transport directly.
   // MessageHandler deals with incoming messages, use call() etc for outgoing.
   clangd::Transport &Transp;
   class MessageHandler;
   std::unique_ptr<MessageHandler> MsgHandler;
   std::atomic<int> NextCallID = {0};
   std::mutex TranspWriter;

clang-tools-extra/clangd/ClangdLSPServer.cpp

 }
 
 void ClangdLSPServer::onDocumentDidOpen(
     const DidOpenTextDocumentParams &Params) {
   PathRef File = Params.textDocument.uri.file();
 
   const std::string &Contents = Params.textDocument.text;
 
   DraftMgr.addDraft(File, Contents);
   Server->addDocument(File, Contents, WantDiagnostics::Yes);
 }
 
 void ClangdLSPServer::onDocumentDidChange(
     const DidChangeTextDocumentParams &Params) {
   auto WantDiags = WantDiagnostics::Auto;
   if (Params.wantDiagnostics.hasValue())
     WantDiags = Params.wantDiagnostics.getValue() ? WantDiagnostics::Yes
                                                   : WantDiagnostics::No;
void ClangdLSPServer::onDocumentDidClose(
     const DidCloseTextDocumentParams &Params) {
   PathRef File = Params.textDocument.uri.file();
   DraftMgr.removeDraft(File);
   Server->removeDocument(File);
 
   {
     std::lock_guard<std::mutex> Lock(FixItsMutex);
     FixItsMap.erase(File);
   }
+  {
+    std::lock_guard<std::mutex> HLock(HighlightingsMutex);
+    FileToHighlightings.erase(File);
+  }
   // clangd will not send updates for this file anymore, so we empty out the
   // list of diagnostics shown on the client (e.g. in the "Problems" pane of
   // VSCode). Note that this cannot race with actual diagnostics responses
   // because removeDocument() guarantees no diagnostic callbacks will be
   // executed after it returns.
   publishDiagnostics(URIForFile::canonicalize(File, /*TUPath=*/File), {});
 }
 
if (Trigger == ">")
     return (*Code)[*Offset - 2] == '-'; // trigger only on '->'.
   if (Trigger == ":")
     return (*Code)[*Offset - 2] == ':'; // trigger only on '::'.
   assert(false && "unhandled trigger character");
   return true;
 }
 
 void ClangdLSPServer::onHighlightingsReady(
-    PathRef File, std::vector<HighlightingToken> Highlightings) {
+    PathRef File, std::vector<HighlightingToken> Highlightings, int NLines) {
+  std::vector<HighlightingToken> Old;
+  {
+    std::lock_guard<std::mutex> Lock(HighlightingsMutex);
+    Old = std::move(FileToHighlightings[File]);
+    FileToHighlightings[File] = Highlightings;
+  }
+  // LSP allows us to send incremental edits of highlightings. Also need to diff
+  // to remove highlightings from tokens that should no longer have them.
+  std::vector<LineHighlightings> Diffed =
+      diffHighlightings(Highlightings, Old, NLines);
   publishSemanticHighlighting(
       {{URIForFile::canonicalize(File, /*TUPath=*/File)},
-       toSemanticHighlightingInformation(Highlightings)});
+       toSemanticHighlightingInformation(Diffed)});
 }
 
 void ClangdLSPServer::onDiagnosticsReady(PathRef File,
                                          std::vector<Diag> Diagnostics) {
   auto URI = URIForFile::canonicalize(File, /*TUPath=*/File);
   std::vector<Diagnostic> LSPDiagnostics;
   DiagnosticToReplacementMap LocalFixIts; // Temporary storage
   for (auto &Diag : Diagnostics) {
     toLSPDiags(Diag, URI, DiagOpts,
                [&](clangd::Diagnostic Diag, llvm::ArrayRef<Fix> Fixes) {
                  auto &FixItsForDiagnostic = LocalFixIts[Diag];

clang-tools-extra/clangd/ClangdServer.h

public:
 
   /// Called by ClangdServer when \p Diagnostics for \p File are ready.
   virtual void onDiagnosticsReady(PathRef File,
                                   std::vector<Diag> Diagnostics) = 0;
   /// Called whenever the file status is updated.
   virtual void onFileUpdated(PathRef File, const TUStatus &Status){};
 
   /// Called by ClangdServer when some \p Highlightings for \p File are ready.
-  virtual void
-  onHighlightingsReady(PathRef File,
-                       std::vector<HighlightingToken> Highlightings) {}
+  /// \p NLines are the number of lines in the file, needed to make sure that
+  /// any diffing does not add lines beyond EOF.
+  virtual void onHighlightingsReady(
+      PathRef File, std::vector<HighlightingToken> Highlightings, int NLines) {}
 };
 
 /// When set, used by ClangdServer to get clang-tidy options for each particular
 /// file. Must be thread-safe. We use this instead of ClangTidyOptionsProvider
 /// to allow reading tidy configs from the VFS used for parsing.
 using ClangTidyOptionsBuilder = std::function<tidy::ClangTidyOptions(
     llvm::vfs::FileSystem &, llvm::StringRef /*File*/)>;
 

clang-tools-extra/clangd/ClangdServer.cpp

void onMainAST(PathRef Path, ParsedAST &AST, PublishFn Publish) override {
     if (FIndex)
       FIndex->updateMain(Path, AST);
 
     std::vector<Diag> Diagnostics = AST.getDiagnostics();
     std::vector<HighlightingToken> Highlightings;
     if (SemanticHighlighting)
       Highlightings = getSemanticHighlightings(AST);
 
+    // FIXME: We need a better way to send the maximum line number to the
+    // differ.
+    // The differ needs the information about the max number of lines
+    // to not send diffs that are outside the file.
+    const SourceManager &SM = AST.getSourceManager();
+    FileID MainFileID = SM.getMainFileID();
+    unsigned FileSize = SM.getFileEntryForID(MainFileID)->getSize();
+    int NLines = AST.getSourceManager().getLineNumber(MainFileID, FileSize);
+
     Publish([&]() {
       DiagConsumer.onDiagnosticsReady(Path, std::move(Diagnostics));
       if (SemanticHighlighting)
-        DiagConsumer.onHighlightingsReady(Path, std::move(Highlightings));
+        DiagConsumer.onHighlightingsReady(Path, std::move(Highlightings),
+                                          NLines);
     });
   }
 
   void onFileUpdated(PathRef File, const TUStatus &Status) override {
     DiagConsumer.onFileUpdated(File, Status);
   }
 
 private:
   FileIndex *FIndex;
   DiagnosticsConsumer &DiagConsumer;
   bool SemanticHighlighting;
 };
 } // namespace
 
 ClangdServer::Options ClangdServer::optsForTest() {
   ClangdServer::Options Opts;
   Opts.UpdateDebounce = std::chrono::steady_clock::duration::zero(); // Faster!
   Opts.StorePreamblesInMemory = true;
   Opts.AsyncThreadsCount = 4; // Consistent!
   Opts.SemanticHighlighting = true;

clang-tools-extra/clangd/SemanticHighlighting.h

 };
 
 // Contains all information needed for the highlighting a token.
 struct HighlightingToken {
   HighlightingKind Kind;
   Range R;
 };
 
-bool operator==(const HighlightingToken &Lhs, const HighlightingToken &Rhs);
+bool operator==(const HighlightingToken &L, const HighlightingToken &R);
+bool operator<(const HighlightingToken &L, const HighlightingToken &R);
+
+/// Contains all information about highlightings on a single line.
+struct LineHighlightings {
+  int Line;
+  std::vector<HighlightingToken> Tokens;
+};
+
+bool operator==(const LineHighlightings &L, const LineHighlightings &R);
 
 // Returns all HighlightingTokens from an AST. Only generates highlights for the
 // main AST.
 std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST);
 
 /// Converts a HighlightingKind to a corresponding TextMate scope
 /// (https://manual.macromates.com/en/language_grammars).
 llvm::StringRef toTextMateScope(HighlightingKind Kind);
 
-// Convert to LSP's semantic highlighting information.
+/// Convert to LSP's semantic highlighting information.
 std::vector<SemanticHighlightingInformation>
-toSemanticHighlightingInformation(llvm::ArrayRef<HighlightingToken> Tokens);
+toSemanticHighlightingInformation(llvm::ArrayRef<LineHighlightings> Tokens);
+
+/// Return a line-by-line diff between two highlightings.
+///  - if the tokens on a line are the same in both hightlightings, this line is
+///  omitted.
+///  - if a line exists in New but not in Old the tokens on this line are
+///  emitted.
+///  - if a line does not exists in New but exists in Old an empty line is
+///  emitted (to tell client to clear the previous highlightings on this line).
+/// \p NewMaxLine is the maximum line number from the new file.
+/// REQUIRED: Old and New are sorted.
+std::vector<LineHighlightings>
+diffHighlightings(ArrayRef<HighlightingToken> New,
+                  ArrayRef<HighlightingToken> Old, int NewMaxLine);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/SemanticHighlighting.cpp

 //===--- SemanticHighlighting.cpp - ------------------------- ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "SemanticHighlighting.h"
 #include "Logger.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include <algorithm>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 // Collects all semantic tokens in an ASTContext.
 class HighlightingTokenCollector
     : public RecursiveASTVisitor<HighlightingTokenCollector> {
   std::vector<HighlightingToken> Tokens;
   ASTContext &Ctx;
   const SourceManager &SM;
 
 public:
   HighlightingTokenCollector(ParsedAST &AST)
       : Ctx(AST.getASTContext()), SM(AST.getSourceManager()) {}
 
   std::vector<HighlightingToken> collectTokens() {
     Tokens.clear();
     TraverseAST(Ctx);
     // Initializer lists can give duplicates of tokens, therefore all tokens
     // must be deduplicated.
-    llvm::sort(Tokens,
-               [](const HighlightingToken &L, const HighlightingToken &R) {
-                 return std::tie(L.R, L.Kind) < std::tie(R.R, R.Kind);
-               });
+    llvm::sort(Tokens);
     auto Last = std::unique(Tokens.begin(), Tokens.end());
     Tokens.erase(Last, Tokens.end());
     return Tokens;
   }
 
   bool VisitNamespaceAliasDecl(NamespaceAliasDecl *NAD) {
     // The target namespace of an alias can not be found in any other way.
     addToken(NAD->getTargetNameLoc(), HighlightingKind::Namespace);

 private:
   void addTypeLoc(SourceLocation Loc, const TypeLoc &TL) {
     if (const Type *TP = TL.getTypePtr())
       if (const TagDecl *TD = TP->getAsTagDecl())
         addToken(Loc, TD);
   }
 
   void addToken(SourceLocation Loc, const NamedDecl *D) {
     if (D->getDeclName().isIdentifier() && D->getName().empty())
       // Don't add symbols that don't have any length.
       return;
     // We highlight class decls, constructor decls and destructor decls as
     // `Class` type. The destructor decls are handled in `VisitTypeLoc` (we will
     // visit a TypeLoc where the underlying Type is a CXXRecordDecl).
     if (isa<ClassTemplateDecl>(D)) {
       addToken(Loc, HighlightingKind::Class);
       return;
void write32be(uint32_t I, llvm::raw_ostream &OS) {
   OS.write(Buf.data(), Buf.size());
 }
 
 void write16be(uint16_t I, llvm::raw_ostream &OS) {
   std::array<char, 2> Buf;
   llvm::support::endian::write16be(Buf.data(), I);
   OS.write(Buf.data(), Buf.size());
 }
+
+// Get the highlightings on \c Line where the first entry of line is at \c
+// StartLineIt. If it is not at \c StartLineIt an empty vector is returned.
+ArrayRef<HighlightingToken>
+takeLine(ArrayRef<HighlightingToken> AllTokens,
+         ArrayRef<HighlightingToken>::iterator StartLineIt, int Line) {
+  return ArrayRef<HighlightingToken>(StartLineIt, AllTokens.end())
+      .take_while([Line](const HighlightingToken &Token) {
+        return Token.R.start.line == Line;
+      });
+}
 } // namespace
 
-bool operator==(const HighlightingToken &Lhs, const HighlightingToken &Rhs) {
-  return Lhs.Kind == Rhs.Kind && Lhs.R == Rhs.R;
+std::vector<LineHighlightings>
+diffHighlightings(ArrayRef<HighlightingToken> New,
+                  ArrayRef<HighlightingToken> Old, int NewMaxLine) {
+  assert(std::is_sorted(New.begin(), New.end()) && "New must be a sorted vector");
+  assert(std::is_sorted(Old.begin(), Old.end()) && "Old must be a sorted vector");
+
+  // FIXME: There's an edge case when tokens span multiple lines. If the first
+  // token on the line started on a line above the current one and the rest of
+  // the line is the equal to the previous one than we will remove all
+  // highlights but the ones for the token spanning multiple lines. This means
+  // that when we get into the LSP layer the only highlights that will be
+  // visible are the ones for the token spanning multiple lines.
+  // Example:
+  // EndOfMultilineToken  Token Token Token
+  // If "Token Token Token" don't differ from previously the line is
+  // incorrectly removed. Suggestion to fix is to separate any multiline tokens
+  // into one token for every line it covers. This requires reading from the
+  // file buffer to figure out the length of each line though.
+  std::vector<LineHighlightings> DiffedLines;
+  // ArrayRefs to the current line in the highlightings.
+  ArrayRef<HighlightingToken> NewLine(New.begin(),
+                                      /*length*/0UL);
+  ArrayRef<HighlightingToken> OldLine(Old.begin(),
+                                      /*length*/ 0UL);
+  auto NewEnd = New.end();
+  auto OldEnd = Old.end();
+  auto NextLineNumber = [&]() {
+    int NextNew = NewLine.end() != NewEnd ? NewLine.end()->R.start.line
+                                             : std::numeric_limits<int>::max();
+    int NextOld = OldLine.end() != OldEnd ? OldLine.end()->R.start.line
+                                             : std::numeric_limits<int>::max();
+    return std::min(NextNew, NextOld);
+  };
+
+  // If the New file has fewer lines than the Old file we don't want to send
+  // highlightings beyond the end of the file.
+  for (int LineNumber = 0; LineNumber < NewMaxLine;
+       LineNumber = NextLineNumber()) {
+    NewLine = takeLine(New, NewLine.end(), LineNumber);
+    OldLine = takeLine(Old, OldLine.end(), LineNumber);
+    if (NewLine != OldLine)
+      DiffedLines.push_back({LineNumber, NewLine});
+  }
+
+  return DiffedLines;
+}
+
+bool operator==(const HighlightingToken &L, const HighlightingToken &R) {
+  return std::tie(L.R, L.Kind) == std::tie(R.R, R.Kind);
+}
+bool operator<(const HighlightingToken &L, const HighlightingToken &R) {
+  return std::tie(L.R, L.Kind) < std::tie(R.R, R.Kind);
+}
+bool operator==(const LineHighlightings &L, const LineHighlightings &R) {
+  return std::tie(L.Line, L.Tokens) == std::tie(R.Line, R.Tokens);
 }
 
 std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST) {
   return HighlightingTokenCollector(AST).collectTokens();
 }
 
 std::vector<SemanticHighlightingInformation>
-toSemanticHighlightingInformation(llvm::ArrayRef<HighlightingToken> Tokens) {
+toSemanticHighlightingInformation(llvm::ArrayRef<LineHighlightings> Tokens) {
   if (Tokens.size() == 0)
     return {};
 
   // FIXME: Tokens might be multiple lines long (block comments) in this case
   // this needs to add multiple lines for those tokens.
-  std::map<int, std::vector<HighlightingToken>> TokenLines;
-  for (const HighlightingToken &Token : Tokens)
-    TokenLines[Token.R.start.line].push_back(Token);
-
   std::vector<SemanticHighlightingInformation> Lines;
-  Lines.reserve(TokenLines.size());
-  for (const auto &Line : TokenLines) {
+  Lines.reserve(Tokens.size());
+  for (const auto &Line : Tokens) {
     llvm::SmallVector<char, 128> LineByteTokens;
     llvm::raw_svector_ostream OS(LineByteTokens);
-    for (const auto &Token : Line.second) {
+    for (const auto &Token : Line.Tokens) {
       // Writes the token to LineByteTokens in the byte format specified by the
       // LSP proposal. Described below.
       // |<---- 4 bytes ---->|<-- 2 bytes -->|<--- 2 bytes -->|
       // |    character      |  length       |    index       |
 
       write32be(Token.R.start.character, OS);
       write16be(Token.R.end.character - Token.R.start.character, OS);
       write16be(static_cast<int>(Token.Kind), OS);
     }
 
-    Lines.push_back({Line.first, encodeBase64(LineByteTokens)});
+    Lines.push_back({Line.Line, encodeBase64(LineByteTokens)});
   }
 
   return Lines;
 }
 
 llvm::StringRef toTextMateScope(HighlightingKind Kind) {
   // FIXME: Add scopes for C and Objective C.
   switch (Kind) {

clang-tools-extra/clangd/test/semantic-highlighting.test

 # CHECK-NEXT:      }
 # CHECK-NEXT:    ],
 # CHECK-NEXT:    "textDocument": {
 # CHECK-NEXT:      "uri": "file:///clangd-test/foo.cpp"
 # CHECK-NEXT:    }
 # CHECK-NEXT:  }
 # CHECK-NEXT:}
 ---
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{"uri":"test:///foo2.cpp","languageId":"cpp","version":1,"text":"int x = 2;\nint y = 2;"}}}
+#      CHECK:  "method": "textDocument/semanticHighlighting",
+# CHECK-NEXT:  "params": {
+# CHECK-NEXT:    "lines": [
+# CHECK-NEXT:      {
+# CHECK-NEXT:        "line": 0,
+# CHECK-NEXT:        "tokens": "AAAABAABAAA="
+# CHECK-NEXT:      }
+# CHECK-NEXT:      {
+# CHECK-NEXT:        "line": 1,
+# CHECK-NEXT:        "tokens": "AAAABAABAAA="
+# CHECK-NEXT:      }
+# CHECK-NEXT:    ],
+# CHECK-NEXT:    "textDocument": {
+# CHECK-NEXT:      "uri": "file:///clangd-test/foo2.cpp"
+# CHECK-NEXT:    }
+# CHECK-NEXT:  }
+# CHECK-NEXT:}
+---
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{"uri":"test:///foo.cpp","version":2},"contentChanges": [{"range":{"start": {"line": 0,"character": 10},"end": {"line": 0,"character": 10}},"rangeLength": 0,"text": "\nint y = 2;"}]}}
+#      CHECK:  "method": "textDocument/semanticHighlighting",
+# CHECK-NEXT:  "params": {
+# CHECK-NEXT:    "lines": [
+# CHECK-NEXT:      {
+# CHECK-NEXT:        "line": 1,
+# CHECK-NEXT:        "tokens": "AAAABAABAAA="
+# CHECK-NEXT:      }
+# CHECK-NEXT:   ],
+# CHECK-NEXT:    "textDocument": {
+# CHECK-NEXT:      "uri": "file:///clangd-test/foo.cpp"
+# CHECK-NEXT:    }
+# CHECK-NEXT:  }
+# CHECK-NEXT:}
+---
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{"uri":"test:///foo.cpp","version":2},"contentChanges": [{"range":{"start": {"line": 0,"character": 10},"end": {"line": 1,"character": 10}},"rangeLength": 11,"text": ""}]}}
+#      CHECK:  "method": "textDocument/semanticHighlighting",
+# CHECK-NEXT:  "params": {
+# CHECK-NEXT:    "lines": [],
+# CHECK-NEXT:    "textDocument": {
+# CHECK-NEXT:      "uri": "file:///clangd-test/foo.cpp"
+# CHECK-NEXT:    }
+# CHECK-NEXT:  }
+# CHECK-NEXT:}
+---
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 ---
 {"jsonrpc":"2.0","method":"exit"}

clang-tools-extra/clangd/unittests/SemanticHighlightingTests.cpp

makeHighlightingTokens(llvm::ArrayRef<Range> Ranges, HighlightingKind Kind) {
   std::vector<HighlightingToken> Tokens(Ranges.size());
   for (int I = 0, End = Ranges.size(); I < End; ++I) {
     Tokens[I].R = Ranges[I];
     Tokens[I].Kind = Kind;
   }
 
   return Tokens;
 }
 
-void checkHighlightings(llvm::StringRef Code) {
-  Annotations Test(Code);
-  auto AST = TestTU::withCode(Test.code()).build();
+std::vector<HighlightingToken> getExpectedTokens(Annotations &Test) {
   static const std::map<HighlightingKind, std::string> KindToString{
       {HighlightingKind::Variable, "Variable"},
       {HighlightingKind::Function, "Function"},
       {HighlightingKind::Class, "Class"},
       {HighlightingKind::Enum, "Enum"},
       {HighlightingKind::Namespace, "Namespace"},
       {HighlightingKind::EnumConstant, "EnumConstant"},
       {HighlightingKind::Field, "Field"},
       {HighlightingKind::Method, "Method"},
       {HighlightingKind::TemplateParameter, "TemplateParameter"}};
   std::vector<HighlightingToken> ExpectedTokens;
   for (const auto &KindString : KindToString) {
     std::vector<HighlightingToken> Toks = makeHighlightingTokens(
         Test.ranges(KindString.second), KindString.first);
     ExpectedTokens.insert(ExpectedTokens.end(), Toks.begin(), Toks.end());
   }
+  llvm::sort(ExpectedTokens);
+  return ExpectedTokens;
+}
 
-  auto ActualTokens = getSemanticHighlightings(AST);
-  EXPECT_THAT(ActualTokens, testing::UnorderedElementsAreArray(ExpectedTokens))
-   << "Inputs is:\n" << Code;
+void checkHighlightings(llvm::StringRef Code) {
+  Annotations Test(Code);
+  auto AST = TestTU::withCode(Test.code()).build();
+  std::vector<HighlightingToken> ActualTokens = getSemanticHighlightings(AST);
+  EXPECT_THAT(ActualTokens, getExpectedTokens(Test));
+}
+
+// Any annotations in OldCode and NewCode are converted into their corresponding
+// HighlightingToken. The tokens are diffed against each other. Any lines where
+// the tokens should diff must be marked with a ^ somewhere on that line in
+// NewCode. If there are diffs that aren't marked with ^ the test fails. The
+// test also fails if there are lines marked with ^ that don't differ.
+void checkDiffedHighlights(llvm::StringRef OldCode, llvm::StringRef NewCode) {
+  Annotations OldTest(OldCode);
+  Annotations NewTest(NewCode);
+  std::vector<HighlightingToken> OldTokens = getExpectedTokens(OldTest);
+  std::vector<HighlightingToken> NewTokens = getExpectedTokens(NewTest);
+
+  llvm::DenseMap<int, std::vector<HighlightingToken>> ExpectedLines;
+  for (const Position &Point : NewTest.points()) {
+    ExpectedLines[Point.line]; // Default initialize to an empty line. Tokens
+                               // are inserted on these lines later.
+  }
+  std::vector<LineHighlightings> ExpectedLinePairHighlighting;
+  for (const HighlightingToken &Token : NewTokens) {
+    auto It = ExpectedLines.find(Token.R.start.line);
+    if (It != ExpectedLines.end())
+      It->second.push_back(Token);
+  }
+  for (auto &LineTokens : ExpectedLines)
+    ExpectedLinePairHighlighting.push_back(
+        {LineTokens.first, LineTokens.second});
+
+  std::vector<LineHighlightings> ActualDiffed =
+      diffHighlightings(NewTokens, OldTokens, NewCode.count('\n'));
+  EXPECT_THAT(ActualDiffed,
+              testing::UnorderedElementsAreArray(ExpectedLinePairHighlighting));
 }
 
 TEST(SemanticHighlighting, GetsCorrectTokens) {
   const char *TestCases[] = {
     R"cpp(
       struct $Class[[AS]] {
         double $Field[[SomeMember]];
       };
       struct {
       } $Variable[[S]];
       void $Function[[foo]](int $Variable[[A]], $Class[[AS]] $Variable[[As]]) {
         auto $Variable[[VeryLongVariableName]] = 12312;
         $Class[[AS]]     $Variable[[AA]];
         auto $Variable[[L]] = $Variable[[AA]].$Field[[SomeMember]] + $Variable[[A]];
         auto $Variable[[FN]] = [ $Variable[[AA]]](int $Variable[[A]]) -> void {};
         $Variable[[FN]](12312);
       }
     )cpp",
     R"cpp(
       void $Function[[foo]](int);
       void $Function[[Gah]]();
       void $Function[[foo]]() {
         auto $Variable[[Bou]] = $Function[[Gah]];
       }
       struct $Class[[A]] {
         void $Method[[abc]]();
       };
     )cpp",
     R"cpp(
       namespace $Namespace[[abc]] {
         template<typename $TemplateParameter[[T]]>
         struct $Class[[A]] {
           $TemplateParameter[[T]] $Field[[t]];
         };
       }
       template<typename $TemplateParameter[[T]]>
       struct $Class[[C]] : $Namespace[[abc]]::$Class[[A]]<$TemplateParameter[[T]]> {
         typename $TemplateParameter[[T]]::A* $Field[[D]];
       };

 }
 
 TEST(SemanticHighlighting, GeneratesHighlightsWhenFileChange) {
   class HighlightingsCounterDiagConsumer : public DiagnosticsConsumer {
   public:
     std::atomic<int> Count = {0};
 
     void onDiagnosticsReady(PathRef, std::vector<Diag>) override {}
-    void onHighlightingsReady(
-        PathRef File, std::vector<HighlightingToken> Highlightings) override {
+    void onHighlightingsReady(PathRef File,
+                              std::vector<HighlightingToken> Highlightings,
+                              int NLines) override {
       ++Count;
     }
   };
 
   auto FooCpp = testPath("foo.cpp");
   MockFSProvider FS;
   FS.Files[FooCpp] = "";
 
   MockCompilationDatabase MCD;
   HighlightingsCounterDiagConsumer DiagConsumer;
   ClangdServer Server(MCD, FS, DiagConsumer, ClangdServer::optsForTest());
   Server.addDocument(FooCpp, "int a;");
   ASSERT_TRUE(Server.blockUntilIdleForTest()) << "Waiting for server";
   ASSERT_EQ(DiagConsumer.Count, 1);
 }
 
 TEST(SemanticHighlighting, toSemanticHighlightingInformation) {
   auto CreatePosition = [](int Line, int Character) -> Position {
     Position Pos;
     Pos.line = Line;
     Pos.character = Character;
     return Pos;
   };
 
-  std::vector<HighlightingToken> Tokens{
-      {HighlightingKind::Variable,
-                        Range{CreatePosition(3, 8), CreatePosition(3, 12)}},
+  std::vector<LineHighlightings> Tokens{
+      {3,
+       {{HighlightingKind::Variable,
         {HighlightingKind::Function,
-                        Range{CreatePosition(3, 4), CreatePosition(3, 7)}},
-      {HighlightingKind::Variable,
-                        Range{CreatePosition(1, 1), CreatePosition(1, 5)}}};
+         Range{CreatePosition(3, 4), CreatePosition(3, 7)}}}},
+      {1,
+       {{HighlightingKind::Variable,
+         Range{CreatePosition(1, 1), CreatePosition(1, 5)}}}}};
   std::vector<SemanticHighlightingInformation> ActualResults =
       toSemanticHighlightingInformation(Tokens);
   std::vector<SemanticHighlightingInformation> ExpectedResults = {
-      {1, "AAAAAQAEAAA="},
-      {3, "AAAACAAEAAAAAAAEAAMAAQ=="}};
+      {3, "AAAACAAEAAAAAAAEAAMAAQ=="}, {1, "AAAAAQAEAAA="}};
   EXPECT_EQ(ActualResults, ExpectedResults);
 }
 
+TEST(SemanticHighlighting, HighlightingDiffer) {
+  struct {
+    llvm::StringRef OldCode;
+    llvm::StringRef NewCode;
+  } TestCases[]{{
+                    R"(
+        $Variable[[A]]
+        $Class[[B]]
+        $Function[[C]]
+      )",
+                    R"(
+        $Variable[[A]]
+        $Class[[D]]
+        $Function[[C]]
+      )"},
+                {
+                    R"(
+        $Class[[C]]
+        $Field[[F]]
+        $Variable[[V]]
+        $Class[[C]] $Variable[[V]] $Field[[F]]
+      )",
+                    R"(
+        $Class[[C]]
+        $Field[[F]]
+       ^$Function[[F]]
+        $Class[[C]] $Variable[[V]] $Field[[F]]
+      )"},
+                {
+                    R"(
+
+        $Class[[A]]
+        $Variable[[A]]
+      )",
+                    R"(
+
+       ^
+       ^$Class[[A]]
+       ^$Variable[[A]]
+      )"},
+                {
+                    R"(
+        $Class[[C]]
+        $Field[[F]]
+        $Variable[[V]]
+        $Class[[C]] $Variable[[V]] $Field[[F]]
+      )",
+                    R"(
+        $Class[[C]]
+       ^
+       ^
+        $Class[[C]] $Variable[[V]] $Field[[F]]
+      )"},
+                {
+                    R"(
+        $Class[[A]]
+        $Variable[[A]]
+        $Variable[[A]]
+      )",
+                    R"(
+        $Class[[A]]
+       ^$Variable[[AA]]
+        $Variable[[A]]
+      )"},
+                {
+                    R"(
+        $Class[[A]]
+        $Variable[[A]]
+        $Class[[A]]
+        $Variable[[A]]
+      )",
+                    R"(
+        $Class[[A]]
+        $Variable[[A]]
+      )"},
+                {
+                    R"(
+        $Class[[A]]
+        $Variable[[A]]
+      )",
+                    R"(
+        $Class[[A]]
+        $Variable[[A]]
+       ^$Class[[A]]
+       ^$Variable[[A]]
+      )"},
+                {
+                    R"(
+        $Variable[[A]]
+        $Variable[[A]]
+        $Variable[[A]]
+      )",
+                    R"(
+       ^$Class[[A]]
+       ^$Class[[A]]
+       ^$Class[[A]]
+      )"}};
+
+  for (const auto &Test : TestCases)
+    checkDiffedHighlights(Test.OldCode, Test.NewCode);
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang