[clangd] Implement findReferences function

[sammccall]

Summary

clangd will use findReferences to provide LSP's reference feature.

[hokein]

More cleanups.

[hokein]

This patch is incomplete, and it is a prototype based on our discussion last week. You can patch it locally, and play around with VSCode.

[sammccall]

Mostly just reviewed the references() implementation.

For the infrastructure I have some high level questions:

• why are we combining SymbolSlab and occurrences?
• what's the motivation for the separate preamble/main-file indexes
• why put more functionality into SymbolCollector?

Let's discuss offline a bit if that works?

clangd/TUScheduler.h

54 ↗

(On Diff #161441)

(we're replacing a std::function, why not just a struct with two std::functions? saves boilerplate in a bunch of places...)

clangd/XRefs.cpp

715

I'm not sure unpacking the reference options into individual bools is going to scale well. I'd suggest passing the whole struct instead.

724

"local" is ambiguous here (this function cares both about file-local references and function-local symbols).
Consider ExternallyVisibleIDs

730

(This saves us hitting the index to look up references for one symbol, not sure if it's worth it at all).

I wouldn't particularly trust the helpers in clang::index::. Indeed the implementation looks wrong here (e.g. it would treat lambda captures as global, I think?)

I'd prefer D->getParentFunctionOrMethod() here.

737

I'm not sure this is the best interpretation of LSP includeDeclaration.

The LSP spec is very vague here, and we can usually assume it's written with typescript in mind :-) where the declaration/definition distinction doesn't really exist.
It appears the distinction they're trying to draw is declaration vs "only" reference, rather than definition vs "only" declaration. So I think if we're going to respect this flag, we should *exclude* definitions when the flag is false.

Alternatively we could punt on this and just ignore the flag for now, and add it in a followup.

743

Reusing symbolcollector here seems odd.

It forces us to go through SymbolID rather than just using the Decl*. This is certainly reliable for global symbols, but I've got no idea how robust USRs are for local symbols. It also ends up building the Symbol objects, which we then throw away.

How much of SymbolCollector are we really reusing, vs a custom IndexDataConsumer? How is this different from document-highlights? Maybe we can reuse the consumer.

748

This comment is surprising as it stands. Maybe

// Look in the AST for references from the current file.

(and hoist the commentto the top of the code that deals with AST, currently line 686)

762

This is also a bit confusing - the symbols are non-local (in the function-scope sense) and the references are non-local (in the file sense).
Consider:

// Consult the index for references in other files.
// We only need to consider symbols visible to other files.

771

in the usual case, SeenURIs is the current file (if there were any local usages), or empty (if there weren't).

This means if there are refs from the current file in the index but not locally (they were deleted), you'll return them when it would be better not to.

Why not just obtain the URI for the main file directly and compare that, instead of extracting it from the local references found?

clangd/index/Index.h

368 ↗

(On Diff #161441)

why are we moving occurrences into the symbol slab?

[ilya-biryukov]

Parts of this patch were landed as D50847 and D50889, happy to discuss the design decisions ("merged" dynamic index, two callbacks), but that's probably out of scope of this patch.

[sammccall]

In D50958#1212179, @ilya-biryukov wrote:

Parts of this patch were landed as D50847 and D50889, happy to discuss the design decisions ("merged" dynamic index, two callbacks), but that's probably out of scope of this patch.

Ah ok, that explains some of the unrelated changes.
I do have some questions there that would be good to discuss. Meanwhile, @hokein can you rebase this patch against head?

[hokein]

Thanks for the comments.

In D50958#1212221, @sammccall wrote:

I do have some questions there that would be good to discuss. Meanwhile, @hokein can you rebase this patch against head?

Yes, I'd love to, but since this patch is quite large, I need split it into smaller patches. I will update this patch (focus on the references implementation), it depends on the following patches

• https://reviews.llvm.org/D50385
• https://reviews.llvm.org/D51279

[hokein]

• address review comments
• rescope the patch only focus on findReferences implementation.

[hokein]

clangd/XRefs.cpp

715

Removed this parameter.

730

I think it is worth the effort here. At least for function-local symbols, traversing AST is sufficient to get complete xref results.

737

Ignore this flag, and always return all kinds.

743

I refined the highlight IndexDataConsumer, and now the common code is shared.

[hokein]

Rebase

[sammccall]

Looking forward to using this!

Unless you object, I'd like to address the remaining comments (and get a review), so you can make the most of your leave!

clangd/XRefs.cpp

334

why is this a map keyed by D? the keys are never used, the result is always flattened into a vector.

362

why is this a subclass rather than just using OccurrenceCollector and putting the finish() code in findReferences()?

379

out of curiosity: how do we actually get duplicates?

389

as above, why does this need to be a subclass

719

(comment just echoes code)

722

visiable -> visible

725

we can check isExternallyVisible, I think?
maybe it's not worth it, but the comment as it stands doesn't seem accurate

unittests/clangd/XRefsTests.cpp

1191

can we use a simple real implementation TU.index() instead of mocking here?
I think this is an artifact of the fact that TestTU doesn't actually expose occurrences in the index, can we fix that?

[hokein]

In D50958#1222249, @sammccall wrote:

Looking forward to using this!

Unless you object, I'd like to address the remaining comments (and get a review), so you can make the most of your leave!

Thanks, feel free to pick it up. The comments make sense to me. Just let you know I have one more xref patch: https://reviews.llvm.org/D50896.

clangd/XRefs.cpp

334

Using map can group results by the Decl, but it doesn't matter here.

379

Some AST nodes will be visited more than once, so we might have duplicated locations.

725

I'm not sure whether isExternallyVisible suits our cases. For example, if a header defines a static const int MAX, and the header is widely included in the project, we want to query it in the index, but the linkage of the symbol MAX is internal.

So I'd prefer to be conservative here (only for function-local symbols).

unittests/clangd/XRefsTests.cpp

1191

Yes, we can. But the test here is only to verify whether the index has been queried. The results from index are not important.

[sammccall]

OK thanks, I'll steal this one.

clangd/XRefs.cpp

725

Agree, this sounds like the right place to start.

unittests/clangd/XRefsTests.cpp

1191

ah, that makes sense. But we should have some actual tests that use the index content too!

[sammccall]

Address comments and tighten up highlights/refs reuse in XRefs.cpp a bit.

Still to do:

• test interaction with index, including actual data
• maybe add the ClangdServer/LSP binding and lit test, if it's just boilerplate

[sammccall]

Oops, and I rebased on head (Occurrences -> Refs) of course.

@ilya-biryukov @ioeric, any interest in taking over review here?

[hokein]

Thanks, looks good from my side.

clangd/XRefs.cpp

329

nit: we can initialize TargetDecls like Targets(TargetDecls.begin(), TargetDecls.end()).

729

nit: use const auto & indicating the variable is immutable.

[sammccall]

Finish tests, fix TestTU::index() to include occurrences.

[sammccall]

clangd/XRefs.cpp

329

In fact SmallSet is missing this constructor :-(

[sammccall]

Oops, I forgot to associate my patch with this review. It landed as rL341458.

clangd/XRefs.h

 
 /// Returns highlights for all usages of a symbol at \p Pos.
 std::vector<DocumentHighlight> findDocumentHighlights(ParsedAST &AST,
                                                       Position Pos);
 
 /// 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
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_XREFS_H

clangd/XRefs.cpp

IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   IndexOpts.IndexFunctionLocals = true;
   indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
                      DeclMacrosFinder, IndexOpts);
 
   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;
 }
 
 } // namespace
 
 std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
                                       const SymbolIndex *Index) {
   const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
std::vector<Location> findDefinitions(ParsedAST &AST, Position Pos,
 
   // Identified symbols at a specific position.
   SourceLocation SourceLocationBeg =
       getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
   auto Symbols = getSymbolAtPosition(AST, SourceLocationBeg);
 
   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);
   }
 
   // Declaration and definition are different terms in C-family languages, and
   // LSP only defines the "GoToDefinition" specification, so we try to perform
   // the "most sensible" GoTo operation:
   //
for (const auto *D : Symbols.Decls) {
     // without USR are rare and unimportant cases, we use the a fake holder to
     // minimize the invasiveness of these cases.
     SymbolID Key("");
     if (auto ID = getSymbolID(D))
       Key = *ID;
 
     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;
     // Populate one of the slots with location for the AST.
     if (!IsDef)
       Candidate.Decl = L;
     else
       Candidate.Def = L;
if (Candidate.Decl &&
       Result.push_back(*Candidate.Decl);
   }
 
   return Result;
 }
 
 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 all occurrences (and related information) in the main file that a
+/// given whitelist of declarations refers to.
+class OccurrenceCollector : public index::IndexDataConsumer {
 public:
-  DocumentHighlightsFinder(raw_ostream &OS, 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);
-  }
-
   bool
   handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
                       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);
+    const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+    SourceLocation FileLoc = SourceMgr.getFileLoc(Loc);
+    // We only collect locations in current main file.
+    if (SourceMgr.getMainFileID() != SourceMgr.getFileID(FileLoc) ||
+        std::find(WhitelistDecls.begin(), WhitelistDecls.end(), D) ==
+            WhitelistDecls.end())
+      return true;
+
+    DeclOccurrences[D].emplace_back(FileLoc, Roles);
 
-    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;
+protected:
+  OccurrenceCollector(ParsedAST &AST,
+                      const std::vector<const Decl *> &WhitelistDecls)
+      : AST(AST), WhitelistDecls(WhitelistDecls) {
 
-    DocumentHighlights.push_back(getDocumentHighlight(SR, Kind));
-    return true;
+    auto &SM = AST.getASTContext().getSourceManager();
+    auto MainFilePath =
+        getRealPath(SM.getFileEntryForID(SM.getMainFileID()), SM);
+    if (!MainFilePath) {
+      log("Fail to get real path!");
+    }
+    MainFileURI = URIForFile(*MainFilePath);
+  }
+
+  ParsedAST &AST;
+  const std::vector<const Decl *> &WhitelistDecls;
+
+  URIForFile MainFileURI;
+
+  using DeclOccurrence = std::pair<SourceLocation, index::SymbolRoleSet>;
+  llvm::DenseMap<const Decl *, std::vector<DeclOccurrence>> DeclOccurrences;
+};
+
+/// Find symbol occurrences that a given whilelist of declarations refers to.
+class OccurrencesFinder : public OccurrenceCollector {
+public:
+  OccurrencesFinder(ParsedAST &AST, std::vector<const Decl *> &WhitelistDecls)
+      : OccurrenceCollector(AST, WhitelistDecls) {}
+
+  std::vector<Location> takeOccurrences() { return std::move(Occurrences); }
+
+  void finish() override {
+    OccurrenceCollector::finish();
+    Occurrences.clear();
+
+    for (auto It : DeclOccurrences) {
+      for (const auto &LocAndRole : It.second) {
+        Occurrences.push_back(
+            {MainFileURI, getTokenRange(AST, LocAndRole.first)});
+      }
+    }
+    // Deduplicate results.
+    std::sort(Occurrences.begin(), Occurrences.end());
+    auto Last = std::unique(Occurrences.begin(), Occurrences.end());
+    Occurrences.erase(Last, Occurrences.end());
   }
 
 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};
+  std::vector<Location> Occurrences;
+};
+
+/// Finds document highlights that a given whitelist of declarations refers to.
+class DocumentHighlightsFinder : public OccurrenceCollector {
+public:
+  DocumentHighlightsFinder(ParsedAST &AST,
+                           const std::vector<const Decl *> &WhitelistDecls)
+      : OccurrenceCollector(AST, WhitelistDecls) {}
+
+  std::vector<DocumentHighlight> takeHighlights() {
+    return std::move(DocumentHighlights);
+  }
+
+  void finish() override {
+    OccurrenceCollector::finish();
+    DocumentHighlights.clear();
+    for (auto It : DeclOccurrences) {
+      for (const auto &LocAndRole : It.second)
+        DocumentHighlights.push_back(
+            createHighlight(LocAndRole.first, LocAndRole.second));
+    }
+    // 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());
+  }
+
+  DocumentHighlight createHighlight(SourceLocation Loc,
+                                    index::SymbolRoleSet Roles) {
     DocumentHighlight DH;
-    DH.range = R;
-    DH.kind = Kind;
+    DH.range = getTokenRange(AST, Loc);
+    DH.kind = DocumentHighlightKind::Text;
+    if (static_cast<index::SymbolRoleSet>(index::SymbolRole::Write) & Roles)
+      DH.kind = DocumentHighlightKind::Write;
+    else if (static_cast<index::SymbolRoleSet>(index::SymbolRole::Read) & Roles)
+      DH.kind = DocumentHighlightKind::Read;
     return DH;
   }
+
+private:
+  std::vector<DocumentHighlight> DocumentHighlights;
 };
 
 } // 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(
-      llvm::errs(), AST.getASTContext(), AST.getPreprocessor(), SelectedDecls);
+  DocumentHighlightsFinder DocHighlightsFinder(AST, SelectedDecls);
 
   index::IndexingOptions IndexOpts;
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::All;
   IndexOpts.IndexFunctionLocals = true;
   indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
                      DocHighlightsFinder, IndexOpts);
 
Optional<Hover> getHover(ParsedAST &AST, Position Pos) {
 
   auto DeducedType = getDeducedType(AST, SourceLocationBeg);
   if (DeducedType && !DeducedType->isNull())
     return getHoverContents(*DeducedType, AST.getASTContext());
 
   return None;
 }
 
+std::vector<Location> findReferences(ParsedAST &AST, Position Pos,
+                                     const SymbolIndex *Index) {
+  const SourceManager &SourceMgr = AST.getASTContext().getSourceManager();
+  SourceLocation SourceLocationBeg =
+      getBeginningOfIdentifier(AST, Pos, SourceMgr.getMainFileID());
+  // Identified symbols at a specific position.
+  auto Symbols = getSymbolAtPosition(AST, SourceLocationBeg);
+
+  // For local symbols (e.g. symbols that are only visiable in main file,
+  // symbols defined in function body), we can get complete references by
+  // traversing the AST, and we don't want to make unnecessary queries to the
+  // index. Howerver, we don't have a a reliable way to distinguish file-local
+  // symbols. We conservatively consider function-local symbols.
+  llvm::DenseSet<SymbolID> PossiblyVisibleIDs;
+  for (const auto *D : Symbols.Decls) {
+    if (auto ID = getSymbolID(D)) {
+      // Ignore if it is a fucntion-scope symbol.
+      if (D->getParentFunctionOrMethod())
+        continue;
+      PossiblyVisibleIDs.insert(*ID);
+    }
+  }
+
+  // Look in the AST for the references from current main file.
+  OccurrencesFinder FindOccurrences(AST, Symbols.Decls);
+  index::IndexingOptions IndexOpts;
+  IndexOpts.SystemSymbolFilter =
+      index::IndexingOptions::SystemSymbolFilterKind::All;
+  IndexOpts.IndexFunctionLocals = true;
+  indexTopLevelDecls(AST.getASTContext(), AST.getLocalTopLevelDecls(),
+                     FindOccurrences, IndexOpts);
+
+  auto Results = FindOccurrences.takeOccurrences();
+
+  // Consult the index for references in other files.
+  // We only need to consider symbols visible to other files.
+  if (Index && !PossiblyVisibleIDs.empty()) {
+    const auto &SM = AST.getASTContext().getSourceManager();
+    auto MainFilePath =
+        getRealPath(SM.getFileEntryForID(SourceMgr.getMainFileID()), SM);
+    if (!MainFilePath) {
+      log("Failed to get real path!");
+      return Results;
+    }
+    auto MainFileURI = URIForFile(*MainFilePath);
+    OccurrencesRequest Request;
+    Request.IDs = std::move(PossiblyVisibleIDs);
+    Request.Filter = SymbolOccurrenceKind::Declaration |
+                     SymbolOccurrenceKind::Reference |
+                     SymbolOccurrenceKind::Definition;
+    Index->findOccurrences(Request, [&](const SymbolOccurrence &O) {
+      if (auto LSPLoc = toLSPLocation(O.Location, /*HintPath=*/*MainFilePath)) {
+        if (MainFileURI != LSPLoc->uri)
+          Results.push_back(*LSPLoc);
+      }
+    });
+  }
+  return Results;
+}
+
 } // namespace clangd
 } // namespace clang

unittests/clangd/XRefsTests.cpp

 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 using namespace llvm;
 
 namespace {
+using testing::_;
 using testing::ElementsAre;
 using testing::Field;
 using testing::IsEmpty;
 using testing::Matcher;
 using testing::UnorderedElementsAreArray;
 
 class IgnoreDiagnostics : public DiagnosticsConsumer {
   void onDiagnosticsReady(PathRef File,
TEST(GoToDefinition, WithPreamble) {
   // Both preamble and AST are built in this request.
   Server.addDocument(FooCpp, FooWithoutHeader.code(), WantDiagnostics::Yes);
   // Use the AST being built in above request.
   EXPECT_THAT(
       cantFail(runFindDefinitions(Server, FooCpp, FooWithoutHeader.point())),
       ElementsAre(Location{FooCppUri, FooWithoutHeader.range()}));
 }
 
+TEST(FindReferences, AllWithoutIndex) {
+  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;
+  }
+}
+
+class MockIndex : public SymbolIndex {
+public:
+  MOCK_CONST_METHOD2(fuzzyFind, bool(const FuzzyFindRequest &,
+                                     llvm::function_ref<void(const Symbol &)>));
+  MOCK_CONST_METHOD2(lookup, void(const LookupRequest &,
+                                  llvm::function_ref<void(const Symbol &)>));
+  MOCK_CONST_METHOD2(findOccurrences,
+                     void(const OccurrencesRequest &,
+                          llvm::function_ref<void(const SymbolOccurrence &)>));
+};
+
+TEST(FindReferences, QueryIndex) {
+  const char *Tests[] = {
+      // Refers to symbols from headers.
+      R"cpp(
+        int main() {
+          F^oo foo;
+        }
+      )cpp",
+      R"cpp(
+        int main() {
+          f^unc();
+        }
+      )cpp",
+      R"cpp(
+        int main() {
+          return I^NT;
+        }
+      )cpp",
+
+      // These are cases of file-local but not function-local symbols, we still
+      // query the index.
+      R"cpp(
+        void MyF^unc() {}
+      )cpp",
+
+      R"cpp(
+        int My^Int = 2;
+      )cpp",
+  };
+
+  TestTU TU;
+  TU.HeaderCode = R"(
+  class Foo {};
+  static const int INT = 3;
+  inline void func() {};
+  )";
+  MockIndex Index;
+  for (const char *Test : Tests) {
+    Annotations T(Test);
+    TU.Code = T.code();
+    auto AST = TU.build();
+    EXPECT_CALL(Index, findOccurrences(_, _));
+    findReferences(AST, T.point(), &Index);
+  }
+}
+
+TEST(FindReferences, DontQueryIndex) {
+  // Don't query index for function-local symbols.
+  const char *Tests[] = {
+      R"cpp(// Local variable in function body
+        int main() {
+          int $foo[[foo]];
+          $foo[[^foo]] = 2;
+        }
+      )cpp",
+
+      R"cpp(// function parameter
+        int f(int fo^o) {
+        }
+      )cpp",
+
+      R"cpp(// function parameter in lambda
+        int f(int foo) {
+          auto func = [](int a, int b) {
+            return ^a = 2;
+          };
+        }
+      )cpp",
+
+      R"cpp(// capture in lambda
+        int f(int foo) {
+          int A;
+          auto func = [&A](int a, int b) {
+            return a = ^A;
+          };
+        }
+      )cpp",
+  };
+
+  MockIndex Index;
+  for (const char *Test : Tests) {
+    Annotations T(Test);
+    auto AST = TestTU::withCode(T.code()).build();
+    EXPECT_CALL(Index, findOccurrences(_, _)).Times(0);
+    findReferences(AST, T.point(), &Index);
+  }
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang