[ObjC][Preprocessor] Handle @import directive as a pp-directive

clang/examples/AnnotateFunctions/AnnotateFunctions.cpp

@@ -65,7 +65,7 @@ class PragmaAnnotateHandler : public PragmaHandler {
     Token Tok;
     PP.LexUnexpandedToken(Tok);
     if (Tok.isNot(tok::eod))
-      PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "pragma";
+      PP.Diag(Tok, diag::ext_pp_extra_tokens_at_eol) << "#pragma";
 
     if (HandledDecl) {
       DiagnosticsEngine &D = PP.getDiagnostics();

clang/include/clang/Basic/DiagnosticLexKinds.td

@@ -465,10 +465,7 @@ def err_pp_embed_device_file : Error<
   "device files are not yet supported by '#embed' directive">;
 
 def ext_pp_extra_tokens_at_eol : ExtWarn<
-  "extra tokens at end of #%0 directive">, InGroup<ExtraTokens>;
-def ext_pp_extra_tokens_at_module_directive_eol
-    : ExtWarn<"extra tokens at end of '%0' directive">,
-      InGroup<ExtraTokens>;
+  "extra tokens at end of %0 directive">, InGroup<ExtraTokens>;
 
 def ext_pp_comma_expr : Extension<"comma operator in operand of #if">;
 def ext_pp_bad_vaargs_use : Extension<

clang/include/clang/Lex/Preprocessor.h

@@ -137,24 +137,6 @@ 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.
 ///
@@ -378,7 +360,7 @@ class Preprocessor {
   bool DeclaringCXXNamedModules = false;
 
   /// Whether the last token we lexed was an 'export' keyword.
-  ExportContextualKeywordInfo LastTokenWasExportKeyword;
+  std::optional<Token> ModuleLikeDirectiveIntroducer;
 
   /// First pp-token source location in current translation unit.
   SourceLocation FirstPPTokenLoc;
@@ -1808,15 +1790,15 @@ class Preprocessor {
   void EnterModuleSuffixTokenStream(ArrayRef<Token> Toks);
   void HandleCXXImportDirective(Token Import);
   void HandleCXXModuleDirective(Token Module);
+  void HandleObjCAtImportDirective(Token &ImportTok);
 
   /// Callback invoked when the lexer sees one of export, import or module token
   /// at the start of a line.
   ///
   /// This consumes the import/module directive, modifies the
   /// lexer/preprocessor state, and advances the lexer(s) so that the next token
   /// read is the correct one.
-  bool HandleModuleContextualKeyword(Token &Result,
-                                     bool TokAtPhysicalStartOfLine);
+  bool HandleModuleContextualKeyword(Token &Result);
 
   /// Get the start location of the first pp-token in main file.
   SourceLocation getMainFileFirstPPTokenLoc() const {
@@ -1825,7 +1807,6 @@ class Preprocessor {
     return FirstPPTokenLoc;
   }
 
-  bool LexAfterModuleImport(Token &Result);
   void CollectPPImportSuffix(SmallVectorImpl<Token> &Toks,
                              bool StopUntilEOD = false);
   bool CollectPPImportSuffixAndEnterStream(SmallVectorImpl<Token> &Toks,
@@ -2356,6 +2337,7 @@ class Preprocessor {
   template <typename... Ts> bool isNextPPTokenOneOf(Ts... Ks) {
     static_assert(sizeof...(Ts) > 0,
                   "requires at least one tok::TokenKind specified");
+
     // Do some quick tests for rejection cases.
     std::optional<Token> Val;
     if (CurLexer)
@@ -3175,9 +3157,6 @@ class Preprocessor {
   static bool CLK_DependencyDirectivesLexer(Preprocessor &P, Token &Result) {
     return P.CurLexer->LexDependencyDirectiveToken(Result);
   }
-  static bool CLK_LexAfterModuleImport(Preprocessor &P, Token &Result) {
-    return P.LexAfterModuleImport(Result);
-  }
 };
 
 /// Abstract base class that describes a handler that will receive

clang/include/clang/Lex/Token.h

@@ -90,8 +90,10 @@ class Token {
                           // re-added, e.g. via EnterTokenStream. Annotation
                           // tokens are *not* reinjected.
     HasSeenNoTrivialPPDirective =
-        0x1000, // Whether we've seen any 'no-trivial' pp-directives before
+        0x1000, // Whether we've seen any 'no-trivial' pp-directives before//
                 // current position.
+    PhysicalStartOfLine =
+        0x2000, // At the physical start of line or only after whitespace.
   };
 
   tok::TokenKind getKind() const { return Kind; }
@@ -277,6 +279,10 @@ class Token {
   ///
   bool isAtStartOfLine() const { return getFlag(StartOfLine); }
 
+  /// isAtPhysicalStartOfLine - Return true if this token is at the physical
+  /// start of a line.
+  bool isAtPhysicalStartOfLine() const { return getFlag(PhysicalStartOfLine); }
+
   /// Return true if this token has whitespace before it.
   ///
   bool hasLeadingSpace() const { return getFlag(LeadingSpace); }

clang/lib/Lex/DependencyDirectivesScanner.cpp

@@ -579,15 +579,12 @@ bool Scanner::lexModuleDirectiveBody(DirectiveKind Kind, const char *&First,
     return false;
   }
 
+  const auto &Tok = lexToken(First, End);
   pushDirective(Kind);
-  skipWhitespace(First, End);
-  if (First == End)
+  if (Tok.is(tok::eof) || Tok.is(tok::eod))
     return false;
-  if (!isVerticalWhitespace(*First))
-    return reportError(
-        DirectiveLoc, diag::err_dep_source_scanner_unexpected_tokens_at_import);
-  skipNewline(First, End);
-  return false;
+  return reportError(DirectiveLoc,
+                     diag::err_dep_source_scanner_unexpected_tokens_at_import);
 }
 
 dependency_directives_scan::Token &Scanner::lexToken(const char *&First,
@@ -931,10 +928,6 @@ bool Scanner::lexPPLine(const char *&First, const char *const End) {
     CurDirToks.clear();
   });
 
-  // FIXME: Shoule we handle @import as a preprocessing directive?
-  if (*First == '@')
-    return lexAt(First, End);
-
   if (*First == '_') {
     if (isNextIdentifierOrSkipLine("_Pragma", First, End))
       return lex_Pragma(First, End);
@@ -947,6 +940,10 @@ bool Scanner::lexPPLine(const char *&First, const char *const End) {
   auto ScEx2 = make_scope_exit(
       [&]() { TheLexer.setParsingPreprocessorDirective(false); });
 
+  // FIXME: Shoule we handle @import as a preprocessing directive?
+  if (*First == '@')
+    return lexAt(First, End);
+
   // Handle module directives for C++20 modules.
   if (*First == 'i' || *First == 'e' || *First == 'm')
     return lexModule(First, End);

clang/lib/Lex/Lexer.cpp

@@ -77,16 +77,13 @@ tok::ObjCKeywordKind Token::getObjCKeywordID() const {
 /// or module).
 bool Token::isModuleContextualKeyword(const LangOptions &LangOpts,
                                       bool AllowExport) const {
-  if (!LangOpts.CPlusPlusModules)
+  if (!LangOpts.CPlusPlusModules || isAnnotation())
     return false;
   if (AllowExport && is(tok::kw_export))
     return true;
-  if (isOneOf(tok::kw_import, tok::kw_module))
-    return true;
-  if (isNot(tok::identifier))
-    return false;
-  const auto *II = getIdentifierInfo();
-  return II->isModulesImport() || II->isModulesDeclaration();
+  if (const auto *II = getIdentifierInfo())
+    return II->isModulesImport() || II->isModulesDeclaration();
+  return false;
 }
 
 /// Determine whether the token kind starts a simple-type-specifier.
@@ -3758,6 +3755,9 @@ bool Lexer::LexTokenInternal(Token &Result, bool TokAtPhysicalStartOfLine) {
   assert(!Result.needsCleaning() && "Result needs cleaning");
   assert(!Result.hasPtrData() && "Result has not been reset");
 
+  if (TokAtPhysicalStartOfLine)
+    Result.setFlag(Token::PhysicalStartOfLine);
+
   // CurPtr - Cache BufferPtr in an automatic variable.
   const char *CurPtr = BufferPtr;
 
@@ -4040,10 +4040,9 @@ bool Lexer::LexTokenInternal(Token &Result, bool TokAtPhysicalStartOfLine) {
     // Notify MIOpt that we read a non-whitespace/non-comment token.
     MIOpt.ReadToken();
     bool returnedToken = LexIdentifierContinue(Result, CurPtr);
-    if (returnedToken && Result.isModuleContextualKeyword(LangOpts) &&
-        !LexingRawMode && !Is_PragmaLexer && !ParsingPreprocessorDirective &&
-        PP &&
-        PP->HandleModuleContextualKeyword(Result, TokAtPhysicalStartOfLine))
+    if (returnedToken && Result.isNot(tok::eof) && !LexingRawMode &&
+        !Is_PragmaLexer && !ParsingPreprocessorDirective && PP &&
+        PP->HandleModuleContextualKeyword(Result))
       goto HandleDirective;
     return returnedToken;
   }
@@ -4454,9 +4453,14 @@ bool Lexer::LexTokenInternal(Token &Result, bool TokAtPhysicalStartOfLine) {
 
   case '@':
     // Objective C support.
-    if (CurPtr[-1] == '@' && LangOpts.ObjC)
+    if (CurPtr[-1] == '@' && LangOpts.ObjC) {
+      if (TokAtPhysicalStartOfLine && !LexingRawMode && !Is_PragmaLexer) {
+        FormTokenWithChars(Result, CurPtr, tok::at);
+        (void)PP->HandleModuleContextualKeyword(Result);
+        return true;
+      }
       Kind = tok::at;
-    else
+    } else
       Kind = tok::unknown;
     break;
 
@@ -4538,10 +4542,6 @@ bool Lexer::LexTokenInternal(Token &Result, bool TokAtPhysicalStartOfLine) {
 HandleDirective:
   PP->HandleDirective(Result);
 
-  if (PP->hadModuleLoaderFatalFailure())
-    // With a fatal failure in the module loader, we abort parsing.
-    return true;
-
   // We parsed the directive; lex a token with the new state.
   return false;
 
@@ -4608,20 +4608,26 @@ bool Lexer::LexDependencyDirectiveToken(Token &Result) {
 
   const char *TokPtr = convertDependencyDirectiveToken(DDTok, Result);
 
-  if (Result.is(tok::hash) && Result.isAtStartOfLine()) {
+  if (Result.is(tok::hash) && Result.isAtStartOfLine() && !isLexingRawMode()) {
     PP->HandleDirective(Result);
     if (PP->hadModuleLoaderFatalFailure())
       // With a fatal failure in the module loader, we abort parsing.
       return true;
     return false;
   }
+  if (Result.is(tok::at) && Result.isAtStartOfLine() && !isLexingRawMode()) {
+    (void)PP->HandleModuleContextualKeyword(Result);
+    return false;
+  }
   if (Result.is(tok::raw_identifier)) {
     Result.setRawIdentifierData(TokPtr);
     if (!isLexingRawMode()) {
       const IdentifierInfo *II = PP->LookUpIdentifierInfo(Result);
-      if (Result.isModuleContextualKeyword(LangOpts) &&
-          PP->HandleModuleContextualKeyword(Result, Result.isAtStartOfLine())) {
+      if (PP->HandleModuleContextualKeyword(Result)) {
         PP->HandleDirective(Result);
+        if (PP->hadModuleLoaderFatalFailure())
+          // With a fatal failure in the module loader, we abort parsing.
+          return true;
         return false;
       }
       if (II->isHandleIdentifierCase())