[clangd] Merge ClangdServer::DynamicIndex into FileIndex. NFC.

Summary:
FileIndex now provides explicit interfaces for preamble and main file updates.
This avoids growing parameter list when preamble and main symbols diverge
further (e.g. D52078). This also gets rid of the hack in `indexAST` that
inferred main file index based on `TopLevelDecls`.

Also separate `indexMainDecls` from `indexAST`.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, kadircet, cfe-commits

Differential Revision: https://reviews.llvm.org/D52222

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@342460 91177308-0d34-0410-b5e6-96231b3b80d8
diff --git a/clangd/ClangdServer.cpp b/clangd/ClangdServer.cpp
index fb4f1be..6f2d13e 100644
--- a/clangd/ClangdServer.cpp
+++ b/clangd/ClangdServer.cpp
@@ -14,6 +14,7 @@
 #include "SourceCode.h"
 #include "Trace.h"
 #include "XRefs.h"
+#include "index/FileIndex.h"
 #include "index/Merge.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
@@ -70,59 +71,23 @@
 };
 } // namespace
 
-/// The dynamic index tracks symbols visible in open files.
-/// For boring reasons, it doesn't implement SymbolIndex directly - use index().
-class ClangdServer::DynamicIndex {
-public:
-  DynamicIndex(std::vector<std::string> URISchemes)
-      : PreambleIdx(URISchemes), MainFileIdx(URISchemes),
-        MergedIndex(mergeIndex(&MainFileIdx, &PreambleIdx)) {}
+// Returns callbacks that can be used to update the FileIndex with new ASTs.
+std::unique_ptr<ParsingCallbacks> makeUpdateCallbacks(FileIndex *FIndex) {
+  struct CB : public ParsingCallbacks {
+    CB(FileIndex *FIndex) : FIndex(FIndex) {}
+    FileIndex *FIndex;
 
-  const SymbolIndex &index() const { return *MergedIndex; }
+    void onPreambleAST(PathRef Path, ASTContext &Ctx,
+                       std::shared_ptr<clang::Preprocessor> PP) override {
+      FIndex->updatePreamble(Path, Ctx, std::move(PP));
+    }
 
-  // Returns callbacks that can be used to update the index with new ASTs.
-  // Index() presents a merged view of the supplied main-file and preamble ASTs.
-  std::unique_ptr<ParsingCallbacks> makeUpdateCallbacks() {
-    struct CB : public ParsingCallbacks {
-      CB(ClangdServer::DynamicIndex *This) : This(This) {}
-      DynamicIndex *This;
-
-      void onPreambleAST(PathRef Path, ASTContext &Ctx,
-                         std::shared_ptr<clang::Preprocessor> PP) override {
-        This->PreambleIdx.update(Path, &Ctx, std::move(PP));
-      }
-
-      void onMainAST(PathRef Path, ParsedAST &AST) override {
-        This->MainFileIdx.update(Path, &AST.getASTContext(),
-                                 AST.getPreprocessorPtr(),
-                                 AST.getLocalTopLevelDecls());
-      }
-    };
-    return llvm::make_unique<CB>(this);
+    void onMainAST(PathRef Path, ParsedAST &AST) override {
+      FIndex->updateMain(Path, AST, AST.getLocalTopLevelDecls());
+    }
   };
-
-private:
-  // Contains information from each file's preamble only.
-  // These are large, but update fairly infrequently (preambles are stable).
-  // 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.
-  FileIndex PreambleIdx;
-  // 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).
-  FileIndex MainFileIdx;
-  std::unique_ptr<SymbolIndex> MergedIndex;
-};
+  return llvm::make_unique<CB>(FIndex);
+}
 
 ClangdServer::Options ClangdServer::optsForTest() {
   ClangdServer::Options Opts;
@@ -139,9 +104,8 @@
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
                                    : getStandardResourceDir()),
-      DynamicIdx(Opts.BuildDynamicSymbolIndex
-                     ? new DynamicIndex(Opts.URISchemes)
-                     : nullptr),
+      DynamicIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex(Opts.URISchemes)
+                                              : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
@@ -149,7 +113,8 @@
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
-                    DynamicIdx ? DynamicIdx->makeUpdateCallbacks() : nullptr,
+                    DynamicIdx ? makeUpdateCallbacks(DynamicIdx.get())
+                               : nullptr,
                     Opts.UpdateDebounce, Opts.RetentionPolicy) {
   if (DynamicIdx && Opts.StaticIndex) {
     MergedIndex = mergeIndex(&DynamicIdx->index(), Opts.StaticIndex);
@@ -162,10 +127,6 @@
     Index = nullptr;
 }
 
-// Destructor has to be in .cpp file to see the definition of
-// ClangdServer::DynamicIndex.
-ClangdServer::~ClangdServer() = default;
-
 const SymbolIndex *ClangdServer::dynamicIndex() const {
   return DynamicIdx ? &DynamicIdx->index() : nullptr;
 }
diff --git a/clangd/ClangdServer.h b/clangd/ClangdServer.h
index afbda4b..1706ebb 100644
--- a/clangd/ClangdServer.h
+++ b/clangd/ClangdServer.h
@@ -111,7 +111,6 @@
   /// synchronize access to shared state.
   ClangdServer(GlobalCompilationDatabase &CDB, FileSystemProvider &FSProvider,
                DiagnosticsConsumer &DiagConsumer, const Options &Opts);
-  ~ClangdServer();
 
   /// Set the root path of the workspace.
   void setRootPath(PathRef RootPath);
@@ -221,7 +220,6 @@
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
-  class DynamicIndex;
   typedef uint64_t DocVersion;
 
   void consumeDiagnostics(PathRef File, DocVersion Version,
@@ -244,7 +242,7 @@
   //   - 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<DynamicIndex> DynamicIdx;
+  std::unique_ptr<FileIndex> DynamicIdx;
   // If present, storage for the merged static/dynamic index. Read via *Index.
   std::unique_ptr<SymbolIndex> MergedIndex;
 
diff --git a/clangd/index/FileIndex.cpp b/clangd/index/FileIndex.cpp
index 6b54674..ad7f634 100644
--- a/clangd/index/FileIndex.cpp
+++ b/clangd/index/FileIndex.cpp
@@ -8,18 +8,22 @@
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
+#include "ClangdUnit.h"
 #include "Logger.h"
 #include "SymbolCollector.h"
+#include "index/Index.h"
+#include "index/Merge.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Preprocessor.h"
+#include <memory>
 
 namespace clang {
 namespace clangd {
 
-std::pair<SymbolSlab, RefSlab>
-indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-         llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls,
-         llvm::ArrayRef<std::string> URISchemes) {
+static std::pair<SymbolSlab, RefSlab>
+indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+             llvm::ArrayRef<Decl *> DeclsToIndex, bool IsIndexMainAST,
+             llvm::ArrayRef<std::string> URISchemes) {
   SymbolCollector::Options CollectorOpts;
   // FIXME(ioeric): we might also want to collect include headers. We would need
   // to make sure all includes are canonicalized (with CanonicalIncludes), which
@@ -28,9 +32,9 @@
   // CommentHandler for IWYU pragma) to canonicalize includes.
   CollectorOpts.CollectIncludePath = false;
   CollectorOpts.CountReferences = false;
+  CollectorOpts.Origin = SymbolOrigin::Dynamic;
   if (!URISchemes.empty())
     CollectorOpts.URISchemes = URISchemes;
-  CollectorOpts.Origin = SymbolOrigin::Dynamic;
 
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
@@ -38,18 +42,8 @@
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
-  std::vector<Decl *> DeclsToIndex;
-  if (TopLevelDecls)
-    DeclsToIndex.assign(TopLevelDecls->begin(), TopLevelDecls->end());
-  else
-    DeclsToIndex.assign(AST.getTranslationUnitDecl()->decls().begin(),
-                        AST.getTranslationUnitDecl()->decls().end());
-
-  // We only collect refs when indexing main AST.
-  // FIXME: this is a hacky way to detect whether we are indexing preamble AST
-  // or main AST, we should make it explicitly.
-  bool IsIndexMainAST = TopLevelDecls.hasValue();
   if (IsIndexMainAST)
+    // We only collect refs when indexing main AST.
     CollectorOpts.RefFilter = RefKind::All;
 
   SymbolCollector Collector(std::move(CollectorOpts));
@@ -62,17 +56,30 @@
 
   auto Syms = Collector.takeSymbols();
   auto Refs = Collector.takeRefs();
-  vlog("index {0}AST for {1}: \n"
+  vlog("index AST for {0} (main={1}): \n"
        "  symbol slab: {2} symbols, {3} bytes\n"
        "  ref slab: {4} symbols, {5} bytes",
-       IsIndexMainAST ? "Main" : "Preamble", FileName, Syms.size(),
-       Syms.bytes(), Refs.size(), Refs.bytes());
+       FileName, IsIndexMainAST, Syms.size(), Syms.bytes(), Refs.size(),
+       Refs.bytes());
   return {std::move(Syms), std::move(Refs)};
 }
 
-FileIndex::FileIndex(std::vector<std::string> URISchemes)
-    : URISchemes(std::move(URISchemes)) {
-  reset(FSymbols.buildMemIndex());
+std::pair<SymbolSlab, RefSlab>
+indexMainDecls(ParsedAST &AST, llvm::ArrayRef<Decl *> TopLevelDecls,
+               llvm::ArrayRef<std::string> URISchemes) {
+  return indexSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
+                      TopLevelDecls,
+                      /*IsIndexMainAST=*/true, URISchemes);
+}
+
+SymbolSlab indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                              llvm::ArrayRef<std::string> URISchemes) {
+  std::vector<Decl *> DeclsToIndex(
+      AST.getTranslationUnitDecl()->decls().begin(),
+      AST.getTranslationUnitDecl()->decls().end());
+  return indexSymbols(AST, std::move(PP), DeclsToIndex,
+                      /*IsIndexMainAST=*/false, URISchemes)
+      .first;
 }
 
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Symbols,
@@ -141,19 +148,28 @@
       StorageSize);
 }
 
-void FileIndex::update(PathRef Path, ASTContext *AST,
-                       std::shared_ptr<Preprocessor> PP,
-                       llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls) {
-  if (!AST) {
-    FSymbols.update(Path, nullptr, nullptr);
-  } else {
-    assert(PP);
-    auto Contents = indexAST(*AST, PP, TopLevelDecls, URISchemes);
-    FSymbols.update(Path,
-                    llvm::make_unique<SymbolSlab>(std::move(Contents.first)),
-                    llvm::make_unique<RefSlab>(std::move(Contents.second)));
-  }
-  reset(FSymbols.buildMemIndex());
+FileIndex::FileIndex(std::vector<std::string> URISchemes)
+    : URISchemes(std::move(URISchemes)),
+      PreambleIndex(PreambleSymbols.buildMemIndex()),
+      MainFileIndex(MainFileSymbols.buildMemIndex()),
+      MergedIndex(mergeIndex(&MainFileIndex, &PreambleIndex)) {}
+
+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());
+}
+
+void FileIndex::updateMain(PathRef Path, ParsedAST &AST,
+                           llvm::ArrayRef<Decl *> TopLevelDecls) {
+  auto Contents = indexMainDecls(AST, TopLevelDecls, URISchemes);
+  MainFileSymbols.update(
+      Path, llvm::make_unique<SymbolSlab>(std::move(Contents.first)),
+      llvm::make_unique<RefSlab>(std::move(Contents.second)));
+  MainFileIndex.reset(MainFileSymbols.buildMemIndex());
 }
 
 } // namespace clangd
diff --git a/clangd/index/FileIndex.h b/clangd/index/FileIndex.h
index 3b0b6fa..421cfa4 100644
--- a/clangd/index/FileIndex.h
+++ b/clangd/index/FileIndex.h
@@ -20,6 +20,7 @@
 #include "Index.h"
 #include "MemIndex.h"
 #include "clang/Lex/Preprocessor.h"
+#include <memory>
 
 namespace clang {
 namespace clangd {
@@ -57,39 +58,66 @@
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
-class FileIndex : public SwapIndex {
+/// FIXME: Expose an interface to remove files that are closed.
+class FileIndex {
 public:
   /// If URISchemes is empty, the default schemes in SymbolCollector will be
   /// used.
   FileIndex(std::vector<std::string> URISchemes = {});
 
-  /// Update symbols in \p Path with symbols in \p AST. If \p AST is
-  /// nullptr, this removes all symbols in the file.
-  /// If \p AST is not null, \p PP cannot be null and it should be the
-  /// preprocessor that was used to build \p AST.
-  /// If \p TopLevelDecls is set, only these decls are indexed. Otherwise, all
-  /// top level decls obtained from \p AST are indexed.
-  void
-  update(PathRef Path, ASTContext *AST, std::shared_ptr<Preprocessor> PP,
-         llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls = llvm::None);
+  // 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 from main file \p Path with symbols in \p TopLevelDecls.
+  void updateMain(PathRef Path, ParsedAST &AST,
+                  llvm::ArrayRef<Decl *> TopLevelDecls);
 
 private:
-  // Only update() should swap the index.
-  using SwapIndex::reset;
-
-  FileSymbols FSymbols;
   std::vector<std::string> URISchemes;
+
+  // Contains information from each file's preamble only.
+  // These are large, but update fairly infrequently (preambles are stable).
+  // 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.
+  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
+  // (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 in \p AST.
+/// Retrieves symbols and refs of \p Decls in \p AST.
 /// Exposed to assist in unit tests.
 /// If URISchemes is empty, the default schemes in SymbolCollector will be used.
-/// If \p TopLevelDecls is set, only these decls are indexed. Otherwise, all top
-/// level decls obtained from \p AST are indexed.
 std::pair<SymbolSlab, RefSlab>
-indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-         llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls = llvm::None,
-         llvm::ArrayRef<std::string> URISchemes = {});
+indexMainDecls(ParsedAST &AST, llvm::ArrayRef<Decl *> Decls,
+               llvm::ArrayRef<std::string> URISchemes = {});
+
+/// Idex declarations from \p AST and macros from \p PP that are declared in
+/// included headers.
+/// If URISchemes is empty, the default schemes in SymbolCollector will be used.
+SymbolSlab indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                              llvm::ArrayRef<std::string> URISchemes = {});
 
 } // namespace clangd
 } // namespace clang
diff --git a/unittests/clangd/FileIndexTests.cpp b/unittests/clangd/FileIndexTests.cpp
index 096ccd9..346560a 100644
--- a/unittests/clangd/FileIndexTests.cpp
+++ b/unittests/clangd/FileIndexTests.cpp
@@ -119,10 +119,10 @@
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 }
 
-std::vector<std::string> match(const SymbolIndex &I,
+std::vector<std::string> match(const FileIndex &I,
                                const FuzzyFindRequest &Req) {
   std::vector<std::string> Matches;
-  I.fuzzyFind(Req, [&](const Symbol &Sym) {
+  I.index().fuzzyFind(Req, [&](const Symbol &Sym) {
     Matches.push_back((Sym.Scope + Sym.Name).str());
   });
   return Matches;
@@ -135,7 +135,8 @@
   File.HeaderFilename = (Basename + ".h").str();
   File.HeaderCode = Code;
   auto AST = File.build();
-  M.update(File.Filename, &AST.getASTContext(), AST.getPreprocessorPtr());
+  M.updatePreamble(File.Filename, AST.getASTContext(),
+                   AST.getPreprocessorPtr());
 }
 
 TEST(FileIndexTest, CustomizedURIScheme) {
@@ -145,7 +146,7 @@
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenSymbol = false;
-  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_EQ(Sym.CanonicalDeclaration.FileURI, "unittest:///f.h");
     SeenSymbol = true;
   });
@@ -182,25 +183,6 @@
   EXPECT_THAT(match(M, Req), UnorderedElementsAre("ns::f", "ns::X", "ns::ff"));
 }
 
-TEST(FileIndexTest, RemoveAST) {
-  FileIndex M;
-  update(M, "f1", "namespace ns { void f() {} class X {}; }");
-
-  FuzzyFindRequest Req;
-  Req.Query = "";
-  Req.Scopes = {"ns::"};
-  EXPECT_THAT(match(M, Req), UnorderedElementsAre("ns::f", "ns::X"));
-
-  M.update("f1.cpp", nullptr, nullptr);
-  EXPECT_THAT(match(M, Req), UnorderedElementsAre());
-}
-
-TEST(FileIndexTest, RemoveNonExisting) {
-  FileIndex M;
-  M.update("no.cpp", nullptr, nullptr);
-  EXPECT_THAT(match(M, FuzzyFindRequest()), UnorderedElementsAre());
-}
-
 TEST(FileIndexTest, ClassMembers) {
   FileIndex M;
   update(M, "f1", "class X { static int m1; int m2; static void f(); };");
@@ -218,7 +200,7 @@
   FuzzyFindRequest Req;
   Req.Query = "";
   bool SeenSymbol = false;
-  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
     EXPECT_TRUE(Sym.IncludeHeaders.empty());
     SeenSymbol = true;
   });
@@ -242,7 +224,7 @@
   Req.Query = "";
   bool SeenVector = false;
   bool SeenMakeVector = false;
-  M.fuzzyFind(Req, [&](const Symbol &Sym) {
+  M.index().fuzzyFind(Req, [&](const Symbol &Sym) {
     if (Sym.Name == "vector") {
       EXPECT_EQ(Sym.Signature, "<class Ty>");
       EXPECT_EQ(Sym.CompletionSnippetSuffix, "<${1:class Ty}>");
@@ -296,7 +278,7 @@
       [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP) {
         EXPECT_FALSE(IndexUpdated) << "Expected only a single index update";
         IndexUpdated = true;
-        Index.update(FooCpp, &Ctx, std::move(PP));
+        Index.updatePreamble(FooCpp, Ctx, std::move(PP));
       });
   ASSERT_TRUE(IndexUpdated);
 
@@ -332,18 +314,16 @@
   Test.Code = MainCode.code();
   Test.Filename = "test.cc";
   auto AST = Test.build();
-  Index.update(Test.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
-               AST.getLocalTopLevelDecls());
+  Index.updateMain(Test.Filename, AST, AST.getLocalTopLevelDecls());
   // Add test2.cc
   TestTU Test2;
   Test2.HeaderCode = HeaderCode;
   Test2.Code = MainCode.code();
   Test2.Filename = "test2.cc";
   AST = Test2.build();
-  Index.update(Test2.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
-               AST.getLocalTopLevelDecls());
+  Index.updateMain(Test2.Filename, AST, AST.getLocalTopLevelDecls());
 
-  EXPECT_THAT(getRefs(Index, Foo.ID),
+  EXPECT_THAT(getRefs(Index.index(), Foo.ID),
               RefsAre({AllOf(RefRange(MainCode.range("foo")),
                              FileURI("unittest:///test.cc")),
                        AllOf(RefRange(MainCode.range("foo")),
diff --git a/unittests/clangd/IndexTests.cpp b/unittests/clangd/IndexTests.cpp
index 187b3c5..1f8129f 100644
--- a/unittests/clangd/IndexTests.cpp
+++ b/unittests/clangd/IndexTests.cpp
@@ -231,7 +231,7 @@
 TEST(MergeIndexTest, Refs) {
   FileIndex Dyn({"unittest"});
   FileIndex StaticIndex({"unittest"});
-  auto MergedIndex = mergeIndex(&Dyn, &StaticIndex);
+  auto MergedIndex = mergeIndex(&Dyn.index(), &StaticIndex.index());
 
   const char *HeaderCode = "class Foo;";
   auto HeaderSymbols = TestTU::withHeaderCode("class Foo;").headerSymbols();
@@ -244,8 +244,7 @@
   Test.Code = Test1Code.code();
   Test.Filename = "test.cc";
   auto AST = Test.build();
-  Dyn.update(Test.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
-             AST.getLocalTopLevelDecls());
+  Dyn.updateMain(Test.Filename, AST, AST.getLocalTopLevelDecls());
 
   // Build static index for test.cc.
   Test.HeaderCode = HeaderCode;
@@ -253,9 +252,8 @@
   Test.Filename = "test.cc";
   auto StaticAST = Test.build();
   // Add stale refs for test.cc.
-  StaticIndex.update(Test.Filename, &StaticAST.getASTContext(),
-                     StaticAST.getPreprocessorPtr(),
-                     StaticAST.getLocalTopLevelDecls());
+  StaticIndex.updateMain(Test.Filename, StaticAST,
+                         StaticAST.getLocalTopLevelDecls());
 
   // Add refs for test2.cc
   Annotations Test2Code(R"(class $Foo[[Foo]] {};)");
@@ -264,9 +262,8 @@
   Test2.Code = Test2Code.code();
   Test2.Filename = "test2.cc";
   StaticAST = Test2.build();
-  StaticIndex.update(Test2.Filename, &StaticAST.getASTContext(),
-                     StaticAST.getPreprocessorPtr(),
-                     StaticAST.getLocalTopLevelDecls());
+  StaticIndex.updateMain(Test2.Filename, StaticAST,
+                         StaticAST.getLocalTopLevelDecls());
 
   RefsRequest Request;
   Request.IDs = {Foo.ID};
diff --git a/unittests/clangd/TestTU.cpp b/unittests/clangd/TestTU.cpp
index e181f68..847f796 100644
--- a/unittests/clangd/TestTU.cpp
+++ b/unittests/clangd/TestTU.cpp
@@ -45,13 +45,13 @@
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return indexAST(AST.getASTContext(), AST.getPreprocessorPtr()).first;
+  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 = indexAST(AST.getASTContext(), AST.getPreprocessorPtr(),
-                          AST.getLocalTopLevelDecls());
+  auto Content = indexMainDecls(AST, AST.getLocalTopLevelDecls());
   return MemIndex::build(std::move(Content.first), std::move(Content.second));
 }