[clangd] Shard preamble symbols in dynamic index

[kadircet]

Summary

This reduces memory usage by dynamic index from more than 400MB to 32MB
when all files in clang-tools-extra/clangd/*.cpp are active in clangd.

[sammccall]

Can't believe we didn't do this before. Nice catch.

clang-tools-extra/clangd/index/FileIndex.cpp

103

This screams out to be shared with BackgroundIndex::update(). Is it very hard to do?

Things that are different:

• no duplication of symbol into defining file (why?)
• refs not gathered here (but I guess empty input -> empty output?)
• include graph not gathered here (same)
• this holds 3 copies of all the data at peak, other impl holds 2

I imagine factoring out a function that returns a StringMap<File> like in BackgroundIndex::Update, but File has a build() function that returns an IndexFileIn or so.
(Ooh, is depending on IndexFileIn a layering violation? Maybe we should move that into another header. Duplicating the struct is probably OK for now)

118

um... I guess that works, because Expected is unchecked in NDEBUG?
I think cantFail is clearer though.

422

nit: main-file

clang-tools-extra/clangd/index/FileIndex.h

129–130

This seems like a caveat rather than a FIXME - we're not really proposing a fix are we?

[kadircet]

clang-tools-extra/clangd/index/FileIndex.cpp

103

In backgroundindex we "prefilter" symbols and only duplicate the ones coming from modified files, whereas in this one we duplicate every symbol no matter what. I didn't want to change backgroundindex into post filtering mode (I suppose that's what you mean by the last item).

as for the first bullet point, duplicating in case of preamble symbols wouldn't matter since we don't merge(rather pick one) while building the index for preamblesymbols. I suppose storing them twice shouldn't be a huge issue.

Depending to IndexFileIn in "Index.h" would be a violation, but depending on it in "FileIndex.h" is fine.

I am happy to refactor a StringMap<File> shardIndexToFiles(IndexFileIn, HintPath) into "FileIndex.h" that can be used by both backgroundindex and this, if you are OK with above mentioned regressions.

[kadircet]

• Unify sharding logic in BackgroundIndex and FileIndex.
• Make sure relations have valid subjects.

[kadircet]

clang-tools-extra/clangd/index/FileIndex.cpp

103

implemented a version of this that only copies data when requested. Hence the pre/post filtering argument is gone.

[kadircet]

clang-tools-extra/clangd/index/SymbolCollector.cpp

286

put this before processRelations to ensure Subject of a relation is a symbol we are going to collect.

[sammccall]

LG, API quibbles.

It would be nice to test this but the good things about it (memory usage) aren't easy to test, and the bad things (sensitivity to header parsing context) seem odd to be the primary test.
Actually I guess one nice thing we get from this is erasure of stale symbols: if we index the preamble {foo.cpp=#include "x.h", x.h=int a;}, and then {bar.cpp=#include "x.h", x.h=int b;}, then I think before this patch we see both a and b in the index, and after we see only b.
That would be a reasonable test.

Or I guess now that the sharding function is public we can test it directly.

clang-tools-extra/clangd/index/FileIndex.h

129–130

nit: header

157

nit: space before parenthetical

159

Hmm, this is a good name for a function, but we want a noun for a class. FileShardedIndex?

165

I do find it a little weird not to expose the map-structure of the vast majority of the data here.

What steers you away from just making this a function
StringMap<X> shardIndexToFile(IndexFileIn, PathRef)
where X could be File with a method to obtain the data, or the more anonymous unique_function<IndexFileIn()>?

172

Can we bundle these parallel functions into one that returns IndexFileIn?
I don't think there's much efficiency in building only a subset, is there? (If there's a lot of data in e.g. refs we're going to ignore, we can strip them from the input)

That would also mean no need to special-case mainfilecmd in the API - the only caller of that function does it in the file loop, so it would be provided for the main-file shard only.

clang-tools-extra/clangd/index/SymbolCollector.cpp

286

The code around this seems unchanged - how does this relate to the rest of the patch?
(Other than I guess you've been inspecting the data and finding bogus shards :-)

This is bit subtle, and hard to reason about since we only have one example. Indeed dangling relations are pointless. However:

• the direction of the relation is fairly arbitrary. Should we be checking shouldCollectSymbol on the object of the relation too?
• shouldCollectSymbols depends somewhat on the file context, and also the indexer options. In principle the subject could be indexable elsewhere but not here. Or this index could see the relation and another could see the symbol itself. Can't think of compelling examples though, just hypotheticals.

I think probably we should give this a comment (about not wanting relations or refs either), add filtering on the relation objects, and add a test (I think that should be pretty easy with inheriting from a main-file-only class?).
Maybe it's a separate patch?

[kadircet]

• Address comments
• Merge slab generations into a single member that returns an IndexFileIn.

[kadircet]

clang-tools-extra/clangd/index/FileIndex.h

165

because the map structure is just pointers into the IndexFileIn, I wanted to make sure those pointers do not outlive the IndexFileIn.

clang-tools-extra/clangd/index/SymbolCollector.cpp

286

The code around this seems unchanged - how does this relate to the rest of the patch?
(Other than I guess you've been inspecting the data and finding bogus shards :-)

Selection of file declaring the Subject seemed arbitrary, and turned out to be broken. We were not hitting this case before because background index was performing pre filtering and skipping any relation that mapped to a "non interesting" file, hence we didn't notice it was broken. Now it become obvious as we assert on the file containing the subject symbol.

the direction of the relation is fairly arbitrary. Should we be checking shouldCollectSymbol on the object of the relation too?
shouldCollectSymbols depends somewhat on the file context, and also the indexer options. In principle the subject could be indexable elsewhere but not here. Or this index could see the relation and another could see the symbol itself. Can't think of compelling examples though, just hypotheticals.
I think probably we should give this a comment (about not wanting relations or refs either), add filtering on the relation objects, and add a test (I think that should be pretty easy with inheriting from a main-file-only class?).
Maybe it's a separate patch?

I also wanted to do this but it got hairy and as you also have suggested I decided to have a separate patch for this. Because there were other problematic stuff, for example IndexerOpts::CollectMainFileSymbols is always true.

[sammccall]

Still LG

clang-tools-extra/clangd/index/FileIndex.h

165

I don't think it's unreasonable to have a non-owning data structure with a comment, but up to you.

[kadircet]

• Add tests for sharding logic and preamble overwrite

clang-tools-extra/clangd/index/Background.cpp

 #include <thread>
 #include <utility>
 #include <vector>
 
 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 Path = URI::resolve(FileURI, 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 "U" and its edges. Any node other than "U" will be
-// empty in the resultant graph.
-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();
-
-  // URIs inside nodes must point into the keys of the same IncludeGraph.
-  for (auto &Include : Node.DirectIncludes) {
-    auto I = IG.try_emplace(Include).first;
-    I->getValue().URI = I->getKey();
-    Include = I->getKey();
-  }
-
-  return IG;
-}
-
 // We cannot use vfs->makeAbsolute because Cmd.FileName is either absolute or
 // relative to Cmd.Directory, which might not be the same as current working
 // directory.
 llvm::SmallString<128> getAbsolutePath(const tooling::CompileCommand &Cmd) {
   llvm::SmallString<128> AbsolutePath;
   if (llvm::sys::path::is_absolute(Cmd.Filename)) {
     AbsolutePath = Cmd.Filename;
   } else {

 
 /// Given index results from a TU, only update symbols coming from files that
 /// are different or missing from than \p ShardVersionsSnapshot. Also stores new
 /// index information on IndexStorage.
 void BackgroundIndex::update(
     llvm::StringRef MainFile, IndexFileIn Index,
     const llvm::StringMap<ShardVersion> &ShardVersionsSnapshot,
     bool HadErrors) {
-  // Partition symbols/references into files.
-  struct File {
-    llvm::DenseSet<const Symbol *> Symbols;
-    llvm::DenseSet<const Ref *> Refs;
-    llvm::DenseSet<const Relation *> Relations;
-    FileDigest Digest;
-  };
-  llvm::StringMap<File> Files;
-  URIToFileCache URICache(MainFile);
+  llvm::StringMap<FileDigest> FilesToUpdate;
   for (const auto &IndexIt : *Index.Sources) {
     const auto &IGN = IndexIt.getValue();
     // Note that sources do not contain any information regarding missing
     // headers, since we don't even know what absolute path they should fall in.
-    const auto AbsPath = URICache.resolve(IGN.URI);
+    auto AbsPath = llvm::cantFail(URI::resolve(IGN.URI, MainFile),
+                                  "Failed to resovle URI");
     const auto DigestIt = ShardVersionsSnapshot.find(AbsPath);
     // File has different contents, or indexing was successful this time.
     if (DigestIt == ShardVersionsSnapshot.end() ||
         DigestIt->getValue().Digest != IGN.Digest ||
         (DigestIt->getValue().HadErrors && !HadErrors))
-      Files.try_emplace(AbsPath).first->getValue().Digest = IGN.Digest;
-  }
-  // This map is used to figure out where to store relations.
-  llvm::DenseMap<SymbolID, File *> SymbolIDToFile;
-  for (const auto &Sym : *Index.Symbols) {
-    if (Sym.CanonicalDeclaration) {
-      auto DeclPath = URICache.resolve(Sym.CanonicalDeclaration.FileURI);
-      const auto FileIt = Files.find(DeclPath);
-      if (FileIt != Files.end()) {
-        FileIt->second.Symbols.insert(&Sym);
-        SymbolIDToFile[Sym.ID] = &FileIt->second;
-      }
-    }
-    // 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
-    // other symbols (e.g. forward declarations).
-    if (Sym.Definition &&
-        Sym.Definition.FileURI != Sym.CanonicalDeclaration.FileURI) {
-      auto DefPath = URICache.resolve(Sym.Definition.FileURI);
-      const auto FileIt = Files.find(DefPath);
-      if (FileIt != Files.end())
-        FileIt->second.Symbols.insert(&Sym);
-    }
-  }
-  llvm::DenseMap<const Ref *, SymbolID> RefToIDs;
-  for (const auto &SymRefs : *Index.Refs) {
-    for (const auto &R : SymRefs.second) {
-      auto Path = URICache.resolve(R.Location.FileURI);
-      const auto FileIt = Files.find(Path);
-      if (FileIt != Files.end()) {
-        auto &F = FileIt->getValue();
-        RefToIDs[&R] = SymRefs.first;
-        F.Refs.insert(&R);
-      }
-    }
-  }
-  for (const auto &Rel : *Index.Relations) {
-    const auto FileIt = SymbolIDToFile.find(Rel.Subject);
-    if (FileIt != SymbolIDToFile.end())
-      FileIt->second->Relations.insert(&Rel);
+      FilesToUpdate[AbsPath] = IGN.Digest;
   }
 
-  // Build and store new slabs for each updated file.
-  for (const auto &FileIt : Files) {
-    llvm::StringRef Path = FileIt.getKey();
-    SymbolSlab::Builder Syms;
-    RefSlab::Builder Refs;
-    RelationSlab::Builder Relations;
-    for (const auto *S : FileIt.second.Symbols)
-      Syms.insert(*S);
-    for (const auto *R : FileIt.second.Refs)
-      Refs.insert(RefToIDs[R], *R);
-    for (const auto *Rel : FileIt.second.Relations)
-      Relations.insert(*Rel);
-    auto SS = std::make_unique<SymbolSlab>(std::move(Syms).build());
-    auto RS = std::make_unique<RefSlab>(std::move(Refs).build());
-    auto RelS = std::make_unique<RelationSlab>(std::move(Relations).build());
-    auto IG = std::make_unique<IncludeGraph>(
-        getSubGraph(URI::create(Path), Index.Sources.getValue()));
+  // Shard slabs into files.
+  FileShardedIndex ShardedIndex(std::move(Index), MainFile);
 
-    // We need to store shards before updating the index, since the latter
-    // consumes slabs.
-    // FIXME: Also skip serializing the shard if it is already up-to-date.
-    BackgroundIndexStorage *IndexStorage = IndexStorageFactory(Path);
-    IndexFileOut Shard;
-    Shard.Symbols = SS.get();
-    Shard.Refs = RS.get();
-    Shard.Relations = RelS.get();
-    Shard.Sources = IG.get();
+  // Build and store new slabs for each updated file.
+  for (const auto &FileIt : FilesToUpdate) {
+    PathRef Path = FileIt.first();
+    auto IF = ShardedIndex.getShard(Path);
 
     // Only store command line hash for main files of the TU, since our
     // current model keeps only one version of a header file.
-    if (Path == MainFile)
-      Shard.Cmd = Index.Cmd.getPointer();
+    if (Path != MainFile)
+      IF.Cmd.reset();
 
-    if (auto Error = IndexStorage->storeShard(Path, Shard))
+    // We need to store shards before updating the index, since the latter
+    // consumes slabs.
+    // FIXME: Also skip serializing the shard if it is already up-to-date.
+    if (auto Error = IndexStorageFactory(Path)->storeShard(Path, IF))
       elog("Failed to write background-index shard for file {0}: {1}", Path,
            std::move(Error));
 
     {
       std::lock_guard<std::mutex> Lock(ShardVersionsMu);
-      auto Hash = FileIt.second.Digest;
+      const auto &Hash = FileIt.getValue();
       auto DigestIt = ShardVersions.try_emplace(Path);
       ShardVersion &SV = DigestIt.first->second;
       // Skip if file is already up to date, unless previous index was broken
       // and this one is not.
       if (!DigestIt.second && SV.Digest == Hash && SV.HadErrors && !HadErrors)
         continue;
       SV.Digest = Hash;
       SV.HadErrors = HadErrors;
 
       // 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.
-      IndexedSymbols.update(Path, std::move(SS), std::move(RS), std::move(RelS),
-                            Path == MainFile);
+      IndexedSymbols.update(
+          Path, std::make_unique<SymbolSlab>(std::move(*IF.Symbols)),
+          std::make_unique<RefSlab>(std::move(*IF.Refs)),
+          std::make_unique<RelationSlab>(std::move(*IF.Relations)),
     }
   }
 }
 
 llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
   trace::Span Tracer("BackgroundIndex");
   SPAN_ATTACH(Tracer, "file", Cmd.Filename);
   auto AbsolutePath = getAbsolutePath(Cmd);

clang-tools-extra/clangd/index/FileIndex.h

 // maintained at source-file granularity (e.g. with ASTs), and files can be
 // updated dynamically.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
+#include "Headers.h"
 #include "Index.h"
 #include "MemIndex.h"
 #include "Merge.h"
 #include "Path.h"
 #include "index/CanonicalIncludes.h"
+#include "index/Ref.h"
+#include "index/Relation.h"
+#include "index/Serialization.h"
 #include "index/Symbol.h"
 #include "clang/Lex/Preprocessor.h"
+#include "clang/Tooling/CompilationDatabase.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
 #include <memory>
+#include <vector>
 
 namespace clang {
 class ASTContext;
 namespace clangd {
 class ParsedAST;
 
 /// Select between in-memory index implementations, which have tradeoffs.
 enum class IndexType {
public:
 
   /// Update symbols and references from main file \p Path with
   /// `indexMainDecls`.
   void updateMain(PathRef Path, ParsedAST &AST);
 
 private:
   bool UseDex; // FIXME: this should be always on.
 
-  // Contains information from each file's preamble only.
-  // These are large, but update fairly infrequently (preambles are stable).
+  // Contains information from each file's preamble only. Symbols and relations
+  // are sharded per declaration file to deduplicate multiple symbols and reduce
+  // memory usage.
   // Missing information:
   //  - symbol refs (these are always "from the main file")
   //  - definition locations in the main file
   //
-  // FIXME: Because the preambles for different TUs have large overlap and
-  // FileIndex doesn't deduplicate, this uses lots of extra RAM.
-  // The biggest obstacle in fixing this: the obvious approach of partitioning
-  // by declaring file (rather than main file) fails if headers provide
-  // different symbols based on preprocessor state.
+  // Note that we store only one version of a header, hence symbols appearing in
+  // different PP states will be missing.
   FileSymbols PreambleSymbols;
   SwapIndex PreambleIndex;
 
   // Contains information from each file's main AST.
   // These are updated frequently (on file change), but are relatively small.
   // Mostly contains:
   //  - refs to symbols declared in the preamble and referenced from main
   //  - symbols declared both in the main file and the preamble

 SlabTuple indexMainDecls(ParsedAST &AST);
 
 /// Index declarations from \p AST and macros from \p PP that are declared in
 /// included headers.
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
                              std::shared_ptr<Preprocessor> PP,
                              const CanonicalIncludes &Includes);
 
+/// Takes slabs coming from a TU (multiple files) and shards them per
+/// declaration location.
+struct FileShardedIndex {
+  /// \p HintPath is used to convert file URIs stored in symbols into absolute
+  /// paths.
+  explicit FileShardedIndex(IndexFileIn Input, PathRef HintPath);
+
+  /// Returns absolute paths for all files that has a shard.
+  std::vector<PathRef> getAllFiles() const;
+
+  /// Generates index shard for the \p File. Note that this function results in
+  /// a copy of all the relevant data.
+  /// Returned index will always have Symbol/Refs/Relation Slabs set, even if
+  /// they are empty.
+  IndexFileIn getShard(PathRef File) const;
+
+private:
+  // Contains all the information that belongs to a single file.
+  struct FileShard {
+    // Either declared or defined in the file.
+    llvm::DenseSet<const Symbol *> Symbols;
+    // Reference occurs in the file.
+    llvm::DenseSet<const Ref *> Refs;
+    // Subject is declared in the file.
+    llvm::DenseSet<const Relation *> Relations;
+    // Contains edges for only the direct includes.
+    IncludeGraph IG;
+  };
+
+  // Keeps all the information alive.
+  const IndexFileIn Index;
+  // Mapping from absolute paths to slab information.
+  llvm::StringMap<FileShard> Shards;
+  // Used to build RefSlabs.
+  llvm::DenseMap<const Ref *, SymbolID> RefToSymID;
+};
+
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clang-tools-extra/clangd/index/FileIndex.cpp

 //===--- FileIndex.cpp - Indexes for files. ------------------------ C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
 #include "CollectMacros.h"
 #include "Logger.h"
 #include "ParsedAST.h"
+#include "Path.h"
 #include "SymbolCollector.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
+#include "index/Ref.h"
+#include "index/Relation.h"
+#include "index/Serialization.h"
+#include "index/Symbol.h"
+#include "index/SymbolID.h"
 #include "index/SymbolOrigin.h"
 #include "index/dex/Dex.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Index/IndexingOptions.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
 #include <memory>
+#include <tuple>
+#include <utility>
+#include <vector>
 
 namespace clang {
 namespace clangd {
+namespace {
 
-static SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-                              llvm::ArrayRef<Decl *> DeclsToIndex,
-                              const MainFileMacros *MacroRefsToIndex,
-                              const CanonicalIncludes &Includes,
-                              bool IsIndexMainAST, llvm::StringRef Version) {
+SlabTuple indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                       const CanonicalIncludes &Includes, bool IsIndexMainAST,
+                       llvm::StringRef Version) {
   SymbolCollector::Options CollectorOpts;
   CollectorOpts.CollectIncludePath = true;
   CollectorOpts.Includes = &Includes;
   CollectorOpts.CountReferences = false;
   CollectorOpts.Origin = SymbolOrigin::Dynamic;
 
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
vlog("indexed {0} AST for {1} version {2}:\n"
        "  relations slab: {8} relations, {9} bytes",
        IsIndexMainAST ? "file" : "preamble", FileName, Version, Syms.size(),
        Syms.bytes(), Refs.size(), Refs.numRefs(), Refs.bytes(),
        Relations.size(), Relations.bytes());
   return std::make_tuple(std::move(Syms), std::move(Refs),
                          std::move(Relations));
 }
 
+// Resolves URI to file paths with cache.
+class URIToFileCache {
+public:
+  URIToFileCache(PathRef HintPath) : HintPath(HintPath) {}
+
+  llvm::StringRef operator[](llvm::StringRef FileURI) {
+    if (FileURI.empty())
+      return "";
+    auto I = URIToPathCache.try_emplace(FileURI);
+    if (I.second) {
+      I.first->second = llvm::cantFail(URI::resolve(FileURI, HintPath),
+                                       "Failed to resolve URI");
+    }
+    return I.first->second;
+  }
+
+private:
+  PathRef HintPath;
+  llvm::StringMap<std::string> URIToPathCache;
+};
+
+// We keep only the node "U" and its edges. Any node other than "U" will be
+// empty in the resultant graph.
+IncludeGraph getSubGraph(llvm::StringRef URI, const IncludeGraph &FullGraph) {
+  IncludeGraph IG;
+
+  auto Entry = IG.try_emplace(URI).first;
+  auto &Node = Entry->getValue();
+  Node = FullGraph.lookup(Entry->getKey());
+  Node.URI = Entry->getKey();
+
+  // URIs inside nodes must point into the keys of the same IncludeGraph.
+  for (auto &Include : Node.DirectIncludes) {
+    auto I = IG.try_emplace(Include).first;
+    I->getValue().URI = I->getKey();
+    Include = I->getKey();
+  }
+  return IG;
+}
+} // namespace
+
+FileShardedIndex::FileShardedIndex(IndexFileIn Input, PathRef HintPath)
+    : Index(std::move(Input)) {
+  URIToFileCache UriToFile(HintPath);
+  // Used to build RelationSlabs.
+  llvm::DenseMap<SymbolID, FileShard *> SymbolIDToFile;
+
+  // Attribute each Symbol to both their declaration and definition locations.
+  if (Index.Symbols) {
+    for (const auto &S : *Index.Symbols) {
+      auto File = UriToFile[S.CanonicalDeclaration.FileURI];
+      auto It = Shards.try_emplace(File);
+      It.first->getValue().Symbols.insert(&S);
+      SymbolIDToFile[S.ID] = &It.first->getValue();
+      // Only bother if definition file is different than declaration file.
+      if (S.Definition &&
+          S.Definition.FileURI != S.CanonicalDeclaration.FileURI) {
+        auto File = UriToFile[S.Definition.FileURI];
+        auto It = Shards.try_emplace(File);
+        It.first->getValue().Symbols.insert(&S);
+      }
+    }
+  }
+  // Attribute references into each file they occured in.
+  if (Index.Refs) {
+    for (const auto &SymRefs : *Index.Refs) {
+      for (const auto &R : SymRefs.second) {
+        auto File = UriToFile[R.Location.FileURI];
+        const auto It = Shards.try_emplace(File);
+        It.first->getValue().Refs.insert(&R);
+        RefToSymID[&R] = SymRefs.first;
+      }
+    }
+  }
+  // Attribute relations to the file declaraing their Subject as Object might
+  // not have been indexed, see SymbolCollector::processRelations for details.
+  if (Index.Relations) {
+    for (const auto &R : *Index.Relations) {
+      auto *File = SymbolIDToFile.lookup(R.Subject);
+      assert(File && "unknown subject in relation");
+      File->Relations.insert(&R);
+    }
+  }
+  // Store only the direct includes of a file in a shard.
+  if (Index.Sources) {
+    const auto &FullGraph = *Index.Sources;
+    for (const auto &It : FullGraph) {
+      auto File = UriToFile[It.first()];
+      auto ShardIt = Shards.try_emplace(File);
+      ShardIt.first->getValue().IG = getSubGraph(It.first(), FullGraph);
+    }
+  }
+}
+std::vector<PathRef> FileShardedIndex::getAllFiles() const {
+  return {Shards.keys().begin(), Shards.keys().end()};
+}
+
+IndexFileIn FileShardedIndex::getShard(PathRef File) const {
+  auto It = Shards.find(File);
+  assert(It != Shards.end() && "received unknown file");
+
+  IndexFileIn IF;
+  IF.Sources = It->getValue().IG;
+  IF.Cmd = Index.Cmd;
+
+  SymbolSlab::Builder SymB;
+  for (const auto *S : It->getValue().Symbols)
+    SymB.insert(*S);
+  IF.Symbols = std::move(SymB).build();
+
+  RefSlab::Builder RefB;
+  for (const auto *Ref : It->getValue().Refs) {
+    auto SID = RefToSymID.lookup(Ref);
+    RefB.insert(SID, *Ref);
+  }
+  IF.Refs = std::move(RefB).build();
+
+  RelationSlab::Builder RelB;
+  for (const auto *Rel : It->getValue().Relations) {
+    RelB.insert(*Rel);
+  }
+  IF.Relations = std::move(RelB).build();
+  return IF;
+}
+
 SlabTuple indexMainDecls(ParsedAST &AST) {
   return indexSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
                       AST.getLocalTopLevelDecls(), &AST.getMacros(),
                       AST.getCanonicalIncludes(),
                       /*IsIndexMainAST=*/true, AST.version());
 }
 
 SlabTuple indexHeaderSymbols(llvm::StringRef Version, ASTContext &AST,
FileIndex::FileIndex(bool UseDex)
     : MergedIndex(&MainFileIndex, &PreambleIndex), UseDex(UseDex),
       PreambleIndex(std::make_unique<MemIndex>()),
       MainFileIndex(std::make_unique<MemIndex>()) {}
 
 void FileIndex::updatePreamble(PathRef Path, llvm::StringRef Version,
                                ASTContext &AST,
                                std::shared_ptr<Preprocessor> PP,
                                const CanonicalIncludes &Includes) {
-  auto Slabs = indexHeaderSymbols(Version, AST, std::move(PP), Includes);
+  IndexFileIn IF;
+  std::tie(IF.Symbols, std::ignore, IF.Relations) =
+      indexHeaderSymbols(Version, AST, std::move(PP), Includes);
+  FileShardedIndex ShardedIndex(std::move(IF), Path);
+  for (PathRef File : ShardedIndex.getAllFiles()) {
+    auto IF = ShardedIndex.getShard(File);
     PreambleSymbols.update(
-      Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Slabs))),
+        File, std::make_unique<SymbolSlab>(std::move(*IF.Symbols)),
         std::make_unique<RefSlab>(),
-      std::make_unique<RelationSlab>(std::move(std::get<2>(Slabs))),
+        std::make_unique<RelationSlab>(std::move(*IF.Relations)),
         /*CountReferences=*/false);
+  }
   PreambleIndex.reset(
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
                                  DuplicateHandling::PickOne));
+  vlog("Build dynamic index for header symbols with estimated memory usage of "
+       "{0} bytes",
+       PreambleIndex.estimateMemoryUsage());
 }
 
 void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
   auto Contents = indexMainDecls(AST);
   MainFileSymbols.update(
       Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Contents))),
       std::make_unique<RefSlab>(std::move(std::get<1>(Contents))),
       std::make_unique<RelationSlab>(std::move(std::get<2>(Contents))),
       /*CountReferences=*/true);
   MainFileIndex.reset(
       MainFileSymbols.buildIndex(IndexType::Light, DuplicateHandling::Merge));
+  vlog("Build dynamic index for main-file symbols with estimated memory usage "
+       "of {0} bytes",
+       MainFileIndex.estimateMemoryUsage());
 }
 
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/index/SymbolCollector.cpp

if (Opts.CountReferences &&
       (Roles & static_cast<unsigned>(index::SymbolRole::Reference)) &&
       SM.getFileID(SM.getSpellingLoc(Loc)) == SM.getMainFileID())
     ReferencedDecls.insert(ND);
 
   auto ID = getSymbolID(ND);
   if (!ID)
     return true;
 
+  // ND is the canonical (i.e. first) declaration. If it's in the main file
+  // (which is not a header), then no public declaration was visible, so assume
+  // it's main-file only.
+  bool IsMainFileOnly =
+      SM.isWrittenInMainFile(SM.getExpansionLoc(ND->getBeginLoc())) &&
+      !isHeaderFile(SM.getFileEntryForID(SM.getMainFileID())->getName(),
+                    ASTCtx->getLangOpts());
+  // In C, printf is a redecl of an implicit builtin! So check OrigD instead.
+  if (ASTNode.OrigD->isImplicit() ||
+      !shouldCollectSymbol(*ND, *ASTCtx, Opts, IsMainFileOnly))
+    return true;
+
   // Note: we need to process relations for all decl occurrences, including
   // refs, because the indexing code only populates relations for specific
   // occurrences. For example, RelationBaseOf is only populated for the
   // occurrence inside the base-specifier.
   processRelations(*ND, *ID, Relations);
 
   bool CollectRef = static_cast<bool>(Opts.RefFilter & toRefKind(Roles));
   bool IsOnlyRef =
       !(Roles & (static_cast<unsigned>(index::SymbolRole::Declaration) |
                  static_cast<unsigned>(index::SymbolRole::Definition)));
 
   if (IsOnlyRef && !CollectRef)
     return true;
 
-  // ND is the canonical (i.e. first) declaration. If it's in the main file
-  // (which is not a header), then no public declaration was visible, so assume
-  // it's main-file only.
-  bool IsMainFileOnly =
-      SM.isWrittenInMainFile(SM.getExpansionLoc(ND->getBeginLoc())) &&
-      !isHeaderFile(SM.getFileEntryForID(SM.getMainFileID())->getName(),
-                    ASTCtx->getLangOpts());
-  // In C, printf is a redecl of an implicit builtin! So check OrigD instead.
-  if (ASTNode.OrigD->isImplicit() ||
-      !shouldCollectSymbol(*ND, *ASTCtx, Opts, IsMainFileOnly))
-    return true;
   // Do not store references to main-file symbols.
   // Unlike other fields, e.g. Symbols (which use spelling locations), we use
   // file locations for references (as it aligns the behavior of clangd's
   // AST-based xref).
   // FIXME: we should try to use the file locations for other fields.
   if (CollectRef && !IsMainFileOnly && !isa<NamespaceDecl>(ND) &&
       (Opts.RefsInHeaders ||
        SM.getFileID(SM.getFileLoc(Loc)) == SM.getMainFileID()))

clang-tools-extra/clangd/unittests/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AST.h"
 #include "Annotations.h"
 #include "Compiler.h"
+#include "Headers.h"
 #include "ParsedAST.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "URI.h"
 #include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
+#include "index/Ref.h"
+#include "index/Relation.h"
+#include "index/Serialization.h"
+#include "index/Symbol.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <utility>
 
 using ::testing::_;
 using ::testing::AllOf;
 using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::IsEmpty;
 using ::testing::Pair;
 using ::testing::UnorderedElementsAre;

 
 // Adds Basename.cpp, which includes Basename.h, which contains Code.
 void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   TestTU File;
   File.Filename = (Basename + ".cpp").str();
   File.HeaderFilename = (Basename + ".h").str();
   File.HeaderCode = std::string(Code);
   auto AST = File.build();
-  M.updatePreamble(File.Filename, /*Version=*/"null", AST.getASTContext(),
-                   AST.getPreprocessorPtr(), AST.getCanonicalIncludes());
+  M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
+                   AST.getASTContext(), AST.getPreprocessorPtr(),
+                   AST.getCanonicalIncludes());
 }
 
 TEST(FileIndexTest, CustomizedURIScheme) {
   FileIndex M;
   update(M, "f", "class string {};");
 
   EXPECT_THAT(runFuzzyFind(M, ""), ElementsAre(DeclURI("unittest:///f.h")));
 }

 
 TEST(FileIndexTest, Relations) {
   TestTU TU;
   TU.Filename = "f.cpp";
   TU.HeaderFilename = "f.h";
   TU.HeaderCode = "class A {}; class B : public A {};";
   auto AST = TU.build();
   FileIndex Index;
-  Index.updatePreamble(TU.Filename, /*Version=*/"null", AST.getASTContext(),
-                       AST.getPreprocessorPtr(), AST.getCanonicalIncludes());
+  Index.updatePreamble(testPath(TU.Filename), /*Version=*/"null",
+                       AST.getASTContext(), AST.getPreprocessorPtr(),
+                       AST.getCanonicalIncludes());
   SymbolID A = findSymbol(TU.headerSymbols(), "A").ID;
   uint32_t Results = 0;
   RelationsRequest Req;
   Req.Subjects.insert(A);
   Req.Predicate = RelationKind::BaseOf;
   Index.relations(Req, [&](const SymbolID &, const Symbol &) { ++Results; });
   EXPECT_EQ(Results, 1u);
 }
FS.update("f3h", numSlab(1, 3), refSlab(SymbolID("3"), "f3.h"), nullptr,
             false);
   EXPECT_THAT(
       runFuzzyFind(*FS.buildIndex(IndexType::Light, DuplicateHandling::Merge),
                    ""),
       UnorderedElementsAre(AllOf(QName("1"), NumReferences(1u)),
                            AllOf(QName("2"), NumReferences(1u)),
                            AllOf(QName("3"), NumReferences(1u))));
 }
+
+TEST(FileIndexTest, StalePreambleSymbolsDeleted) {
+  FileIndex M;
+  TestTU File;
+  File.HeaderFilename = "a.h";
+
+  File.Filename = "f1.cpp";
+  File.HeaderCode = "int a;";
+  auto AST = File.build();
+  M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
+                   AST.getASTContext(), AST.getPreprocessorPtr(),
+                   AST.getCanonicalIncludes());
+  EXPECT_THAT(runFuzzyFind(M, ""), UnorderedElementsAre(QName("a")));
+
+  File.Filename = "f2.cpp";
+  File.HeaderCode = "int b;";
+  AST = File.build();
+  M.updatePreamble(testPath(File.Filename), /*Version=*/"null",
+                   AST.getASTContext(), AST.getPreprocessorPtr(),
+                   AST.getCanonicalIncludes());
+  EXPECT_THAT(runFuzzyFind(M, ""), UnorderedElementsAre(QName("b")));
+}
+
+TEST(FileShardedIndexTest, Sharding) {
+  auto AHeaderUri = URI::create(testPath("a.h")).toString();
+  auto BHeaderUri = URI::create(testPath("b.h")).toString();
+  auto BSourceUri = URI::create(testPath("b.cc")).toString();
+
+  auto Sym1 = symbol("1");
+  Sym1.CanonicalDeclaration.FileURI = AHeaderUri.c_str();
+
+  auto Sym2 = symbol("2");
+  Sym2.CanonicalDeclaration.FileURI = BHeaderUri.c_str();
+  Sym2.Definition.FileURI = BSourceUri.c_str();
+
+  IndexFileIn IF;
+  {
+    SymbolSlab::Builder B;
+    // Should be stored in only a.h
+    B.insert(Sym1);
+    // Should be stored in both b.h and b.cc
+    B.insert(Sym2);
+    IF.Symbols = std::move(B).build();
+  }
+  {
+    // Should be stored in b.cc
+    IF.Refs = std::move(*refSlab(Sym1.ID, BSourceUri.c_str()).release());
+  }
+  {
+    RelationSlab::Builder B;
+    // Should be stored in a.h
+    B.insert(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID});
+    // Should be stored in b.h
+    B.insert(Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID});
+    IF.Relations = std::move(B).build();
+  }
+
+  IF.Sources.emplace();
+  IncludeGraph &IG = *IF.Sources;
+  {
+    // b.cc includes b.h
+    auto &Node = IG[BSourceUri];
+    Node.DirectIncludes = {BHeaderUri};
+    Node.URI = BSourceUri;
+  }
+  {
+    // b.h includes a.h
+    auto &Node = IG[BHeaderUri];
+    Node.DirectIncludes = {AHeaderUri};
+    Node.URI = BHeaderUri;
+  }
+  {
+    // a.h includes nothing.
+    auto &Node = IG[AHeaderUri];
+    Node.DirectIncludes = {};
+    Node.URI = AHeaderUri;
+  }
+
+  IF.Cmd = tooling::CompileCommand(testRoot(), "b.cc", {"clang"}, "out");
+
+  FileShardedIndex ShardedIndex(std::move(IF), testPath("b.cc"));
+  ASSERT_THAT(
+      ShardedIndex.getAllFiles(),
+      UnorderedElementsAre(testPath("a.h"), testPath("b.h"), testPath("b.cc")));
+
+  {
+    auto Shard = ShardedIndex.getShard(testPath("a.h"));
+    EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("1")));
+    EXPECT_THAT(Shard.Refs.getValue(), IsEmpty());
+    EXPECT_THAT(
+        Shard.Relations.getValue(),
+        UnorderedElementsAre(Relation{Sym1.ID, RelationKind::BaseOf, Sym2.ID}));
+    ASSERT_THAT(Shard.Sources.getValue().keys(),
+                UnorderedElementsAre(AHeaderUri));
+    EXPECT_THAT(Shard.Sources->lookup(AHeaderUri).DirectIncludes, IsEmpty());
+    EXPECT_TRUE(Shard.Cmd.hasValue());
+  }
+  {
+    auto Shard = ShardedIndex.getShard(testPath("b.h"));
+    EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("2")));
+    EXPECT_THAT(Shard.Refs.getValue(), IsEmpty());
+    EXPECT_THAT(
+        Shard.Relations.getValue(),
+        UnorderedElementsAre(Relation{Sym2.ID, RelationKind::BaseOf, Sym1.ID}));
+    ASSERT_THAT(Shard.Sources.getValue().keys(),
+                UnorderedElementsAre(BHeaderUri, AHeaderUri));
+    EXPECT_THAT(Shard.Sources->lookup(BHeaderUri).DirectIncludes,
+                UnorderedElementsAre(AHeaderUri));
+    EXPECT_TRUE(Shard.Cmd.hasValue());
+  }
+  {
+    auto Shard = ShardedIndex.getShard(testPath("b.cc"));
+    EXPECT_THAT(Shard.Symbols.getValue(), UnorderedElementsAre(QName("2")));
+    EXPECT_THAT(Shard.Refs.getValue(), UnorderedElementsAre(Pair(Sym1.ID, _)));
+    EXPECT_THAT(Shard.Relations.getValue(), IsEmpty());
+    ASSERT_THAT(Shard.Sources.getValue().keys(),
+                UnorderedElementsAre(BSourceUri, BHeaderUri));
+    EXPECT_THAT(Shard.Sources->lookup(BSourceUri).DirectIncludes,
+                UnorderedElementsAre(BHeaderUri));
+    EXPECT_TRUE(Shard.Cmd.hasValue());
+  }
+}
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/TestTU.cpp

SymbolSlab TestTU::headerSymbols() const {
   return std::get<0>(indexHeaderSymbols(/*Version=*/"null", AST.getASTContext(),
                                         AST.getPreprocessorPtr(),
                                         AST.getCanonicalIncludes()));
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
   auto AST = build();
   auto Idx = std::make_unique<FileIndex>(/*UseDex=*/true);
-  Idx->updatePreamble(Filename, /*Version=*/"null", AST.getASTContext(),
-                      AST.getPreprocessorPtr(), AST.getCanonicalIncludes());
-  Idx->updateMain(Filename, AST);
+  Idx->updatePreamble(testPath(Filename), /*Version=*/"null",
+                      AST.getASTContext(), AST.getPreprocessorPtr(),
+                      AST.getCanonicalIncludes());
+  Idx->updateMain(testPath(Filename), AST);
   return std::move(Idx);
 }
 
 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;