[clangd] Support multiline semantic tokens

[kadircet]

Summary

Per LSP, multiline tokens should be handled as if they end at the end
of the line starting the token (there's also a capability to enable them, but
that's an adventure for a different day).

Fixes https://github.com/clangd/clangd/issues/1145

[sammccall]

(My first reaction was that this belongs in semanticHighlights() rather than toSemanticTokens, but I think you got it right - the single-line restriction is pretty unnatural from clang's point of view, and cuts across tokens generated in different places).

clang-tools-extra/clangd/SemanticHighlighting.cpp

951

This wording is hard for me to follow. I think it's saying:

"If the token spans a line break, truncate it to avoid this"

951

It seems it would be better to split into one token per line, rather than simply truncating.

Is truncation for simplicity or is there a reason to prefer it?

FWIW I think this wouldn't be too hard to implement if you reordered the tokenType/tokenModifiers above so this part is the last step, and just copied the whole SemanticToken object from the back of the Result. But on the other hand it's not terribly important, either.

At least I think we should have a comment for the truncate/split tradeoff.

clang-tools-extra/clangd/SemanticHighlighting.h

90

I think this deserves a comment like:

// This is clang's token bounds, which may span multiple lines.
// This may be split/truncated to a single line when required by LSP.

[kadircet]

Split highlights into multiple tokens rather than trimming

clang-tools-extra/clangd/SemanticHighlighting.cpp

951

This wording is hard for me to follow. I think it's saying:

It is actually from LSP 😅 changed into a simpler version, I don't think we need references to LSP in here.

Is truncation for simplicity or is there a reason to prefer it?

Yeah I didn't want to do a loop, as in theory there can be many lines not just two. Also we need to take care of the Last and delta calculations for the following tokens.
But it isn't as bad, i suppose. Changed into splitting.

[sammccall]

clang-tools-extra/clangd/SemanticHighlighting.cpp

950

this copy feels gratuitous, can we just use Tokens.back(), or do the copy in the rare case?

956

This isn't literally true - files with multiline raw strings aren't rare. maybe "multiline tokens are rare"?
(even so I wonder if there's a real performance problem here...)

962

can we use < rather than != here? (or an assert)
I realize something has to have gone horribly wrong to infloop here, but it's something nonlocal...

973

reassigning fields one by one seems fragile, copy the last token and then overwrite the fields needed?

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

1007

please have the first token start somewhere other than column 0

[kadircet]

Get rid of the copy in common case

clang-tools-extra/clangd/SemanticHighlighting.cpp

950

unfortunately copy is needed, as the delta needs to be from start of the last token :/

changing the logic to do the copy only in the rare case by using a scratch element, LMK if you have better ideas here.

clang-tools-extra/clangd/ClangdLSPServer.cpp

for (char &C : llvm::reverse(S)) {
     }
     C = '0';
   }
   S.insert(S.begin(), '1');
 }
 
 void ClangdLSPServer::onSemanticTokens(const SemanticTokensParams &Params,
                                        Callback<SemanticTokens> CB) {
+  auto File = Params.textDocument.uri.file();
   Server->semanticHighlights(
       Params.textDocument.uri.file(),
-      [this, File(Params.textDocument.uri.file().str()), CB(std::move(CB))](
+      [this, File(File.str()), CB(std::move(CB)), Code(Server->getDraft(File))](
           llvm::Expected<std::vector<HighlightingToken>> HT) mutable {
         if (!HT)
           return CB(HT.takeError());
         SemanticTokens Result;
-        Result.tokens = toSemanticTokens(*HT);
+        Result.tokens = toSemanticTokens(*HT, *Code);
         {
           std::lock_guard<std::mutex> Lock(SemanticTokensMutex);
           auto &Last = LastSemanticTokens[File];
 
           Last.tokens = Result.tokens;
           increment(Last.resultId);
           Result.resultId = Last.resultId;
         }
         CB(std::move(Result));
       });
 }
 
 void ClangdLSPServer::onSemanticTokensDelta(
     const SemanticTokensDeltaParams &Params,
     Callback<SemanticTokensOrDelta> CB) {
+  auto File = Params.textDocument.uri.file();
   Server->semanticHighlights(
       Params.textDocument.uri.file(),
-      [this, PrevResultID(Params.previousResultId),
-       File(Params.textDocument.uri.file().str()), CB(std::move(CB))](
+      [this, PrevResultID(Params.previousResultId), File(File.str()),
+       CB(std::move(CB)), Code(Server->getDraft(File))](
           llvm::Expected<std::vector<HighlightingToken>> HT) mutable {
         if (!HT)
           return CB(HT.takeError());
-        std::vector<SemanticToken> Toks = toSemanticTokens(*HT);
+        std::vector<SemanticToken> Toks = toSemanticTokens(*HT, *Code);
 
         SemanticTokensOrDelta Result;
         {
           std::lock_guard<std::mutex> Lock(SemanticTokensMutex);
           auto &Last = LastSemanticTokens[File];
 
           if (PrevResultID == Last.resultId) {
             Result.edits = diffTokens(Last.tokens, Toks);

clang-tools-extra/clangd/SemanticHighlighting.h

 // calls etc.).
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_SEMANTICHIGHLIGHTING_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_SEMANTICHIGHLIGHTING_H
 
 #include "Protocol.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 class ParsedAST;
 
 enum class HighlightingKind {
   Variable = 0,

 static_assert(static_cast<unsigned>(HighlightingModifier::LastModifier) < 32,
               "Increase width of modifiers bitfield!");
 llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, HighlightingModifier K);
 
 // Contains all information needed for the highlighting a token.
 struct HighlightingToken {
   HighlightingKind Kind;
   uint32_t Modifiers = 0;
   Range R;
 
   HighlightingToken &addModifier(HighlightingModifier M) {
     Modifiers |= 1 << static_cast<unsigned>(M);
     return *this;
   }
 };
 
 bool operator==(const HighlightingToken &L, const HighlightingToken &R);
 bool operator<(const HighlightingToken &L, const HighlightingToken &R);
 
 // Returns all HighlightingTokens from an AST. Only generates highlights for the
 // main AST.
 std::vector<HighlightingToken> getSemanticHighlightings(ParsedAST &AST);
 
-std::vector<SemanticToken> toSemanticTokens(llvm::ArrayRef<HighlightingToken>);
+std::vector<SemanticToken> toSemanticTokens(llvm::ArrayRef<HighlightingToken>,
+                                            llvm::StringRef Code);
 llvm::StringRef toSemanticTokenType(HighlightingKind Kind);
 llvm::StringRef toSemanticTokenModifier(HighlightingModifier Modifier);
 std::vector<SemanticTokensEdit> diffTokens(llvm::ArrayRef<SemanticToken> Before,
                                            llvm::ArrayRef<SemanticToken> After);
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/SemanticHighlighting.cpp

 #include "clang/AST/TypeLoc.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/Error.h"
 #include <algorithm>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 /// Some names are not written in the source code and cannot be highlighted,
 /// e.g. anonymous classes. This function detects those cases.
return std::tie(L.R, L.Kind, L.Modifiers) ==
          std::tie(R.R, R.Kind, R.Modifiers);
 }
 bool operator<(const HighlightingToken &L, const HighlightingToken &R) {
   return std::tie(L.R, L.Kind, L.Modifiers) <
          std::tie(R.R, R.Kind, R.Modifiers);
 }
 
 std::vector<SemanticToken>
-toSemanticTokens(llvm::ArrayRef<HighlightingToken> Tokens) {
+toSemanticTokens(llvm::ArrayRef<HighlightingToken> Tokens,
+                 llvm::StringRef Code) {
   assert(std::is_sorted(Tokens.begin(), Tokens.end()));
   std::vector<SemanticToken> Result;
+  // In case we split a HighlightingToken into multiple tokens (e.g. because it
+  // was spanning multiple lines), this tracks the last one. This prevents
+  // having a copy all the time.
+  HighlightingToken Scratch;
   const HighlightingToken *Last = nullptr;
   for (const HighlightingToken &Tok : Tokens) {
     Result.emplace_back();
-    SemanticToken &Out = Result.back();
+    SemanticToken *Out = &Result.back();
     // deltaStart/deltaLine are relative if possible.
     if (Last) {
-      assert(Tok.R.start.line >= Last->R.start.line);
-      Out.deltaLine = Tok.R.start.line - Last->R.start.line;
-      if (Out.deltaLine == 0) {
+      assert(Tok.R.start.line >= Last->R.end.line);
+      Out->deltaLine = Tok.R.start.line - Last->R.end.line;
+      if (Out->deltaLine == 0) {
         assert(Tok.R.start.character >= Last->R.start.character);
-        Out.deltaStart = Tok.R.start.character - Last->R.start.character;
+        Out->deltaStart = Tok.R.start.character - Last->R.start.character;
       } else {
-        Out.deltaStart = Tok.R.start.character;
+        Out->deltaStart = Tok.R.start.character;
       }
     } else {
-      Out.deltaLine = Tok.R.start.line;
-      Out.deltaStart = Tok.R.start.character;
+      Out->deltaLine = Tok.R.start.line;
+      Out->deltaStart = Tok.R.start.character;
     }
-    assert(Tok.R.end.line == Tok.R.start.line);
-    Out.length = Tok.R.end.character - Tok.R.start.character;
-    Out.tokenType = static_cast<unsigned>(Tok.Kind);
-    Out.tokenModifiers = Tok.Modifiers;
-
+    Out->tokenType = static_cast<unsigned>(Tok.Kind);
+    Out->tokenModifiers = Tok.Modifiers;
     Last = &Tok;
+
+    if (Tok.R.end.line == Tok.R.start.line) {
+      Out->length = Tok.R.end.character - Tok.R.start.character;
+    } else {
+      // If the token spans a line break, split it into multiple pieces for each
+      // line.
+      // This is slow, but multiline tokens are rare.
+      // FIXME: There's a client capability for supporting multiline tokens,
+      // respect that.
+      auto TokStartOffset = llvm::cantFail(positionToOffset(Code, Tok.R.start));
+      // Note that the loop doesn't cover the last line, which has a special
+      // length.
+      for (int I = Tok.R.start.line; I < Tok.R.end.line; ++I) {
+        auto LineEnd = Code.find('\n', TokStartOffset);
+        assert(LineEnd != Code.npos);
+        Out->length = LineEnd - TokStartOffset;
+        // Token continues on next line, right after the line break.
+        TokStartOffset = LineEnd + 1;
+        Result.emplace_back();
+        Out = &Result.back();
+        *Out = Result[Result.size() - 2];
+        // New token starts at the first column of the next line.
+        Out->deltaLine = 1;
+        Out->deltaStart = 0;
+      }
+      // This is the token on last line.
+      Out->length = Tok.R.end.character;
+      // Update the start location for last token, as that's used in the
+      // relative delta calculation for following tokens.
+      Scratch = *Last;
+      Scratch.R.start.line = Tok.R.end.line;
+      Scratch.R.start.character = 0;
+      Last = &Scratch;
+    }
   }
   return Result;
 }
 llvm::StringRef toSemanticTokenType(HighlightingKind Kind) {
   switch (Kind) {
   case HighlightingKind::Variable:
   case HighlightingKind::LocalVariable:
   case HighlightingKind::StaticField:

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

 TEST(SemanticHighlighting, toSemanticTokens) {
   auto Tokens = tokens(R"(
  [[blah]]
 
     $Function[[big]] [[bang]]
   )");
   Tokens.front().Modifiers |= unsigned(HighlightingModifier::Declaration);
   Tokens.front().Modifiers |= unsigned(HighlightingModifier::Readonly);
-  auto Results = toSemanticTokens(Tokens);
+  auto Results = toSemanticTokens(Tokens, /*Code=*/"");
 
   ASSERT_THAT(Results, SizeIs(3));
   EXPECT_EQ(Results[0].tokenType, unsigned(HighlightingKind::Variable));
   EXPECT_EQ(Results[0].tokenModifiers,
             unsigned(HighlightingModifier::Declaration) |
                 unsigned(HighlightingModifier::Readonly));
   EXPECT_EQ(Results[0].deltaLine, 1u);
   EXPECT_EQ(Results[0].deltaStart, 1u);
[[blah]]
   EXPECT_EQ(Results[2].deltaStart, 4u);
   EXPECT_EQ(Results[2].length, 4u);
 }
 
 TEST(SemanticHighlighting, diffSemanticTokens) {
   auto Before = toSemanticTokens(tokens(R"(
     [[foo]] [[bar]] [[baz]]
     [[one]] [[two]] [[three]]
-  )"));
+  )"),
+                                 /*Code=*/"");
   EXPECT_THAT(diffTokens(Before, Before), IsEmpty());
 
   auto After = toSemanticTokens(tokens(R"(
     [[foo]] [[hello]] [[world]] [[baz]]
     [[one]] [[two]] [[three]]
-  )"));
+  )"),
+                                /*Code=*/"");
 
   // Replace [bar, baz] with [hello, world, baz]
   auto Diff = diffTokens(Before, After);
   ASSERT_THAT(Diff, SizeIs(1));
   EXPECT_EQ(1u, Diff.front().startToken);
   EXPECT_EQ(2u, Diff.front().deleteTokens);
   ASSERT_THAT(Diff.front().tokens, SizeIs(3));
   // hello
   EXPECT_EQ(0u, Diff.front().tokens[0].deltaLine);
   EXPECT_EQ(4u, Diff.front().tokens[0].deltaStart);
   EXPECT_EQ(5u, Diff.front().tokens[0].length);
   // world
   EXPECT_EQ(0u, Diff.front().tokens[1].deltaLine);
   EXPECT_EQ(6u, Diff.front().tokens[1].deltaStart);
   EXPECT_EQ(5u, Diff.front().tokens[1].length);
   // baz
   EXPECT_EQ(0u, Diff.front().tokens[2].deltaLine);
   EXPECT_EQ(6u, Diff.front().tokens[2].deltaStart);
   EXPECT_EQ(3u, Diff.front().tokens[2].length);
 }
 
+TEST(SemanticHighlighting, MultilineTokens) {
+  llvm::StringRef AnnotatedCode = R"cpp(
+  [[fo
+o
+o]] [[bar]])cpp";
+  auto Toks = toSemanticTokens(tokens(AnnotatedCode),
+                               Annotations(AnnotatedCode).code());
+  ASSERT_THAT(Toks, SizeIs(4));
+  // foo
+  EXPECT_EQ(Toks[0].deltaLine, 1u);
+  EXPECT_EQ(Toks[0].deltaStart, 2u);
+  EXPECT_EQ(Toks[0].length, 2u);
+  EXPECT_EQ(Toks[1].deltaLine, 1u);
+  EXPECT_EQ(Toks[1].deltaStart, 0u);
+  EXPECT_EQ(Toks[1].length, 1u);
+  EXPECT_EQ(Toks[2].deltaLine, 1u);
+  EXPECT_EQ(Toks[2].deltaStart, 0u);
+  EXPECT_EQ(Toks[2].length, 1u);
+
+  // bar
+  EXPECT_EQ(Toks[3].deltaLine, 0u);
+  EXPECT_EQ(Toks[3].deltaStart, 2u);
+  EXPECT_EQ(Toks[3].length, 3u);
+}
 } // namespace
 } // namespace clangd
 } // namespace clang