[clangd] Add beforeExecute() callback to FeatureModules.

It runs immediatelly before FrontendAction::Execute() with a mutable
CompilerInstance, allowing FeatureModules to register callbacks, remap
files, etc.

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

clang-tools-extra/clangd/FeatureModule.h

@@ -20,6 +20,7 @@
 #include <vector>
 
 namespace clang {
+class CompilerInstance;
 namespace clangd {
 struct Diag;
 class LSPBinder;
@@ -105,6 +106,11 @@ public:
     /// Listeners are destroyed once the AST is built.
     virtual ~ASTListener() = default;
 
+    /// Called before every AST build, both for main file and preamble. The call
+    /// happens immediately before FrontendAction::Execute(), with Preprocessor
+    /// set up already and after BeginSourceFile() on main file was called.
+    virtual void beforeExecute(CompilerInstance &CI) {}
+
     /// Called everytime a diagnostic is encountered. Modules can use this
     /// modify the final diagnostic, or store some information to surface code
     /// actions later on.

clang-tools-extra/clangd/ParsedAST.cpp

@@ -550,6 +550,12 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // Collect tokens of the main file.
   syntax::TokenCollector CollectTokens(Clang->getPreprocessor());
 
+  // To remain consistent with preamble builds, these callbacks must be called
+  // exactly here, after preprocessor is initialized and BeginSourceFile() was
+  // called already.
+  for (const auto &L : ASTListeners)
+    L->beforeExecute(*Clang);
+
   if (llvm::Error Err = Action->Execute())
     log("Execute() failed when building AST for {0}: {1}", MainInput.getFile(),
         toString(std::move(Err)));

clang-tools-extra/clangd/Preamble.cpp

@@ -64,9 +64,12 @@ bool compileCommandsAreEqual(const tooling::CompileCommand &LHS,
 
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
-  CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback,
-                           PreambleBuildStats *Stats)
-      : File(File), ParsedCallback(ParsedCallback), Stats(Stats) {}
+  CppFilePreambleCallbacks(
+      PathRef File, PreambleParsedCallback ParsedCallback,
+      PreambleBuildStats *Stats,
+      std::function<void(CompilerInstance &)> BeforeExecuteCallback)
+      : File(File), ParsedCallback(ParsedCallback), Stats(Stats),
+        BeforeExecuteCallback(std::move(BeforeExecuteCallback)) {}
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
@@ -115,6 +118,8 @@ public:
     LangOpts = &CI.getLangOpts();
     SourceMgr = &CI.getSourceManager();
     Includes.collect(CI);
+    if (BeforeExecuteCallback)
+      BeforeExecuteCallback(CI);
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
@@ -156,6 +161,7 @@ private:
   const clang::LangOptions *LangOpts = nullptr;
   const SourceManager *SourceMgr = nullptr;
   PreambleBuildStats *Stats;
+  std::function<void(CompilerInstance &)> BeforeExecuteCallback;
 };
 
 // Represents directives other than includes, where basic textual information is
@@ -477,7 +483,11 @@ buildPreamble(PathRef FileName, CompilerInvocation CI,
   // to read back. We rely on dynamic index for the comments instead.
   CI.getPreprocessorOpts().WriteCommentListToPCH = false;
 
-  CppFilePreambleCallbacks CapturedInfo(FileName, PreambleCallback, Stats);
+  CppFilePreambleCallbacks CapturedInfo(FileName, PreambleCallback, Stats,
+                                        [&ASTListeners](CompilerInstance &CI) {
+                                          for (const auto &L : ASTListeners)
+                                            L->beforeExecute(CI);
+                                        });
   auto VFS = Inputs.TFS->view(Inputs.CompileCommand.Directory);
   llvm::SmallString<32> AbsFileName(FileName);
   VFS->makeAbsolute(AbsFileName);
@@ -716,5 +726,6 @@ SourceLocation translatePreamblePatchLocation(SourceLocation Loc,
   }
   return Loc;
 }
+
 } // namespace clangd
 } // namespace clang

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

@@ -12,6 +12,7 @@
 #include "TestTU.h"
 #include "refactor/Tweak.h"
 #include "support/Logger.h"
+#include "clang/Lex/PreprocessorOptions.h"
 #include "llvm/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -85,6 +86,41 @@ TEST(FeatureModulesTest, SuppressDiags) {
   }
 }
 
+TEST(FeatureModulesTest, BeforeExecute) {
+  struct BeforeExecuteModule final : public FeatureModule {
+    struct Listener : public FeatureModule::ASTListener {
+      void beforeExecute(CompilerInstance &CI) override {
+        CI.getPreprocessor().SetSuppressIncludeNotFoundError(true);
+      }
+    };
+    std::unique_ptr<ASTListener> astListeners() override {
+      return std::make_unique<Listener>();
+    };
+  };
+  FeatureModuleSet FMS;
+  FMS.add(std::make_unique<BeforeExecuteModule>());
+
+  TestTU TU = TestTU::withCode(R"cpp(
+    /*error-ok*/
+    #include "not_found.h"
+
+    void foo() {
+      #include "not_found_not_preamble.h"
+    }
+  )cpp");
+
+  {
+    auto AST = TU.build();
+    EXPECT_THAT(*AST.getDiagnostics(), testing::Not(testing::IsEmpty()));
+  }
+
+  TU.FeatureModules = &FMS;
+  {
+    auto AST = TU.build();
+    EXPECT_THAT(*AST.getDiagnostics(), testing::IsEmpty());
+  }
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang