[clang] Fix the behavior of __COUNT__ macros when PCH is enabled (PR #105591)

Giuliano Belinassi via cfe-commits cfe-commits at lists.llvm.org
Wed Aug 21 15:34:20 PDT 2024


https://github.com/giulianobelinassi updated https://github.com/llvm/llvm-project/pull/105591

>From 914427af9b83848ac395bbff663133aadf387161 Mon Sep 17 00:00:00 2001
From: Giuliano Belinassi <gbelinassi at suse.de>
Date: Wed, 21 Aug 2024 18:31:36 -0300
Subject: [PATCH] Fix the behavior of __COUNT__ macros when PCH is enabled

Previoulsy, calling `ASTUnit::LoadFromCompilerInvocation` with
`PrecompilePreambleAfterNParses > 0` caused the `__COUNT__` macro value
to be restarted.  This commit fixes this by remembering the value of
this macro when the PCH finished parsing.

Signed-off-by: Giuliano Belinassi <gbelinassi at suse.de>
---
 clang/include/clang/Frontend/ASTUnit.h        |  6 +-
 clang/include/clang/Frontend/FrontendAction.h |  3 +
 .../clang/Frontend/PrecompiledPreamble.h      |  7 ++-
 clang/lib/Frontend/ASTUnit.cpp                | 39 ++++++++----
 clang/lib/Frontend/FrontendAction.cpp         |  3 +-
 clang/lib/Frontend/PrecompiledPreamble.cpp    |  8 ++-
 clang/unittests/Frontend/ASTUnitTest.cpp      | 60 +++++++++++++++++++
 7 files changed, 109 insertions(+), 17 deletions(-)

diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index 080844893c13c9..2762abedaa2927 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -372,12 +372,14 @@ class ASTUnit {
 
   bool Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
              std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
-             IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS);
+             IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+             unsigned PCHCountValue = 0);
 
   std::unique_ptr<llvm::MemoryBuffer> getMainBufferWithPrecompiledPreamble(
       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
       CompilerInvocation &PreambleInvocationIn,
-      IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS, bool AllowRebuild = true,
+      IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+      unsigned *CountValueAtFinish = nullptr, bool AllowRebuild = true,
       unsigned MaxLines = 0);
   void RealizeTopLevelDeclsFromPreamble();
 
diff --git a/clang/include/clang/Frontend/FrontendAction.h b/clang/include/clang/Frontend/FrontendAction.h
index 039f6f247b6d8c..b8112cb7e21469 100644
--- a/clang/include/clang/Frontend/FrontendAction.h
+++ b/clang/include/clang/Frontend/FrontendAction.h
@@ -178,6 +178,9 @@ class FrontendAction {
   /// details.
   virtual bool isModelParsingAction() const { return false; }
 
+  /// Do we need to reuse existing preprocessor?
+  virtual bool reuseExistingPreprocessor() const { return false; }
+
   /// Does this action only use the preprocessor?
   ///
   /// If so no AST context will be created and this action will be invalid
diff --git a/clang/include/clang/Frontend/PrecompiledPreamble.h b/clang/include/clang/Frontend/PrecompiledPreamble.h
index 624df004bf89e4..ef3302b82a30b0 100644
--- a/clang/include/clang/Frontend/PrecompiledPreamble.h
+++ b/clang/include/clang/Frontend/PrecompiledPreamble.h
@@ -102,6 +102,9 @@ class PrecompiledPreamble {
   /// be used for logging and debugging purposes only.
   std::size_t getSize() const;
 
+  /// Returns the value of __COUNT__ macro when the process finished.
+  unsigned getCountValue() const { return CountValueAtFinish; }
+
   /// Returned string is not null-terminated.
   llvm::StringRef getContents() const {
     return {PreambleBytes.data(), PreambleBytes.size()};
@@ -137,7 +140,7 @@ class PrecompiledPreamble {
                       std::vector<char> PreambleBytes,
                       bool PreambleEndsAtStartOfLine,
                       llvm::StringMap<PreambleFileHash> FilesInPreamble,
-                      llvm::StringSet<> MissingFiles);
+                      llvm::StringSet<> MissingFiles, unsigned CountPreamble);
 
   /// Data used to determine if a file used in the preamble has been changed.
   struct PreambleFileHash {
@@ -205,6 +208,8 @@ class PrecompiledPreamble {
   std::vector<char> PreambleBytes;
   /// See PreambleBounds::PreambleEndsAtStartOfLine
   bool PreambleEndsAtStartOfLine;
+  /// __COUNT__ macro value when the preamble finished building.
+  unsigned CountValueAtFinish;
 };
 
 /// A set of callbacks to gather useful information while building a preamble.
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index 84e273a3949ef0..bd6e4f902fc38e 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -1046,6 +1046,7 @@ class TopLevelDeclTrackerConsumer : public ASTConsumer {
 class TopLevelDeclTrackerAction : public ASTFrontendAction {
 public:
   ASTUnit &Unit;
+  bool ReusePreprocessor;
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override {
@@ -1056,11 +1057,13 @@ class TopLevelDeclTrackerAction : public ASTFrontendAction {
         Unit, Unit.getCurrentTopLevelHashValue());
   }
 
-public:
-  TopLevelDeclTrackerAction(ASTUnit &_Unit) : Unit(_Unit) {}
+  TopLevelDeclTrackerAction(ASTUnit &_Unit, bool _ReusePreprocessor = false)
+    : Unit(_Unit), ReusePreprocessor(_ReusePreprocessor) {}
 
   bool hasCodeCompletionSupport() const override { return false; }
 
+  bool reuseExistingPreprocessor() const override { return ReusePreprocessor; }
+
   TranslationUnitKind getTranslationUnitKind() override {
     return Unit.getTranslationUnitKind();
   }
@@ -1146,7 +1149,8 @@ static void checkAndSanitizeDiags(SmallVectorImpl<StoredDiagnostic> &
 /// contain any translation-unit information, false otherwise.
 bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
                     std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer,
-                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS) {
+                    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
+                    unsigned PCHCountValue) {
   if (!Invocation)
     return true;
 
@@ -1243,12 +1247,19 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   }
 
   std::unique_ptr<TopLevelDeclTrackerAction> Act(
-      new TopLevelDeclTrackerAction(*this));
+      new TopLevelDeclTrackerAction(*this, true));
 
   // Recover resources if we crash before exiting this method.
   llvm::CrashRecoveryContextCleanupRegistrar<TopLevelDeclTrackerAction>
     ActCleanup(Act.get());
 
+  if (!Clang->hasPreprocessor()) {
+    // Create the Preprocessor.
+    Clang->createPreprocessor(TUKind);
+  }
+
+  Clang->getPreprocessor().setCounterValue(PCHCountValue);
+
   if (!Act->BeginSourceFile(*Clang.get(), Clang->getFrontendOpts().Inputs[0]))
     return true;
 
@@ -1343,8 +1354,8 @@ std::unique_ptr<llvm::MemoryBuffer>
 ASTUnit::getMainBufferWithPrecompiledPreamble(
     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
     CompilerInvocation &PreambleInvocationIn,
-    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS, bool AllowRebuild,
-    unsigned MaxLines) {
+    IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS, unsigned *CountValueAtFinish,
+    bool AllowRebuild, unsigned MaxLines) {
   auto MainFilePath =
       PreambleInvocationIn.getFrontendOpts().Inputs[0].getFile();
   std::unique_ptr<llvm::MemoryBuffer> MainFileBuffer =
@@ -1420,6 +1431,10 @@ ASTUnit::getMainBufferWithPrecompiledPreamble(
         PCHContainerOps, StorePreamblesInMemory, PreambleStoragePath,
         Callbacks);
 
+    if (NewPreamble && CountValueAtFinish) {
+      *CountValueAtFinish = NewPreamble->getCountValue();
+    }
+
     PreambleInvocationIn.getFrontendOpts().SkipFunctionBodies =
         PreviousSkipFunctionBodies;
 
@@ -1697,6 +1712,8 @@ bool ASTUnit::LoadFromCompilerInvocation(
 
   assert(VFS && "VFS is null");
 
+  unsigned PCHCountValue = 0;
+
   // We'll manage file buffers ourselves.
   Invocation->getPreprocessorOpts().RetainRemappedFileBuffers = true;
   Invocation->getFrontendOpts().DisableFree = false;
@@ -1706,8 +1723,8 @@ bool ASTUnit::LoadFromCompilerInvocation(
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   if (PrecompilePreambleAfterNParses > 0) {
     PreambleRebuildCountdown = PrecompilePreambleAfterNParses;
-    OverrideMainBuffer =
-        getMainBufferWithPrecompiledPreamble(PCHContainerOps, *Invocation, VFS);
+    OverrideMainBuffer = getMainBufferWithPrecompiledPreamble(
+        PCHContainerOps, *Invocation, VFS, &PCHCountValue);
     getDiagnostics().Reset();
     ProcessWarningOptions(getDiagnostics(), Invocation->getDiagnosticOpts());
   }
@@ -1719,7 +1736,8 @@ bool ASTUnit::LoadFromCompilerInvocation(
   llvm::CrashRecoveryContextCleanupRegistrar<llvm::MemoryBuffer>
     MemBufferCleanup(OverrideMainBuffer.get());
 
-  return Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer), VFS);
+  return Parse(std::move(PCHContainerOps), std::move(OverrideMainBuffer), VFS,
+               PCHCountValue);
 }
 
 std::unique_ptr<ASTUnit> ASTUnit::LoadFromCompilerInvocation(
@@ -2299,7 +2317,8 @@ void ASTUnit::CodeComplete(
   std::unique_ptr<llvm::MemoryBuffer> OverrideMainBuffer;
   if (Preamble && Line > 1 && hasSameUniqueID(File, OriginalSourceFile)) {
     OverrideMainBuffer = getMainBufferWithPrecompiledPreamble(
-        PCHContainerOps, Inv, &FileMgr.getVirtualFileSystem(), false, Line - 1);
+        PCHContainerOps, Inv, &FileMgr.getVirtualFileSystem(), nullptr, false,
+        Line - 1);
   }
 
   // If the main file has been overridden due to the use of a preamble,
diff --git a/clang/lib/Frontend/FrontendAction.cpp b/clang/lib/Frontend/FrontendAction.cpp
index a9c45e525c696c..9e0fe1031e060e 100644
--- a/clang/lib/Frontend/FrontendAction.cpp
+++ b/clang/lib/Frontend/FrontendAction.cpp
@@ -816,7 +816,8 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
   // Set up the preprocessor if needed. When parsing model files the
   // preprocessor of the original source is reused.
-  if (!isModelParsingAction())
+  if (!CI.hasPreprocessor() ||
+      (!isModelParsingAction() && !reuseExistingPreprocessor()))
     CI.createPreprocessor(getTranslationUnitKind());
 
   // Inform the diagnostic client we are processing a source file.
diff --git a/clang/lib/Frontend/PrecompiledPreamble.cpp b/clang/lib/Frontend/PrecompiledPreamble.cpp
index cab5838fceb24d..c7cd108301f162 100644
--- a/clang/lib/Frontend/PrecompiledPreamble.cpp
+++ b/clang/lib/Frontend/PrecompiledPreamble.cpp
@@ -571,11 +571,12 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
   // Shrinking the storage requires extra temporary memory.
   // Destroying clang first reduces peak memory usage.
   CICleanup.unregister();
+  unsigned CountValue = Clang->getPreprocessor().getCounterValue();
   Clang.reset();
   Storage->shrink();
   return PrecompiledPreamble(
       std::move(Storage), std::move(PreambleBytes), PreambleEndsAtStartOfLine,
-      std::move(FilesInPreamble), std::move(MissingFiles));
+      std::move(FilesInPreamble), std::move(MissingFiles), CountValue);
 }
 
 PreambleBounds PrecompiledPreamble::getBounds() const {
@@ -728,11 +729,12 @@ PrecompiledPreamble::PrecompiledPreamble(
     std::unique_ptr<PCHStorage> Storage, std::vector<char> PreambleBytes,
     bool PreambleEndsAtStartOfLine,
     llvm::StringMap<PreambleFileHash> FilesInPreamble,
-    llvm::StringSet<> MissingFiles)
+    llvm::StringSet<> MissingFiles, unsigned CountValue)
     : Storage(std::move(Storage)), FilesInPreamble(std::move(FilesInPreamble)),
       MissingFiles(std::move(MissingFiles)),
       PreambleBytes(std::move(PreambleBytes)),
-      PreambleEndsAtStartOfLine(PreambleEndsAtStartOfLine) {
+      PreambleEndsAtStartOfLine(PreambleEndsAtStartOfLine),
+      CountValueAtFinish(CountValue) {
   assert(this->Storage != nullptr);
 }
 
diff --git a/clang/unittests/Frontend/ASTUnitTest.cpp b/clang/unittests/Frontend/ASTUnitTest.cpp
index 30d2731897e7f3..e0fabbf0f124a8 100644
--- a/clang/unittests/Frontend/ASTUnitTest.cpp
+++ b/clang/unittests/Frontend/ASTUnitTest.cpp
@@ -210,4 +210,64 @@ TEST_F(ASTUnitTest, LoadFromCommandLineWorkingDirectory) {
   ASSERT_EQ(FM.getFileSystemOpts().WorkingDir, WorkingDir.str());
 }
 
+TEST_F(ASTUnitTest, PCHCount) {
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> InMemoryFs =
+      new llvm::vfs::InMemoryFileSystem();
+
+  InMemoryFs->addFile("header.h", 0, llvm::MemoryBuffer::getMemBuffer(R"cpp(
+      #define __DEFINE_FUNCTION(i) \
+        int function_ ## i (void) { return 0; }
+      #define _DEFINE_FUNCTION(i) __DEFINE_FUNCTION(i)
+      #define DEFINE_FUNCTION      _DEFINE_FUNCTION(__COUNTER__)
+      DEFINE_FUNCTION;\n";
+    )cpp"));
+
+  InMemoryFs->addFile("test.cpp", 0, llvm::MemoryBuffer::getMemBuffer(R"cpp(
+      #include "header.h"
+      DEFINE_FUNCTION;
+      int main(void)
+      {
+        function_0();
+        function_1();
+        return 0;
+      }
+    )cpp"));
+
+  const char *Args[] = {"clang", "test.cpp"};
+  Diags = CompilerInstance::createDiagnostics(new DiagnosticOptions());
+  CreateInvocationOptions CIOpts;
+  CIOpts.Diags = Diags;
+  CInvok = createInvocation(Args, std::move(CIOpts));
+  ASSERT_TRUE(CInvok);
+
+  FileManager *FileMgr = new FileManager(FileSystemOptions(), InMemoryFs);
+  PCHContainerOps = std::make_shared<PCHContainerOperations>();
+
+  auto AU = ASTUnit::LoadFromCompilerInvocation(
+      CInvok, PCHContainerOps, Diags, FileMgr, false, CaptureDiagsKind::None, 1,
+      TU_Complete, false, false, false);
+  ASSERT_TRUE(AU);
+
+  /* Check if the AU is free of errors.  */
+  const DiagnosticsEngine &de = AU->getDiagnostics();
+  ASSERT_TRUE(de.hasErrorOccurred());
+
+  /* Make sure we have the two functions.  */
+  bool functions[2] = {false, false};
+  for (auto it = AU->top_level_begin(); it != AU->top_level_end(); ++it) {
+    if (FunctionDecl *fdecl = dyn_cast<FunctionDecl>(*it)) {
+      if (fdecl->getName() == "function_0") {
+        functions[0] = true;
+      }
+
+      if (fdecl->getName() == "function_1") {
+        functions[1] = true;
+      }
+    }
+  }
+
+  ASSERT_TRUE(functions[0]);
+  ASSERT_TRUE(functions[1]);
+}
+
 } // anonymous namespace



More information about the cfe-commits mailing list