A module directive may only appear as the first preprocessing tokens in a file

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

@@ -137,6 +137,24 @@ struct CXXStandardLibraryVersionInfo {
   std::uint64_t Version;
 };
 
+class ExportContextualKeywordInfo {
+  Token ExportTok;
+  bool AtPhysicalStartOfLine = false;
+
+public:
+  ExportContextualKeywordInfo() = default;
+  ExportContextualKeywordInfo(const Token &Tok, bool AtPhysicalStartOfLine)
+      : ExportTok(Tok), AtPhysicalStartOfLine(AtPhysicalStartOfLine) {}
+
+  bool isValid() const { return ExportTok.is(tok::kw_export); }
+  bool isAtPhysicalStartOfLine() const { return AtPhysicalStartOfLine; }
+  Token getExportTok() const { return ExportTok; }
+  void reset() {
+    ExportTok.startToken();
+    AtPhysicalStartOfLine = false;
+  }
+};
+
 /// Engages in a tight little dance with the lexer to efficiently
 /// preprocess tokens.
 ///
@@ -359,14 +377,11 @@ class Preprocessor {
   /// Whether we're declaring a standard C++20 named Modules.
   bool DeclaringCXXNamedModules = false;
 
-  struct ExportContextualKeywordInfo {
-    Token ExportTok;
-    bool TokAtPhysicalStartOfLine;
-  };
-
   /// Whether the last token we lexed was an 'export' keyword.
-  std::optional<ExportContextualKeywordInfo> LastTokenWasExportKeyword =
-      std::nullopt;
+  ExportContextualKeywordInfo LastTokenWasExportKeyword;
+
+  /// First pp-token in current translation unit.
+  Token FirstPPToken;
 
   /// A position within a C++20 import-seq.
   class StdCXXImportSeq {
@@ -1774,6 +1789,9 @@ class Preprocessor {
   void HandleCXXImportDirective(Token Import);
   void HandleCXXModuleDirective(Token Module);
 
+  void setFirstPPToken(const Token &Tok) { FirstPPToken = Tok; }
+  Token getFirstPPToken() const { return FirstPPToken; }
+
   /// Callback invoked when the lexer sees one of export, import or module token
   /// at the start of a line.
   ///

clang/include/clang/Lex/Token.h

@@ -89,6 +89,8 @@ class Token {
     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; }
@@ -325,6 +327,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

@@ -9937,9 +9937,9 @@ class Sema final : public SemaBase {
   /// of a module interface or implementation.
   DeclGroupPtrTy ActOnModuleDecl(SourceLocation StartLoc,
                                  SourceLocation ModuleLoc, ModuleDeclKind MDK,
-                                 ModuleNameLoc *PathLoc,
-                                 ModuleNameLoc *PartitionLoc,
-                                 ModuleImportState &ImportState);
+                                 ModuleIdPath Path, ModuleIdPath Partition,
+                                 ModuleImportState &ImportState,
+                                 bool AtStartOfTU);
 
   /// 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

@@ -187,6 +187,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
@@ -3739,11 +3741,20 @@ 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);
+
+  if (returnedToken && Result.isFirstPPToken() && PP)
+    PP->setFirstPPToken(Result);
+
   // (After the LexTokenInternal call, the lexer might be destroyed.)
   assert((returnedToken || !isRawLex) && "Raw lex must succeed");
   return returnedToken;

clang/lib/Lex/PPDirectives.cpp

@@ -4089,7 +4089,7 @@ void Preprocessor::HandleCXXImportDirective(Token ImportTok) {
       this->ImportingCXXNamedModules);
   ImportingCXXNamedModules = true;
 
-  if (LastTokenWasExportKeyword)
+  if (LastTokenWasExportKeyword.isValid())
     LastTokenWasExportKeyword.reset();
 
   Token Tok;
@@ -4216,13 +4216,26 @@ void Preprocessor::HandleCXXImportDirective(Token ImportTok) {
 
 void Preprocessor::HandleCXXModuleDirective(Token ModuleTok) {
   assert(getLangOpts().CPlusPlusModules && ModuleTok.is(tok::kw_module));
-  SourceLocation StartLoc = ModuleTok.getLocation();
-  if (LastTokenWasExportKeyword) {
-    StartLoc = LastTokenWasExportKeyword->ExportTok.getLocation();
+  Token Introducer = ModuleTok;
+  if (LastTokenWasExportKeyword.isValid()) {
+    Introducer = LastTokenWasExportKeyword.getExportTok();
     LastTokenWasExportKeyword.reset();
   }
-  bool IsInHeaderInclusion = !IncludeMacroStack.empty();
-  bool IsInConditionBlock = CurPPLexer->getConditionalStackDepth() != 0;
+
+  SourceLocation StartLoc = Introducer.getLocation();
+  if (!IncludeMacroStack.empty()) {
+    SourceLocation End = DiscardUntilEndOfDirective().getEnd();
+    Diag(StartLoc, diag::err_module_decl_in_header)
+        << SourceRange(StartLoc, End);
+    return;
+  }
+
+  if (CurPPLexer->getConditionalStackDepth() != 0) {
+    SourceLocation End = DiscardUntilEndOfDirective().getEnd();
+    Diag(StartLoc, diag::err_pp_cond_span_module_decl)
+        << SourceRange(StartLoc, End);
+    return;
+  }
 
   Token Tok;
   SourceLocation UseLoc = ModuleTok.getLocation();
@@ -4305,13 +4318,6 @@ void Preprocessor::HandleCXXModuleDirective(Token ModuleTok) {
     End = CheckEndOfDirective(ModuleTok.getIdentifierInfo()->getName());
   else
     End = DirToks.pop_back_val().getLocation();
-  if (IsInHeaderInclusion)
-    Diag(StartLoc, diag::err_module_decl_in_header)
-        << SourceRange(StartLoc, End);
-
-  if (IsInConditionBlock)
-    Diag(StartLoc, diag::err_pp_cond_span_module_decl)
-        << SourceRange(StartLoc, End);
 
   EnterModuleSuffixTokenStream(DirToks);
 }

clang/lib/Lex/Preprocessor.cpp

@@ -1185,11 +1185,10 @@ bool Preprocessor::HandleModuleContextualKeyword(
     return false;
   }
 
-  if (LastTokenWasExportKeyword) {
-    auto Export = *LastTokenWasExportKeyword;
+  if (LastTokenWasExportKeyword.isValid()) {
     // The export keyword was not at the start of line, it's not a
     // directive-introducing token.
-    if (!Export.TokAtPhysicalStartOfLine)
+    if (!LastTokenWasExportKeyword.isAtPhysicalStartOfLine())
       return false;
     // [cpp.pre]/1.4
     // export                  // not a preprocessing directive

clang/lib/Parse/Parser.cpp

@@ -17,6 +17,7 @@
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/DiagnosticParse.h"
 #include "clang/Basic/StackExhaustionHandler.h"
+#include "clang/Lex/ModuleLoader.h"
 #include "clang/Parse/RAIIObjectsForParser.h"
 #include "clang/Sema/DeclSpec.h"
 #include "clang/Sema/EnterExpressionEvaluationContext.h"
@@ -2303,7 +2304,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
@@ -2322,7 +2324,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;
@@ -2378,8 +2380,10 @@ Parser::ParseModuleDecl(Sema::ModuleImportState &ImportState) {
 
   ExpectAndConsumeSemi(diag::err_module_expected_semi);
 
-  return Actions.ActOnModuleDecl(StartLoc, ModuleLoc, MDK, Path, Partition,
-                                 ImportState);
+  return Actions.ActOnModuleDecl(
+      StartLoc, ModuleLoc, MDK, Path->getModuleIdPath(),
+      Partition ? Partition->getModuleIdPath() : ModuleIdPath{}, ImportState,
+      Introducer.isFirstPPToken());
 }
 
 Decl *Parser::ParseModuleImport(SourceLocation AtLoc,

clang/lib/Sema/SemaModule.cpp

@@ -13,7 +13,9 @@
 
 #include "clang/AST/ASTConsumer.h"
 #include "clang/AST/ASTMutationListener.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Lex/HeaderSearch.h"
+#include "clang/Lex/ModuleLoader.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Sema/ParsedAttr.h"
 #include "clang/Sema/SemaInternal.h"
@@ -238,22 +240,21 @@ static bool DiagReservedModuleName(Sema &S, const IdentifierInfo *II,
   llvm_unreachable("fell off a fully covered switch");
 }
 
-Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
-                                           SourceLocation ModuleLoc,
-                                           ModuleDeclKind MDK,
-                                           ModuleNameLoc *PathLoc,
-                                           ModuleNameLoc *PartitionLoc,
-                                           ModuleImportState &ImportState) {
+Sema::DeclGroupPtrTy
+Sema::ActOnModuleDecl(SourceLocation StartLoc, SourceLocation ModuleLoc,
+                      ModuleDeclKind MDK, ModuleIdPath Path,
+                      ModuleIdPath Partition, ModuleImportState &ImportState,
+                      bool AtStartOfTU) {
   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;
   ImportState = ModuleImportState::NotACXX20Module;
 
-  if (PartitionLoc)
+  bool IsPartition = !Partition.empty();
+  if (IsPartition)
     switch (MDK) {
     case ModuleDeclKind::Implementation:
       MDK = ModuleDeclKind::PartitionImplementation;
@@ -282,7 +283,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
     // We were asked to compile a module interface unit but this is a module
     // implementation unit.
     Diag(ModuleLoc, diag::err_module_interface_implementation_mismatch)
-      << FixItHint::CreateInsertion(ModuleLoc, "export ");
+        << FixItHint::CreateInsertion(ModuleLoc, "export ");
     MDK = ModuleDeclKind::Interface;
     break;
 
@@ -314,7 +315,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
 
   // In C++20, the module-declaration must be the first declaration if there
   // is no global module fragment.
-  if (getLangOpts().CPlusPlusModules && !IsFirstDecl && !SeenGMF) {
+  if (getLangOpts().CPlusPlusModules && !AtStartOfTU && !SeenGMF) {
     Diag(ModuleLoc, diag::err_module_decl_not_at_start);
     SourceLocation BeginLoc =
         ModuleScopes.empty()
@@ -335,39 +336,41 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
 
   // Test the first part of the path to see if it's std[0-9]+ but allow the
   // name in a system header.
-  StringRef FirstComponentName =
-      PathLoc->getModuleIdPath()[0].getIdentifierInfo()->getName();
-  if (!getSourceManager().isInSystemHeader(PathLoc->getBeginLoc()) &&
+  StringRef FirstComponentName = Path[0].getIdentifierInfo()->getName();
+  if (!getSourceManager().isInSystemHeader(Path[0].getLoc()) &&
       (FirstComponentName == "std" ||
        (FirstComponentName.starts_with("std") &&
         llvm::all_of(FirstComponentName.drop_front(3), &llvm::isDigit))))
-    Diag(PathLoc->getBeginLoc(), diag::warn_reserved_module_name)
-        << PathLoc->getModuleIdPath()[0].getIdentifierInfo();
+    Diag(Path[0].getLoc(), diag::warn_reserved_module_name)
+        << Path[0].getIdentifierInfo();
 
   // Then test all of the components in the path to see if any of them are
   // using another kind of reserved or invalid identifier.
-  for (auto Part : PathLoc->getModuleIdPath()) {
+  for (auto Part : Path) {
     if (DiagReservedModuleName(*this, Part.getIdentifierInfo(), Part.getLoc()))
       return nullptr;
   }
 
   // Flatten the dots in a module name. Unlike Clang's hierarchical module map
   // modules, the dots here are just another character that can appear in a
   // module name.
-  std::string ModuleName = PathLoc->str();
-  if (PartitionLoc) {
+  std::string ModuleName = ModuleNameLoc::stringFromModuleIdPath(Path);
+  if (IsPartition) {
     ModuleName += ":";
-    ModuleName += PartitionLoc->str();
+    ModuleName += ModuleNameLoc::stringFromModuleIdPath(Partition);
   }
   // If a module name was explicitly specified on the command line, it must be
   // correct.
   if (!getLangOpts().CurrentModule.empty() &&
       getLangOpts().CurrentModule != ModuleName) {
-    Diag(PathLoc->getBeginLoc(), diag::err_current_module_name_mismatch)
-        << PathLoc->getRange() << getLangOpts().CurrentModule;
+    Diag(Path.front().getLoc(), diag::err_current_module_name_mismatch)
+        << SourceRange(Path.front().getLoc(), IsPartition
+                                                  ? Partition.back().getLoc()
+                                                  : Path.back().getLoc())
+        << getLangOpts().CurrentModule;
     return nullptr;
   }
-  const_cast<LangOptions&>(getLangOpts()).CurrentModule = ModuleName;
+  const_cast<LangOptions &>(getLangOpts()).CurrentModule = ModuleName;
 
   auto &Map = PP.getHeaderSearchInfo().getModuleMap();
   Module *Mod;                 // The module we are creating.
@@ -378,7 +381,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
     // We can't have parsed or imported a definition of this module or parsed a
     // module map defining it already.
     if (auto *M = Map.findOrLoadModule(ModuleName)) {
-      Diag(PathLoc->getBeginLoc(), diag::err_module_redefinition) << ModuleName;
+      Diag(Path[0].getLoc(), diag::err_module_redefinition) << ModuleName;
       if (M->DefinitionLoc.isValid())
         Diag(M->DefinitionLoc, diag::note_prev_module_definition);
       else if (OptionalFileEntryRef FE = M->getASTFile())
@@ -401,7 +404,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
     // keyword nor a module-partition implicitly imports the primary
     // module interface unit of the module as if by a module-import-
     // declaration.
-    IdentifierLoc ModuleNameLoc(PathLoc->getBeginLoc(),
+    IdentifierLoc ModuleNameLoc(Path[0].getLoc(),
                                 PP.getIdentifierInfo(ModuleName));
 
     // The module loader will assume we're trying to import the module that
@@ -412,7 +415,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
     Interface = getModuleLoader().loadModule(ModuleLoc, {ModuleNameLoc},
                                              Module::AllVisible,
                                              /*IsInclusionDirective=*/false);
-    const_cast<LangOptions&>(getLangOpts()).CurrentModule = ModuleName;
+    const_cast<LangOptions &>(getLangOpts()).CurrentModule = ModuleName;
 
     if (!Interface) {
       Diag(ModuleLoc, diag::err_module_not_defined) << ModuleName;
@@ -474,7 +477,7 @@ Sema::DeclGroupPtrTy Sema::ActOnModuleDecl(SourceLocation StartLoc,
 
     // Make the import decl for the interface in the impl module.
     ImportDecl *Import = ImportDecl::Create(Context, CurContext, ModuleLoc,
-                                            Interface, PathLoc->getBeginLoc());
+                                            Interface, Path[0].getLoc());
     CurContext->addDecl(Import);
 
     // Sequence initialization of the imported module before that of the current