[clangd] Add beforeExecute() callback to FeatureModules.

[adamcz]

Summary

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

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/FeatureModule.h

109

nit: triple slashes

109

comment still says preamble build. I'd either talk more about Allows customizing the compiler instance before starting the execution on input or get rid of CompilerInvocation in that list. as in theory modifications to only parts of it will be respected (since we've already issued beginsourcefile on compiler instance and created some structs with whatever we had in compiler invocation prior to this call).

clang-tools-extra/clangd/ParsedAST.cpp

553

could you have a comment that says we must perform this after Action::BeginSourceFile to ensure PP and rest of the helper structs are initialized and closer to the end so that other modifications we do in clangd (like adjusting diag severities/marking main file as include-guarded) is visible to modules?

to make sure it doesn't get moved around without these in mind.

clang-tools-extra/clangd/Preamble.cpp

72

nit: std::move

117

could we move this to the bottom, to make it closer to the way parsedast calls it?

488

nit: drop braces

[adamcz]

addressed review comments

clang-tools-extra/clangd/FeatureModule.h

 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/JSON.h"
 #include <memory>
 #include <type_traits>
 #include <vector>
 
 namespace clang {
+class CompilerInstance;
 namespace clangd {
 struct Diag;
 class LSPBinder;
 class SymbolIndex;
 class ThreadsafeFS;
 class TUScheduler;
 class Tweak;
 
public:
   /// Extension point that allows modules to observe and modify an AST build.
   /// One instance is created each time clangd produces a ParsedAST or
   /// PrecompiledPreamble. For a given instance, lifecycle methods are always
   /// called on a single thread.
   struct ASTListener {
     /// 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.
     virtual void sawDiagnostic(const clang::Diagnostic &, clangd::Diag &) {}
   };
   /// Can be called asynchronously before building an AST.
   virtual std::unique_ptr<ASTListener> astListeners() { return nullptr; }
 

clang-tools-extra/clangd/ParsedAST.cpp

else
     CanonIncludes.addSystemHeadersMapping(Clang->getLangOpts());
   std::unique_ptr<CommentHandler> IWYUHandler =
       collectIWYUHeaderMaps(&CanonIncludes);
   Clang->getPreprocessor().addCommentHandler(IWYUHandler.get());
 
   // 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)));
 
   // We have to consume the tokens before running clang-tidy to avoid collecting
   // tokens from running the preprocessor inside the checks (only
   // modernize-use-trailing-return-type does that today).
   syntax::TokenBuffer Tokens = std::move(CollectTokens).consume();

clang-tools-extra/clangd/Preamble.cpp

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);
 }
 
 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); }
 
   MainFileMacros takeMacros() { return std::move(Macros); }
 
   std::vector<PragmaMark> takeMarks() { return std::move(Marks); }
 
   CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
if (Stats) {
       Stats->BuildSize += PP.getTotalMemory();
       if (PreprocessingRecord *PRec = PP.getPreprocessingRecord())
         Stats->BuildSize += PRec->getTotalMemory();
       Stats->BuildSize += PP.getHeaderSearchInfo().getTotalMemory();
     }
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     CanonIncludes.addSystemHeadersMapping(CI.getLangOpts());
     LangOpts = &CI.getLangOpts();
     SourceMgr = &CI.getSourceManager();
     Includes.collect(CI);
+    if (BeforeExecuteCallback)
+      BeforeExecuteCallback(CI);
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
     assert(SourceMgr && LangOpts &&
            "SourceMgr and LangOpts must be set at this point");
 
     return std::make_unique<PPChainedCallbacks>(
         std::make_unique<CollectMainFileMacros>(*SourceMgr, Macros),
private:
   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;
   PreambleBuildStats *Stats;
+  std::function<void(CompilerInstance &)> BeforeExecuteCallback;
 };
 
 // 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, 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);
   auto StatCache = std::make_unique<PreambleFileStatusCache>(AbsFileName);
   auto StatCacheFS = StatCache->getProducingFS(VFS);
   llvm::IntrusiveRefCntPtr<TimerFS> TimedFS(new TimerFS(StatCacheFS));
 
   WallTimer PreambleTimer;
if (IncludeLoc.isValid() && SM.isWrittenInBuiltinFile(IncludeLoc) &&
           isMainFile(Presumed.getFilename(), SM)) {
         Loc = SM.translateLineCol(SM.getMainFileID(), Presumed.getLine(),
                                   Presumed.getColumn());
       }
     }
   }
   return Loc;
 }
+
 } // namespace clangd
 } // namespace clang

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

 //===--- FeatureModulesTests.cpp  -------------------------------*- 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 "Annotations.h"
 #include "FeatureModule.h"
 #include "Selection.h"
 #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"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 namespace {
TEST(FeatureModulesTest, SuppressDiags) {
 
   TU.FeatureModules = &FMS;
   {
     auto AST = TU.build();
     EXPECT_THAT(*AST.getDiagnostics(), testing::IsEmpty());
   }
 }
 
+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