[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

97

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)

112

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

331

nit: main-file

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

119

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

97

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

97

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 ↗

(On Diff #257003)

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

119–120

nit: header

147

nit: space before parenthetical

149

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

155

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()>?

162

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 ↗

(On Diff #257003)

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

155

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 ↗

(On Diff #257003)

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

155

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/FileIndex.h

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.
+  // FIXME: We ignore symbol resulting from different PP states while parsing a
+  // file and store only the last one.
   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);
 
 } // 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/Relation.h"
+#include "index/Symbol.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 <memory>
+#include <tuple>
 
 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));
 }
 
+/// Takes slabs coming from a TU(multiple files) and shards them per declaration
+/// location. Afterwards updates \p FileSyms.
+/// \p HintPath is used to convert file uris stored in symbols into absolute
+/// paths used by \p FileSyms.
+void shardSlabToFiles(const SymbolSlab &Syms, const RelationSlab &Rels,
+                      FileSymbols &FileSyms, llvm::StringRef HintPath) {
+  llvm::StringMap<SymbolSlab::Builder> FileToSyms;
+  llvm::StringMap<RelationSlab::Builder> FileToRels;
+  llvm::DenseMap<SymbolID, llvm::StringRef> SymbolIDToFile;
+  llvm::StringMap<std::string> URIToPathCache;
+  auto FileForSymbol = [&](const Symbol &S) -> llvm::StringRef {
+    llvm::StringRef FileURI = S.CanonicalDeclaration.FileURI;
+    // Symbols might contain empty URIs when they are not coming from a file,
+    // e.g. command line macros.
+    if (FileURI.empty())
+      return "";
+    auto I = URIToPathCache.try_emplace(FileURI);
+    if (I.second) {
+      auto URIPath = URI::resolve(FileURI, HintPath);
+      assert(URIPath);
+      I.first->second = *URIPath;
+    }
+    return I.first->second;
+  };
+  for (const auto &S : Syms) {
+    auto File = FileForSymbol(S);
+    FileToSyms[File].insert(S);
+    SymbolIDToFile[S.ID] = File;
+  }
+  for (const auto &R : Rels) {
+    auto File = SymbolIDToFile.lookup(R.Subject);
+    assert(!File.empty() && "unknown symbol in relation slab");
+    FileToRels[File].insert(R);
+  }
+  for (auto &Entry : FileToSyms) {
+    auto File = Entry.first();
+    auto &SymSlab = Entry.getValue();
+    auto &RelSlab = FileToRels[File];
+    FileSyms.update(File,
+                    std::make_unique<SymbolSlab>(std::move(SymSlab).build()),
+                    std::make_unique<RefSlab>(),
+                    std::make_unique<RelationSlab>(std::move(RelSlab).build()),
+                    /*CountReferences=*/false);
+  }
+}
+} // namespace
+
 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);
-  PreambleSymbols.update(
-      Path, std::make_unique<SymbolSlab>(std::move(std::get<0>(Slabs))),
-      std::make_unique<RefSlab>(),
-      std::make_unique<RelationSlab>(std::move(std::get<2>(Slabs))),
-      /*CountReferences=*/false);
+  SymbolSlab Syms;
+  RelationSlab Rels;
+  std::tie(Syms, std::ignore, Rels) =
+      indexHeaderSymbols(Version, AST, std::move(PP), Includes);
+  shardSlabToFiles(Syms, Rels, PreambleSymbols, Path);
   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 mainfile symbols with estimated memory usage "
+       "of {0} bytes",
+       MainFileIndex.estimateMemoryUsage());
 }
 
 } // namespace clangd
 } // namespace clang

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

 
 // 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);
 }

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;