Index: clangd/XRefs.h
===================================================================
--- clangd/XRefs.h
+++ clangd/XRefs.h
@@ -34,6 +34,10 @@
 /// Get the hover information when hovering at \p Pos.
 llvm::Optional<Hover> getHover(ParsedAST &AST, Position Pos);
 
+/// Returns reference locations of the symbol at a specified \p Pos.
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     const SymbolIndex *Index = nullptr);
+
 } // namespace clangd
 } // namespace clang
 
Index: clangd/XRefs.cpp
===================================================================
--- clangd/XRefs.cpp
+++ clangd/XRefs.cpp
@@ -174,30 +174,27 @@
   return {DeclMacrosFinder.takeDecls(), DeclMacrosFinder.takeMacroInfos()};
 }
 
-llvm::Optional<Location>
-makeLocation(ParsedAST &AST, const SourceRange &ValSourceRange) {
+Range getTokenRange(ParsedAST &AST, SourceLocation TokLoc) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-  const LangOptions &LangOpts = AST.getASTContext().getLangOpts();
-  SourceLocation LocStart = ValSourceRange.getBegin();
+  SourceLocation LocEnd = Lexer::getLocForEndOfToken(
+      TokLoc, 0, SourceMgr, AST.getASTContext().getLangOpts());
+  return {sourceLocToPosition(SourceMgr, TokLoc),
+          sourceLocToPosition(SourceMgr, LocEnd)};
+}
 
-  const FileEntry *F =
-      SourceMgr.getFileEntryForID(SourceMgr.getFileID(LocStart));
+llvm::Optional<Location> makeLocation(ParsedAST &AST, SourceLocation TokLoc) {
+  const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+  const FileEntry *F = SourceMgr.getFileEntryForID(SourceMgr.getFileID(TokLoc));
   if (!F)
     return llvm::None;
-  SourceLocation LocEnd = Lexer::getLocForEndOfToken(ValSourceRange.getEnd(), 0,
-                                                     SourceMgr, LangOpts);
-  Position Begin = sourceLocToPosition(SourceMgr, LocStart);
-  Position End = sourceLocToPosition(SourceMgr, LocEnd);
-  Range R = {Begin, End};
-  Location L;
-
   auto FilePath = getRealPath(F, SourceMgr);
   if (!FilePath) {
     log("failed to get path!");
     return llvm::None;
   }
+  Location L;
   L.uri = URIForFile(*FilePath);
-  L.range = R;
+  L.range = getTokenRange(AST, TokLoc);
   return L;
 }
 
@@ -223,7 +220,7 @@
 
   for (auto Item : Symbols.Macros) {
     auto Loc = Item.Info->getDefinitionLoc();
-    auto L = makeLocation(AST, SourceRange(Loc, Loc));
+    auto L = makeLocation(AST, Loc);
     if (L)
       Result.push_back(*L);
   }
@@ -266,7 +263,7 @@
 
     auto &Candidate = ResultCandidates[Key];
     auto Loc = findNameLoc(D);
-    auto L = makeLocation(AST, SourceRange(Loc, Loc));
+    auto L = makeLocation(AST, Loc);
     // The declaration in the identified symbols is a definition if possible
     // otherwise it is declaration.
     bool IsDef = getDefinition(D) == D;
@@ -316,24 +313,36 @@
 
 namespace {
 
-/// Finds document highlights that a given list of declarations refers to.
-class DocumentHighlightsFinder : public index::IndexDataConsumer {
-  std::vector<const Decl *> &Decls;
-  std::vector<DocumentHighlight> DocumentHighlights;
-  const ASTContext &AST;
-
+/// Collects references to symbols within the main file.
+class ReferenceFinder : public index::IndexDataConsumer {
 public:
-  DocumentHighlightsFinder(ASTContext &AST, Preprocessor &PP,
-                           std::vector<const Decl *> &Decls)
-      : Decls(Decls), AST(AST) {}
-  std::vector<DocumentHighlight> takeHighlights() {
-    // Don't keep the same highlight multiple times.
-    // This can happen when nodes in the AST are visited twice.
-    std::sort(DocumentHighlights.begin(), DocumentHighlights.end());
-    auto Last =
-        std::unique(DocumentHighlights.begin(), DocumentHighlights.end());
-    DocumentHighlights.erase(Last, DocumentHighlights.end());
-    return std::move(DocumentHighlights);
+  struct Reference {
+    const Decl *Target;
+    SourceLocation Loc;
+    index::SymbolRoleSet Role;
+  };
+
+  ReferenceFinder(ASTContext &AST, Preprocessor &PP,
+                  const std::vector<const Decl *> &TargetDecls)
+      : AST(AST) {
+    for (const Decl *D : TargetDecls)
+      Targets.insert(D);
+  }
+
+  std::vector<Reference> take() && {
+    std::sort(References.begin(), References.end(),
+              [](const Reference &L, const Reference &R) {
+                return std::tie(L.Loc, L.Target, L.Role) <
+                       std::tie(R.Loc, R.Target, R.Role);
+              });
+    // We sometimes see duplicates when parts of the AST get traversed twice.
+    References.erase(std::unique(References.begin(), References.end(),
+                                 [](const Reference &L, const Reference &R) {
+                                   return std::tie(L.Target, L.Loc, L.Role) ==
+                                          std::tie(R.Target, R.Loc, R.Role);
+                                 }),
+                     References.end());
+    return std::move(References);
   }
 
   bool
@@ -341,63 +350,53 @@
                       ArrayRef<index::SymbolRelation> Relations,
                       SourceLocation Loc,
                       index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
-    const SourceManager &SourceMgr = AST.getSourceManager();
-    SourceLocation HighlightStartLoc = SourceMgr.getFileLoc(Loc);
-    if (SourceMgr.getMainFileID() != SourceMgr.getFileID(HighlightStartLoc) ||
-        std::find(Decls.begin(), Decls.end(), D) == Decls.end()) {
-      return true;
-    }
-    SourceLocation End;
-    const LangOptions &LangOpts = AST.getLangOpts();
-    End = Lexer::getLocForEndOfToken(HighlightStartLoc, 0, SourceMgr, LangOpts);
-    SourceRange SR(HighlightStartLoc, End);
-
-    DocumentHighlightKind Kind = DocumentHighlightKind::Text;
-    if (static_cast<index::SymbolRoleSet>(index::SymbolRole::Write) & Roles)
-      Kind = DocumentHighlightKind::Write;
-    else if (static_cast<index::SymbolRoleSet>(index::SymbolRole::Read) & Roles)
-      Kind = DocumentHighlightKind::Read;
-
-    DocumentHighlights.push_back(getDocumentHighlight(SR, Kind));
+    const SourceManager &SM = AST.getSourceManager();
+    Loc = SM.getFileLoc(Loc);
+    if (SM.isWrittenInMainFile(Loc) && Targets.count(D))
+      References.push_back({D, Loc, Roles});
     return true;
   }
 
 private:
-  DocumentHighlight getDocumentHighlight(SourceRange SR,
-                                         DocumentHighlightKind Kind) {
-    const SourceManager &SourceMgr = AST.getSourceManager();
-    Position Begin = sourceLocToPosition(SourceMgr, SR.getBegin());
-    Position End = sourceLocToPosition(SourceMgr, SR.getEnd());
-    Range R = {Begin, End};
-    DocumentHighlight DH;
-    DH.range = R;
-    DH.kind = Kind;
-    return DH;
-  }
+  llvm::SmallSet<const Decl *, 4> Targets;
+  std::vector<Reference> References;
+  const ASTContext &AST;
 };
 
-} // namespace
-
-std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
-                                                      Position Pos) {
-  const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
-  SourceLocation SourceLocationBeg =
-      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
-
-  auto Symbols = getSymbolAtPosition(AST, SourceLocationBeg);
-  std::vector<const Decl *> SelectedDecls = Symbols.Decls;
-
-  DocumentHighlightsFinder DocHighlightsFinder(
-      AST.getASTContext(), AST.getPreprocessor(), SelectedDecls);
-
+std::vector<ReferenceFinder::Reference>
+findRefs(const std::vector<const Decl *> &Decls, ParsedAST &AST) {
+  ReferenceFinder RefFinder(AST.getASTContext(), AST.getPreprocessor(), Decls);
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   IndexOpts.IndexFunctionLocals = true;
   indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
-                     DocHighlightsFinder, IndexOpts);
+                     RefFinder, IndexOpts);
+  return std::move(RefFinder).take();
+}
+
+} // namespace
+
+std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
+                                                      Position Pos) {
+  const SourceManager &SM = AST.getASTContext().getSourceManager();
+  auto Symbols = getSymbolAtPosition(
+      AST, getBeginningOfIdentifier(AST, Pos, SM.getMainFileID()));
+  auto References = findRefs(Symbols.Decls, AST);
 
-  return DocHighlightsFinder.takeHighlights();
+  std::vector<DocumentHighlight> Result;
+  for (const auto &Ref : References) {
+    DocumentHighlight DH;
+    DH.range = getTokenRange(AST, 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));
+  }
+  return Result;
 }
 
 static PrintingPolicy printingPolicyForDecls(PrintingPolicy Base) {
@@ -659,5 +658,51 @@
   return None;
 }
 
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     const SymbolIndex *Index) {
+  std::vector<Location> Results;
+  const SourceManager &SM = AST.getASTContext().getSourceManager();
+  auto MainFilePath = getRealPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
+  if (!MainFilePath) {
+    elog("Failed to get a path for the main file, so no references");
+    return Results;
+  }
+  auto Loc = getBeginningOfIdentifier(AST, Pos, SM.getMainFileID());
+  auto Symbols = getSymbolAtPosition(AST, Loc);
+
+  // We traverse the AST to find references in the main file.
+  // TODO: should we handle macros, too?
+  auto MainFileRefs = findRefs(Symbols.Decls, AST);
+  for (auto &Ref : MainFileRefs) {
+    Location Result;
+    Result.range = getTokenRange(AST, Ref.Loc);
+    Result.uri = URIForFile(*MainFilePath);
+    Results.push_back(std::move(Result));
+  }
+
+  // Now query the index for references from other files.
+  if (!Index)
+    return Results;
+  RefsRequest Req;
+  for (const Decl *D : Symbols.Decls) {
+    // Not all symbols can be referenced from outside (e.g. function-locals).
+    // TODO: we could skip TU-scoped symbols here (e.g. static functions) if
+    // we know this file isn't a header. The details might be tricky.
+    if (D->getParentFunctionOrMethod())
+      continue;
+    if (auto ID = getSymbolID(D))
+      Req.IDs.insert(*ID);
+  }
+  if (Req.IDs.empty())
+    return Results;
+  Index->refs(Req, [&](const Ref &R) {
+    auto LSPLoc = toLSPLocation(R.Location, /*HintPath=*/*MainFilePath);
+    // Avoid indexed results for the main file - the AST is authoritative.
+    if (LSPLoc && LSPLoc->uri.file() != *MainFilePath)
+      Results.push_back(std::move(*LSPLoc));
+  });
+  return Results;
+}
+
 } // namespace clangd
 } // namespace clang
Index: unittests/clangd/XRefsTests.cpp
===================================================================
--- unittests/clangd/XRefsTests.cpp
+++ unittests/clangd/XRefsTests.cpp
@@ -1068,6 +1068,81 @@
       ElementsAre(Location{FooCppUri, FooWithoutHeader.range()}));
 }
 
+TEST(FindReferences, WithinAST) {
+  const char *Tests[] = {
+      R"cpp(// Local variable
+        int main() {
+          int $foo[[foo]];
+          $foo[[^foo]] = 2;
+          int test1 = $foo[[foo]];
+        }
+      )cpp",
+
+      R"cpp(// Struct
+        namespace ns1 {
+        struct $foo[[Foo]] {};
+        } // namespace ns1
+        int main() {
+          ns1::$foo[[Fo^o]]* Params;
+        }
+      )cpp",
+
+      R"cpp(// Function
+        int $foo[[foo]](int) {}
+        int main() {
+          auto *X = &$foo[[^foo]];
+          $foo[[foo]](42)
+        }
+      )cpp",
+
+      R"cpp(// Field
+        struct Foo {
+          int $foo[[foo]];
+          Foo() : $foo[[foo]](0) {}
+        };
+        int main() {
+          Foo f;
+          f.$foo[[f^oo]] = 1;
+        }
+      )cpp",
+
+      R"cpp(// Method call
+        struct Foo { int [[foo]](); };
+        int Foo::[[foo]]() {}
+        int main() {
+          Foo f;
+          f.^foo();
+        }
+      )cpp",
+
+      R"cpp(// Typedef
+        typedef int $foo[[Foo]];
+        int main() {
+          $foo[[^Foo]] bar;
+        }
+      )cpp",
+
+      R"cpp(// Namespace
+        namespace $foo[[ns]] {
+        struct Foo {};
+        } // namespace ns
+        int main() { $foo[[^ns]]::Foo foo; }
+      )cpp",
+  };
+  for (const char *Test : Tests) {
+    Annotations T(Test);
+    auto AST = TestTU::withCode(T.code()).build();
+    std::vector<Matcher<Location>> ExpectedLocations;
+    for (const auto &R : T.ranges("foo"))
+      ExpectedLocations.push_back(RangeIs(R));
+    EXPECT_THAT(findReferences(AST, T.point()),
+                ElementsAreArray(ExpectedLocations))
+        << Test;
+  }
+}
+
+// TODO: add tests that rely on the index, and verify whether it was queried.
+
 } // namespace
 } // namespace clangd
 } // namespace clang