[clangd] expose MergedIndex class

[sammccall]

Summary

This allows inheriting from it, so index() can ga away and allowing
TestTU::index) to be fixed.

[ioeric]

lgtm

[sammccall]

Rebase

clangd/ClangdServer.h

public:
   /// here, as this metric does not account (at least) for:
   ///   - memory occupied by static and dynamic index,
   ///   - memory required for in-flight requests,
   /// FIXME: those metrics might be useful too, we should add them.
   std::vector<std::pair<Path, std::size_t>> getUsedBytesPerFile() const;
 
   /// Returns the active dynamic index if one was built.
   /// This can be useful for testing, debugging, or observing memory usage.
-  const SymbolIndex *dynamicIndex() const;
+  const SymbolIndex *dynamicIndex() const { return DynamicIdx.get(); }
 
   // Blocks the main thread until the server is idle. Only for use in tests.
   // Returns false if the timeout expires.
   LLVM_NODISCARD bool
   blockUntilIdleForTest(llvm::Optional<double> TimeoutSeconds = 10);
 
 private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
private:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (DynamicIdx)
   //   - the static index passed to the constructor
   //   - a merged view of a static and dynamic index (MergedIndex)
   const SymbolIndex *Index;
   // If present, an index of symbols in open files. Read via *Index.
   std::unique_ptr<FileIndex> DynamicIdx;
   // If present, storage for the merged static/dynamic index. Read via *Index.
-  std::unique_ptr<SymbolIndex> MergedIndex;
+  std::unique_ptr<SymbolIndex> MergedIdx;
 
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   // If set, this represents the workspace path.
   llvm::Optional<std::string> RootPath;

clangd/ClangdServer.cpp

: CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       // is parsed.
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
                     DynamicIdx ? makeUpdateCallbacks(DynamicIdx.get())
                                : nullptr,
                     Opts.UpdateDebounce, Opts.RetentionPolicy) {
   if (DynamicIdx && Opts.StaticIndex) {
-    MergedIndex = mergeIndex(&DynamicIdx->index(), Opts.StaticIndex);
-    Index = MergedIndex.get();
+    MergedIdx =
+        llvm::make_unique<MergedIndex>(DynamicIdx.get(), Opts.StaticIndex);
+    Index = MergedIdx.get();
   } else if (DynamicIdx)
-    Index = &DynamicIdx->index();
+    Index = DynamicIdx.get();
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else
     Index = nullptr;
 }
 
-const SymbolIndex *ClangdServer::dynamicIndex() const {
-  return DynamicIdx ? &DynamicIdx->index() : nullptr;
-}
-
 void ClangdServer::setRootPath(PathRef RootPath) {
   auto FS = FSProvider.getFileSystem();
   auto Status = FS->status(RootPath);
   if (!Status)
     elog("Failed to get status for RootPath {0}: {1}", RootPath,
          Status.getError().message());
   else if (Status->isDirectory())
     this->RootPath = RootPath;

clangd/index/FileIndex.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "ClangdUnit.h"
 #include "Index.h"
 #include "MemIndex.h"
+#include "Merge.h"
 #include "clang/Lex/Preprocessor.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
 /// A container of Symbols from several source files. It can be updated
 /// at source-file granularity, replacing all symbols from one file with a new
private:
   /// Stores the latest symbol snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSymbols;
   /// Stores the latest ref snapshots for all active files.
   llvm::StringMap<std::shared_ptr<RefSlab>> FileToRefs;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
 /// FIXME: Expose an interface to remove files that are closed.
-class FileIndex {
+class FileIndex : public MergedIndex {
 public:
   /// If URISchemes is empty, the default schemes in SymbolCollector will be
   /// used.
   FileIndex(std::vector<std::string> URISchemes = {});
 
-  // Presents a merged view of the supplied main-file and preamble ASTs.
-  const SymbolIndex &index() const { return *MergedIndex; }
-
   /// Update preamble symbols of file \p Path with all declarations in \p AST
   /// and macros in \p PP.
   void updatePreamble(PathRef Path, ASTContext &AST,
                       std::shared_ptr<Preprocessor> PP);
 
   /// Update symbols and references from main file \p Path with
   /// `indexMainDecls`.
   void updateMain(PathRef Path, ParsedAST &AST);
private:
   // 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
   // (Note that symbols *only* in the main file are not indexed).
   FileSymbols MainFileSymbols;
   SwapIndex MainFileIndex;
-
-  std::unique_ptr<SymbolIndex> MergedIndex;  // Merge preamble and main index.
 };
 
 /// Retrieves symbols and refs of local top level decls in \p AST (i.e.
 /// `AST.getLocalTopLevelDecls()`).
 /// Exposed to assist in unit tests.
 /// If URISchemes is empty, the default schemes in SymbolCollector will be used.
 std::pair<SymbolSlab, RefSlab>
 indexMainDecls(ParsedAST &AST, llvm::ArrayRef<std::string> URISchemes = {});

clangd/index/FileIndex.cpp

std::unique_ptr<SymbolIndex> FileSymbols::buildMemIndex() {
   return llvm::make_unique<MemIndex>(
       llvm::make_pointee_range(AllSymbols), std::move(AllRefs),
       std::make_tuple(std::move(SymbolSlabs), std::move(RefSlabs),
                       std::move(RefsStorage)),
       StorageSize);
 }
 
 FileIndex::FileIndex(std::vector<std::string> URISchemes)
-    : URISchemes(std::move(URISchemes)),
+    : MergedIndex(&MainFileIndex, &PreambleIndex),
+      URISchemes(std::move(URISchemes)),
       PreambleIndex(PreambleSymbols.buildMemIndex()),
-      MainFileIndex(MainFileSymbols.buildMemIndex()),
-      MergedIndex(mergeIndex(&MainFileIndex, &PreambleIndex)) {}
+      MainFileIndex(MainFileSymbols.buildMemIndex()) {}
 
 void FileIndex::updatePreamble(PathRef Path, ASTContext &AST,
                                std::shared_ptr<Preprocessor> PP) {
   auto Symbols = indexHeaderSymbols(AST, std::move(PP), URISchemes);
   PreambleSymbols.update(Path,
                          llvm::make_unique<SymbolSlab>(std::move(Symbols)),
                          llvm::make_unique<RefSlab>());
   PreambleIndex.reset(PreambleSymbols.buildMemIndex());

clangd/index/Merge.h

 namespace clang {
 namespace clangd {
 
 // Merge symbols L and R, preferring data from L in case of conflict.
 // The two symbols must have the same ID.
 // Returned symbol may contain data owned by either source.
 Symbol mergeSymbol(const Symbol &L, const Symbol &R);
 
-// mergeIndex returns a composite index based on two provided Indexes:
+// MergedIndex is a composite index based on two provided Indexes:
 //  - the Dynamic index covers few files, but is relatively up-to-date.
 //  - the Static index covers a bigger set of files, but is relatively stale.
 // The returned index attempts to combine results, and avoid duplicates.
 //
 // FIXME: We don't have a mechanism in Index to track deleted symbols and
 // refs in dirty files, so the merged index may return stale symbols
 // and refs from Static index.
-std::unique_ptr<SymbolIndex> mergeIndex(const SymbolIndex *Dynamic,
-                                        const SymbolIndex *Static);
+class MergedIndex : public SymbolIndex {
+  const SymbolIndex *Dynamic, *Static;
+
+public:
+  // The constructor does not access the symbols.
+  // It's safe to inherit from this class and pass pointers to derived members.
+  MergedIndex(const SymbolIndex *Dynamic, const SymbolIndex *Static)
+      : Dynamic(Dynamic), Static(Static) {}
+
+  bool fuzzyFind(const FuzzyFindRequest &,
+                 llvm::function_ref<void(const Symbol &)>) const override;
+  void lookup(const LookupRequest &,
+              llvm::function_ref<void(const Symbol &)>) const override;
+  void refs(const RefsRequest &,
+            llvm::function_ref<void(const Ref &)>) const override;
+  size_t estimateMemoryUsage() const override {
+    return Dynamic->estimateMemoryUsage() + Static->estimateMemoryUsage();
+  }
+};
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_MERGE_H

clangd/index/Merge.cpp

 #include "Trace.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/raw_ostream.h"
 #include <set>
 
 namespace clang {
 namespace clangd {
-namespace {
 
 using namespace llvm;
 
-class MergedIndex : public SymbolIndex {
- public:
-   MergedIndex(const SymbolIndex *Dynamic, const SymbolIndex *Static)
-       : Dynamic(Dynamic), Static(Static) {}
-
-   // FIXME: Deleted symbols in dirty files are still returned (from Static).
-   //        To identify these eliminate these, we should:
-   //          - find the generating file from each Symbol which is Static-only
-   //          - ask Dynamic if it has that file (needs new SymbolIndex method)
-   //          - if so, drop the Symbol.
-   bool fuzzyFind(const FuzzyFindRequest &Req,
-                  function_ref<void(const Symbol &)> Callback) const override {
-     // We can't step through both sources in parallel. So:
-     //  1) query all dynamic symbols, slurping results into a slab
-     //  2) query the static symbols, for each one:
-     //    a) if it's not in the dynamic slab, yield it directly
-     //    b) if it's in the dynamic slab, merge it and yield the result
-     //  3) now yield all the dynamic symbols we haven't processed.
-     trace::Span Tracer("MergedIndex fuzzyFind");
-     bool More = false; // We'll be incomplete if either source was.
-     SymbolSlab::Builder DynB;
-     unsigned DynamicCount = 0;
-     unsigned StaticCount = 0;
-     unsigned MergedCount = 0;
-     More |= Dynamic->fuzzyFind(Req, [&](const Symbol &S) {
-       ++DynamicCount;
-       DynB.insert(S);
-     });
-     SymbolSlab Dyn = std::move(DynB).build();
+bool MergedIndex::fuzzyFind(const FuzzyFindRequest &Req,
+                            function_ref<void(const Symbol &)> Callback) const {
 
-     DenseSet<SymbolID> SeenDynamicSymbols;
-     More |= Static->fuzzyFind(Req, [&](const Symbol &S) {
-       auto DynS = Dyn.find(S.ID);
-       ++StaticCount;
-       if (DynS == Dyn.end())
-         return Callback(S);
-       ++MergedCount;
-       SeenDynamicSymbols.insert(S.ID);
-       Callback(mergeSymbol(*DynS, S));
-     });
-     SPAN_ATTACH(Tracer, "dynamic", DynamicCount);
-     SPAN_ATTACH(Tracer, "static", StaticCount);
-     SPAN_ATTACH(Tracer, "merged", MergedCount);
-     for (const Symbol &S : Dyn)
-       if (!SeenDynamicSymbols.count(S.ID))
-         Callback(S);
-     return More;
-  }
 
-  void
-  lookup(const LookupRequest &Req,
-         llvm::function_ref<void(const Symbol &)> Callback) const override {
-    trace::Span Tracer("MergedIndex lookup");
-    SymbolSlab::Builder B;
+void MergedIndex::lookup(
+    const LookupRequest &Req,
+    llvm::function_ref<void(const Symbol &)> Callback) const {
 
-    Dynamic->lookup(Req, [&](const Symbol &S) { B.insert(S); });
 
-    auto RemainingIDs = Req.IDs;
-    Static->lookup(Req, [&](const Symbol &S) {
-      const Symbol *Sym = B.find(S.ID);
-      RemainingIDs.erase(S.ID);
-      if (!Sym)
-        Callback(S);
-      else
-        Callback(mergeSymbol(*Sym, S));
-    });
-    for (const auto &ID : RemainingIDs)
-      if (const Symbol *Sym = B.find(ID))
-        Callback(*Sym);
-  }
 
-  void refs(const RefsRequest &Req,
-            llvm::function_ref<void(const Ref &)> Callback) const override {
-    trace::Span Tracer("MergedIndex refs");
-    // We don't want duplicated refs from the static/dynamic indexes,
-    // and we can't reliably duplicate them because offsets may differ slightly.
-    // We consider the dynamic index authoritative and report all its refs,
-    // and only report static index refs from other files.
-    //
-    // FIXME: The heuristic fails if the dynamic index contains a file, but all
-    // refs were removed (we will report stale ones from the static index).
-    // Ultimately we should explicit check which index has the file instead.
-    llvm::StringSet<> DynamicIndexFileURIs;
-    Dynamic->refs(Req, [&](const Ref &O) {
-      DynamicIndexFileURIs.insert(O.Location.FileURI);
-      Callback(O);
-    });
-    Static->refs(Req, [&](const Ref &O) {
-      if (!DynamicIndexFileURIs.count(O.Location.FileURI))
-        Callback(O);
-    });
-  }
+void MergedIndex::refs(const RefsRequest &Req,
+                       llvm::function_ref<void(const Ref &)> Callback) const {
 
-  size_t estimateMemoryUsage() const override {
-    return Dynamic->estimateMemoryUsage() + Static->estimateMemoryUsage();
-  }
-
-private:
-  const SymbolIndex *Dynamic, *Static;
-};
-} // namespace
-
 Symbol mergeSymbol(const Symbol &L, const Symbol &R) {
   assert(L.ID == R.ID);
   // We prefer information from TUs that saw the definition.
   // Classes: this is the def itself. Functions: hopefully the header decl.
   // If both did (or both didn't), continue to prefer L over R.
   bool PreferR = R.Definition && !L.Definition;
   // Merge include headers only if both have definitions or both have no
   // definition; otherwise, only accumulate references of common includes.
if (!Found && MergeIncludes)
       S.IncludeHeaders.emplace_back(OI.IncludeHeader, OI.References);
   }
 
   S.Origin |= O.Origin | SymbolOrigin::Merge;
   S.Flags |= O.Flags;
   return S;
 }
 
-std::unique_ptr<SymbolIndex> mergeIndex(const SymbolIndex *Dynamic,
-                                        const SymbolIndex *Static) {
-  return llvm::make_unique<MergedIndex>(Dynamic, Static);
-}
-
 } // namespace clangd
 } // namespace clang

unittests/clangd/FileIndexTests.cpp

TEST(FileSymbolsTest, SnapshotAliveAfterRemove) {
   auto Empty = FS.buildMemIndex();
   EXPECT_THAT(getSymbolNames(*Empty), UnorderedElementsAre());
   EXPECT_THAT(getRefs(*Empty, ID), ElementsAre());
 
   EXPECT_THAT(getSymbolNames(*Symbols), UnorderedElementsAre("1", "2", "3"));
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 }
 
-std::vector<std::string> match(const FileIndex &I,
+std::vector<std::string> match(const SymbolIndex &I,
                                const FuzzyFindRequest &Req) {
   std::vector<std::string> Matches;
-  I.index().fuzzyFind(Req, [&](const Symbol &Sym) {
+  I.fuzzyFind(Req, [&](const Symbol &Sym) {
     Matches.push_back((Sym.Scope + Sym.Name).str());
   });
   return Matches;
 }
 
 // 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 = Code;
   auto AST = File.build();
   M.updatePreamble(File.Filename, AST.getASTContext(),
                    AST.getPreprocessorPtr());
 }
 
 TEST(FileIndexTest, CustomizedURIScheme) {
   FileIndex M({"unittest"});
   update(M, "f", "class string {};");
 
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenSymbol = false;
-  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_EQ(Sym.CanonicalDeclaration.FileURI, "unittest:///f.h");
     SeenSymbol = true;
   });
   EXPECT_TRUE(SeenSymbol);
 }
 
 TEST(FileIndexTest, IndexAST) {
   FileIndex M;

 
 TEST(FileIndexTest, NoIncludeCollected) {
   FileIndex M;
   update(M, "f", "class string {};");
 
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenSymbol = false;
-  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_TRUE(Sym.IncludeHeaders.empty());
     SeenSymbol = true;
   });
   EXPECT_TRUE(SeenSymbol);
 }
 
 TEST(FileIndexTest, TemplateParamsInLabel) {
   auto Source = R"cpp(
 template <class Ty>
 class vector {
 };
 
 template <class Ty, class Arg>
 vector<Ty> make_vector(Arg A) {}
 )cpp";
 
   FileIndex M;
   update(M, "f", Source);
 
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenVector = false;
   bool SeenMakeVector = false;
-  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
     if (Sym.Name == "vector") {
       EXPECT_EQ(Sym.Signature, "<class Ty>");
       EXPECT_EQ(Sym.CompletionSnippetSuffix, "<${1:class Ty}>");
       SeenVector = true;
       return;
     }
 
     if (Sym.Name == "make_vector") {
TEST(FileIndexTest, Refs) {
   // Add test2.cc
   TestTU Test2;
   Test2.HeaderCode = HeaderCode;
   Test2.Code = MainCode.code();
   Test2.Filename = "test2.cc";
   AST = Test2.build();
   Index.updateMain(Test2.Filename, AST);
 
-  EXPECT_THAT(getRefs(Index.index(), Foo.ID),
+  EXPECT_THAT(getRefs(Index, Foo.ID),
               RefsAre({AllOf(RefRange(MainCode.range("foo")),
                              FileURI("unittest:///test.cc")),
                        AllOf(RefRange(MainCode.range("foo")),
                              FileURI("unittest:///test2.cc"))}));
 }
 
 TEST(FileIndexTest, CollectMacros) {
   FileIndex M;
   update(M, "f", "#define CLANGD 1");
 
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenSymbol = false;
-  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_EQ(Sym.Name, "CLANGD");
     EXPECT_EQ(Sym.SymInfo.Kind, index::SymbolKind::Macro);
     SeenSymbol = true;
   });
   EXPECT_TRUE(SeenSymbol);
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

unittests/clangd/IndexTests.cpp

TEST(MemIndexTest, Lookup) {
   EXPECT_THAT(lookup(*I, {SymbolID("ns::nonono"), SymbolID("ns::xyz")}),
               UnorderedElementsAre("ns::xyz"));
   EXPECT_THAT(lookup(*I, SymbolID("ns::nonono")), UnorderedElementsAre());
 }
 
 TEST(MergeIndexTest, Lookup) {
   auto I = MemIndex::build(generateSymbols({"ns::A", "ns::B"}), RefSlab()),
        J = MemIndex::build(generateSymbols({"ns::B", "ns::C"}), RefSlab());
-  auto M = mergeIndex(I.get(), J.get());
-  EXPECT_THAT(lookup(*M, SymbolID("ns::A")), UnorderedElementsAre("ns::A"));
-  EXPECT_THAT(lookup(*M, SymbolID("ns::B")), UnorderedElementsAre("ns::B"));
-  EXPECT_THAT(lookup(*M, SymbolID("ns::C")), UnorderedElementsAre("ns::C"));
-  EXPECT_THAT(lookup(*M, {SymbolID("ns::A"), SymbolID("ns::B")}),
+  MergedIndex M(I.get(), J.get());
+  EXPECT_THAT(lookup(M, SymbolID("ns::A")), UnorderedElementsAre("ns::A"));
+  EXPECT_THAT(lookup(M, SymbolID("ns::B")), UnorderedElementsAre("ns::B"));
+  EXPECT_THAT(lookup(M, SymbolID("ns::C")), UnorderedElementsAre("ns::C"));
+  EXPECT_THAT(lookup(M, {SymbolID("ns::A"), SymbolID("ns::B")}),
               UnorderedElementsAre("ns::A", "ns::B"));
-  EXPECT_THAT(lookup(*M, {SymbolID("ns::A"), SymbolID("ns::C")}),
+  EXPECT_THAT(lookup(M, {SymbolID("ns::A"), SymbolID("ns::C")}),
               UnorderedElementsAre("ns::A", "ns::C"));
-  EXPECT_THAT(lookup(*M, SymbolID("ns::D")), UnorderedElementsAre());
-  EXPECT_THAT(lookup(*M, {}), UnorderedElementsAre());
+  EXPECT_THAT(lookup(M, SymbolID("ns::D")), UnorderedElementsAre());
+  EXPECT_THAT(lookup(M, {}), UnorderedElementsAre());
 }
 
 TEST(MergeIndexTest, FuzzyFind) {
   auto I = MemIndex::build(generateSymbols({"ns::A", "ns::B"}), RefSlab()),
        J = MemIndex::build(generateSymbols({"ns::B", "ns::C"}), RefSlab());
   FuzzyFindRequest Req;
   Req.Scopes = {"ns::"};
-  EXPECT_THAT(match(*mergeIndex(I.get(), J.get()), Req),
+  EXPECT_THAT(match(MergedIndex(I.get(), J.get()), Req),
               UnorderedElementsAre("ns::A", "ns::B", "ns::C"));
 }
 
 TEST(MergeTest, Merge) {
   Symbol L, R;
   L.ID = R.ID = SymbolID("hello");
   L.Name = R.Name = "Foo";                           // same in both
   L.CanonicalDeclaration.FileURI = "file:///left.h"; // differs
TEST(MergeTest, PreferSymbolWithDefn) {
   EXPECT_EQ(M.CanonicalDeclaration.FileURI, "file:/right.h");
   EXPECT_EQ(M.Definition.FileURI, "file:/right.cpp");
   EXPECT_EQ(M.Name, "right");
 }
 
 TEST(MergeIndexTest, Refs) {
   FileIndex Dyn({"unittest"});
   FileIndex StaticIndex({"unittest"});
-  auto MergedIndex = mergeIndex(&Dyn.index(), &StaticIndex.index());
+  MergedIndex Merge(&Dyn, &StaticIndex);
 
   const char *HeaderCode = "class Foo;";
   auto HeaderSymbols = TestTU::withHeaderCode("class Foo;").headerSymbols();
   auto Foo = findSymbol(HeaderSymbols, "Foo");
 
   // Build dynamic index for test.cc.
   Annotations Test1Code(R"(class $Foo[[Foo]];)");
   TestTU Test;
TEST(MergeIndexTest, Refs) {
   Test2.Code = Test2Code.code();
   Test2.Filename = "test2.cc";
   StaticAST = Test2.build();
   StaticIndex.updateMain(Test2.Filename, StaticAST);
 
   RefsRequest Request;
   Request.IDs = {Foo.ID};
   RefSlab::Builder Results;
-  MergedIndex->refs(Request, [&](const Ref &O) { Results.insert(Foo.ID, O); });
+  Merge.refs(Request, [&](const Ref &O) { Results.insert(Foo.ID, O); });
 
   EXPECT_THAT(
       std::move(Results).build(),
       ElementsAre(Pair(
           _, UnorderedElementsAre(AllOf(RefRange(Test1Code.range("Foo")),
                                         FileURI("unittest:///test.cc")),
                                   AllOf(RefRange(Test2Code.range("Foo")),
                                         FileURI("unittest:///test2.cc"))))));

unittests/clangd/TestTU.cpp

ParsedAST TestTU::build() const {
   return std::move(*AST);
 }
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
   return indexHeaderSymbols(AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
-// FIXME: This should return a FileIndex with both preamble and main index.
 std::unique_ptr<SymbolIndex> TestTU::index() const {
   auto AST = build();
-  auto Content = indexMainDecls(AST);
-  return MemIndex::build(std::move(Content.first), std::move(Content.second));
+  auto Idx = llvm::make_unique<FileIndex>();
+  Idx->updatePreamble(Filename, AST.getASTContext(), AST.getPreprocessorPtr());
+  Idx->updateMain(Filename, AST);
+  return 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;
     if (Result) {