[clangd] Fix unicode handling, using UTF-16 where LSP requires it.

Summary:
The Language Server Protocol unfortunately mandates that locations in files
be represented by line/column pairs, where the "column" is actually an index
into the UTF-16-encoded text of the line.
(This is because VSCode is written in JavaScript, which is UTF-16-native).

Internally clangd treats source files at UTF-8, the One True Encoding, and
generally deals with byte offsets (though there are exceptions).

Before this patch, conversions between offsets and LSP Position pretended
that Position.character was UTF-8 bytes, which is only true for ASCII lines.
Now we examine the text to convert correctly (but don't actually need to
transcode it, due to some nice details of the encodings).

The updated functions in SourceCode are the blessed way to interact with
the Position.character field, and anything else is likely to be wrong.
So I also updated the other accesses:
 - CodeComplete needs a "clang-style" line/column, with column in utf-8 bytes.
   This is now converted via Position -> offset -> clang line/column
   (a new function is added to SourceCode.h for the second conversion).
 - getBeginningOfIdentifier skipped backwards in UTF-16 space, which is will
   behave badly when it splits a surrogate pair. Skipping backwards in UTF-8
   coordinates gives the lexer a fighting chance of getting this right.
   While here, I clarified(?) the logic comments, fixed a bug with identifiers
   containing digits, simplified the signature slightly and added a test.

This seems likely to cause problems with editors that have the same bug, and
treat the protocol as if columns are UTF-8 bytes. But we can find and fix those.

Reviewers: hokein

Subscribers: klimek, ilya-biryukov, ioeric, MaskRay, jkorous, cfe-commits

Differential Revision: https://reviews.llvm.org/D46035

8 files changed, 147 insertions, 88 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index e72cf40159c..cd8b3f4a859 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -232,14 +232,8 @@ void ClangdServer::rename(PathRef File, Position Pos, llvm::StringRef NewName,
 
     RefactoringResultCollector ResultCollector;
     const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-    const FileEntry *FE =
-        SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
-    if (!FE)
-      return CB(llvm::make_error<llvm::StringError>(
-          "rename called for non-added document",
-          llvm::errc::invalid_argument));
     SourceLocation SourceLocationBeg =
-        clangd::getBeginningOfIdentifier(AST, Pos, FE);
+        clangd::getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
     tooling::RefactoringRuleContext Context(
         AST.getASTContext().getSourceManager());
     Context.setASTContext(AST.getASTContext());
diff --git a/clang-tools-extra/clangd/ClangdUnit.cpp b/clang-tools-extra/clangd/ClangdUnit.cpp
index 7309e1ad49b..c750f578e65 100644
--- a/clang-tools-extra/clangd/ClangdUnit.cpp
+++ b/clang-tools-extra/clangd/ClangdUnit.cpp
@@ -215,19 +215,6 @@ ParsedAST::Build(std::unique_ptr<clang::CompilerInvocation> CI,
                    std::move(IncLocations));
 }
 
-namespace {
-
-SourceLocation getMacroArgExpandedLocation(const SourceManager &Mgr,
-                                           const FileEntry *FE, Position Pos) {
-  // The language server protocol uses zero-based line and column numbers.
-  // Clang uses one-based numbers.
-  SourceLocation InputLoc =
-      Mgr.translateFileLineCol(FE, Pos.line + 1, Pos.character + 1);
-  return Mgr.getMacroArgExpandedLocation(InputLoc);
-}
-
-} // namespace
-
 void ParsedAST::ensurePreambleDeclsDeserialized() {
   if (PreambleDeclsDeserialized || !Preamble)
     return;
@@ -470,40 +457,34 @@ CppFile::rebuildPreamble(CompilerInvocation &CI,
 
 SourceLocation clangd::getBeginningOfIdentifier(ParsedAST &Unit,
                                                 const Position &Pos,
-                                                const FileEntry *FE) {
+                                                const FileID FID) {
   const ASTContext &AST = Unit.getASTContext();
   const SourceManager &SourceMgr = AST.getSourceManager();
-
-  SourceLocation InputLocation =
-      getMacroArgExpandedLocation(SourceMgr, FE, Pos);
-  if (Pos.character == 0) {
-    return InputLocation;
-  }
-
-  // This handle cases where the position is in the middle of a token or right
-  // after the end of a token. In theory we could just use GetBeginningOfToken
-  // to find the start of the token at the input position, but this doesn't
-  // work when right after the end, i.e. foo|.
-  // So try to go back by one and see if we're still inside an identifier
-  // token. If so, Take the beginning of this token.
-  // (It should be the same identifier because you can't have two adjacent
-  // identifiers without another token in between.)
-  Position PosCharBehind = Pos;
-  --PosCharBehind.character;
-
-  SourceLocation PeekBeforeLocation =
-      getMacroArgExpandedLocation(SourceMgr, FE, PosCharBehind);
-  Token Result;
-  if (Lexer::getRawToken(PeekBeforeLocation, Result, SourceMgr,
-                         AST.getLangOpts(), false)) {
-    // getRawToken failed, just use InputLocation.
-    return InputLocation;
+  auto Offset = positionToOffset(SourceMgr.getBufferData(FID), Pos);
+  if (!Offset) {
+    log("getBeginningOfIdentifier: " + toString(Offset.takeError()));
+    return SourceLocation();
   }
-
-  if (Result.is(tok::raw_identifier)) {
-    return Lexer::GetBeginningOfToken(PeekBeforeLocation, SourceMgr,
-                                      AST.getLangOpts());
-  }
-
-  return InputLocation;
+  SourceLocation InputLoc = SourceMgr.getComposedLoc(FID, *Offset);
+
+  // GetBeginningOfToken(pos) is almost what we want, but does the wrong thing
+  // if the cursor is at the end of the identifier.
+  // Instead, we lex at GetBeginningOfToken(pos - 1). The cases are:
+  //  1) at the beginning of an identifier, we'll be looking at something
+  //  that isn't an identifier.
+  //  2) at the middle or end of an identifier, we get the identifier.
+  //  3) anywhere outside an identifier, we'll get some non-identifier thing.
+  // We can't actually distinguish cases 1 and 3, but returning the original
+  // location is correct for both!
+  if (*Offset == 0) // Case 1 or 3.
+    return SourceMgr.getMacroArgExpandedLocation(InputLoc);
+  SourceLocation Before =
+      SourceMgr.getMacroArgExpandedLocation(InputLoc.getLocWithOffset(-1));
+  Before = Lexer::GetBeginningOfToken(Before, SourceMgr, AST.getLangOpts());
+  Token Tok;
+  if (Before.isValid() &&
+      !Lexer::getRawToken(Before, Tok, SourceMgr, AST.getLangOpts(), false) &&
+      Tok.is(tok::raw_identifier))
+    return Before;                                        // Case 2.
+  return SourceMgr.getMacroArgExpandedLocation(InputLoc); // Case 1 or 3.
 }
diff --git a/clang-tools-extra/clangd/ClangdUnit.h b/clang-tools-extra/clangd/ClangdUnit.h
index b3535b5a536..84b5acca76e 100644
--- a/clang-tools-extra/clangd/ClangdUnit.h
+++ b/clang-tools-extra/clangd/ClangdUnit.h
@@ -173,8 +173,9 @@ private:
 };
 
 /// Get the beginning SourceLocation at a specified \p Pos.
+/// May be invalid if Pos is, or if there's no identifier.
 SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
-                                        const FileEntry *FE);
+                                        const FileID FID);
 
 /// For testing/debugging purposes. Note that this method deserializes all
 /// unserialized Decls, so use with care.
diff --git a/clang-tools-extra/clangd/CodeComplete.cpp b/clang-tools-extra/clangd/CodeComplete.cpp
index 87536459db1..552f37e8030 100644
--- a/clang-tools-extra/clangd/CodeComplete.cpp
+++ b/clang-tools-extra/clangd/CodeComplete.cpp
@@ -729,8 +729,15 @@ bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   FrontendOpts.SkipFunctionBodies = true;
   FrontendOpts.CodeCompleteOpts = Options;
   FrontendOpts.CodeCompletionAt.FileName = Input.FileName;
-  FrontendOpts.CodeCompletionAt.Line = Input.Pos.line + 1;
-  FrontendOpts.CodeCompletionAt.Column = Input.Pos.character + 1;
+  auto Offset = positionToOffset(Input.Contents, Input.Pos);
+  if (!Offset) {
+    log("Code completion position was invalid " +
+        llvm::toString(Offset.takeError()));
+    return false;
+  }
+  std::tie(FrontendOpts.CodeCompletionAt.Line,
+           FrontendOpts.CodeCompletionAt.Column) =
+      offsetToClangLineColumn(Input.Contents, *Offset);
 
   Clang->setCodeCompletionConsumer(Consumer.release());
 
diff --git a/clang-tools-extra/clangd/Protocol.h b/clang-tools-extra/clangd/Protocol.h
index a54179ca3ee..2e2c4a29692 100644
--- a/clang-tools-extra/clangd/Protocol.h
+++ b/clang-tools-extra/clangd/Protocol.h
@@ -91,6 +91,8 @@ struct Position {
   int line = 0;
 
   /// Character offset on a line in a document (zero-based).
+  /// WARNING: this is in UTF-16 codepoints, not bytes or characters!
+  /// Use the functions in SourceCode.h to construct/interpret Positions.
   int character = 0;
 
   friend bool operator==(const Position &LHS, const Position &RHS) {
diff --git a/clang-tools-extra/clangd/SourceCode.cpp b/clang-tools-extra/clangd/SourceCode.cpp
index 65a7ca987c1..ddfaec888e2 100644
--- a/clang-tools-extra/clangd/SourceCode.cpp
+++ b/clang-tools-extra/clangd/SourceCode.cpp
@@ -16,6 +16,66 @@ namespace clang {
 namespace clangd {
 using namespace llvm;
 
+// Here be dragons. LSP positions use columns measured in *UTF-16 code units*!
+// Clangd uses UTF-8 and byte-offsets internally, so conversion is nontrivial.
+
+// Iterates over unicode codepoints in the (UTF-8) string. For each,
+// invokes CB(UTF-8 length, UTF-16 length), and breaks if it returns true.
+// Returns true if CB returned true, false if we hit the end of string.
+template <typename Callback>
+static bool iterateCodepoints(StringRef U8, const Callback &CB) {
+  for (size_t I = 0; I < U8.size();) {
+    unsigned char C = static_cast<unsigned char>(U8[I]);
+    if (LLVM_LIKELY(!(C & 0x80))) { // ASCII character.
+      if (CB(1, 1))
+        return true;
+      ++I;
+      continue;
+    }
+    // This convenient property of UTF-8 holds for all non-ASCII characters.
+    size_t UTF8Length = countLeadingOnes(C);
+    // 0xxx is ASCII, handled above. 10xxx is a trailing byte, invalid here.
+    // 11111xxx is not valid UTF-8 at all. Assert because it's probably our bug.
+    assert((UTF8Length >= 2 && UTF8Length <= 4) &&
+           "Invalid UTF-8, or transcoding bug?");
+    I += UTF8Length; // Skip over all trailing bytes.
+    // A codepoint takes two UTF-16 code unit if it's astral (outside BMP).
+    // Astral codepoints are encoded as 4 bytes in UTF-8 (11110xxx ...)
+    if (CB(UTF8Length, UTF8Length == 4 ? 2 : 1))
+      return true;
+  }
+  return false;
+}
+
+// Returns the offset into the string that matches \p Units UTF-16 code units.
+// Conceptually, this converts to UTF-16, truncates to CodeUnits, converts back
+// to UTF-8, and returns the length in bytes.
+static size_t measureUTF16(StringRef U8, int U16Units, bool &Valid) {
+  size_t Result = 0;
+  Valid = U16Units == 0 || iterateCodepoints(U8, [&](int U8Len, int U16Len) {
+            Result += U8Len;
+            U16Units -= U16Len;
+            return U16Units <= 0;
+          });
+  if (U16Units < 0) // Offset was into the middle of a surrogate pair.
+    Valid = false;
+  // Don't return an out-of-range index if we overran.
+  return std::min(Result, U8.size());
+}
+
+// Counts the number of UTF-16 code units needed to represent a string.
+// Like most strings in clangd, the input is UTF-8 encoded.
+static size_t utf16Len(StringRef U8) {
+  // A codepoint takes two UTF-16 code unit if it's astral (outside BMP).
+  // Astral codepoints are encoded as 4 bytes in UTF-8, starting with 11110xxx.
+  size_t Count = 0;
+  iterateCodepoints(U8, [&](int U8Len, int U16Len) {
+    Count += U16Len;
+    return false;
+  });
+  return Count;
+}
+
 llvm::Expected<size_t> positionToOffset(StringRef Code, Position P,
                                         bool AllowColumnsBeyondLineLength) {
   if (P.line < 0)
@@ -40,12 +100,15 @@ llvm::Expected<size_t> positionToOffset(StringRef Code, Position P,
   if (NextNL == StringRef::npos)
     NextNL = Code.size();
 
-  if (StartOfLine + P.character > NextNL && !AllowColumnsBeyondLineLength)
+  bool Valid;
+  size_t ByteOffsetInLine = measureUTF16(
+      Code.substr(StartOfLine, NextNL - StartOfLine), P.character, Valid);
+  if (!Valid && !AllowColumnsBeyondLineLength)
     return llvm::make_error<llvm::StringError>(
-        llvm::formatv("Character value is out of range ({0})", P.character),
+        llvm::formatv("UTF-16 offset {0} is invalid for line {1}", P.character,
+                      P.line),
         llvm::errc::invalid_argument);
-  // FIXME: officially P.character counts UTF-16 code units, not UTF-8 bytes!
-  return std::min(NextNL, StartOfLine + P.character);
+  return StartOfLine + ByteOffsetInLine;
 }
 
 Position offsetToPosition(StringRef Code, size_t Offset) {
@@ -54,17 +117,26 @@ Position offsetToPosition(StringRef Code, size_t Offset) {
   int Lines = Before.count('\n');
   size_t PrevNL = Before.rfind('\n');
   size_t StartOfLine = (PrevNL == StringRef::npos) ? 0 : (PrevNL + 1);
-  // FIXME: officially character counts UTF-16 code units, not UTF-8 bytes!
   Position Pos;
   Pos.line = Lines;
-  Pos.character = static_cast<int>(Offset - StartOfLine);
+  Pos.character = utf16Len(Before.substr(StartOfLine));
   return Pos;
 }
 
 Position sourceLocToPosition(const SourceManager &SM, SourceLocation Loc) {
+  // We use the SourceManager's line tables, but its column number is in bytes.
+  FileID FID;
+  unsigned Offset;
+  std::tie(FID, Offset) = SM.getDecomposedSpellingLoc(Loc);
   Position P;
-  P.line = static_cast<int>(SM.getSpellingLineNumber(Loc)) - 1;
-  P.character = static_cast<int>(SM.getSpellingColumnNumber(Loc)) - 1;
+  P.line = static_cast<int>(SM.getLineNumber(FID, Offset)) - 1;
+  bool Invalid = false;
+  StringRef Code = SM.getBufferData(FID, &Invalid);
+  if (!Invalid) {
+    auto ColumnInBytes = SM.getColumnNumber(FID, Offset) - 1;
+    auto LineSoFar = Code.substr(Offset - ColumnInBytes, ColumnInBytes);
+    P.character = utf16Len(LineSoFar);
+  }
   return P;
 }
 
@@ -76,6 +148,16 @@ Range halfOpenToRange(const SourceManager &SM, CharSourceRange R) {
   return {Begin, End};
 }
 
+std::pair<size_t, size_t> offsetToClangLineColumn(StringRef Code,
+                                                  size_t Offset) {
+  Offset = std::min(Code.size(), Offset);
+  StringRef Before = Code.substr(0, Offset);
+  int Lines = Before.count('\n');
+  size_t PrevNL = Before.rfind('\n');
+  size_t StartOfLine = (PrevNL == StringRef::npos) ? 0 : (PrevNL + 1);
+  return {Lines + 1, Offset - StartOfLine + 1};
+}
+
 std::pair<llvm::StringRef, llvm::StringRef>
 splitQualifiedName(llvm::StringRef QName) {
   size_t Pos = QName.rfind("::");
diff --git a/clang-tools-extra/clangd/SourceCode.h b/clang-tools-extra/clangd/SourceCode.h
index 43fd07015ec..8867b90c3c7 100644
--- a/clang-tools-extra/clangd/SourceCode.h
+++ b/clang-tools-extra/clangd/SourceCode.h
@@ -23,14 +23,9 @@ namespace clangd {
 
 /// Turn a [line, column] pair into an offset in Code.
 ///
-/// If the character value is greater than the line length, the behavior depends
-/// on AllowColumnsBeyondLineLength:
-///
-/// - if true: default back to the end of the line
-/// - if false: return an error
-///
-/// If the line number is greater than the number of lines in the document,
-/// always return an error.
+/// If P.character exceeds the line length, returns the offset at end-of-line.
+/// (If !AllowColumnsBeyondLineLength, then returns an error instead).
+/// If the line number is out of range, returns an error.
 ///
 /// The returned value is in the range [0, Code.size()].
 llvm::Expected<size_t>
@@ -38,7 +33,7 @@ positionToOffset(llvm::StringRef Code, Position P,
                  bool AllowColumnsBeyondLineLength = true);
 
 /// Turn an offset in Code into a [line, column] pair.
-/// FIXME: This should return an error if the offset is invalid.
+/// The offset must be in range [0, Code.size()].
 Position offsetToPosition(llvm::StringRef Code, size_t Offset);
 
 /// Turn a SourceLocation into a [line, column] pair.
@@ -49,6 +44,12 @@ Position sourceLocToPosition(const SourceManager &SM, SourceLocation Loc);
 // Note that clang also uses closed source ranges, which this can't handle!
 Range halfOpenToRange(const SourceManager &SM, CharSourceRange R);
 
+// Converts an offset to a clang line/column (1-based, columns are bytes).
+// The offset must be in range [0, Code.size()].
+// Prefer to use SourceManager if one is available.
+std::pair<size_t, size_t> offsetToClangLineColumn(llvm::StringRef Code,
+                                                  size_t Offset);
+
 /// From "a::b::c", return {"a::b::", "c"}. Scope is empty if there's no
 /// qualifier.
 std::pair<llvm::StringRef, llvm::StringRef>
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 01d4e952355..45241b7b291 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -164,11 +164,8 @@ makeLocation(ParsedAST &AST, const SourceRange &ValSourceRange) {
 
 std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
-  if (!FE)
-    return {};
-
-  SourceLocation SourceLocationBeg = getBeginningOfIdentifier(AST, Pos, FE);
+  SourceLocation SourceLocationBeg =
+      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
 
   std::vector<Location> Result;
   // Handle goto definition for #include.
@@ -280,11 +277,8 @@ private:
 std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
                                                       Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
-  if (!FE)
-    return {};
-
-  SourceLocation SourceLocationBeg = getBeginningOfIdentifier(AST, Pos, FE);
+  SourceLocation SourceLocationBeg =
+      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
 
   DeclarationAndMacrosFinder DeclMacrosFinder(llvm::errs(), SourceLocationBeg,
                                               AST.getASTContext(),
@@ -413,11 +407,8 @@ static Hover getHoverContents(StringRef MacroName) {
 
 Hover getHover(ParsedAST &AST, Position Pos) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-  const FileEntry *FE = SourceMgr.getFileEntryForID(SourceMgr.getMainFileID());
-  if (FE == nullptr)
-    return Hover();
-
-  SourceLocation SourceLocationBeg = getBeginningOfIdentifier(AST, Pos, FE);
+  SourceLocation SourceLocationBeg =
+      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
   DeclarationAndMacrosFinder DeclMacrosFinder(llvm::errs(), SourceLocationBeg,
                                               AST.getASTContext(),
                                               AST.getPreprocessor());