[clangd] Migrate Lexer usages in TypeHierarchy to TokenBuffers

[kadircet]

Summary

Also fixes a bug, resulting from directly using ND.getEndLoc() for end
location of the range. As ND.getEndLoc() points to the begining of the last
token, whereas it should point one past the end, since LSP ranges are half open
(exclusive on the end).

[sammccall]

clang-tools-extra/clangd/XRefs.cpp

613

I think you're trying to test if NameTok has the location you're searching for, but you'd need to check the partition point for equality.

Couldn't you just do NameTokens = spelledForExpanded(expandedTokens(NameLoc))?

[kadircet]

• Address comments

clang-tools-extra/clangd/XRefs.cpp

 #include "clang/Basic/LLVM.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Index/IndexDataConsumer.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Index/USRGeneration.h"
+#include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
llvm::raw_ostream &operator<<(llvm::raw_ostream &OS, const LocatedSymbol &S) {
   OS << S.Name << ": " << S.PreferredDeclaration;
   if (S.Definition)
     OS << " def=" << *S.Definition;
   return OS;
 }
 
 // FIXME(nridge): Reduce duplication between this function and declToSym().
 static llvm::Optional<TypeHierarchyItem>
-declToTypeHierarchyItem(ASTContext &Ctx, const NamedDecl &ND) {
+declToTypeHierarchyItem(ASTContext &Ctx, const NamedDecl &ND,
+                        const syntax::TokenBuffer &TB) {
   auto &SM = Ctx.getSourceManager();
-
   SourceLocation NameLoc = nameLocation(ND, Ctx.getSourceManager());
-  // getFileLoc is a good choice for us, but we also need to make sure
-  // sourceLocToPosition won't switch files, so we call getSpellingLoc on top of
-  // that to make sure it does not switch files.
-  // FIXME: sourceLocToPosition should not switch files!
-  SourceLocation BeginLoc = SM.getSpellingLoc(SM.getFileLoc(ND.getBeginLoc()));
-  SourceLocation EndLoc = SM.getSpellingLoc(SM.getFileLoc(ND.getEndLoc()));
-  if (NameLoc.isInvalid() || BeginLoc.isInvalid() || EndLoc.isInvalid())
-    return llvm::None;
+  auto FilePath =
+      getCanonicalPath(SM.getFileEntryForID(SM.getFileID(NameLoc)), SM);
+  auto TUPath = getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
+  if (!FilePath || !TUPath)
+    return llvm::None; // Not useful without a uri.
 
-  Position NameBegin = sourceLocToPosition(SM, NameLoc);
-  Position NameEnd = sourceLocToPosition(
-      SM, Lexer::getLocForEndOfToken(NameLoc, 0, SM, Ctx.getLangOpts()));
+  auto Tokens = TB.spelledForExpanded(TB.expandedTokens(ND.getSourceRange()));
+  if (!Tokens || Tokens->empty())
+    return llvm::None;
+  const auto *NameTok =
+      llvm::partition_point(*Tokens, [&SM, &NameLoc](const syntax::Token &T) {
+        return SM.isBeforeInTranslationUnit(T.location(), NameLoc);
+      });
+  if (!NameTok)
+    return llvm::None;
 
   index::SymbolInfo SymInfo = index::getSymbolInfo(&ND);
   // FIXME: this is not classifying constructors, destructors and operators
   //        correctly (they're all "methods").
   SymbolKind SK = indexSymbolKindToSymbolKind(SymInfo.Kind);
 
   TypeHierarchyItem THI;
   THI.name = printName(Ctx, ND);
   THI.kind = SK;
   THI.deprecated = ND.isDeprecated();
-  THI.range =
-      Range{sourceLocToPosition(SM, BeginLoc), sourceLocToPosition(SM, EndLoc)};
-  THI.selectionRange = Range{NameBegin, NameEnd};
+  THI.range = halfOpenToRange(
+      SM, syntax::Token::range(SM, Tokens->front(), Tokens->back())
+              .toCharRange(SM));
+  THI.selectionRange = halfOpenToRange(SM, NameTok->range(SM).toCharRange(SM));
   if (!THI.range.contains(THI.selectionRange)) {
     // 'selectionRange' must be contained in 'range', so in cases where clang
     // reports unrelated ranges we need to reconcile somehow.
     THI.range = THI.selectionRange;
   }
 
-  auto FilePath =
-      getCanonicalPath(SM.getFileEntryForID(SM.getFileID(BeginLoc)), SM);
-  auto TUPath = getCanonicalPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
-  if (!FilePath || !TUPath)
-    return llvm::None; // Not useful without a uri.
   THI.uri = URIForFile::canonicalize(*FilePath, *TUPath);
 
   return THI;
 }
 
 static Optional<TypeHierarchyItem>
 symbolToTypeHierarchyItem(const Symbol &S, const SymbolIndex *Index,
                           PathRef TUPath) {
Index->relations(Req, [&](const SymbolID &Subject, const Symbol &Object) {
     }
   });
 }
 
 using RecursionProtectionSet = llvm::SmallSet<const CXXRecordDecl *, 4>;
 
 static void fillSuperTypes(const CXXRecordDecl &CXXRD, ASTContext &ASTCtx,
                            std::vector<TypeHierarchyItem> &SuperTypes,
-                           RecursionProtectionSet &RPSet) {
+                           RecursionProtectionSet &RPSet,
+                           const syntax::TokenBuffer &TB) {
   // typeParents() will replace dependent template specializations
   // with their class template, so to avoid infinite recursion for
   // certain types of hierarchies, keep the templates encountered
   // along the parent chain in a set, and stop the recursion if one
   // starts to repeat.
   auto *Pattern = CXXRD.getDescribedTemplate() ? &CXXRD : nullptr;
   if (Pattern) {
     if (!RPSet.insert(Pattern).second) {
       return;
     }
   }
 
   for (const CXXRecordDecl *ParentDecl : typeParents(&CXXRD)) {
     if (Optional<TypeHierarchyItem> ParentSym =
-            declToTypeHierarchyItem(ASTCtx, *ParentDecl)) {
+            declToTypeHierarchyItem(ASTCtx, *ParentDecl, TB)) {
       ParentSym->parents.emplace();
-      fillSuperTypes(*ParentDecl, ASTCtx, *ParentSym->parents, RPSet);
+      fillSuperTypes(*ParentDecl, ASTCtx, *ParentSym->parents, RPSet, TB);
       SuperTypes.emplace_back(std::move(*ParentSym));
     }
   }
 
   if (Pattern) {
     RPSet.erase(Pattern);
   }
 }

 getTypeHierarchy(ParsedAST &AST, Position Pos, int ResolveLevels,
                  TypeHierarchyDirection Direction, const SymbolIndex *Index,
                  PathRef TUPath) {
   const CXXRecordDecl *CXXRD = findRecordTypeAt(AST, Pos);
   if (!CXXRD)
     return llvm::None;
 
   Optional<TypeHierarchyItem> Result =
-      declToTypeHierarchyItem(AST.getASTContext(), *CXXRD);
+      declToTypeHierarchyItem(AST.getASTContext(), *CXXRD, AST.getTokens());
   if (!Result)
     return Result;
 
   if (Direction == TypeHierarchyDirection::Parents ||
       Direction == TypeHierarchyDirection::Both) {
     Result->parents.emplace();
 
     RecursionProtectionSet RPSet;
-    fillSuperTypes(*CXXRD, AST.getASTContext(), *Result->parents, RPSet);
+    fillSuperTypes(*CXXRD, AST.getASTContext(), *Result->parents, RPSet,
+                   AST.getTokens());
   }
 
   if ((Direction == TypeHierarchyDirection::Children ||
        Direction == TypeHierarchyDirection::Both) &&
       ResolveLevels > 0) {
     Result->children.emplace();
 
     if (Index) {

clang-tools-extra/clangd/test/type-hierarchy.test

 # CHECK-NEXT:        "name": "Child1",
 # CHECK-NEXT:        "parents": [
 # CHECK-NEXT:          {
 # CHECK-NEXT:            "kind": 23,
 # CHECK-NEXT:            "name": "Parent",
 # CHECK-NEXT:            "parents": [],
 # CHECK-NEXT:            "range": {
 # CHECK-NEXT:              "end": {
-# CHECK-NEXT:                "character": 15,
+# CHECK-NEXT:                "character": 16,
 # CHECK-NEXT:                "line": 0
 # CHECK-NEXT:              },
 # CHECK-NEXT:              "start": {
 # CHECK-NEXT:                "character": 0,
 # CHECK-NEXT:                "line": 0
 # CHECK-NEXT:              }
 # CHECK-NEXT:            },
 # CHECK-NEXT:            "selectionRange": {
 # CHECK-NEXT:              "end": {
 # CHECK-NEXT:                "character": 13,
 # CHECK-NEXT:                "line": 0
 # CHECK-NEXT:              },
 # CHECK-NEXT:              "start": {
 # CHECK-NEXT:                "character": 7,
 # CHECK-NEXT:                "line": 0
 # CHECK-NEXT:              }
 # CHECK-NEXT:            },
 # CHECK-NEXT:            "uri": "file://{{.*}}/clangd-test/main.cpp"
 # CHECK-NEXT:          }
 # CHECK-NEXT:        ],
 # CHECK-NEXT:        "range": {
 # CHECK-NEXT:          "end": {
-# CHECK-NEXT:            "character": 24,
+# CHECK-NEXT:            "character": 25,
 # CHECK-NEXT:            "line": 1
 # CHECK-NEXT:          },
 # CHECK-NEXT:          "start": {
 # CHECK-NEXT:            "character": 0,
 # CHECK-NEXT:            "line": 1
 # CHECK-NEXT:          }
 # CHECK-NEXT:        },
 # CHECK-NEXT:        "selectionRange": {
 # CHECK-NEXT:          "end": {
 # CHECK-NEXT:            "character": 13,
 # CHECK-NEXT:            "line": 1
 # CHECK-NEXT:          },
 # CHECK-NEXT:          "start": {
 # CHECK-NEXT:            "character": 7,
 # CHECK-NEXT:            "line": 1
 # CHECK-NEXT:          }
 # CHECK-NEXT:        },
 # CHECK-NEXT:        "uri": "file://{{.*}}/clangd-test/main.cpp"
 # CHECK-NEXT:      }
 # CHECK-NEXT:    ],
 # CHECK-NEXT:    "range": {
 # CHECK-NEXT:      "end": {
-# CHECK-NEXT:        "character": 24,
+# CHECK-NEXT:        "character": 25,
 # CHECK-NEXT:        "line": 2
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "start": {
 # CHECK-NEXT:        "character": 0,
 # CHECK-NEXT:        "line": 2
 # CHECK-NEXT:      }
 # CHECK-NEXT:    },
 # CHECK-NEXT:    "selectionRange": {