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

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

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())

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");
+  ASSERT_TRUE(!Buffer.expandedTokens().empty());
+  ASSERT_EQ(Buffer.expandedTokens().back().kind(), tok::eof);
+  // Expanded range including `eof` is handled gracefully (`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