From f4feb15d81a87656029dc5a9da5391665c255283 Mon Sep 17 00:00:00 2001
From: Adam Balogh <adam.balogh@ericsson.com>
Date: Fri, 15 Jun 2018 06:45:39 +0000
Subject: [PATCH] [ASTImporter] Corrected diagnostic client handling in tests.
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

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



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@334804 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/clang/Frontend/ASTUnit.h  |  9 +++++++++
 lib/Frontend/ASTUnit.cpp          |  6 ++++++
 unittests/AST/ASTImporterTest.cpp | 10 +++++++++-
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/include/clang/Frontend/ASTUnit.h b/include/clang/Frontend/ASTUnit.h
index 35cacee57f..9f529ba0e0 100644
--- a/include/clang/Frontend/ASTUnit.h
+++ b/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/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index d7bbf6fe4a..e4c313fed3 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/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/unittests/AST/ASTImporterTest.cpp b/unittests/AST/ASTImporterTest.cpp
index 5e2dcf73fc..c81b4b338b 100644
--- a/unittests/AST/ASTImporterTest.cpp
+++ b/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
-- 
2.50.1