[clangd] Proof of concept: indexing after the preamble is built

[sammccall]

Summary

We can use this technique to take preamble indexing off the critical path on startup.

[sammccall]

rebase

[sammccall]

@kadircet FYI this demonstrates how to extend the lifetime of the preamble AST beyond building the preamble, which would allow us to take preamble indexing off the critical path.
It's pretty hacky, but the doc comments on CompilerInstance more or less suggest this sort of thing!

In reality I imagine this looking more like: pass the whole CompilerInstance to the PreambleCallback, and have ClangdServer do the lifetime extension and throw the AliveAST struct onto a queue.
This ends up resembling https://reviews.llvm.org/D115232 a little bit, and they should share queues I think. So maybe I shouldn't develop this further until that one lands in some form.

[nridge]

When you say "take preamble indexing off the critical path on startup", do you mean the first time a file is opened in general? Or is this specific to startup somehow?

clang-tools-extra/clangd/Preamble.cpp

 #include "support/ThreadsafeFS.h"
 #include "support/Trace.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticLex.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TokenKinds.h"
+#include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"

 
 bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
                              const tooling::CompileCommand &RHS) {
   // We don't check for Output, it should not matter to clangd.
   return LHS.Directory == RHS.Directory && LHS.Filename == RHS.Filename &&
          llvm::makeArrayRef(LHS.CommandLine).equals(RHS.CommandLine);
 }
 
+struct AliveAST {
+public:
+  AliveAST(CompilerInstance &Clang)
+      : Invocation(Clang.getInvocationPtr()),
+        Diagnostics(&Clang.getDiagnostics()), FileMgr(&Clang.getFileManager()),
+        SourceMgr(&Clang.getSourceManager()), PP(Clang.getPreprocessorPtr()),
+        Context(&Clang.getASTContext()) {}
+
+  ASTContext &getASTContext() { return *Context; }
+  const std::shared_ptr<Preprocessor> &getPreprocessor() { return PP; }
+
+private:
+  std::shared_ptr<CompilerInvocation> Invocation;
+  IntrusiveRefCntPtr<DiagnosticsEngine> Diagnostics;
+  IntrusiveRefCntPtr<TargetInfo> Target;
+  IntrusiveRefCntPtr<TargetInfo> AuxTarget;
+  IntrusiveRefCntPtr<FileManager> FileMgr;
+  IntrusiveRefCntPtr<SourceManager> SourceMgr;
+  std::shared_ptr<Preprocessor> PP;
+  IntrusiveRefCntPtr<ASTContext> Context;
+};
+
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
-  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
-      : File(File), ParsedCallback(ParsedCallback) {}
+  CppFilePreambleCallbacks(PathRef File) : File(File) {}
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
   MainFileMacros takeMacros() { return std::move(Macros); }
 
   std::vector<PragmaMark> takeMarks() { return std::move(Marks); }
 
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
+  AliveAST takeLife() { return std::move(*Alive); }
+
   bool isMainFileIncludeGuarded() const { return IsMainFileIncludeGuarded; }
 
   void AfterExecute(CompilerInstance &CI) override {
-    if (ParsedCallback) {
-      trace::Span Tracer("Running PreambleCallback");
-      ParsedCallback(CI.getASTContext(), CI.getPreprocessorPtr(),
-                     CanonIncludes);
-    }
+    Alive.emplace(CI);
 
     const SourceManager &SM = CI.getSourceManager();
     const FileEntry *MainFE = SM.getFileEntryForID(SM.getMainFileID());
     IsMainFileIncludeGuarded =
         CI.getPreprocessor().getHeaderSearchInfo().isFileMultipleIncludeGuarded(
             MainFE);
   }
 
if (const auto *FT = llvm::dyn_cast<clang::FunctionTemplateDecl>(D)) {
         if (II->isStr("make_unique") && FT->isInStdNamespace())
           return false;
     }
     return true;
   }
 
 private:
   PathRef File;
-  PreambleParsedCallback ParsedCallback;
   IncludeStructure Includes;
   CanonicalIncludes CanonIncludes;
   MainFileMacros Macros;
   std::vector<PragmaMark> Marks;
   bool IsMainFileIncludeGuarded = false;
   std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   const clang::LangOptions *LangOpts = nullptr;
   const SourceManager *SourceMgr = nullptr;
+
+  llvm::Optional<AliveAST> Alive;
 };
 
 // Represents directives other than includes, where basic textual information is
 // enough.
 struct TextualPPDirective {
   unsigned DirectiveLine;
   // Full text that's representing the directive, including the `#`.
   std::string Text;
buildPreamble(PathRef FileName, CompilerInvocation CI,
   // Skip function bodies when building the preamble to speed up building
   // the preamble and make it smaller.
   assert(!CI.getFrontendOpts().SkipFunctionBodies);
   CI.getFrontendOpts().SkipFunctionBodies = true;
   // We don't want to write comment locations into PCH. They are racy and slow
   // to read back. We rely on dynamic index for the comments instead.
   CI.getPreprocessorOpts().WriteCommentListToPCH = false;
 
-  CppFilePreambleCallbacks CapturedInfo(FileName, PreambleCallback);
+  CppFilePreambleCallbacks CapturedInfo(FileName);
   auto VFS = Inputs.TFS->view(Inputs.CompileCommand.Directory);
   llvm::SmallString<32> AbsFileName(FileName);
   VFS->makeAbsolute(AbsFileName);
   auto StatCache = std::make_unique<PreambleFileStatusCache>(AbsFileName);
   auto BuiltPreamble = PrecompiledPreamble::Build(
       CI, ContentsBuffer.get(), Bounds, *PreambleDiagsEngine,
       StatCache->getProducingFS(VFS),
       std::make_shared<PCHContainerOperations>(), StoreInMemory, CapturedInfo);
if (BuiltPreamble) {
     Result->CompileCommand = Inputs.CompileCommand;
     Result->Diags = std::move(Diags);
     Result->Includes = CapturedInfo.takeIncludes();
     Result->Macros = CapturedInfo.takeMacros();
     Result->Marks = CapturedInfo.takeMarks();
     Result->CanonIncludes = CapturedInfo.takeCanonicalIncludes();
     Result->StatCache = std::move(StatCache);
     Result->MainIsIncludeGuarded = CapturedInfo.isMainFileIncludeGuarded();
+
+    if (PreambleCallback) {
+      trace::Span Tracer("Running PreambleCallback");
+      auto Alive = CapturedInfo.takeLife();
+      PreambleCallback(Alive.getASTContext(), Alive.getPreprocessor(),
+                       Result->CanonIncludes);
+    }
+
     return Result;
   }
 
   elog("Could not build a preamble for file {0} version {1}: {2}", FileName,
        Inputs.Version, BuiltPreamble.getError().message());
   return nullptr;
 }
 

clang/include/clang/Frontend/CompilerInstance.h

public:
 
   bool hasInvocation() const { return Invocation != nullptr; }
 
   CompilerInvocation &getInvocation() {
     assert(Invocation && "Compiler instance has no invocation!");
     return *Invocation;
   }
 
+  std::shared_ptr<CompilerInvocation> getInvocationPtr() { return Invocation; }
+
   /// setInvocation - Replace the current invocation.
   void setInvocation(std::shared_ptr<CompilerInvocation> Value);
 
   /// Indicates whether we should (re)build the global module index.
   bool shouldBuildGlobalModuleIndex() const;
 
   /// Set the flag indicating whether we should (re)build the global
   /// module index.