[clangd] Fix some references missing in dynamic index.

[hokein]

Summary

Previously, SymbolCollector postfilters all references at the end to
find all references of interesting symbols.
It was incorrect when indxing main AST where we don't see locations
of symbol declarations and definitions in the main AST (as those are in
preamble AST).

The fix is to do earily check during collecting references.

[sammccall]

Nice fix, just performance concerns.

(I do think this might be significant, but measuring dynamic index time before/after for a big TU might show this to be unimportant)

clangd/index/SymbolCollector.cpp

348

This seems better for the main-AST case, but substantially more expensive for indexing preambles: we may not want references at all, so paying for shouldCollectSymbol for every reference seems wasteful.

what about

bool RefMayBeInteresting = RefFilter & Roles;
bool IsOnlyRef = !(Roles & (Declaration | Definition));
if (IsOnlyRef && !RefMayBeInteresting)
  return;
if (!shouldCollectSymbol())
  return;
if (RefMayBeInteresting && File == getMainFileID())
  // add ref
if (IsOnlyRef) // don't continue if mere reference
  return;

[hokein]

avoid calling shouldCollectSymbol every time during indexing.

clangd/index/SymbolCollector.cpp

348

I thought shouldCollectSymbol is not as expensive as before (since we have removed AST matchers stuff there), but I'm definitely +1 on the idea of avoid calling it for every reference.

clangd/index/SymbolCollector.cpp

bool SymbolCollector::handleDeclOccurence(
   // D may not be an interesting symbol, but it's cheaper to check at the end.
   auto &SM = ASTCtx->getSourceManager();
   auto SpellingLoc = SM.getSpellingLoc(Loc);
   if (Opts.CountReferences &&
       (Roles & static_cast<unsigned>(index::SymbolRole::Reference)) &&
       SM.getFileID(SpellingLoc) == SM.getMainFileID())
     ReferencedDecls.insert(ND);
 
-  if ((static_cast<unsigned>(Opts.RefFilter) & Roles) &&
-      SM.getFileID(SpellingLoc) == SM.getMainFileID())
-    DeclRefs[ND].emplace_back(SpellingLoc, Roles);
+  bool CollectRef = static_cast<unsigned>(Opts.RefFilter) & Roles;
+  bool IsOnlyRef =
+      !(Roles & (static_cast<unsigned>(index::SymbolRole::Declaration) |
+                 static_cast<unsigned>(index::SymbolRole::Definition)));
 
-  // Don't continue indexing if this is a mere reference.
-  if (!(Roles & static_cast<unsigned>(index::SymbolRole::Declaration) ||
-        Roles & static_cast<unsigned>(index::SymbolRole::Definition)))
+  if (IsOnlyRef && !CollectRef)
     return true;
   if (!shouldCollectSymbol(*ND, *ASTCtx, Opts))
     return true;
+  if (CollectRef && SM.getFileID(SpellingLoc) == SM.getMainFileID())
+    DeclRefs[ND].emplace_back(SpellingLoc, Roles);
+  // Don't continue indexing if this is a mere reference.
+  if (IsOnlyRef)
+    return true;
 
   auto ID = getSymbolID(ND);
   if (!ID)
     return true;
 
   const NamedDecl &OriginalDecl = *cast<NamedDecl>(ASTNode.OrigD);
   const Symbol *BasicSymbol = Symbols.find(*ID);
   if (!BasicSymbol) // Regardless of role, ND is the canonical declaration.
void SymbolCollector::finish() {
 
   const auto &SM = ASTCtx->getSourceManager();
   auto* MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
 
   if (auto MainFileURI = toURI(SM, MainFileEntry->getName(), Opts)) {
     std::string MainURI = *MainFileURI;
     for (const auto &It : DeclRefs) {
       if (auto ID = getSymbolID(It.first)) {
-        if (Symbols.find(*ID)) {
-          for (const auto &LocAndRole : It.second) {
-            Ref R;
-            auto Range =
-                getTokenRange(LocAndRole.first, SM, ASTCtx->getLangOpts());
-            R.Location.Start = Range.first;
-            R.Location.End = Range.second;
-            R.Location.FileURI = MainURI;
-            R.Kind = toRefKind(LocAndRole.second);
-            Refs.insert(*ID, R);
-          }
-        }
-      }
-    }
   } else {
     log("Failed to create URI for main file: {0}", MainFileEntry->getName());
   }
 
   ReferencedDecls.clear();
   ReferencedMacros.clear();
   DeclRefs.clear();
 }

unittests/clangd/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
+#include "AST.h"
 #include "ClangdUnit.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "gmock/gmock.h"
 #include "index/FileIndex.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
+#include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "gtest/gtest.h"
 
 using testing::_;
 using testing::AllOf;
 using testing::ElementsAre;
TEST(FileIndexTest, CollectMacros) {
   M.fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_EQ(Sym.Name, "CLANGD");
     EXPECT_EQ(Sym.SymInfo.Kind, index::SymbolKind::Macro);
     SeenSymbol = true;
   });
   EXPECT_TRUE(SeenSymbol);
 }
 
+TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
+  const std::string Header = R"cpp(
+    class Foo {};
+  )cpp";
+  Annotations Main(R"cpp(
+    #include "foo.h"
+    void f() {
+      [[Foo]] foo;
+    }
+  )cpp");
+  auto MainFile = testPath("foo.cpp");
+  auto HeaderFile = testPath("foo.h");
+  std::vector<const char*> Cmd = {"clang", "-xc++", MainFile.c_str()};
+  // Preparse ParseInputs.
+  ParseInputs PI;
+  PI.CompileCommand.Directory = testRoot();
+  PI.CompileCommand.Filename = MainFile;
+  PI.CompileCommand.CommandLine = {Cmd.begin(), Cmd.end()};
+  PI.Contents = Main.code();
+  PI.FS = buildTestFS({{MainFile, Main.code()}, {HeaderFile, Header}});
+
+  // Prepare preamble.
+  auto CI = buildCompilerInvocation(PI);
+  auto PreambleData = buildPreamble(
+      MainFile,
+      *buildCompilerInvocation(PI), /*OldPreamble=*/nullptr,
+      tooling::CompileCommand(), PI,
+      std::make_shared<PCHContainerOperations>(), /*StoreInMemory=*/true,
+      [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP) {});
+  // Build AST for main file with preamble.
+  auto AST = ParsedAST::build(
+      createInvocationFromCommandLine(Cmd), PreambleData,
+      llvm::MemoryBuffer::getMemBufferCopy(Main.code()),
+      std::make_shared<PCHContainerOperations>(),
+      PI.FS);
+  ASSERT_TRUE(AST);
+  FileIndex Index;
+  Index.updateMain(MainFile, *AST);
+
+  auto Foo =
+      findSymbol(TestTU::withHeaderCode(Header).headerSymbols(), "Foo");
+  RefsRequest Request;
+  Request.IDs.insert(Foo.ID);
+
+  // Expect to see references in main file, references in headers are excluded
+  // because we only index main AST.
+  EXPECT_THAT(getRefs(Index, Foo.ID), RefsAre({RefRange(Main.range())}));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang