[Frontend] Add ASTConsumerInjector

[jkorous]

Summary

Follow-up to recent index-related FrontendActions refactoring.
Prerequisite for [WIP] Index-while-building https://reviews.llvm.org/D64384

clang/include/clang/Frontend/FrontendAction.h

public:
   bool usesPreprocessorOnly() const override;
   TranslationUnitKind getTranslationUnitKind() override;
   bool hasPCHSupport() const override;
   bool hasASTFileSupport() const override;
   bool hasIRSupport() const override;
   bool hasCodeCompletionSupport() const override;
 };
 
+/// Class for adding custom ASTConsumer to a FrontendAction.
+/// Wraps the base action and forwards all method call except CreateASTConsumer.
+class ASTConsumerInjector : public WrapperFrontendAction {
+  std::function<std::unique_ptr<ASTConsumer>(CompilerInstance &CI,
+                                             StringRef InFile)>
+      ASTConsumerFactory;
+
+public:
+  /// \p BaseActionInput is the base action to which all method calls are
+  /// forwarded. \p ASTConsumerFactory is factory to produce the custom
+  /// ASTConsumer in CreateASTConsumer method.
+  ASTConsumerInjector(std::unique_ptr<FrontendAction> BaseActionInput,
+                      std::function<std::unique_ptr<ASTConsumer>(
+                          CompilerInstance &CI, StringRef InFile)>
+                          ASTConsumerFactory);
+
+  /// Returns MultiplexConsumer consumer with first element produced by
+  /// ASTConsumerFactory, followed by the ASTConsumer from the base action.
+  std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
+                                                 StringRef InFile) override;
+};
+
 }  // end namespace clang
 
 #endif

clang/lib/Frontend/FrontendAction.cpp

 bool WrapperFrontendAction::hasCodeCompletionSupport() const {
   return WrappedAction->hasCodeCompletionSupport();
 }
 
 WrapperFrontendAction::WrapperFrontendAction(
     std::unique_ptr<FrontendAction> WrappedAction)
   : WrappedAction(std::move(WrappedAction)) {}
 
+ASTConsumerInjector::ASTConsumerInjector(
+    std::unique_ptr<FrontendAction> BaseActionInput,
+    std::function<std::unique_ptr<ASTConsumer>(CompilerInstance &CI,
+                                               StringRef InFile)>
+        ASTConsumerFactory)
+    : WrapperFrontendAction(std::move(BaseActionInput)),
+      ASTConsumerFactory(ASTConsumerFactory) {}
+
+std::unique_ptr<ASTConsumer>
+ASTConsumerInjector::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
+  std::vector<std::unique_ptr<ASTConsumer>> Consumers;
+  Consumers.emplace_back(ASTConsumerFactory(CI, InFile));
+  Consumers.emplace_back(WrapperFrontendAction::CreateASTConsumer(CI, InFile));
+  return std::make_unique<MultiplexConsumer>(std::move(Consumers));
+}

clang/unittests/Frontend/FrontendActionTest.cpp

 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Serialization/InMemoryModuleCache.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "gtest/gtest.h"
+#include <memory>
 
 using namespace llvm;
 using namespace clang;
 
 namespace {
 
 class TestASTFrontendAction : public ASTFrontendAction {
 public:
if (ShouldCache)
       EXPECT_EQ(InMemoryModuleCache::Final,
                 Compiler.getModuleCache().getPCMState(PCHFilename));
     else
       EXPECT_EQ(InMemoryModuleCache::Unknown,
                 Compiler.getModuleCache().getPCMState(PCHFilename));
   }
 }
 
+TEST(ASTFrontendAction, ASTConsumerInjector) {
+  class InjectTestASTConsumer : public ASTConsumer {
+    std::string& TestLog;
+    const std::string LogValue;
+  public:
+    InjectTestASTConsumer(std::string& TestLog, const std::string& LogValue)
+    : TestLog(TestLog), LogValue(LogValue)
+    { }
+
+    bool HandleTopLevelDecl(DeclGroupRef D) override {
+      TestLog += LogValue;
+      return true;
+    }
+  };
+
+  class InjectTestASTFrontendAction : public ASTFrontendAction {
+    std::string& TestLog;
+    const std::string LogValue;
+  public:
+    InjectTestASTFrontendAction(
+      std::string& TestLog,
+      const std::string& LogValue
+    ) : TestLog(TestLog),
+        LogValue(LogValue) { }
+
+    bool BeginSourceFileAction(CompilerInstance &ci) override {
+      return ASTFrontendAction::BeginSourceFileAction(ci);
+    }
+
+    std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
+                                                  StringRef) override {
+      return std::make_unique<InjectTestASTConsumer>(TestLog, LogValue);
+    }
+  };
+
+  auto invocation = std::make_shared<CompilerInvocation>();
+  invocation->getLangOpts()->CPlusPlus = true;
+  invocation->getLangOpts()->DelayedTemplateParsing = true;
+  invocation->getPreprocessorOpts().addRemappedFile(
+    "test.cpp", MemoryBuffer::getMemBuffer(
+      "int main() { return 0;}\n").release());
+  invocation->getFrontendOpts().Inputs.push_back(
+      FrontendInputFile("test.cpp", Language::CXX));
+  invocation->getFrontendOpts().ProgramAction = frontend::ParseSyntaxOnly;
+  invocation->getTargetOpts().Triple = "i386-unknown-linux-gnu";
+  CompilerInstance compiler;
+  compiler.setInvocation(std::move(invocation));
+  compiler.createDiagnostics();
+
+  std::string TestLog;
+
+  ASTConsumerInjector WrappedTestAction(
+    std::make_unique<InjectTestASTFrontendAction>(TestLog, "A"),
+    [&](CompilerInstance&, StringRef) {
+      return std::make_unique<InjectTestASTConsumer>(TestLog, "B");
+    }
+  );
+
+  ASSERT_TRUE(compiler.ExecuteAction(WrappedTestAction));
+  ASSERT_EQ(TestLog, "BA");
+}
+
 } // anonymous namespace