Lex: add a callback for `#pragma mark`

Allow a preprocessor observer to be notified of mark pragmas.  Although
this does not impact code generation in any way, it is useful for other
clients, such as clangd, to be able to identify any marked regions.

Reviewed By: dgoldman

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

clang/include/clang/Lex/PPCallbacks.h

@@ -191,6 +191,10 @@ public:
                              StringRef Str) {
   }
 
+  /// Callback invoked when a \#pragma mark comment is read.
+  virtual void PragmaMark(SourceLocation Loc, StringRef Trivia) {
+  }
+
   /// Callback invoked when a \#pragma detect_mismatch directive is
   /// read.
   virtual void PragmaDetectMismatch(SourceLocation Loc, StringRef Name,

clang/include/clang/Lex/Preprocessor.h

@@ -2365,7 +2365,7 @@ private:
 
 public:
   void HandlePragmaOnce(Token &OnceTok);
-  void HandlePragmaMark();
+  void HandlePragmaMark(Token &MarkTok);
   void HandlePragmaPoison();
   void HandlePragmaSystemHeader(Token &SysHeaderTok);
   void HandlePragmaDependency(Token &DependencyTok);

clang/lib/Lex/Pragma.cpp

@@ -412,9 +412,13 @@ void Preprocessor::HandlePragmaOnce(Token &OnceTok) {
   HeaderInfo.MarkFileIncludeOnce(getCurrentFileLexer()->getFileEntry());
 }
 
-void Preprocessor::HandlePragmaMark() {
+void Preprocessor::HandlePragmaMark(Token &MarkTok) {
   assert(CurPPLexer && "No current lexer?");
-  CurLexer->ReadToEndOfLine();
+
+  SmallString<64> Buffer;
+  CurLexer->ReadToEndOfLine(&Buffer);
+  if (Callbacks)
+    Callbacks->PragmaMark(MarkTok.getLocation(), Buffer);
 }
 
 /// HandlePragmaPoison - Handle \#pragma GCC poison.  PoisonTok is the 'poison'.
@@ -992,7 +996,7 @@ struct PragmaMarkHandler : public PragmaHandler {
 
   void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,
                     Token &MarkTok) override {
-    PP.HandlePragmaMark();
+    PP.HandlePragmaMark(MarkTok);
   }
 };
 

clang/unittests/Lex/PPCallbacksTest.cpp

@@ -112,6 +112,20 @@ public:
   unsigned State;
 };
 
+class PragmaMarkCallbacks : public PPCallbacks {
+public:
+  struct Mark {
+    SourceLocation Location;
+    std::string Trivia;
+  };
+
+  std::vector<Mark> Marks;
+
+  void PragmaMark(SourceLocation Loc, StringRef Trivia) override {
+    Marks.emplace_back(Mark{Loc, Trivia.str()});
+  }
+};
+
 // PPCallbacks test fixture.
 class PPCallbacksTest : public ::testing::Test {
 protected:
@@ -256,6 +270,36 @@ protected:
     return Callbacks->Results;
   }
 
+  std::vector<PragmaMarkCallbacks::Mark>
+  PragmaMarkCall(const char *SourceText) {
+    std::unique_ptr<llvm::MemoryBuffer> SourceBuf =
+        llvm::MemoryBuffer::getMemBuffer(SourceText, "test.c");
+    SourceMgr.setMainFileID(SourceMgr.createFileID(std::move(SourceBuf)));
+
+    HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                            Diags, LangOpts, Target.get());
+    TrivialModuleLoader ModLoader;
+
+    Preprocessor PP(std::make_shared<PreprocessorOptions>(), Diags, LangOpts,
+                    SourceMgr, HeaderInfo, ModLoader, /*IILookup=*/nullptr,
+                    /*OwnsHeaderSearch=*/false);
+    PP.Initialize(*Target);
+
+    auto *Callbacks = new PragmaMarkCallbacks;
+    PP.addPPCallbacks(std::unique_ptr<PPCallbacks>(Callbacks));
+
+    // Lex source text.
+    PP.EnterMainSourceFile();
+    while (true) {
+      Token Tok;
+      PP.Lex(Tok);
+      if (Tok.is(tok::eof))
+        break;
+    }
+
+    return Callbacks->Marks;
+  }
+
   PragmaOpenCLExtensionCallbacks::CallbackParameters
   PragmaOpenCLExtensionCall(const char *SourceText) {
     LangOptions OpenCLLangOpts;
@@ -424,6 +468,24 @@ TEST_F(PPCallbacksTest, OpenCLExtensionPragmaDisabled) {
   ASSERT_EQ(ExpectedState, Parameters.State);
 }
 
+TEST_F(PPCallbacksTest, CollectMarks) {
+  const char *Source =
+    "#pragma mark\n"
+    "#pragma mark\r\n"
+    "#pragma mark - trivia\n"
+    "#pragma mark - trivia\r\n";
+
+  auto Marks = PragmaMarkCall(Source);
+
+  ASSERT_EQ(4u, Marks.size());
+  ASSERT_TRUE(Marks[0].Trivia.empty());
+  ASSERT_TRUE(Marks[1].Trivia.empty());
+  ASSERT_FALSE(Marks[2].Trivia.empty());
+  ASSERT_FALSE(Marks[3].Trivia.empty());
+  ASSERT_EQ(" - trivia", Marks[2].Trivia);
+  ASSERT_EQ(" - trivia", Marks[3].Trivia);
+}
+
 TEST_F(PPCallbacksTest, DirectiveExprRanges) {
   const auto &Results1 = DirectiveExprRange("#if FLUZZY_FLOOF\n#endif\n");
   EXPECT_EQ(Results1.size(), 1U);