[clang] [clang-tools-extra] [clangd] Fix crash with null check for Token at Loc (PR #94528)

Utkarsh Saxena via cfe-commits cfe-commits at lists.llvm.org
Fri Jun 7 02:06:57 PDT 2024


https://github.com/usx95 updated https://github.com/llvm/llvm-project/pull/94528

>From b13a663ae347649a3bcf9d6e381a5fbbfdc9ea4b Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Wed, 5 Jun 2024 19:48:48 +0000
Subject: [PATCH 1/6] [clangd] Fix crash with null check for Token at Loc

---
 clang-tools-extra/clangd/XRefs.cpp            | 32 +++++++++++--------
 .../clangd/unittests/XRefsTests.cpp           | 14 +++++++-
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index cd909266489a8..f52228e599591 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -1354,6 +1354,8 @@ maybeFindIncludeReferences(ParsedAST &AST, Position Pos,
 
         ReferencesResult::Reference Result;
         const auto *Token = AST.getTokens().spelledTokenAt(Loc);
+        if (!Token)
+          return;
         Result.Loc.range = Range{sourceLocToPosition(SM, Token->location()),
                                  sourceLocToPosition(SM, Token->endLocation())};
         Result.Loc.uri = URIMainFile;
@@ -2012,15 +2014,15 @@ static QualType typeForNode(const SelectionTree::Node *N) {
   return QualType();
 }
 
-// Given a type targeted by the cursor, return one or more types that are more interesting
-// to target.
-static void unwrapFindType(
-    QualType T, const HeuristicResolver* H, llvm::SmallVector<QualType>& Out) {
+// Given a type targeted by the cursor, return one or more types that are more
+// interesting to target.
+static void unwrapFindType(QualType T, const HeuristicResolver *H,
+                           llvm::SmallVector<QualType> &Out) {
   if (T.isNull())
     return;
 
   // If there's a specific type alias, point at that rather than unwrapping.
-  if (const auto* TDT = T->getAs<TypedefType>())
+  if (const auto *TDT = T->getAs<TypedefType>())
     return Out.push_back(QualType(TDT, 0));
 
   // Pointers etc => pointee type.
@@ -2044,17 +2046,18 @@ static void unwrapFindType(
 
   // For smart pointer types, add the underlying type
   if (H)
-    if (const auto* PointeeType = H->getPointeeType(T.getNonReferenceType().getTypePtr())) {
-        unwrapFindType(QualType(PointeeType, 0), H, Out);
-        return Out.push_back(T);
+    if (const auto *PointeeType =
+            H->getPointeeType(T.getNonReferenceType().getTypePtr())) {
+      unwrapFindType(QualType(PointeeType, 0), H, Out);
+      return Out.push_back(T);
     }
 
   return Out.push_back(T);
 }
 
 // Convenience overload, to allow calling this without the out-parameter
-static llvm::SmallVector<QualType> unwrapFindType(
-    QualType T, const HeuristicResolver* H) {
+static llvm::SmallVector<QualType> unwrapFindType(QualType T,
+                                                  const HeuristicResolver *H) {
   llvm::SmallVector<QualType> Result;
   unwrapFindType(T, H, Result);
   return Result;
@@ -2076,10 +2079,11 @@ std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos,
     std::vector<LocatedSymbol> LocatedSymbols;
 
     // NOTE: unwrapFindType might return duplicates for something like
-    // unique_ptr<unique_ptr<T>>. Let's *not* remove them, because it gives you some
-    // information about the type you may have not known before
-    // (since unique_ptr<unique_ptr<T>> != unique_ptr<T>).
-    for (const QualType& Type : unwrapFindType(typeForNode(N), AST.getHeuristicResolver()))
+    // unique_ptr<unique_ptr<T>>. Let's *not* remove them, because it gives you
+    // some information about the type you may have not known before (since
+    // unique_ptr<unique_ptr<T>> != unique_ptr<T>).
+    for (const QualType &Type :
+         unwrapFindType(typeForNode(N), AST.getHeuristicResolver()))
       llvm::copy(locateSymbolForType(AST, Type, Index),
                  std::back_inserter(LocatedSymbols));
 
diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
index f53cbf01b7992..c4def624a2fcc 100644
--- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
@@ -2358,7 +2358,14 @@ TEST(FindReferences, UsedSymbolsFromInclude) {
 
       R"cpp([[#in^clude <vector>]]
         std::[[vector]]<int> vec;
-      )cpp"};
+      )cpp",
+
+      R"cpp(
+        [[#include ^"operator_qoutes.h"]]
+        using ::operator_qoutes::[[operator]]"" _b;
+        auto x = 1_b;
+      )cpp",
+  };
   for (const char *Test : Tests) {
     Annotations T(Test);
     auto TU = TestTU::withCode(T.code());
@@ -2375,6 +2382,11 @@ TEST(FindReferences, UsedSymbolsFromInclude) {
         class vector{};
       }
     )cpp");
+    TU.AdditionalFiles["operator_qoutes.h"] = guard(R"cpp(
+      namespace operator_qoutes {
+        bool operator"" _b(unsigned long long value);
+      }
+    )cpp");
     TU.ExtraArgs.push_back("-isystem" + testPath("system"));
 
     auto AST = TU.build();

>From 3fe19589a332342abb069316224bdf9a6150c660 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Wed, 5 Jun 2024 19:59:44 +0000
Subject: [PATCH 2/6] remove unintentional format

---
 clang-tools-extra/clangd/XRefs.cpp | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index f52228e599591..033c75f02d561 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -2014,15 +2014,15 @@ static QualType typeForNode(const SelectionTree::Node *N) {
   return QualType();
 }
 
-// Given a type targeted by the cursor, return one or more types that are more
-// interesting to target.
-static void unwrapFindType(QualType T, const HeuristicResolver *H,
-                           llvm::SmallVector<QualType> &Out) {
+// Given a type targeted by the cursor, return one or more types that are more interesting
+// to target.
+static void unwrapFindType(
+    QualType T, const HeuristicResolver* H, llvm::SmallVector<QualType>& Out) {
   if (T.isNull())
     return;
 
   // If there's a specific type alias, point at that rather than unwrapping.
-  if (const auto *TDT = T->getAs<TypedefType>())
+  if (const auto* TDT = T->getAs<TypedefType>())
     return Out.push_back(QualType(TDT, 0));
 
   // Pointers etc => pointee type.
@@ -2046,18 +2046,17 @@ static void unwrapFindType(QualType T, const HeuristicResolver *H,
 
   // For smart pointer types, add the underlying type
   if (H)
-    if (const auto *PointeeType =
-            H->getPointeeType(T.getNonReferenceType().getTypePtr())) {
-      unwrapFindType(QualType(PointeeType, 0), H, Out);
-      return Out.push_back(T);
+    if (const auto* PointeeType = H->getPointeeType(T.getNonReferenceType().getTypePtr())) {
+        unwrapFindType(QualType(PointeeType, 0), H, Out);
+        return Out.push_back(T);
     }
 
   return Out.push_back(T);
 }
 
 // Convenience overload, to allow calling this without the out-parameter
-static llvm::SmallVector<QualType> unwrapFindType(QualType T,
-                                                  const HeuristicResolver *H) {
+static llvm::SmallVector<QualType> unwrapFindType(
+    QualType T, const HeuristicResolver* H) {
   llvm::SmallVector<QualType> Result;
   unwrapFindType(T, H, Result);
   return Result;
@@ -2079,11 +2078,10 @@ std::vector<LocatedSymbol> findType(ParsedAST &AST, Position Pos,
     std::vector<LocatedSymbol> LocatedSymbols;
 
     // NOTE: unwrapFindType might return duplicates for something like
-    // unique_ptr<unique_ptr<T>>. Let's *not* remove them, because it gives you
-    // some information about the type you may have not known before (since
-    // unique_ptr<unique_ptr<T>> != unique_ptr<T>).
-    for (const QualType &Type :
-         unwrapFindType(typeForNode(N), AST.getHeuristicResolver()))
+    // unique_ptr<unique_ptr<T>>. Let's *not* remove them, because it gives you some
+    // information about the type you may have not known before
+    // (since unique_ptr<unique_ptr<T>> != unique_ptr<T>).
+    for (const QualType& Type : unwrapFindType(typeForNode(N), AST.getHeuristicResolver()))
       llvm::copy(locateSymbolForType(AST, Type, Index),
                  std::back_inserter(LocatedSymbols));
 

>From 1a2e3e4d7f381553f855389c4ff471dac5fa751d Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Thu, 6 Jun 2024 10:53:49 +0000
Subject: [PATCH 3/6] Modify spelledTokenAt to support locations in middle of a
 spelled token

---
 .../clangd/unittests/XRefsTests.cpp              | 16 +++++++++-------
 clang/lib/Tooling/Syntax/Tokens.cpp              |  4 ++--
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/clang-tools-extra/clangd/unittests/XRefsTests.cpp b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
index c4def624a2fcc..cbceb9a343f87 100644
--- a/clang-tools-extra/clangd/unittests/XRefsTests.cpp
+++ b/clang-tools-extra/clangd/unittests/XRefsTests.cpp
@@ -2173,6 +2173,11 @@ TEST(FindReferences, WithinAST) {
         using $def[[MyTypeD^ef]] = int;
         enum MyEnum : $(MyEnum)[[MyTy^peDef]] { };
       )cpp",
+      // UDL
+      R"cpp(
+        bool $decl[[operator]]"" _u^dl(unsigned long long value);
+        bool x = $(x)[[1_udl]];
+      )cpp",
   };
   for (const char *Test : Tests)
     checkFindRefs(Test);
@@ -2361,9 +2366,8 @@ TEST(FindReferences, UsedSymbolsFromInclude) {
       )cpp",
 
       R"cpp(
-        [[#include ^"operator_qoutes.h"]]
-        using ::operator_qoutes::[[operator]]"" _b;
-        auto x = 1_b;
+        [[#include ^"udl_header.h"]]
+        auto x = [[1_b]];
       )cpp",
   };
   for (const char *Test : Tests) {
@@ -2382,10 +2386,8 @@ TEST(FindReferences, UsedSymbolsFromInclude) {
         class vector{};
       }
     )cpp");
-    TU.AdditionalFiles["operator_qoutes.h"] = guard(R"cpp(
-      namespace operator_qoutes {
-        bool operator"" _b(unsigned long long value);
-      }
+    TU.AdditionalFiles["udl_header.h"] = guard(R"cpp(
+      bool operator"" _b(unsigned long long value);
     )cpp");
     TU.ExtraArgs.push_back("-isystem" + testPath("system"));
 
diff --git a/clang/lib/Tooling/Syntax/Tokens.cpp b/clang/lib/Tooling/Syntax/Tokens.cpp
index 8d32c45a4a70c..f26e556d762c7 100644
--- a/clang/lib/Tooling/Syntax/Tokens.cpp
+++ b/clang/lib/Tooling/Syntax/Tokens.cpp
@@ -387,8 +387,8 @@ const syntax::Token *TokenBuffer::spelledTokenAt(SourceLocation Loc) const {
   assert(Loc.isFileID());
   const auto *Tok = llvm::partition_point(
       spelledTokens(SourceMgr->getFileID(Loc)),
-      [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
-  if (!Tok || Tok->location() != Loc)
+      [&](const syntax::Token &Tok) { return Tok.endLocation() <= Loc; });
+  if (!Tok || Tok->location() > Loc || Loc >= Tok->endLocation())
     return nullptr;
   return Tok;
 }

>From c8467210ad82399482da5c9294d302e3695de288 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Thu, 6 Jun 2024 15:33:36 +0000
Subject: [PATCH 4/6] Rename and add more tests

---
 clang-tools-extra/clangd/FindSymbols.cpp        |  2 +-
 clang-tools-extra/clangd/IncludeCleaner.cpp     |  2 +-
 .../clangd/SemanticHighlighting.cpp             |  4 ++--
 clang-tools-extra/clangd/XRefs.cpp              | 11 +++++------
 clang-tools-extra/clangd/refactor/Rename.cpp    |  2 +-
 .../clangd/unittests/PreambleTests.cpp          |  6 +++---
 clang/include/clang/Tooling/Syntax/Tokens.h     |  4 ++--
 clang/lib/Tooling/Syntax/Tokens.cpp             |  5 +++--
 clang/unittests/Tooling/Syntax/TokensTest.cpp   | 17 +++++++++++++++--
 9 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/clang-tools-extra/clangd/FindSymbols.cpp b/clang-tools-extra/clangd/FindSymbols.cpp
index 5244a4e893769..55f16b7085a6f 100644
--- a/clang-tools-extra/clangd/FindSymbols.cpp
+++ b/clang-tools-extra/clangd/FindSymbols.cpp
@@ -454,7 +454,7 @@ class DocumentOutline {
       if (!MacroName.isValid() || !MacroName.isFileID())
         continue;
       // All conditions satisfied, add the macro.
-      if (auto *Tok = AST.getTokens().spelledTokenAt(MacroName))
+      if (auto *Tok = AST.getTokens().spelledTokenContaining(MacroName))
         CurParent = &CurParent->inMacro(
             *Tok, SM, AST.getTokens().expansionStartingAt(Tok));
     }
diff --git a/clang-tools-extra/clangd/IncludeCleaner.cpp b/clang-tools-extra/clangd/IncludeCleaner.cpp
index 01b47679790f1..dc5b7ec95db5f 100644
--- a/clang-tools-extra/clangd/IncludeCleaner.cpp
+++ b/clang-tools-extra/clangd/IncludeCleaner.cpp
@@ -303,7 +303,7 @@ collectMacroReferences(ParsedAST &AST) {
   for (const auto &[_, Refs] : AST.getMacros().MacroRefs) {
     for (const auto &Ref : Refs) {
       auto Loc = SM.getComposedLoc(SM.getMainFileID(), Ref.StartOffset);
-      const auto *Tok = AST.getTokens().spelledTokenAt(Loc);
+      const auto *Tok = AST.getTokens().spelledTokenContaining(Loc);
       if (!Tok)
         continue;
       auto Macro = locateMacroAt(*Tok, PP);
diff --git a/clang-tools-extra/clangd/SemanticHighlighting.cpp b/clang-tools-extra/clangd/SemanticHighlighting.cpp
index eb025f21f3616..81327c4090cc4 100644
--- a/clang-tools-extra/clangd/SemanticHighlighting.cpp
+++ b/clang-tools-extra/clangd/SemanticHighlighting.cpp
@@ -447,7 +447,7 @@ class HighlightingsBuilder {
     if (!RLoc.isValid())
       return;
 
-    const auto *RTok = TB.spelledTokenAt(RLoc);
+    const auto *RTok = TB.spelledTokenContaining(RLoc);
     // Handle `>>`. RLoc is always pointing at the right location, just change
     // the end to be offset by 1.
     // We'll either point at the beginning of `>>`, hence get a proper spelled
@@ -577,7 +577,7 @@ class HighlightingsBuilder {
       return std::nullopt;
     // We might have offsets in the main file that don't correspond to any
     // spelled tokens.
-    const auto *Tok = TB.spelledTokenAt(Loc);
+    const auto *Tok = TB.spelledTokenContaining(Loc);
     if (!Tok)
       return std::nullopt;
     return halfOpenToRange(SourceMgr,
diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 033c75f02d561..f94cadeffaa29 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -844,7 +844,7 @@ std::vector<DocumentLink> getDocumentLinks(ParsedAST &AST) {
     if (Inc.Resolved.empty())
       continue;
     auto HashLoc = SM.getComposedLoc(SM.getMainFileID(), Inc.HashOffset);
-    const auto *HashTok = AST.getTokens().spelledTokenAt(HashLoc);
+    const auto *HashTok = AST.getTokens().spelledTokenContaining(HashLoc);
     assert(HashTok && "got inclusion at wrong offset");
     const auto *IncludeTok = std::next(HashTok);
     const auto *FileTok = std::next(IncludeTok);
@@ -938,7 +938,7 @@ class ReferenceFinder : public index::IndexDataConsumer {
     CollectorOpts.CollectMainFileSymbols = true;
     for (SourceLocation L : Locs) {
       L = SM.getFileLoc(L);
-      if (const auto *Tok = TB.spelledTokenAt(L))
+      if (const auto *Tok = TB.spelledTokenContaining(L))
         References.push_back(
             {*Tok, Roles,
              SymbolCollector::getRefContainer(ASTNode.Parent, CollectorOpts)});
@@ -1216,7 +1216,7 @@ DocumentHighlight toHighlight(const ReferenceFinder::Reference &Ref,
 std::optional<DocumentHighlight> toHighlight(SourceLocation Loc,
                                              const syntax::TokenBuffer &TB) {
   Loc = TB.sourceManager().getFileLoc(Loc);
-  if (const auto *Tok = TB.spelledTokenAt(Loc)) {
+  if (const auto *Tok = TB.spelledTokenContaining(Loc)) {
     DocumentHighlight Result;
     Result.range = halfOpenToRange(
         TB.sourceManager(),
@@ -1353,9 +1353,8 @@ maybeFindIncludeReferences(ParsedAST &AST, Position Pos,
           Loc = SM.getIncludeLoc(SM.getFileID(Loc));
 
         ReferencesResult::Reference Result;
-        const auto *Token = AST.getTokens().spelledTokenAt(Loc);
-        if (!Token)
-          return;
+        const auto *Token = AST.getTokens().spelledTokenContaining(Loc);
+        assert(Token && "references expected token here");
         Result.Loc.range = Range{sourceLocToPosition(SM, Token->location()),
                                  sourceLocToPosition(SM, Token->endLocation())};
         Result.Loc.uri = URIMainFile;
diff --git a/clang-tools-extra/clangd/refactor/Rename.cpp b/clang-tools-extra/clangd/refactor/Rename.cpp
index c0fc4453a3fcc..c85e13dbdfe97 100644
--- a/clang-tools-extra/clangd/refactor/Rename.cpp
+++ b/clang-tools-extra/clangd/refactor/Rename.cpp
@@ -748,7 +748,7 @@ std::vector<SymbolRange> collectRenameIdentifierRanges(
 clangd::Range tokenRangeForLoc(ParsedAST &AST, SourceLocation TokLoc,
                                const SourceManager &SM,
                                const LangOptions &LangOpts) {
-  const auto *Token = AST.getTokens().spelledTokenAt(TokLoc);
+  const auto *Token = AST.getTokens().spelledTokenContaining(TokLoc);
   assert(Token && "rename expects spelled tokens");
   clangd::Range Result;
   Result.start = sourceLocToPosition(SM, Token->location());
diff --git a/clang-tools-extra/clangd/unittests/PreambleTests.cpp b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
index 6420516e78557..240d9bfc7254c 100644
--- a/clang-tools-extra/clangd/unittests/PreambleTests.cpp
+++ b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
@@ -417,7 +417,7 @@ TEST(PreamblePatchTest, LocateMacroAtWorks) {
     ASSERT_TRUE(AST);
 
     const auto &SM = AST->getSourceManager();
-    auto *MacroTok = AST->getTokens().spelledTokenAt(
+    auto *MacroTok = AST->getTokens().spelledTokenContaining(
         SM.getComposedLoc(SM.getMainFileID(), Modified.point("use")));
     ASSERT_TRUE(MacroTok);
 
@@ -441,7 +441,7 @@ TEST(PreamblePatchTest, LocateMacroAtDeletion) {
     ASSERT_TRUE(AST);
 
     const auto &SM = AST->getSourceManager();
-    auto *MacroTok = AST->getTokens().spelledTokenAt(
+    auto *MacroTok = AST->getTokens().spelledTokenContaining(
         SM.getComposedLoc(SM.getMainFileID(), Modified.point()));
     ASSERT_TRUE(MacroTok);
 
@@ -512,7 +512,7 @@ TEST(PreamblePatchTest, RefsToMacros) {
       ExpectedLocations.push_back(referenceRangeIs(R));
 
     for (const auto &P : Modified.points()) {
-      auto *MacroTok = AST->getTokens().spelledTokenAt(SM.getComposedLoc(
+      auto *MacroTok = AST->getTokens().spelledTokenContaining(SM.getComposedLoc(
           SM.getMainFileID(),
           llvm::cantFail(positionToOffset(Modified.code(), P))));
       ASSERT_TRUE(MacroTok);
diff --git a/clang/include/clang/Tooling/Syntax/Tokens.h b/clang/include/clang/Tooling/Syntax/Tokens.h
index b1bdefed7c97f..f71b8d67bfea4 100644
--- a/clang/include/clang/Tooling/Syntax/Tokens.h
+++ b/clang/include/clang/Tooling/Syntax/Tokens.h
@@ -292,9 +292,9 @@ class TokenBuffer {
   ///     "DECL", "(", "a", ")", ";"}
   llvm::ArrayRef<syntax::Token> spelledTokens(FileID FID) const;
 
-  /// Returns the spelled Token starting at Loc, if there are no such tokens
+  /// Returns the spelled Token containing the Loc, if there are no such tokens
   /// returns nullptr.
-  const syntax::Token *spelledTokenAt(SourceLocation Loc) const;
+  const syntax::Token *spelledTokenContaining(SourceLocation Loc) const;
 
   /// Get all tokens that expand a macro in \p FID. For the following input
   ///     #define FOO B
diff --git a/clang/lib/Tooling/Syntax/Tokens.cpp b/clang/lib/Tooling/Syntax/Tokens.cpp
index f26e556d762c7..0a656dff38421 100644
--- a/clang/lib/Tooling/Syntax/Tokens.cpp
+++ b/clang/lib/Tooling/Syntax/Tokens.cpp
@@ -383,12 +383,13 @@ llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens(FileID FID) const {
   return It->second.SpelledTokens;
 }
 
-const syntax::Token *TokenBuffer::spelledTokenAt(SourceLocation Loc) const {
+const syntax::Token *
+TokenBuffer::spelledTokenContaining(SourceLocation Loc) const {
   assert(Loc.isFileID());
   const auto *Tok = llvm::partition_point(
       spelledTokens(SourceMgr->getFileID(Loc)),
       [&](const syntax::Token &Tok) { return Tok.endLocation() <= Loc; });
-  if (!Tok || Tok->location() > Loc || Loc >= Tok->endLocation())
+  if (!Tok || Loc < Tok->location())
     return nullptr;
   return Tok;
 }
diff --git a/clang/unittests/Tooling/Syntax/TokensTest.cpp b/clang/unittests/Tooling/Syntax/TokensTest.cpp
index 42f5169713965..dc8a11dbc345c 100644
--- a/clang/unittests/Tooling/Syntax/TokensTest.cpp
+++ b/clang/unittests/Tooling/Syntax/TokensTest.cpp
@@ -374,11 +374,24 @@ TEST_F(TokenCollectorTest, Locations) {
 
   auto StartLoc = SourceMgr->getLocForStartOfFile(SourceMgr->getMainFileID());
   for (auto &R : Code.ranges()) {
-    EXPECT_THAT(Buffer.spelledTokenAt(StartLoc.getLocWithOffset(R.Begin)),
-                Pointee(RangeIs(R)));
+    EXPECT_THAT(
+        Buffer.spelledTokenContaining(StartLoc.getLocWithOffset(R.Begin)),
+        Pointee(RangeIs(R)));
   }
 }
 
+TEST_F(TokenCollectorTest, LocationInMiddleOfSpelledToken) {
+  llvm::Annotations Code(R"cpp(
+    int foo = [[baa^aar]];
+  )cpp");
+  recordTokens(Code.code());
+  // Check spelled tokens.
+  auto StartLoc = SourceMgr->getLocForStartOfFile(SourceMgr->getMainFileID());
+  EXPECT_THAT(
+      Buffer.spelledTokenContaining(StartLoc.getLocWithOffset(Code.point())),
+      Pointee(RangeIs(Code.range())));
+}
+
 TEST_F(TokenCollectorTest, MacroDirectives) {
   // Macro directives are not stored anywhere at the moment.
   std::string Code = R"cpp(

>From 9f957dafddd275e4579389a72aa50f2dbf2994f4 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Thu, 6 Jun 2024 15:39:23 +0000
Subject: [PATCH 5/6] format

---
 clang-tools-extra/clangd/unittests/PreambleTests.cpp | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/clang-tools-extra/clangd/unittests/PreambleTests.cpp b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
index 240d9bfc7254c..16a2f9448b1ec 100644
--- a/clang-tools-extra/clangd/unittests/PreambleTests.cpp
+++ b/clang-tools-extra/clangd/unittests/PreambleTests.cpp
@@ -512,9 +512,10 @@ TEST(PreamblePatchTest, RefsToMacros) {
       ExpectedLocations.push_back(referenceRangeIs(R));
 
     for (const auto &P : Modified.points()) {
-      auto *MacroTok = AST->getTokens().spelledTokenContaining(SM.getComposedLoc(
-          SM.getMainFileID(),
-          llvm::cantFail(positionToOffset(Modified.code(), P))));
+      auto *MacroTok =
+          AST->getTokens().spelledTokenContaining(SM.getComposedLoc(
+              SM.getMainFileID(),
+              llvm::cantFail(positionToOffset(Modified.code(), P))));
       ASSERT_TRUE(MacroTok);
       EXPECT_THAT(findReferences(*AST, P, 0).References,
                   testing::ElementsAreArray(ExpectedLocations));

>From 7e316eb4b5b632c05bf2d9b1cb7c8d30f73d7bf2 Mon Sep 17 00:00:00 2001
From: Utkarsh Saxena <usx at google.com>
Date: Fri, 7 Jun 2024 09:06:42 +0000
Subject: [PATCH 6/6] Fix handle angleBracketTokens

---
 clang-tools-extra/clangd/SemanticHighlighting.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/clang-tools-extra/clangd/SemanticHighlighting.cpp b/clang-tools-extra/clangd/SemanticHighlighting.cpp
index 81327c4090cc4..a366f1331c2d3 100644
--- a/clang-tools-extra/clangd/SemanticHighlighting.cpp
+++ b/clang-tools-extra/clangd/SemanticHighlighting.cpp
@@ -448,10 +448,9 @@ class HighlightingsBuilder {
       return;
 
     const auto *RTok = TB.spelledTokenContaining(RLoc);
-    // Handle `>>`. RLoc is always pointing at the right location, just change
-    // the end to be offset by 1.
-    // We'll either point at the beginning of `>>`, hence get a proper spelled
-    // or point in the middle of `>>` hence get no spelled tok.
+    // Handle `>>`. RLoc is either part of `>>` or a spelled token on its own
+    // `>`. If it's the former, slice to have length of 1, if latter use the
+    // token as-is.
     if (!RTok || RTok->kind() == tok::greatergreater) {
       Position Begin = sourceLocToPosition(SourceMgr, RLoc);
       Position End = sourceLocToPosition(SourceMgr, RLoc.getLocWithOffset(1));



More information about the cfe-commits mailing list