[Preprocessor] Preserve #pragma clang assume_nonnull in preprocessed output

Patch by Zbigniew Sarbinowski!

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

llvm-svn: 314364

2 files changed, 23 insertions, 0 deletions

diff --git a/clang/lib/Frontend/PrintPreprocessedOutput.cpp b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
index 914039ad5bb..8ae9038316c 100644
--- a/clang/lib/Frontend/PrintPreprocessedOutput.cpp
+++ b/clang/lib/Frontend/PrintPreprocessedOutput.cpp
@@ -143,6 +143,8 @@ public:
                      ArrayRef<int> Ids) override;
   void PragmaWarningPush(SourceLocation Loc, int Level) override;
   void PragmaWarningPop(SourceLocation Loc) override;
+  void PragmaAssumeNonNullBegin(SourceLocation Loc) override;
+  void PragmaAssumeNonNullEnd(SourceLocation Loc) override;
 
   bool HandleFirstTokOnLine(Token &Tok);
 
@@ -549,6 +551,22 @@ void PrintPPOutputPPCallbacks::PragmaWarningPop(SourceLocation Loc) {
   setEmittedDirectiveOnThisLine();
 }
 
+void PrintPPOutputPPCallbacks::
+PragmaAssumeNonNullBegin(SourceLocation Loc) {
+  startNewLineIfNeeded();
+  MoveToLine(Loc);
+  OS << "#pragma clang assume_nonnull begin";
+  setEmittedDirectiveOnThisLine();
+}
+
+void PrintPPOutputPPCallbacks::
+PragmaAssumeNonNullEnd(SourceLocation Loc) {
+  startNewLineIfNeeded();
+  MoveToLine(Loc);
+  OS << "#pragma clang assume_nonnull end";
+  setEmittedDirectiveOnThisLine();
+}
+
 /// HandleFirstTokOnLine - When emitting a preprocessed file in -E mode, this
 /// is called for the first token on each new line.  If this really is the start
 /// of a new logical line, handle it and return true, otherwise return false.
diff --git a/clang/lib/Lex/Pragma.cpp b/clang/lib/Lex/Pragma.cpp
index bf2363a0a6f..b5d0f939549 100644
--- a/clang/lib/Lex/Pragma.cpp
+++ b/clang/lib/Lex/Pragma.cpp
@@ -1725,6 +1725,7 @@ struct PragmaAssumeNonNullHandler : public PragmaHandler {
 
     // The start location we want after processing this.
     SourceLocation NewLoc;
+    PPCallbacks *Callbacks = PP.getPPCallbacks();
 
     if (IsBegin) {
       // Complain about attempts to re-enter an audit.
@@ -1733,6 +1734,8 @@ struct PragmaAssumeNonNullHandler : public PragmaHandler {
         PP.Diag(BeginLoc, diag::note_pragma_entered_here);
       }
       NewLoc = Loc;
+      if (Callbacks)
+        Callbacks->PragmaAssumeNonNullBegin(NewLoc);
     } else {
       // Complain about attempts to leave an audit that doesn't exist.
       if (!BeginLoc.isValid()) {
@@ -1740,6 +1743,8 @@ struct PragmaAssumeNonNullHandler : public PragmaHandler {
         return;
       }
       NewLoc = SourceLocation();
+      if (Callbacks)
+        Callbacks->PragmaAssumeNonNullEnd(NewLoc);
     }
 
     PP.setPragmaAssumeNonNullLoc(NewLoc);