[C++][Modules] A module directive may only appear as the first preprocessing tokens in a file (#144233)

This PR is 2nd part of
[P1857R3](#107168)
implementation, and mainly implement the restriction `A module directive
may only appear as the first preprocessing tokens in a file (excluding
the global module fragment.)`:
[cpp.pre](https://eel.is/c++draft/cpp.pre):
```
module-file:
    pp-global-module-fragment[opt] pp-module group[opt] pp-private-module-fragment[opt]
```

We also refine tests use `split-file` instead of conditional macro.

Signed-off-by: yronglin <[email protected]>

Author: yronglin

clang/include/clang/Lex/Lexer.h

@@ -143,6 +143,9 @@ class Lexer : public PreprocessorLexer {
   /// True if this is the first time we're lexing the input file.
   bool IsFirstTimeLexingFile;
 
+  /// True if current lexing token is the first pp-token.
+  bool IsFirstPPToken;
+
   // NewLinePtr - A pointer to new line character '\n' being lexed. For '\r\n',
   // it also points to '\n.'
   const char *NewLinePtr;

clang/include/clang/Lex/Preprocessor.h

@@ -350,6 +350,9 @@ class Preprocessor {
   /// Whether the last token we lexed was an '@'.
   bool LastTokenWasAt = false;
 
+  /// First pp-token in current translation unit.
+  std::optional<Token> FirstPPToken;
+
   /// A position within a C++20 import-seq.
   class StdCXXImportSeq {
   public:
@@ -1766,6 +1769,20 @@ class Preprocessor {
   std::optional<LexEmbedParametersResult> LexEmbedParameters(Token &Current,
                                                              bool ForHasEmbed);
 
+  /// Whether the preprocessor already seen the first pp-token in main file.
+  bool hasSeenMainFileFirstPPToken() const { return FirstPPToken.has_value(); }
+
+  /// Record first pp-token and check if it has a Token::FirstPPToken flag.
+  void HandleMainFileFirstPPToken(const Token &Tok) {
+    if (!hasSeenMainFileFirstPPToken() && Tok.isFirstPPToken() &&
+        SourceMgr.isWrittenInMainFile(Tok.getLocation()))
+      FirstPPToken = Tok;
+  }
+
+  Token getMainFileFirstPPToken() const {
+    assert(FirstPPToken && "First main file pp-token doesn't exists");
+    return *FirstPPToken;
+  }
   bool LexAfterModuleImport(Token &Result);
   void CollectPpImportSuffix(SmallVectorImpl<Token> &Toks);
 

clang/include/clang/Lex/Token.h

@@ -86,9 +86,12 @@ class Token {
                                 // macro stringizing or charizing operator.
     CommaAfterElided = 0x200, // The comma following this token was elided (MS).
     IsEditorPlaceholder = 0x400, // This identifier is a placeholder.
-    IsReinjected = 0x800, // A phase 4 token that was produced before and
-                          // re-added, e.g. via EnterTokenStream. Annotation
-                          // tokens are *not* reinjected.
+
+    IsReinjected = 0x800,  // A phase 4 token that was produced before and
+                           // re-added, e.g. via EnterTokenStream. Annotation
+                           // tokens are *not* reinjected.
+    FirstPPToken = 0x1000, // This token is the first pp token in the
+                           // translation unit.
   };
 
   tok::TokenKind getKind() const { return Kind; }
@@ -318,6 +321,9 @@ class Token {
   /// represented as characters between '<#' and '#>' in the source code. The
   /// lexer uses identifier tokens to represent placeholders.
   bool isEditorPlaceholder() const { return getFlag(IsEditorPlaceholder); }
+
+  /// Returns true if this token is the first pp-token.
+  bool isFirstPPToken() const { return getFlag(FirstPPToken); }
 };
 
 /// Information about the conditional stack (\#if directives)

clang/include/clang/Sema/Sema.h

@@ -9822,7 +9822,8 @@ class Sema final : public SemaBase {
   DeclGroupPtrTy ActOnModuleDecl(SourceLocation StartLoc,
                                  SourceLocation ModuleLoc, ModuleDeclKind MDK,
                                  ModuleIdPath Path, ModuleIdPath Partition,
-                                 ModuleImportState &ImportState);
+                                 ModuleImportState &ImportState,
+                                 bool IntroducerIsFirstPPToken);
 
   /// The parser has processed a global-module-fragment declaration that begins
   /// the definition of the global module fragment of the current module unit.

clang/lib/Lex/Lexer.cpp

@@ -174,6 +174,8 @@ void Lexer::InitLexer(const char *BufStart, const char *BufPtr,
   ExtendedTokenMode = 0;
 
   NewLinePtr = nullptr;
+
+  IsFirstPPToken = true;
 }
 
 /// Lexer constructor - Create a new lexer object for the specified buffer
@@ -3725,13 +3727,22 @@ bool Lexer::Lex(Token &Result) {
     HasLeadingEmptyMacro = false;
   }
 
+  if (IsFirstPPToken) {
+    Result.setFlag(Token::FirstPPToken);
+    IsFirstPPToken = false;
+  }
+
   bool atPhysicalStartOfLine = IsAtPhysicalStartOfLine;
   IsAtPhysicalStartOfLine = false;
   bool isRawLex = isLexingRawMode();
   (void) isRawLex;
   bool returnedToken = LexTokenInternal(Result, atPhysicalStartOfLine);
   // (After the LexTokenInternal call, the lexer might be destroyed.)
   assert((returnedToken || !isRawLex) && "Raw lex must succeed");
+
+  if (returnedToken && Result.isFirstPPToken() && PP &&
+      !PP->hasSeenMainFileFirstPPToken())
+    PP->HandleMainFileFirstPPToken(Result);
   return returnedToken;
 }
 
@@ -4535,6 +4546,8 @@ const char *Lexer::convertDependencyDirectiveToken(
   Result.setFlag((Token::TokenFlags)DDTok.Flags);
   Result.setLength(DDTok.Length);
   BufferPtr = TokPtr + DDTok.Length;
+  if (PP && !PP->hasSeenMainFileFirstPPToken() && Result.isFirstPPToken())
+    PP->HandleMainFileFirstPPToken(Result);
   return TokPtr;
 }
 

clang/lib/Lex/PPDirectives.cpp

@@ -1242,6 +1242,9 @@ void Preprocessor::HandleDirective(Token &Result) {
   // pp-directive.
   bool ReadAnyTokensBeforeDirective =CurPPLexer->MIOpt.getHasReadAnyTokensVal();
 
+  if (!hasSeenMainFileFirstPPToken())
+    HandleMainFileFirstPPToken(Result);
+
   // Save the '#' token in case we need to return it later.
   Token SavedHash = Result;
 

clang/lib/Lex/PPMacroExpansion.cpp

@@ -469,6 +469,9 @@ bool Preprocessor::HandleMacroExpandedIdentifier(Token &Identifier,
   // to disable the optimization in this case.
   if (CurPPLexer) CurPPLexer->MIOpt.ExpandedMacro();
 
+  if (!hasSeenMainFileFirstPPToken())
+    HandleMainFileFirstPPToken(Identifier);
+
   // If this is a builtin macro, like __LINE__ or _Pragma, handle it specially.
   if (MI->isBuiltinMacro()) {
     if (Callbacks)

clang/lib/Lex/Preprocessor.cpp

@@ -247,6 +247,8 @@ void Preprocessor::DumpToken(const Token &Tok, bool DumpFlags) const {
     llvm::errs() << " [LeadingSpace]";
   if (Tok.isExpandDisabled())
     llvm::errs() << " [ExpandDisabled]";
+  if (Tok.isFirstPPToken())
+    llvm::errs() << " [First pp-token]";
   if (Tok.needsCleaning()) {
     const char *Start = SourceMgr.getCharacterData(Tok.getLocation());
     llvm::errs() << " [UnClean='" << StringRef(Start, Tok.getLength())

clang/lib/Parse/Parser.cpp

@@ -2340,7 +2340,8 @@ void Parser::ParseMicrosoftIfExistsExternalDeclaration() {
 
 Parser::DeclGroupPtrTy
 Parser::ParseModuleDecl(Sema::ModuleImportState &ImportState) {
-  SourceLocation StartLoc = Tok.getLocation();
+  Token Introducer = Tok;
+  SourceLocation StartLoc = Introducer.getLocation();
 
   Sema::ModuleDeclKind MDK = TryConsumeToken(tok::kw_export)
                                  ? Sema::ModuleDeclKind::Interface
@@ -2359,7 +2360,7 @@ Parser::ParseModuleDecl(Sema::ModuleImportState &ImportState) {
   // Parse a global-module-fragment, if present.
   if (getLangOpts().CPlusPlusModules && Tok.is(tok::semi)) {
     SourceLocation SemiLoc = ConsumeToken();
-    if (ImportState != Sema::ModuleImportState::FirstDecl) {
+    if (!Introducer.isFirstPPToken()) {
       Diag(StartLoc, diag::err_global_module_introducer_not_at_start)
         << SourceRange(StartLoc, SemiLoc);
       return nullptr;
@@ -2416,7 +2417,7 @@ Parser::ParseModuleDecl(Sema::ModuleImportState &ImportState) {
   ExpectAndConsumeSemi(diag::err_module_expected_semi);
 
   return Actions.ActOnModuleDecl(StartLoc, ModuleLoc, MDK, Path, Partition,
-                                 ImportState);
+                                 ImportState, Introducer.isFirstPPToken());
 }
 
 Decl *Parser::ParseModuleImport(SourceLocation AtLoc,

clang/lib/Sema/SemaModule.cpp

@@ -263,11 +263,11 @@ static bool DiagReservedModuleName(Sema &S, const IdentifierInfo *II,
 Sema::DeclGroupPtrTy
 Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
                       ModuleDeclKind MDK, ModuleIdPath Path,
-                      ModuleIdPath Partition, ModuleImportState &ImportState) {
+                      ModuleIdPath Partition, ModuleImportState &ImportState,
+                      bool IntroducerIsFirstPPToken) {
   assert(getLangOpts().CPlusPlusModules &&
          "should only have module decl in standard C++ modules");
 
-  bool IsFirstDecl = ImportState == ModuleImportState::FirstDecl;
   bool SeenGMF = ImportState == ModuleImportState::GlobalFragment;
   // If any of the steps here fail, we count that as invalidating C++20
   // module state;
@@ -333,14 +333,11 @@ Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
           SeenGMF == (bool)this->TheGlobalModuleFragment) &&
          "mismatched global module state");
 
-  // In C++20, the module-declaration must be the first declaration if there
-  // is no global module fragment.
-  if (getLangOpts().CPlusPlusModules && !IsFirstDecl && !SeenGMF) {
+  // In C++20, A module directive may only appear as the first preprocessing
+  // tokens in a file (excluding the global module fragment.).
+  if (getLangOpts().CPlusPlusModules && !IntroducerIsFirstPPToken && !SeenGMF) {
     Diag(ModuleLoc, diag::err_module_decl_not_at_start);
-    SourceLocation BeginLoc =
-        ModuleScopes.empty()
-            ? SourceMgr.getLocForStartOfFile(SourceMgr.getMainFileID())
-            : ModuleScopes.back().BeginLoc;
+    SourceLocation BeginLoc = PP.getMainFileFirstPPToken().getLocation();
     if (BeginLoc.isValid()) {
       Diag(BeginLoc, diag::note_global_module_introducer_missing)
           << FixItHint::CreateInsertion(BeginLoc, "module;\n");