[clang] Reset FileID based diag state mappings (#143695)

When sharing same compiler instance for multiple compilations, we reset
source manager's file id tables in between runs. Diagnostics engine
keeps a cache based on these file ids, that became dangling references
across compilations.

This patch makes sure we reset those whenever sourcemanager is trashing
its FileIDs.

4 files changed, 67 insertions, 4 deletions

diff --git a/clang/include/clang/Basic/Diagnostic.h b/clang/include/clang/Basic/Diagnostic.h
index efee830..7ae4ef7 100644
--- a/clang/include/clang/Basic/Diagnostic.h
+++ b/clang/include/clang/Basic/Diagnostic.h
@@ -424,10 +424,13 @@ private:
     bool empty() const { return Files.empty(); }
 
     /// Clear out this map.
-    void clear() {
+    void clear(bool Soft) {
+      // Just clear the cache when in soft mode.
       Files.clear();
-      FirstDiagState = CurDiagState = nullptr;
-      CurDiagStateLoc = SourceLocation();
+      if (!Soft) {
+        FirstDiagState = CurDiagState = nullptr;
+        CurDiagStateLoc = SourceLocation();
+      }
     }
 
     /// Produce a debugging dump of the diagnostic state.
@@ -920,6 +923,10 @@ public:
   /// Reset the state of the diagnostic object to its initial configuration.
   /// \param[in] soft - if true, doesn't reset the diagnostic mappings and state
   void Reset(bool soft = false);
+  /// We keep a cache of FileIDs for diagnostics mapped by pragmas. These might
+  /// get invalidated when diagnostics engine is shared across different
+  /// compilations. Provide users with a way to reset that.
+  void ResetPragmas();
 
   //===--------------------------------------------------------------------===//
   // DiagnosticsEngine classification and reporting interfaces.
diff --git a/clang/lib/Basic/Diagnostic.cpp b/clang/lib/Basic/Diagnostic.cpp
index 95d86cb..a30bfa2 100644
--- a/clang/lib/Basic/Diagnostic.cpp
+++ b/clang/lib/Basic/Diagnostic.cpp
@@ -119,6 +119,8 @@ bool DiagnosticsEngine::popMappings(SourceLocation Loc) {
   return true;
 }
 
+void DiagnosticsEngine::ResetPragmas() { DiagStatesByLoc.clear(/*Soft=*/true); }
+
 void DiagnosticsEngine::Reset(bool soft /*=false*/) {
   ErrorOccurred = false;
   UncompilableErrorOccurred = false;
@@ -135,7 +137,7 @@ void DiagnosticsEngine::Reset(bool soft /*=false*/) {
   if (!soft) {
     // Clear state related to #pragma diagnostic.
     DiagStates.clear();
-    DiagStatesByLoc.clear();
+    DiagStatesByLoc.clear(false);
     DiagStateOnPushStack.clear();
 
     // Create a DiagState and DiagStatePoint representing diagnostic changes
diff --git a/clang/lib/Basic/SourceManager.cpp b/clang/lib/Basic/SourceManager.cpp
index 09e5c65..053e826 100644
--- a/clang/lib/Basic/SourceManager.cpp
+++ b/clang/lib/Basic/SourceManager.cpp
@@ -344,6 +344,9 @@ void SourceManager::clearIDTables() {
   NextLocalOffset = 0;
   CurrentLoadedOffset = MaxLoadedOffset;
   createExpansionLoc(SourceLocation(), SourceLocation(), SourceLocation(), 1);
+  // Diagnostics engine keeps some references to fileids, mostly for dealing
+  // with diagnostic pragmas, make sure they're reset as well.
+  Diag.ResetPragmas();
 }
 
 bool SourceManager::isMainFile(const FileEntry &SourceFile) {
diff --git a/clang/unittests/Frontend/CompilerInstanceTest.cpp b/clang/unittests/Frontend/CompilerInstanceTest.cpp
index a7b258d..459a386 100644
--- a/clang/unittests/Frontend/CompilerInstanceTest.cpp
+++ b/clang/unittests/Frontend/CompilerInstanceTest.cpp
@@ -9,9 +9,12 @@
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Frontend/CompilerInvocation.h"
+#include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/TextDiagnosticPrinter.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "gtest/gtest.h"
@@ -97,4 +100,52 @@ TEST(CompilerInstance, AllowDiagnosticLogWithUnownedDiagnosticConsumer) {
   ASSERT_EQ(DiagnosticOutput, "error: expected no crash\n");
 }
 
+TEST(CompilerInstance, MultipleInputsCleansFileIDs) {
+  auto VFS = makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  VFS->addFile("a.cc", /*ModificationTime=*/{},
+               MemoryBuffer::getMemBuffer(R"cpp(
+      #include "a.h"
+      )cpp"));
+  // Paddings of `void foo();` in the sources below are "important". We're
+  // testing against source locations from previous compilations colliding.
+  // Hence the `unused` variable in `b.h` needs to be within `#pragma clang
+  // diagnostic` block from `a.h`.
+  VFS->addFile("a.h", /*ModificationTime=*/{}, MemoryBuffer::getMemBuffer(R"cpp(
+      #include "b.h"
+      #pragma clang diagnostic push
+      #pragma clang diagnostic warning "-Wunused"
+      void foo();
+      #pragma clang diagnostic pop
+      )cpp"));
+  VFS->addFile("b.h", /*ModificationTime=*/{}, MemoryBuffer::getMemBuffer(R"cpp(
+      void foo(); void foo(); void foo(); void foo();
+      inline void foo() { int unused = 2; }
+      )cpp"));
+
+  DiagnosticOptions DiagOpts;
+  IntrusiveRefCntPtr<DiagnosticsEngine> Diags =
+      CompilerInstance::createDiagnostics(*VFS, DiagOpts);
+
+  CreateInvocationOptions CIOpts;
+  CIOpts.Diags = Diags;
+
+  const char *Args[] = {"clang", "-xc++", "a.cc"};
+  std::shared_ptr<CompilerInvocation> CInvok =
+      createInvocation(Args, std::move(CIOpts));
+  ASSERT_TRUE(CInvok) << "could not create compiler invocation";
+
+  CompilerInstance Instance(std::move(CInvok));
+  Instance.setDiagnostics(Diags.get());
+  Instance.createFileManager(VFS);
+
+  // Run once for `a.cc` and then for `a.h`. This makes sure we get the same
+  // file ID for `b.h` in the second run as `a.h` from first run.
+  const auto &OrigInputKind = Instance.getFrontendOpts().Inputs[0].getKind();
+  Instance.getFrontendOpts().Inputs.emplace_back("a.h", OrigInputKind);
+
+  SyntaxOnlyAction Act;
+  EXPECT_TRUE(Instance.ExecuteAction(Act)) << "Failed to execute action";
+  EXPECT_FALSE(Diags->hasErrorOccurred());
+  EXPECT_EQ(Diags->getNumWarnings(), 0u);
+}
 } // anonymous namespace