[clangd] Handle an expanded token range that ends in the eof token in TokenBuffer::spelledForExpanded()

[HighCommander4]

Such ranges can legitimately arise in the case of invalid code, such as a declaration missing an ending brace.

Fixes clangd/clangd#1559

[HighCommander4]

Based on previous discussion and analysis in #69849

[llvmbot]

@llvm/pr-subscribers-clangd

@llvm/pr-subscribers-clang

Author: Nathan Ridge (HighCommander4)

Changes

Such ranges can legitimately arise in the case of invalid code, such as a declaration missing an ending brace.

---

Full diff: https://github.com/llvm/llvm-project/pull/78092.diff

3 Files Affected:

• (modified) clang-tools-extra/clangd/unittests/DumpASTTests.cpp (+11)
• (modified) clang/lib/Tooling/Syntax/Tokens.cpp (+6)
• (modified) clang/unittests/Tooling/Syntax/TokensTest.cpp (+12)

diff --git a/clang-tools-extra/clangd/unittests/DumpASTTests.cpp b/clang-tools-extra/clangd/unittests/DumpASTTests.cpp
index d1b8f21b82c65a..304682118c871d 100644
--- a/clang-tools-extra/clangd/unittests/DumpASTTests.cpp
+++ b/clang-tools-extra/clangd/unittests/DumpASTTests.cpp
@@ -186,6 +186,17 @@ TEST(DumpASTTests, Arcana) {
   EXPECT_THAT(Node.children.front().arcana, testing::StartsWith("QualType "));
 }
 
+TEST(DumpASTTests, UnbalancedBraces) {
+  // Test that we don't crash while trying to compute a source range for the
+  // node whose ending brace is missing, and also that the source range is
+  // not empty.
+  Annotations Case("/*error-ok*/ $func[[int main() {]]");
+  ParsedAST AST = TestTU::withCode(Case.code()).build();
+  auto Node = dumpAST(DynTypedNode::create(findDecl(AST, "main")),
+                      AST.getTokens(), AST.getASTContext());
+  ASSERT_EQ(Node.range, Case.range("func"));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
diff --git a/clang/lib/Tooling/Syntax/Tokens.cpp b/clang/lib/Tooling/Syntax/Tokens.cpp
index 2f28b9cf286a63..8d32c45a4a70cf 100644
--- a/clang/lib/Tooling/Syntax/Tokens.cpp
+++ b/clang/lib/Tooling/Syntax/Tokens.cpp
@@ -401,6 +401,12 @@ std::string TokenBuffer::Mapping::str() const {
 
 std::optional<llvm::ArrayRef<syntax::Token>>
 TokenBuffer::spelledForExpanded(llvm::ArrayRef<syntax::Token> Expanded) const {
+  // In cases of invalid code, AST nodes can have source ranges that include
+  // the `eof` token. As there's no spelling for this token, exclude it from
+  // the range.
+  if (!Expanded.empty() && Expanded.back().kind() == tok::eof) {
+    Expanded = Expanded.drop_back();
+  }
   // Mapping an empty range is ambiguous in case of empty mappings at either end
   // of the range, bail out in that case.
   if (Expanded.empty())
diff --git a/clang/unittests/Tooling/Syntax/TokensTest.cpp b/clang/unittests/Tooling/Syntax/TokensTest.cpp
index 0c08318a637c0b..45b2f0de4eabc1 100644
--- a/clang/unittests/Tooling/Syntax/TokensTest.cpp
+++ b/clang/unittests/Tooling/Syntax/TokensTest.cpp
@@ -816,6 +816,18 @@ TEST_F(TokenBufferTest, SpelledByExpanded) {
   EXPECT_EQ(Buffer.spelledForExpanded(findExpanded("prev good")), std::nullopt);
 }
 
+TEST_F(TokenBufferTest, NoCrashForEofToken) {
+  recordTokens(R"cpp(
+    int main() {
+  )cpp");
+  // Calling spelledForExpanded() on the entire range of expanded tokens (which
+  // includes the `eof` token at the end) produces the range of all the spelled
+  // tokens (the `eof` is ignored).
+  EXPECT_THAT(
+      Buffer.spelledForExpanded(Buffer.expandedTokens()),
+      ValueIs(SameRange(Buffer.spelledTokens(SourceMgr->getMainFileID()))));
+}
+
 TEST_F(TokenBufferTest, ExpandedTokensForRange) {
   recordTokens(R"cpp(
     #define SIGN(X) X##_washere

[ilya-biryukov]

Thanks, LGTM! I only have a minor (and optional) suggestion.