[ASTImporter] Corrected diagnostic client handling in tests.

ASTImporter tests may produce source file related warnings, the diagnostic
client should be in correct state to handle it. Added 'beginSourceFile' to set
the client state.

Patch by: Balázs Kéri

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

llvm-svn: 334804

3 files changed, 24 insertions, 1 deletions

diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index 35cacee57f7..9f529ba0e0e 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -438,6 +438,15 @@ public:
   void setASTContext(ASTContext *ctx) { Ctx = ctx; }
   void setPreprocessor(std::shared_ptr<Preprocessor> pp);
 
+  /// Enable source-range based diagnostic messages.
+  ///
+  /// If diagnostic messages with source-range information are to be expected
+  /// and AST comes not from file (e.g. after LoadFromCompilerInvocation) this
+  /// function has to be called.
+  /// The function is to be called only once and the AST should be associated
+  /// with the same source file afterwards.
+  void enableSourceFileDiagnostics();
+
   bool hasSema() const { return (bool)TheSema; }
 
   Sema &getSema() const { 
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index d7bbf6fe4ae..e4c313fed30 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -275,6 +275,12 @@ void ASTUnit::setPreprocessor(std::shared_ptr<Preprocessor> PP) {
   this->PP = std::move(PP);
 }
 
+void ASTUnit::enableSourceFileDiagnostics() {
+  assert(getDiagnostics().getClient() && Ctx &&
+      "Bad context for source file");
+  getDiagnostics().getClient()->BeginSourceFile(Ctx->getLangOpts(), PP.get());
+}
+
 /// Determine the set of code-completion contexts in which this
 /// declaration should be shown.
 static unsigned getDeclShowContexts(const NamedDecl *ND,
diff --git a/clang/unittests/AST/ASTImporterTest.cpp b/clang/unittests/AST/ASTImporterTest.cpp
index 5e2dcf73fc3..c81b4b338b7 100644
--- a/clang/unittests/AST/ASTImporterTest.cpp
+++ b/clang/unittests/AST/ASTImporterTest.cpp
@@ -98,6 +98,9 @@ testImport(const std::string &FromCode, const ArgVector &FromArgs,
   ASTContext &FromCtx = FromAST->getASTContext(),
       &ToCtx = ToAST->getASTContext();
 
+  FromAST->enableSourceFileDiagnostics();
+  ToAST->enableSourceFileDiagnostics();
+
   ASTImporter Importer(ToCtx, ToAST->getFileManager(),
                        FromCtx, FromAST->getFileManager(), false);
 
@@ -172,7 +175,9 @@ class ASTImporterTestBase : public ::testing::TestWithParam<ArgVector> {
         : Code(Code), FileName(FileName),
           Unit(tooling::buildASTFromCodeWithArgs(this->Code, Args,
                                                  this->FileName)),
-          TUDecl(Unit->getASTContext().getTranslationUnitDecl()) {}
+          TUDecl(Unit->getASTContext().getTranslationUnitDecl()) {
+      Unit->enableSourceFileDiagnostics();
+    }
   };
 
   // We may have several From contexts and related translation units. In each
@@ -214,6 +219,7 @@ public:
     ToCode = ToSrcCode;
     assert(!ToAST);
     ToAST = tooling::buildASTFromCodeWithArgs(ToCode, ToArgs, OutputFileName);
+    ToAST->enableSourceFileDiagnostics();
 
     ASTContext &FromCtx = FromTU.Unit->getASTContext(),
                &ToCtx = ToAST->getASTContext();
@@ -261,6 +267,7 @@ public:
     ToCode = ToSrcCode;
     assert(!ToAST);
     ToAST = tooling::buildASTFromCodeWithArgs(ToCode, ToArgs, OutputFileName);
+    ToAST->enableSourceFileDiagnostics();
 
     return ToAST->getASTContext().getTranslationUnitDecl();
   }
@@ -274,6 +281,7 @@ public:
       // Build the AST from an empty file.
       ToAST =
           tooling::buildASTFromCodeWithArgs(/*Code=*/"", ToArgs, "empty.cc");
+      ToAST->enableSourceFileDiagnostics();
     }
 
     // Create a virtual file in the To Ctx which corresponds to the file from