[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

662

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

671

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

677

(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.

684

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.

690

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.

695

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)

709

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.

718

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

662

Removed this parameter.

677

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

684

Ignore this flag, and always return all kinds.

690

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

354

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

363

as above, why does this need to be a subclass

364

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

381

out of curiosity: how do we actually get duplicates?

666

(comment just echoes code)

669

visiable -> visible

672

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

1192

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

354

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

381

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

672

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

1192

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

672

Agree, this sounds like the right place to start.

unittests/clangd/XRefsTests.cpp

1192

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

328

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

676

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

[sammccall]

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

[sammccall]

clangd/XRefs.cpp

328

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 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
   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);
-
-    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
 
-  return DocHighlightsFinder.takeHighlights();
+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);
+
+  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) {
   PrintingPolicy Policy(Base);
 
   Policy.AnonymousTagLocations = false;
   Policy.TerseOutput = true;
   Policy.PolishForDeclaration = true;
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) {
+  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 (const 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

unittests/clangd/TestTU.cpp

 }
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
   return indexAST(AST.getASTContext(), AST.getPreprocessorPtr()).first;
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
-  // FIXME: we should generate proper refs for TestTU.
-  return MemIndex::build(headerSymbols(), RefSlab());
+  auto AST = build();
+  auto Content = indexAST(AST.getASTContext(), AST.getPreprocessorPtr(),
+                          AST.getLocalTopLevelDecls());
+  return MemIndex::build(std::move(Content.first), std::move(Content.second));
 }
 
 const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {
   const Symbol *Result = nullptr;
   for (const Symbol &S : Slab) {
     if (QName != (S.Scope + S.Name).str())
       continue;
     if (Result) {

unittests/clangd/XRefsTests.cpp

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, 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;
+  }
+}
+
+TEST(FindReferences, NeedsIndex) {
+  const char *Header = "int foo();";
+  Annotations Main("int main() { [[f^oo]](); }");
+  TestTU TU;
+  TU.Code = Main.code();
+  TU.HeaderCode = Header;
+  auto AST = TU.build();
+
+  // References in main file are returned without index.
+  EXPECT_THAT(findReferences(AST, Main.point(), /*Index=*/nullptr),
+              ElementsAre(RangeIs(Main.range())));
+  Annotations IndexedMain(R"cpp(
+    int main() { [[f^oo]](); }
+  )cpp");
+
+  // References from indexed files are included.
+  TestTU IndexedTU;
+  IndexedTU.Code = IndexedMain.code();
+  IndexedTU.Filename = "Indexed.cpp";
+  IndexedTU.HeaderCode = Header;
+  EXPECT_THAT(findReferences(AST, Main.point(), IndexedTU.index().get()),
+              ElementsAre(RangeIs(Main.range()), RangeIs(IndexedMain.range())));
+
+  // If the main file is in the index, we don't return duplicates.
+  // (even if the references are in a different location)
+  TU.Code = ("\n\n" + Main.code()).str();
+  EXPECT_THAT(findReferences(AST, Main.point(), TU.index().get()),
+              ElementsAre(RangeIs(Main.range())));
+};
+
+TEST(FindReferences, NoQueryForLocalSymbols) {
+  struct RecordingIndex : public MemIndex {
+    mutable Optional<DenseSet<SymbolID>> RefIDs;
+    void refs(const RefsRequest &Req,
+              llvm::function_ref<void(const Ref &)>) const override {
+      RefIDs = Req.IDs;
+    }
+  };
+
+  struct Test {
+    StringRef AnnotatedCode;
+    bool WantQuery;
+  } Tests[] = {
+      {"int ^x;", true},
+      // For now we don't assume header structure which would allow skipping.
+      {"namespace { int ^x; }", true},
+      {"static int ^x;", true},
+      // Anything in a function certainly can't be referenced though.
+      {"void foo() { int ^x; }", false},
+      {"void foo() { struct ^x{}; }", false},
+      {"auto lambda = []{ int ^x; };", false},
+  };
+  for (Test T : Tests) {
+    Annotations File(T.AnnotatedCode);
+    RecordingIndex Rec;
+    auto AST = TestTU::withCode(File.code()).build();
+    findReferences(AST, File.point(), &Rec);
+    if (T.WantQuery)
+      EXPECT_NE(Rec.RefIDs, llvm::None) << T.AnnotatedCode;
+    else
+      EXPECT_EQ(Rec.RefIDs, llvm::None) << T.AnnotatedCode;
+  }
+};
+
 } // namespace
 } // namespace clangd
 } // namespace clang