[clangd] Partition include graph on auto-index.

[kadircet]

Summary

Partitions include graphs in auto-index so that each shards contains
only part of the include graph related to itself.

[ilya-biryukov]

clangd/index/Background.cpp

207

We should make the function static or put it into the anonymous namespace.
While here the same should be done for URIToFileCache.

210

NIT: Maybe accept a URI directly? E.g. if we ever have clients that already store a URI they won't need to do double conversions.

217

Maybe add a comment mentioning we do this merely to intern the strings?
This might be hard to follow on the first read.

385–386

Maybe also log the stats from the include graph?
It's probably less useful than symbols and refs count, but still...

388

Same here, maybe also attach some stats from the include graph?

unittests/clangd/BackgroundIndexTests.cpp

192

This sets a compile command for a directory. Should it be "root/A.cc" or am I missing something?

[kadircet]

• Address comments

[kadircet]

unittests/clangd/BackgroundIndexTests.cpp

192

WOW! Thanks for noticing it :D It was working correctly before because we were deducing the file's absolute path from the compile command itself rather than what we've received by watching the CDB. Changed all references in the file.

[ilya-biryukov]

clangd/index/Background.cpp

70

Naming: technically the variable name should be URI, but to avoid clashing with the type name we could use U

259

Maybe use auto? The type is spelled explicitly in the rhs

260

What are the cases when it happens in practice? Do we ever run the indexer without producing the include graph?

386–391

add a comma before {3} sources?

[kadircet]

• Address comments

clangd/index/Background.cpp

260

You are right that's not the case any more.

[kadircet]

• Rebase

[ilya-biryukov]

Everything looks good, just a single important request about testing the included files do not have any edges in the resulting graph

clangd/index/Background.cpp

79

NIT: the comment could probably be simplified to something like URIs inside nodes must point into the keys of the same IncludeGraph

386

"sources" might be a bit confusing. I'd probably use "files" instead, but that doesn't look ideal too.

unittests/clangd/BackgroundIndexTests.cpp

197

Maybe check the full shape of the graph? E.g. also add a check that A.h does not have any includes and its digest is populated.

[kadircet]

• Addresss comments
• Make sure there are no uninitialized values in IncludeGraphNode.

[ilya-biryukov]

LGTM

unittests/clangd/BackgroundIndexTests.cpp

199

NIT: maybe also check the hash of A.cc is initialized (the value doesn't matter, probably, just to make sure it's not 0)?

207

NIT: maybe add a comma? A.h, B.h

[kadircet]

• Address comments

clangd/index/Background.cpp

 #include <numeric>
 #include <queue>
 #include <random>
 #include <string>
 
 using namespace llvm;
 namespace clang {
 namespace clangd {
+namespace {
+// Resolves URI to file paths with cache.
+class URIToFileCache {
+public:
+  URIToFileCache(llvm::StringRef HintPath) : HintPath(HintPath) {}
+
+  llvm::StringRef resolve(llvm::StringRef FileURI) {
+    auto I = URIToPathCache.try_emplace(FileURI);
+    if (I.second) {
+      auto U = URI::parse(FileURI);
+      if (!U) {
+        elog("Failed to parse URI {0}: {1}", FileURI, U.takeError());
+        assert(false && "Failed to parse URI");
+        return "";
+      }
+      auto Path = URI::resolve(*U, HintPath);
+      if (!Path) {
+        elog("Failed to resolve URI {0}: {1}", FileURI, Path.takeError());
+        assert(false && "Failed to resolve URI");
+        return "";
+      }
+      I.first->second = *Path;
+    }
+    return I.first->second;
+  }
+
+private:
+  std::string HintPath;
+  llvm::StringMap<std::string> URIToPathCache;
+};
+
+// We keep only the node "Path" and its edges.
+IncludeGraph getSubGraph(const URI &U, const IncludeGraph &FullGraph) {
+  IncludeGraph IG;
+
+  std::string FileURI = U.toString();
+  auto Entry = IG.try_emplace(FileURI).first;
+  auto &Node = Entry->getValue();
+  Node = FullGraph.lookup(Entry->getKey());
+  Node.URI = Entry->getKey();
+
+  // Since the strings in direct includes are references to the keys of the
+  // fullgraph, we need to create entries in our new sub-graph for those
+  // includes and change references to point to the new graph's keys.
+  for (auto &Include : Node.DirectIncludes) {
+    auto I = IG.try_emplace(Include).first;
+    I->getValue().URI = I->getKey();
+    Include = I->getKey();
+  }
+
+  return IG;
+}
+} // namespace
 
 BackgroundIndex::BackgroundIndex(
     Context BackgroundContext, StringRef ResourceDir,
     const FileSystemProvider &FSProvider, const GlobalCompilationDatabase &CDB,
     BackgroundIndexStorage::Factory IndexStorageFactory, size_t ThreadPoolSize)
     : SwapIndex(make_unique<MemIndex>()), ResourceDir(ResourceDir),
       FSProvider(FSProvider), CDB(CDB),
       BackgroundContext(std::move(BackgroundContext)),

 void BackgroundIndex::enqueueTask(Task T) {
   {
     std::lock_guard<std::mutex> Lock(QueueMu);
     Queue.push_back(std::move(T));
   }
   QueueCV.notify_all();
 }
 
-// Resolves URI to file paths with cache.
-class URIToFileCache {
-public:
-  URIToFileCache(llvm::StringRef HintPath) : HintPath(HintPath) {}
-
-  llvm::StringRef resolve(llvm::StringRef FileURI) {
-    auto I = URIToPathCache.try_emplace(FileURI);
-    if (I.second) {
-      auto U = URI::parse(FileURI);
-      if (!U) {
-        elog("Failed to parse URI {0}: {1}", FileURI, U.takeError());
-        assert(false && "Failed to parse URI");
-        return "";
-      }
-      auto Path = URI::resolve(*U, HintPath);
-      if (!Path) {
-        elog("Failed to resolve URI {0}: {1}", FileURI, Path.takeError());
-        assert(false && "Failed to resolve URI");
-        return "";
-      }
-      I.first->second = *Path;
-    }
-    return I.first->second;
-  }
-
-private:
-  std::string HintPath;
-  llvm::StringMap<std::string> URIToPathCache;
-};
-
 /// Given index results from a TU, only update files in \p FilesToUpdate.
 void BackgroundIndex::update(StringRef MainFile, IndexFileIn Index,
                              const StringMap<FileDigest> &FilesToUpdate,
                              BackgroundIndexStorage *IndexStorage) {
   // Partition symbols/references into files.
   struct File {
     DenseSet<const Symbol *> Symbols;
     DenseSet<const Ref *> Refs;
   };
   StringMap<File> Files;
   URIToFileCache URICache(MainFile);
   for (const auto &Sym : *Index.Symbols) {
     if (Sym.CanonicalDeclaration) {
       auto DeclPath = URICache.resolve(Sym.CanonicalDeclaration.FileURI);
       if (FilesToUpdate.count(DeclPath) != 0)
         Files[DeclPath].Symbols.insert(&Sym);
     }
     // For symbols with different declaration and definition locations, we store
     // the full symbol in both the header file and the implementation file, so
     // that merging can tell the preferred symbols (from canonical headers) from
for (const auto &F : Files) {
     RefSlab::Builder Refs;
     for (const auto *S : F.second.Symbols)
       Syms.insert(*S);
     for (const auto *R : F.second.Refs)
       Refs.insert(RefToIDs[R], *R);
 
     auto SS = llvm::make_unique<SymbolSlab>(std::move(Syms).build());
     auto RS = llvm::make_unique<RefSlab>(std::move(Refs).build());
+    auto IG = llvm::make_unique<IncludeGraph>(
+        getSubGraph(URI::create(Path), Index.Sources.getValue()));
 
     auto Hash = FilesToUpdate.lookup(Path);
     // We need to store shards before updating the index, since the latter
     // consumes slabs.
     if (IndexStorage) {
       IndexFileOut Shard;
       Shard.Symbols = SS.get();
       Shard.Refs = RS.get();
+      Shard.Sources = IG.get();
 
       if (auto Error = IndexStorage->storeShard(Path, Shard))
         elog("Failed to write background-index shard for file {0}: {1}", Path,
              std::move(Error));
     }
 
     std::lock_guard<std::mutex> Lock(DigestsMu);
     // This can override a newer version that is added in another thread,
     // if this thread sees the older version but finishes later. This should be
     // rare in practice.
     IndexedFileDigests[Path] = Hash;
     IndexedSymbols.update(Path, std::move(SS), std::move(RS));
   }
 }
 
 // Creates a filter to not collect index results from files with unchanged
 // digests.
 // \p FileDigests contains file digests for the current indexed files, and all
 // changed files will be added to \p FilesToUpdate.
-decltype(SymbolCollector::Options::FileFilter) createFileFilter(
-    const llvm::StringMap<FileDigest> &FileDigests,
+decltype(SymbolCollector::Options::FileFilter)
+createFileFilter(const llvm::StringMap<FileDigest> &FileDigests,
                  llvm::StringMap<FileDigest> &FilesToUpdate) {
   return [&FileDigests, &FilesToUpdate](const SourceManager &SM, FileID FID) {
     StringRef Path;
     if (const auto *F = SM.getFileEntryForID(FID))
       Path = F->getName();
     if (Path.empty())
       return false; // Skip invalid files.
     SmallString<128> AbsPath(Path);
     if (std::error_code EC =
auto Clang = prepareCompilerInstance(
       std::make_shared<PCHContainerOperations>(), Inputs.FS, IgnoreDiags);
   if (!Clang)
     return createStringError(inconvertibleErrorCode(),
                              "Couldn't build compiler instance");
 
   SymbolCollector::Options IndexOpts;
   StringMap<FileDigest> FilesToUpdate;
   IndexOpts.FileFilter = createFileFilter(DigestsSnapshot, FilesToUpdate);
-  SymbolSlab Symbols;
-  RefSlab Refs;
+  IndexFileIn Index;
   auto Action = createStaticIndexingAction(
-      IndexOpts, [&](SymbolSlab S) { Symbols = std::move(S); },
-      [&](RefSlab R) { Refs = std::move(R); },
-      /*IncludeGraphCallback=*/nullptr);
+      IndexOpts, [&](SymbolSlab S) { Index.Symbols = std::move(S); },
+      [&](RefSlab R) { Index.Refs = std::move(R); },
+      [&](IncludeGraph IG) { Index.Sources = std::move(IG); });
 
   // We're going to run clang here, and it could potentially crash.
   // We could use CrashRecoveryContext to try to make indexing crashes nonfatal,
   // but the leaky "recovery" is pretty scary too in a long-running process.
   // If crashes are a real problem, maybe we should fork a child process.
 
   const FrontendInputFile &Input = Clang->getFrontendOpts().Inputs.front();
   if (!Action->BeginSourceFile(*Clang, Input))
     return createStringError(inconvertibleErrorCode(),
                              "BeginSourceFile() failed");
   if (!Action->Execute())
     return createStringError(inconvertibleErrorCode(), "Execute() failed");
   Action->EndSourceFile();
 
-  log("Indexed {0} ({1} symbols, {2} refs)", Inputs.CompileCommand.Filename,
-      Symbols.size(), Refs.numRefs());
-  SPAN_ATTACH(Tracer, "symbols", int(Symbols.size()));
-  SPAN_ATTACH(Tracer, "refs", int(Refs.numRefs()));
-  IndexFileIn Index;
-  Index.Symbols = std::move(Symbols);
-  Index.Refs = std::move(Refs);
+  log("Indexed {0} ({1} symbols, {2} refs, {3} sources)",
+      Inputs.CompileCommand.Filename, Index.Symbols->size(),
+      Index.Refs->numRefs(), Index.Sources->size());
+  SPAN_ATTACH(Tracer, "symbols", int(Index.Symbols->size()));
+  SPAN_ATTACH(Tracer, "refs", int(Index.Refs->numRefs()));
+  SPAN_ATTACH(Tracer, "sources", int(Index.Sources->size()));
 
   update(AbsolutePath, std::move(Index), FilesToUpdate, IndexStorage);
   {
     // Make sure hash for the main file is always updated even if there is no
     // index data in it.
     std::lock_guard<std::mutex> Lock(DigestsMu);
     IndexedFileDigests[AbsolutePath] = Hash;
   }

unittests/clangd/BackgroundIndexTests.cpp

TEST_F(BackgroundIndexTest, IndexTwoFiles) {
   OverlayCDB CDB(/*Base=*/nullptr);
   BackgroundIndex Idx(Context::empty(), "", FS, CDB,
                       [&](llvm::StringRef) { return &MSS; });
 
   tooling::CompileCommand Cmd;
   Cmd.Filename = testPath("root/A.cc");
   Cmd.Directory = testPath("root");
   Cmd.CommandLine = {"clang++", "-DA=1", testPath("root/A.cc")};
-  CDB.setCompileCommand(testPath("root"), Cmd);
+  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
 
   ASSERT_TRUE(Idx.blockUntilIdleForTest());
   EXPECT_THAT(
       runFuzzyFind(Idx, ""),
       UnorderedElementsAre(Named("common"), Named("A_CC"),
                            AllOf(Named("f_b"), Declared(), Not(Defined()))));
 
   Cmd.Filename = testPath("root/B.cc");
   Cmd.CommandLine = {"clang++", Cmd.Filename};
-  CDB.setCompileCommand(testPath("root"), Cmd);
+  CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
 
   ASSERT_TRUE(Idx.blockUntilIdleForTest());
   // B_CC is dropped as we don't collect symbols from A.h in this compilation.
   EXPECT_THAT(runFuzzyFind(Idx, ""),
               UnorderedElementsAre(Named("common"), Named("A_CC"),
                                    AllOf(Named("f_b"), Declared(), Defined())));
 
   auto Syms = runFuzzyFind(Idx, "common");
TEST_F(BackgroundIndexTest, ShardStorageWriteTest) {
   Cmd.Filename = testPath("root/A.cc");
   Cmd.Directory = testPath("root");
   Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
   // Check nothing is loaded from Storage, but A.cc and A.h has been stored.
   {
     OverlayCDB CDB(/*Base=*/nullptr);
     BackgroundIndex Idx(Context::empty(), "", FS, CDB,
                         [&](llvm::StringRef) { return &MSS; });
-    CDB.setCompileCommand(testPath("root"), Cmd);
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
     ASSERT_TRUE(Idx.blockUntilIdleForTest());
   }
   EXPECT_EQ(CacheHits, 0U);
   EXPECT_EQ(Storage.size(), 2U);
 
   auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
   EXPECT_NE(ShardHeader, nullptr);
   EXPECT_THAT(
       *ShardHeader->Symbols,
       UnorderedElementsAre(Named("common"), Named("A_CC"),
                            AllOf(Named("f_b"), Declared(), Not(Defined()))));
   for (const auto &Ref : *ShardHeader->Refs)
     EXPECT_THAT(Ref.second,
                 UnorderedElementsAre(FileURI("unittest:///root/A.h")));
 
   auto ShardSource = MSS.loadShard(testPath("root/A.cc"));
   EXPECT_NE(ShardSource, nullptr);
   EXPECT_THAT(*ShardSource->Symbols, UnorderedElementsAre());
   EXPECT_THAT(*ShardSource->Refs, RefsAre({FileURI("unittest:///root/A.cc")}));
 }
 
+TEST_F(BackgroundIndexTest, DirectIncludesTest) {
+  MockFSProvider FS;
+  FS.Files[testPath("root/B.h")] = "";
+  FS.Files[testPath("root/A.h")] = R"cpp(
+      #include "B.h"
+      void common();
+      void f_b();
+      class A_CC {};
+      )cpp";
+  std::string A_CC = "#include \"A.h\"\nvoid g() { (void)common; }";
+  FS.Files[testPath("root/A.cc")] = A_CC;
+
+  llvm::StringMap<std::string> Storage;
+  size_t CacheHits = 0;
+  MemoryShardStorage MSS(Storage, CacheHits);
+
+  tooling::CompileCommand Cmd;
+  Cmd.Filename = testPath("root/A.cc");
+  Cmd.Directory = testPath("root");
+  Cmd.CommandLine = {"clang++", testPath("root/A.cc")};
+  {
+    OverlayCDB CDB(/*Base=*/nullptr);
+    BackgroundIndex Idx(Context::empty(), "", FS, CDB,
+                        [&](llvm::StringRef) { return &MSS; });
+    CDB.setCompileCommand(testPath("root/A.cc"), Cmd);
+    ASSERT_TRUE(Idx.blockUntilIdleForTest());
+  }
+
+  auto ShardSource = MSS.loadShard(testPath("root/A.cc"));
+  EXPECT_TRUE(ShardSource->Sources);
+  EXPECT_EQ(ShardSource->Sources->size(), 2U); // A.cc, A.h
+  EXPECT_THAT(
+      ShardSource->Sources->lookup("unittest:///root/A.cc").DirectIncludes,
+      UnorderedElementsAre("unittest:///root/A.h"));
+
+  auto ShardHeader = MSS.loadShard(testPath("root/A.h"));
+  EXPECT_TRUE(ShardHeader->Sources);
+  EXPECT_EQ(ShardHeader->Sources->size(), 2U); // A.h B.h
+  EXPECT_THAT(
+      ShardHeader->Sources->lookup("unittest:///root/A.h").DirectIncludes,
+      UnorderedElementsAre("unittest:///root/B.h"));
+}
+
 } // namespace clangd
 } // namespace clang