[clang] 24f4c3e - Lex: add a callback for `#pragma mark`

[Saleem Abdulrasool via cfe-commits]

Author: Saleem Abdulrasool
Date: 2021-07-02T15:44:01-07:00
New Revision: 24f4c3ebef63c7d2553132d7d9d75ea19338bcb7

URL: https://github.com/llvm/llvm-project/commit/24f4c3ebef63c7d2553132d7d9d75ea19338bcb7
DIFF: https://github.com/llvm/llvm-project/commit/24f4c3ebef63c7d2553132d7d9d75ea19338bcb7.diff

LOG: 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

Added: 
    

Modified: 
    clang/include/clang/Lex/PPCallbacks.h
    clang/include/clang/Lex/Preprocessor.h
    clang/lib/Lex/Pragma.cpp
    clang/unittests/Lex/PPCallbacksTest.cpp

Removed: 
    


################################################################################
diff  --git a/clang/include/clang/Lex/PPCallbacks.h b/clang/include/clang/Lex/PPCallbacks.h
index d57be1990caf..bcf49c577735 100644
--- a/clang/include/clang/Lex/PPCallbacks.h
+++ b/clang/include/clang/Lex/PPCallbacks.h
@@ -191,6 +191,10 @@ class PPCallbacks {
                              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,

diff  --git a/clang/include/clang/Lex/Preprocessor.h b/clang/include/clang/Lex/Preprocessor.h
index 16fbf5ea5a5b..2d6335471383 100644
--- a/clang/include/clang/Lex/Preprocessor.h
+++ b/clang/include/clang/Lex/Preprocessor.h
@@ -2365,7 +2365,7 @@ class Preprocessor {
 
 public:
   void HandlePragmaOnce(Token &OnceTok);
-  void HandlePragmaMark();
+  void HandlePragmaMark(Token &MarkTok);
   void HandlePragmaPoison();
   void HandlePragmaSystemHeader(Token &SysHeaderTok);
   void HandlePragmaDependency(Token &DependencyTok);

diff  --git a/clang/lib/Lex/Pragma.cpp b/clang/lib/Lex/Pragma.cpp
index 081b92ac21d9..c89061ba6d02 100644
--- a/clang/lib/Lex/Pragma.cpp
+++ b/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);
   }
 };
 

diff  --git a/clang/unittests/Lex/PPCallbacksTest.cpp b/clang/unittests/Lex/PPCallbacksTest.cpp
index 5581f9fb82f3..f92587af8dc5 100644
--- a/clang/unittests/Lex/PPCallbacksTest.cpp
+++ b/clang/unittests/Lex/PPCallbacksTest.cpp
@@ -112,6 +112,20 @@ class PragmaOpenCLExtensionCallbacks : public PPCallbacks {
   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 @@ class PPCallbacksTest : public ::testing::Test {
     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);