Serialize PragmaAssumeNonNullLoc to support preambles

[dgoldman]

Summary

Previously, if a #pragma clang assume_nonnull begin was at the
end of a premable with a #pragma clang assume_nonnull end at the
end of the main file, clang would diagnose an unterminated begin in
the preamble and an unbalanced end in the main file.

With this change, those errors no longer occur and the case above is
now properly handled. I've added a corresponding test to clangd,
which makes use of preambles, in order to verify this works as
expected.

[dgoldman]

LMK if there are other folks more suited to review this, I'm not sure...

[sammccall]

clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

750

This test seems to have a lot of extraneous elements.

The following seems sufficient, with no header:

#pragma clang assume_nonnull begin
void foo(int *x);
#pragma clang assume_nonnull end

750

this should be a clang test, rather than a clangd test (or in addition)

probably under clang/PCH/, of the form

// RUN: %clang_cc1 -emit-pch ...
// RUN: %clang_cc1 -include-pch -fsyntax-only -verify -ast-dump ...
// and optionally without PCH:
// RUN: %clang_cc1 -fsyntax-only -verify -ast-dump ...

Then you can verify both the lack of diagnostics and the AST that gets generated.

777

you verify no diagnostics, but likely also want to verify that the parameter has the nonnull attribute.

This is probably simple enough, e.g

ParsedAST AST = TU.build();
ParmVarDecl *X = cast<FunctionDecl>(findDecl(AST, "foo")).getParamDecl(0);
ASSERT_TRUE(X->hasAttr<NonnullAttr>());

(or by matching the AST dump)

clang/include/clang/Lex/PreprocessorOptions.h

131–132

nit: "Similarly, we track an unterminated #pragma assume_nonnull"?

Currently it's not completely clear what the connection is.

clang/lib/Lex/PPLexerChange.cpp

430

unless we're hitting _the end_ of the preamble

This is a special case though, pull it out of the one-sentence summary (as the other special case is)

433

It's not valid *anywhere* we exit a file when building a preamble, only when we fall off the end of the preamble region itself within the main file.

434

what's the PredefinesFileID special case about?

clang/lib/Serialization/ASTWriter.cpp

2303

nit: #pragma assume_nonnull

2304

this "might" is doing a lot of spooky action at a distance.

At minimum we should assert this is a preamble, right? Or move this block into the if() below?

[dgoldman]

Fixes for review

clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

750

Hmm I tried this earlier, maybe I did it wrong, but I don't think it works since it doesn't set the GeneratePreamble bit

clang/lib/Lex/PPLexerChange.cpp

430

Done, but isn't it still 2 places - generating the preamble and actually using it both need the exception?

434

See ExitedFromPredefinesFile below, this is how I check if we're falling off the end of the preamble region into the main file (same as done for the conditional stack), LMK if there's a better way.

clang/lib/Serialization/ASTWriter.cpp

2304

Yeah that makes added, added an assert.

[sammccall]

clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

750

(please don't mark comment threads done unless they've been addressed)

You're right, regular PCHes are not preambles and so don't handle these cases.
Probably the best place then is something using c-index-test similar to clang/test/Index/preamble-conditionals.cpp

756

nit: not necessary, this pragma isn't tied to objc

clang/lib/Lex/PPLexerChange.cpp

434

I don't understand how this checks if we're falling off the preamble region, and the code around ExitedFromPredefinesFile doesn't clarify this for me. Can you explain?

The if (ExitedFromPredefinesFile) appears to be handling the logical *insertion* of preamble PP events when *consuming* a preamble, which is not what you're trying to do here.

The condition here is of the form "if we have an open pragma, and we're not generating a preamble, and some other stuff, then diagnose". So if the baseline case is "diagnose" and the preamble case is an exception, the "other stuff" clauses don't limit the preamble exception, they add extra exceptions!

clang/lib/Serialization/ASTWriter.cpp

2303

it should be -> we should be
the preamble -> a preamble
?

[dgoldman]

Minor fixes

[dgoldman]

clang/lib/Lex/PPLexerChange.cpp

434

!(CurLexer && CurLexer->getFileID() == PredefinesFileID) makes sure that if we're consuming the preamble, we don't emit the warning. AFAICT this is the way to tell that the preamble is terminated, since the current file ID being processed is the predefines file and the file is now terminated as of this method call (since !isEndOfMacro && !(CurLexer && CurLexer->Is_PragmaLexer)). !this->PPOpts->GeneratePreamble makes sure that if we're generating the preamble, we don't emit the warning. We need to special case both cases otherwise we get an error when generating the preamble or when we load the preamble before even processing the rest of the main file.

Does that makes sense?

[sammccall]

clang/lib/Lex/PPLexerChange.cpp

434

!(CurLexer && CurLexer->getFileID() == PredefinesFileID) makes sure that if we're consuming the preamble, we don't emit the warning

1. if we're consuming the preamble, what exactly is the sequence of events that would otherwise lead us to emit the warning?
2. what if we're in the predefines file for some other reason instead?

I'm hoping you'll explain to me what you think the predefines file is, and what its relationship to the preamble is, and so why this condition is correct :-)

My understanding (which is pretty shaky!) is that the predefines file is a kind of catchall used to inject text into the preprocessor that doesn't appear in the source file - that it contains definitions of builtin integer types, macros defined by -D on the command line, and so on. If it has any relationship to the preamble, it's something subtle that's to do with the relative order in which the preprocessor sees entities like the predefines, the preprocesor, and the main file.
So according to my understanding, interpreting "reaching EOF in the predefines file" as "consuming a preamble" is either wrong or something very subtle requiring a significant comment.

The code you refer to below is doing something very different: it's saying that if we have a preamble, then reaching the end of the predefines file is the *trigger* to inject state from it!

!this->PPOpts->GeneratePreamble makes sure that if we're generating the preamble, we don't emit the warning

Sure, but that doesn't sound correct at all! A preamble mostly consists of parsing a lot of headers, and if any of those headers have an unpaired pragma, we should be warning on that at EOF. It's only if we hit the "pretend" EOF from the truncated main file that we want to suppress the warning.

[sammccall]

clang/lib/Lex/PPLexerChange.cpp

434

the preprocessor sees entities like the predefines, the preprocesor, and the main file.

Oops... sees the predefines, any injected events from the preamble, and the main file.

[dgoldman]

clang/lib/Lex/PPLexerChange.cpp

434

I think that definition is correct, but given that there is only one predefines file, when it ends, it must go into the main file, no? And if we have a pragma assume nonnull loc, it shouldn't be from the clang builtins, it should be from the preamble (I'd certainly hope!)

re: headers, I think there might be a misunderstanding, as soon as an #include is seen the pragma is ended here: https://github.com/llvm/llvm-project/blob/7631c366c8589dda488cb7ff1df26cc134002208/clang/lib/Lex/PPDirectives.cpp#L2006, so if PragmaAssumeNonNullLoc is valid here we haven't had any includes since the pragma was seen.

[sammccall]

clang/lib/Lex/PPLexerChange.cpp

434

I think that definition is correct, but given that there is only one predefines file, when it ends, it must go into the main file, no?

I'm sorry, but i really can't follow what you're asking, or how it relates to the original problem, or what the problem even is. If *you're* not sure, it's probably worth tracing execution of the preprocessor and describing the precise sequence of events that cause a problem. And then trying to come up with a *direct* way of solving them.

Here's my *guess* as to what's going on (but please check!):

• the assume_nonnull state is very sensitive to shifts between files - it's supposed to be off before reaching EOF and before #include. (This is why it can appear to be file-local syntax but be implemented as a global flag).
• therefore the AssumeNonnullLoc is not just some property of the preprocessor, it's a transient state that should be set at precisely the right point in the program. However you're setting it at the very beginning, as a side effect of ReadASTBlock().
• when consuming a preamble, the preprocessor simulates a structure like:

// AssumeNonnullLoc is being set here
#include <predefines> // implicit
#inject pp conditional stack
// AssumeNonnnullLoc *should* be set here
// now continue with the post-preamble part of main file

• because AssumeNonnullLoc is set too early, the PP events in between these locations mess with it. Fortunately that is probably only the predefines, although that's a brittle assumption to make. And we don't run into problems when entering the predefines file, because HandleHeaderIncludeOrImport isn't called, we just enter the file directly. However you do run into issues when exiting that file again, and so a special case works around those.

If this is right, it's not a good solution. There are other observable effects of setting AssumeNonnullLoc in the wrong place. All the predefines will be parsed in "assume nonnull mode". And if any files are #included from the predefines file, then the enter/exit events will cause problems. Predefines may contain #include with the -include and -imacros clang flags, and due to clangd's preamble patching.

In this case the right fix is to set AssumeNonnullLoc at the correct point in the file: triggered by exiting the predefines, along with the PP conditional stack. The way to do this is probably to have a second property on preprocessor which stores the loc that will be restored when the trigger occurs. (Similar to how PP distinguishes between the *current* conditional stack and the replayable one set by ASTReader when it calls PP.setReplayablePreambleConditionalStack(ConditionalStack, SkipInfo);

re: headers, I think there might be a misunderstanding, as soon as an #include is seen the pragma is ended here

The scenario is:

We are building a preamble from main.c
main.c:

#include "foo.h"
int x;

foo.h:

 #pragma clang assume_nonnull begin
void foo(int*);
// missing end pragma

We should issue a diagnostic when we hit the end of foo.h, but with the code you have here it will be supposed instead. (There's no #include in scope of the pragma here, so I don't think that's relevant)

[dgoldman]

Properly handle generation and restoring

[dgoldman]

clang/lib/Lex/PPLexerChange.cpp

434

🤦 Thanks, that example is super helpful, not sure why I wasn't able to understand it earlier. Yeah, this doesn't properly handle EOF for headers during preamble generation and we should only set the assume_nonnull loc only when transitioning, not earlier. Like you said, we need to basically mirror the replay/record logic for the preamble conditional stack for it to work properly. I went ahead and implemented the approach that you mentioned to restore the value when exiting from the predefines as well as added a test to clangd for the edge case you gave. PTAL, I kept the storage separate but I wonder if it would be better to merge in with the existing replay logic.

[dgoldman]

Minor lint fixes

[dgoldman]

if/else formatting

[dgoldman]

Another formatting fix

[sammccall]

Thanks, this looks really solid now!

clang/lib/Lex/PPLexerChange.cpp

430

revert comment change? this is still specific to assume_nonnull

[dgoldman]

Revert comment change

[dgoldman]

Rebase

[sammccall]

(Just checking you've seen this is accepted and you can land when you're ready!)

clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

TEST(DiagnosticsTest, RecursivePreambleIfndefGuard) {
   )cpp");
   TU.Filename = "foo.h";
   // FIXME: should be no errors here.
   EXPECT_THAT(*TU.build().getDiagnostics(),
               ElementsAre(diagName("pp_including_mainfile_in_preamble")));
   EXPECT_THAT(TU.build().getLocalTopLevelDecls(), SizeIs(1));
 }
 
+TEST(DiagnosticsTest, PreambleWithPragmaAssumeNonnull) {
+  auto TU = TestTU::withCode(R"cpp(
+#pragma clang assume_nonnull begin
+void foo(int *x);
+#pragma clang assume_nonnull end
+)cpp");
+  auto AST = TU.build();
+  EXPECT_THAT(*AST.getDiagnostics(), IsEmpty());
+  const auto *X = cast<FunctionDecl>(findDecl(AST, "foo")).getParamDecl(0);
+  ASSERT_TRUE(X->getOriginalType()->getNullability(X->getASTContext()) ==
+              NullabilityKind::NonNull);
+}
+
+TEST(DiagnosticsTest, PreambleHeaderWithBadPragmaAssumeNonnull) {
+  Annotations Header(R"cpp(
+#pragma clang assume_nonnull begin  // error-ok
+void foo(int *X);
+)cpp");
+  auto TU = TestTU::withCode(R"cpp(
+#include "foo.h"  // unterminated assume_nonnull should not affect bar.
+void bar(int *Y);
+)cpp");
+  TU.AdditionalFiles = {{"foo.h", std::string(Header.code())}};
+  auto AST = TU.build();
+  EXPECT_THAT(*AST.getDiagnostics(),
+              ElementsAre(diagName("pp_eof_in_assume_nonnull")));
+  const auto *X = cast<FunctionDecl>(findDecl(AST, "foo")).getParamDecl(0);
+  ASSERT_TRUE(X->getOriginalType()->getNullability(X->getASTContext()) ==
+              NullabilityKind::NonNull);
+  const auto *Y = cast<FunctionDecl>(findDecl(AST, "bar")).getParamDecl(0);
+  ASSERT_FALSE(
+      Y->getOriginalType()->getNullability(X->getASTContext()).hasValue());
+}
+
 TEST(DiagnosticsTest, InsideMacros) {
   Annotations Test(R"cpp(
     #define TEN 10
     #define RET(x) return x + 10
 
     int* foo() {
       RET($foo[[0]]); // error-ok
     }

clang/include/clang/Lex/Preprocessor.h

private:
   /// The identifier and source location of the currently-active
   /// \#pragma clang arc_cf_code_audited begin.
   std::pair<IdentifierInfo *, SourceLocation> PragmaARCCFCodeAuditedInfo;
 
   /// The source location of the currently-active
   /// \#pragma clang assume_nonnull begin.
   SourceLocation PragmaAssumeNonNullLoc;
 
+  /// Set only for preambles which end with an active
+  /// \#pragma clang assume_nonnull begin.
+  ///
+  /// When the preamble is loaded into the main file,
+  /// `PragmaAssumeNonNullLoc` will be set to this to
+  /// replay the unterminated assume_nonnull.
+  SourceLocation PreambleRecordedPragmaAssumeNonNullLoc;
+
   /// True if we hit the code-completion point.
   bool CodeCompletionReached = false;
 
   /// The code completion token containing the information
   /// on the stem that is to be code completed.
   IdentifierInfo *CodeCompletionII = nullptr;
 
   /// Range for the code completion token.
public:
   }
 
   /// Set the location of the currently-active \#pragma clang
   /// assume_nonnull begin.  An invalid location ends the pragma.
   void setPragmaAssumeNonNullLoc(SourceLocation Loc) {
     PragmaAssumeNonNullLoc = Loc;
   }
 
+  /// Get the location of the recorded unterminated \#pragma clang
+  /// assume_nonnull begin in the preamble, if one exists.
+  ///
+  /// Returns an invalid location if the premable did not end with
+  /// such a pragma active or if there is no recorded preamble.
+  SourceLocation getPreambleRecordedPragmaAssumeNonNullLoc() const {
+    return PreambleRecordedPragmaAssumeNonNullLoc;
+  }
+
+  /// Record the location of the unterminated \#pragma clang
+  /// assume_nonnull begin in the preamble.
+  void setPreambleRecordedPragmaAssumeNonNullLoc(SourceLocation Loc) {
+    PreambleRecordedPragmaAssumeNonNullLoc = Loc;
+  }
+
   /// Set the directory in which the main file should be considered
   /// to have been found, if it is not a real file.
   void setMainFileDir(const DirectoryEntry *Dir) {
     MainFileDir = Dir;
   }
 
   /// Instruct the preprocessor to skip part of the main source file.
   ///

clang/include/clang/Lex/PreprocessorOptions.h

public:
   /// The boolean indicates whether the preamble ends at the start of a new
   /// line.
   std::pair<unsigned, bool> PrecompiledPreambleBytes;
 
   /// True indicates that a preamble is being generated.
   ///
   /// When the lexer is done, one of the things that need to be preserved is the
   /// conditional #if stack, so the ASTWriter/ASTReader can save/restore it when
-  /// processing the rest of the file.
+  /// processing the rest of the file. Similarly, we track an unterminated
+  /// #pragma assume_nonnull.
   bool GeneratePreamble = false;
 
   /// Whether to write comment locations into the PCH when building it.
   /// Reading the comments from the PCH can be a performance hit even if the
   /// clients don't use them.
   bool WriteCommentListToPCH = true;
 
   /// When enabled, preprocessor is in a mode for parsing a single file only.

clang/include/clang/Serialization/ASTBitCodes.h

enum ASTRecordTypes {
   /// Record code for the Decls to be checked for deferred diags.
   DECLS_TO_CHECK_FOR_DEFERRED_DIAGS = 64,
 
   /// Record code for \#pragma float_control options.
   FLOAT_CONTROL_PRAGMA_OPTIONS = 65,
 
   /// Record code for included files.
   PP_INCLUDED_FILES = 66,
+
+  /// Record code for an unterminated \#pragma clang assume_nonnull begin
+  /// recorded in a preamble.
+  PP_ASSUME_NONNULL_LOC = 67,
 };
 
 /// Record types used within a source manager block.
 enum SourceManagerRecordTypes {
   /// Describes a source location entry (SLocEntry) for a
   /// file.
   SM_SLOC_FILE_ENTRY = 1,
 

clang/lib/Lex/PPLexerChange.cpp

if (PragmaARCCFCodeAuditedInfo.second.isValid() && !isEndOfMacro &&
       !(CurLexer && CurLexer->Is_PragmaLexer)) {
     Diag(PragmaARCCFCodeAuditedInfo.second,
          diag::err_pp_eof_in_arc_cf_code_audited);
 
     // Recover by leaving immediately.
     PragmaARCCFCodeAuditedInfo = {nullptr, SourceLocation()};
   }
 
   // Complain about reaching a true EOF within assume_nonnull.
   // We don't want to complain about reaching the end of a macro
   // instantiation or a _Pragma.
   if (PragmaAssumeNonNullLoc.isValid() &&
       !isEndOfMacro && !(CurLexer && CurLexer->Is_PragmaLexer)) {
+    // If we're at the end of generating a preamble, we should record the
+    // unterminated \#pragma clang assume_nonnull so we can restore it later
+    // when the preamble is loaded into the main file.
+    if (isRecordingPreamble() && isInPrimaryFile())
+      PreambleRecordedPragmaAssumeNonNullLoc = PragmaAssumeNonNullLoc;
+    else
       Diag(PragmaAssumeNonNullLoc, diag::err_pp_eof_in_assume_nonnull);
-
     // Recover by leaving immediately.
     PragmaAssumeNonNullLoc = SourceLocation();
   }
 
   bool LeavingPCHThroughHeader = false;
 
   // If this is a #include'd file, pop it off the include stack and continue
   // lexing the #includer file.
if (!IncludeMacroStack.empty()) {
     // Notify the client, if desired, that we are in a new source file.
     if (Callbacks && !isEndOfMacro && CurPPLexer) {
       SrcMgr::CharacteristicKind FileType =
         SourceMgr.getFileCharacteristic(CurPPLexer->getSourceLocation());
       Callbacks->FileChanged(CurPPLexer->getSourceLocation(),
                              PPCallbacks::ExitFile, FileType, ExitedFID);
     }
 
-    // Restore conditional stack from the preamble right after exiting from the
-    // predefines file.
-    if (ExitedFromPredefinesFile)
+    // Restore conditional stack as well as the recorded
+    // \#pragma clang assume_nonnull from the preamble right after exiting
+    // from the predefines file.
+    if (ExitedFromPredefinesFile) {
       replayPreambleConditionalStack();
+      if (PreambleRecordedPragmaAssumeNonNullLoc.isValid())
+        PragmaAssumeNonNullLoc = PreambleRecordedPragmaAssumeNonNullLoc;
+    }
 
     if (!isEndOfMacro && CurPPLexer && FoundPCHThroughHeader &&
         (isInPrimaryFile() ||
          CurPPLexer->getFileID() == getPredefinesFileID())) {
       // Leaving the through header. Continue directly to end of main file
       // processing.
       LeavingPCHThroughHeader = true;
     } else {

clang/lib/Serialization/ASTReader.cpp

while (true) {
 
     // If we're not loading an AST context, we don't care about most records.
     if (!ContextObj) {
       switch (RecordType) {
       case IDENTIFIER_TABLE:
       case IDENTIFIER_OFFSET:
       case INTERESTING_IDENTIFIERS:
       case STATISTICS:
+      case PP_ASSUME_NONNULL_LOC:
       case PP_CONDITIONAL_STACK:
       case PP_COUNTER_VALUE:
       case SOURCE_LOCATION_OFFSETS:
       case MODULE_OFFSET_MAP:
       case SOURCE_MANAGER_LINE_TABLE:
       case SOURCE_LOCATION_PRELOADS:
       case PPD_ENTITIES_OFFSETS:
       case HEADER_SEARCH_TABLE:
case REFERENCED_SELECTOR_POOL:
           ReferencedSelectorsData.push_back(getGlobalSelectorID(F,
                                                                 Record[Idx++]));
           ReferencedSelectorsData.push_back(ReadSourceLocation(F, Record, Idx).
                                               getRawEncoding());
         }
       }
       break;
 
+    case PP_ASSUME_NONNULL_LOC: {
+      unsigned Idx = 0;
+      if (!Record.empty())
+        PP.setPreambleRecordedPragmaAssumeNonNullLoc(
+            ReadSourceLocation(F, Record, Idx));
+      break;
+    }
+
     case PP_CONDITIONAL_STACK:
       if (!Record.empty()) {
         unsigned Idx = 0, End = Record.size() - 1;
         bool ReachedEOFWhileSkipping = Record[Idx++];
         llvm::Optional<Preprocessor::PreambleSkipInfo> SkipInfo;
         if (ReachedEOFWhileSkipping) {
           SourceLocation HashToken = ReadSourceLocation(F, Record, Idx);
           SourceLocation IfTokenLoc = ReadSourceLocation(F, Record, Idx);

clang/lib/Serialization/ASTWriter.cpp

#define RECORD(X) EmitRecordID(X, #X, Stream, Record)
   RECORD(MSSTRUCT_PRAGMA_OPTIONS);
   RECORD(POINTERS_TO_MEMBERS_PRAGMA_OPTIONS);
   RECORD(UNUSED_LOCAL_TYPEDEF_NAME_CANDIDATES);
   RECORD(DELETE_EXPRS_TO_ANALYZE);
   RECORD(CUDA_PRAGMA_FORCE_HOST_DEVICE_DEPTH);
   RECORD(PP_CONDITIONAL_STACK);
   RECORD(DECLS_TO_CHECK_FOR_DEFERRED_DIAGS);
   RECORD(PP_INCLUDED_FILES);
+  RECORD(PP_ASSUME_NONNULL_LOC);
 
   // SourceManager Block.
   BLOCK(SOURCE_MANAGER_BLOCK);
   RECORD(SM_SLOC_FILE_ENTRY);
   RECORD(SM_SLOC_BUFFER_ENTRY);
   RECORD(SM_SLOC_BUFFER_BLOB);
   RECORD(SM_SLOC_BUFFER_BLOB_COMPRESSED);
   RECORD(SM_SLOC_EXPANSION_ENTRY);
void ASTWriter::WritePreprocessor(const Preprocessor &PP, bool IsModule) {
   RecordData ModuleMacroRecord;
 
   // If the preprocessor __COUNTER__ value has been bumped, remember it.
   if (PP.getCounterValue() != 0) {
     RecordData::value_type Record[] = {PP.getCounterValue()};
     Stream.EmitRecord(PP_COUNTER_VALUE, Record);
   }
 
+  // If we have a recorded #pragma assume_nonnull, remember it so it can be
+  // replayed when the preamble terminates into the main file.
+  SourceLocation AssumeNonNullLoc =
+      PP.getPreambleRecordedPragmaAssumeNonNullLoc();
+  if (AssumeNonNullLoc.isValid()) {
+    assert(PP.isRecordingPreamble());
+    AddSourceLocation(AssumeNonNullLoc, Record);
+    Stream.EmitRecord(PP_ASSUME_NONNULL_LOC, Record);
+    Record.clear();
+  }
+
   if (PP.isRecordingPreamble() && PP.hasRecordedPreamble()) {
     assert(!IsModule);
     auto SkipInfo = PP.getPreambleSkipInfo();
     if (SkipInfo.hasValue()) {
       Record.push_back(true);
       AddSourceLocation(SkipInfo->HashTokenLoc, Record);
       AddSourceLocation(SkipInfo->IfTokenLoc, Record);
       Record.push_back(SkipInfo->FoundNonSkipPortion);

clang/test/Index/preamble-assume-nonnull.c

+// RUN: env CINDEXTEST_EDITING=1 c-index-test -test-load-source local %s 2>&1 \
+// RUN: | FileCheck %s --implicit-check-not "error:"
+
+#pragma clang assume_nonnull begin
+void foo(int *x);
+#pragma clang assume_nonnull end