1 files changed, 32 insertions, 38 deletions

diff --git a/clang-tools-extra/clangd/XRefs.cpp b/clang-tools-extra/clangd/XRefs.cpp
index 2dc3ac03095..687652ca78f 100644
--- a/clang-tools-extra/clangd/XRefs.cpp
+++ b/clang-tools-extra/clangd/XRefs.cpp
@@ -260,14 +260,6 @@ IdentifiedSymbol getSymbolAtPosition(ParsedAST &AST, SourceLocation Pos) {
   return {DeclMacrosFinder.getFoundDecls(), DeclMacrosFinder.takeMacroInfos()};
 }
 
-Range getTokenRange(ASTContext &AST, SourceLocation TokLoc) {
-  const SourceManager &SourceMgr = AST.getSourceManager();
-  SourceLocation LocEnd =
-      Lexer::getLocForEndOfToken(TokLoc, 0, SourceMgr, AST.getLangOpts());
-  return {sourceLocToPosition(SourceMgr, TokLoc),
-          sourceLocToPosition(SourceMgr, LocEnd)};
-}
-
 llvm::Optional<Location> makeLocation(ASTContext &AST, SourceLocation TokLoc,
                                       llvm::StringRef TUPath) {
   const SourceManager &SourceMgr = AST.getSourceManager();
@@ -279,10 +271,14 @@ llvm::Optional<Location> makeLocation(ASTContext &AST, SourceLocation TokLoc,
     log("failed to get path!");
     return None;
   }
-  Location L;
-  L.uri = URIForFile::canonicalize(*FilePath, TUPath);
-  L.range = getTokenRange(AST, TokLoc);
-  return L;
+  if (auto Range =
+          getTokenRange(AST.getSourceManager(), AST.getLangOpts(), TokLoc)) {
+    Location L;
+    L.uri = URIForFile::canonicalize(*FilePath, TUPath);
+    L.range = *Range;
+    return L;
+  }
+  return None;
 }
 
 } // namespace
@@ -471,15 +467,19 @@ std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
 
   std::vector<DocumentHighlight> Result;
   for (const auto &Ref : References) {
-    DocumentHighlight DH;
-    DH.range = getTokenRange(AST.getASTContext(), Ref.Loc);
-    if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Write))
-      DH.kind = DocumentHighlightKind::Write;
-    else if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Read))
-      DH.kind = DocumentHighlightKind::Read;
-    else
-      DH.kind = DocumentHighlightKind::Text;
-    Result.push_back(std::move(DH));
+    if (auto Range =
+            getTokenRange(AST.getASTContext().getSourceManager(),
+                          AST.getASTContext().getLangOpts(), Ref.Loc)) {
+      DocumentHighlight DH;
+      DH.range = *Range;
+      if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Write))
+        DH.kind = DocumentHighlightKind::Write;
+      else if (Ref.Role & index::SymbolRoleSet(index::SymbolRole::Read))
+        DH.kind = DocumentHighlightKind::Read;
+      else
+        DH.kind = DocumentHighlightKind::Text;
+      Result.push_back(std::move(DH));
+    }
   }
   return Result;
 }
@@ -610,18 +610,6 @@ fetchTemplateParameters(const TemplateParameterList *Params,
   return TempParameters;
 }
 
-static llvm::Optional<Range> getTokenRange(SourceLocation Loc,
-                                           const ASTContext &Ctx) {
-  if (!Loc.isValid())
-    return llvm::None;
-  SourceLocation End = Lexer::getLocForEndOfToken(
-      Loc, 0, Ctx.getSourceManager(), Ctx.getLangOpts());
-  if (!End.isValid())
-    return llvm::None;
-  return halfOpenToRange(Ctx.getSourceManager(),
-                         CharSourceRange::getCharRange(Loc, End));
-}
-
 static const FunctionDecl *getUnderlyingFunction(const Decl *D) {
   // Extract lambda from variables.
   if (const VarDecl *VD = llvm::dyn_cast<VarDecl>(D)) {
@@ -910,7 +898,9 @@ llvm::Optional<HoverInfo> getHover(ParsedAST &AST, Position Pos,
           tooling::applyAllReplacements(HI->Definition, Replacements))
     HI->Definition = *Formatted;
 
-  HI->SymRange = getTokenRange(SourceLocationBeg, AST.getASTContext());
+  HI->SymRange =
+      getTokenRange(AST.getASTContext().getSourceManager(),
+                    AST.getASTContext().getLangOpts(), SourceLocationBeg);
   return HI;
 }
 
@@ -933,10 +923,14 @@ std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
   // TODO: should we handle macros, too?
   auto MainFileRefs = findRefs(Symbols.Decls, AST);
   for (const auto &Ref : MainFileRefs) {
-    Location Result;
-    Result.range = getTokenRange(AST.getASTContext(), Ref.Loc);
-    Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
-    Results.push_back(std::move(Result));
+    if (auto Range =
+            getTokenRange(AST.getASTContext().getSourceManager(),
+                          AST.getASTContext().getLangOpts(), Ref.Loc)) {
+      Location Result;
+      Result.range = *Range;
+      Result.uri = URIForFile::canonicalize(*MainFilePath, *MainFilePath);
+      Results.push_back(std::move(Result));
+    }
   }
 
   // Now query the index for references from other files.