1 files changed, 23 insertions, 9 deletions

diff --git a/tools/clang-import-test/clang-import-test.cpp b/tools/clang-import-test/clang-import-test.cpp
index 567a4bb4f0a26..db4dc76eded54 100644
--- a/tools/clang-import-test/clang-import-test.cpp
+++ b/tools/clang-import-test/clang-import-test.cpp
@@ -17,7 +17,9 @@
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/TargetOptions.h"
 #include "clang/CodeGen/ModuleBuilder.h"
+#include "clang/Frontend/ASTConsumers.h"
 #include "clang/Frontend/CompilerInstance.h"
+#include "clang/Frontend/MultiplexConsumer.h"
 #include "clang/Frontend/TextDiagnosticBuffer.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/Preprocessor.h"
@@ -51,6 +53,10 @@ static llvm::cl::list<std::string>
               llvm::cl::desc("Argument to pass to the CompilerInvocation"),
               llvm::cl::CommaSeparated);
 
+static llvm::cl::opt<bool>
+DumpAST("dump-ast", llvm::cl::init(false),
+        llvm::cl::desc("Dump combined AST"));
+
 namespace init_convenience {
 class TestDiagnosticConsumer : public DiagnosticConsumer {
 private:
@@ -233,7 +239,7 @@ std::unique_ptr<CompilerInstance> BuildIndirect(std::unique_ptr<CompilerInstance
 }
 
 llvm::Error ParseSource(const std::string &Path, CompilerInstance &CI,
-                        CodeGenerator &CG) {
+                        ASTConsumer &Consumer) {
   SourceManager &SM = CI.getSourceManager();
   const FileEntry *FE = CI.getFileManager().getFile(Path);
   if (!FE) {
@@ -241,13 +247,14 @@ llvm::Error ParseSource(const std::string &Path, CompilerInstance &CI,
         llvm::Twine("Couldn't open ", Path), std::error_code());
   }
   SM.setMainFileID(SM.createFileID(FE, SourceLocation(), SrcMgr::C_User));
-  ParseAST(CI.getPreprocessor(), &CG, CI.getASTContext());
+  ParseAST(CI.getPreprocessor(), &Consumer, CI.getASTContext());
   return llvm::Error::success();
 }
 
 llvm::Expected<std::unique_ptr<CompilerInstance>>
 Parse(const std::string &Path,
-      llvm::ArrayRef<std::unique_ptr<CompilerInstance>> Imports) {
+      llvm::ArrayRef<std::unique_ptr<CompilerInstance>> Imports,
+      bool ShouldDumpAST) {
   std::vector<const char *> ClangArgv(ClangArgs.size());
   std::transform(ClangArgs.begin(), ClangArgs.end(), ClangArgv.begin(),
                  [](const std::string &s) -> const char * { return s.data(); });
@@ -261,14 +268,20 @@ Parse(const std::string &Path,
   if (Imports.size())
     AddExternalSource(*CI, Imports);
 
+  std::vector<std::unique_ptr<ASTConsumer>> ASTConsumers;
+
   auto LLVMCtx = llvm::make_unique<llvm::LLVMContext>();
-  std::unique_ptr<CodeGenerator> CG =
-      init_convenience::BuildCodeGen(*CI, *LLVMCtx);
-  CG->Initialize(CI->getASTContext());
+  ASTConsumers.push_back(init_convenience::BuildCodeGen(*CI, *LLVMCtx));
+
+  if (ShouldDumpAST)
+    ASTConsumers.push_back(CreateASTDumper("", true, false, false));
 
   CI->getDiagnosticClient().BeginSourceFile(CI->getLangOpts(),
                                             &CI->getPreprocessor());
-  if (llvm::Error PE = ParseSource(Path, *CI, *CG)) {
+  MultiplexConsumer Consumers(std::move(ASTConsumers));
+  Consumers.Initialize(CI->getASTContext());
+
+  if (llvm::Error PE = ParseSource(Path, *CI, Consumers)) {
     return std::move(PE);
   }
   CI->getDiagnosticClient().EndSourceFile();
@@ -288,7 +301,8 @@ int main(int argc, const char **argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv);
   std::vector<std::unique_ptr<CompilerInstance>> ImportCIs;
   for (auto I : Imports) {
-    llvm::Expected<std::unique_ptr<CompilerInstance>> ImportCI = Parse(I, {});
+    llvm::Expected<std::unique_ptr<CompilerInstance>> ImportCI =
+      Parse(I, {}, false);
     if (auto E = ImportCI.takeError()) {
       llvm::errs() << llvm::toString(std::move(E));
       exit(-1);
@@ -310,7 +324,7 @@ int main(int argc, const char **argv) {
     }
   }
   llvm::Expected<std::unique_ptr<CompilerInstance>> ExpressionCI =
-      Parse(Expression, Direct ? ImportCIs : IndirectCIs);
+      Parse(Expression, Direct ? ImportCIs : IndirectCIs, DumpAST);
   if (auto E = ExpressionCI.takeError()) {
     llvm::errs() << llvm::toString(std::move(E));
     exit(-1);