[refactor] allow the use of refactoring diagnostics

This commit allows the refactoring library to use its own set of
refactoring-specific diagnostics to reports things like initiation errors.

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

llvm-svn: 315924
diff --git a/clang/tools/clang-refactor/ClangRefactor.cpp b/clang/tools/clang-refactor/ClangRefactor.cpp
index 8a5c8e6..ed9bf5e 100644
--- a/clang/tools/clang-refactor/ClangRefactor.cpp
+++ b/clang/tools/clang-refactor/ClangRefactor.cpp
@@ -15,6 +15,7 @@
 
 #include "TestSupport.h"
 #include "clang/Frontend/CommandLineSourceLoc.h"
+#include "clang/Frontend/TextDiagnosticPrinter.h"
 #include "clang/Rewrite/Core/Rewriter.h"
 #include "clang/Tooling/CommonOptionsParser.h"
 #include "clang/Tooling/Refactoring.h"
@@ -65,7 +66,8 @@
   /// logic into the refactoring operation. The test-specific consumer
   /// ensures that the individual results in a particular test group are
   /// identical.
-  virtual std::unique_ptr<RefactoringResultConsumer> createCustomConsumer() {
+  virtual std::unique_ptr<ClangRefactorToolConsumerInterface>
+  createCustomConsumer() {
     return nullptr;
   }
 
@@ -85,7 +87,8 @@
 
   void print(raw_ostream &OS) override { TestSelections.dump(OS); }
 
-  std::unique_ptr<RefactoringResultConsumer> createCustomConsumer() override {
+  std::unique_ptr<ClangRefactorToolConsumerInterface>
+  createCustomConsumer() override {
     return TestSelections.createConsumer();
   }
 
@@ -304,10 +307,20 @@
   RefactoringActionCommandLineOptions Options;
 };
 
-class ClangRefactorConsumer : public RefactoringResultConsumer {
+class ClangRefactorConsumer final : public ClangRefactorToolConsumerInterface {
 public:
+  ClangRefactorConsumer() {}
+
   void handleError(llvm::Error Err) override {
-    llvm::errs() << llvm::toString(std::move(Err)) << "\n";
+    Optional<PartialDiagnosticAt> Diag = DiagnosticError::take(Err);
+    if (!Diag) {
+      llvm::errs() << llvm::toString(std::move(Err)) << "\n";
+      return;
+    }
+    llvm::cantFail(std::move(Err)); // This is a success.
+    DiagnosticBuilder DB(
+        getDiags().Report(Diag->first, Diag->second.getDiagID()));
+    Diag->second.Emit(DB);
   }
 
   void handle(AtomicChanges Changes) override {
@@ -468,8 +481,8 @@
       return true;
     }
 
-    bool HasFailed = false;
     ClangRefactorConsumer Consumer;
+    bool HasFailed = false;
     if (foreachTranslationUnit(DB, Sources, [&](ASTContext &AST) {
           RefactoringRuleContext Context(AST.getSourceManager());
           Context.setASTContext(AST);
@@ -488,21 +501,27 @@
                 "The action must have at least one selection rule");
           };
 
+          std::unique_ptr<ClangRefactorToolConsumerInterface> CustomConsumer;
+          if (HasSelection)
+            CustomConsumer = Subcommand.getSelection()->createCustomConsumer();
+          ClangRefactorToolConsumerInterface &ActiveConsumer =
+              CustomConsumer ? *CustomConsumer : Consumer;
+          ActiveConsumer.beginTU(AST);
           if (HasSelection) {
             assert(Subcommand.getSelection() && "Missing selection argument?");
             if (opts::Verbose)
               Subcommand.getSelection()->print(llvm::outs());
-            auto CustomConsumer =
-                Subcommand.getSelection()->createCustomConsumer();
             if (Subcommand.getSelection()->forAllRanges(
                     Context.getSources(), [&](SourceRange R) {
                       Context.setSelectionRange(R);
-                      InvokeRule(CustomConsumer ? *CustomConsumer : Consumer);
+                      InvokeRule(ActiveConsumer);
                     }))
               HasFailed = true;
+            ActiveConsumer.endTU();
             return;
           }
           // FIXME (Alex L): Implement non-selection based invocation path.
+          ActiveConsumer.endTU();
         }))
       return true;
     return HasFailed || applySourceChanges(Consumer.getSourceChanges());