[clang] Reset FileID based diag state mappings

clang/include/clang/Basic/Diagnostic.h

@@ -420,10 +420,13 @@ class DiagnosticsEngine : public RefCountedBase<DiagnosticsEngine> {
     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.
@@ -918,6 +921,10 @@ class DiagnosticsEngine : public RefCountedBase<DiagnosticsEngine> {
   /// 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.

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

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) {

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