3 files changed, 53 insertions, 19 deletions

diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index 4ba0534fb5..48578569d5 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -2512,11 +2512,6 @@ bool ASTReader::ReadASTBlock(ModuleFile &F) {
     }
 
     case DIAG_PRAGMA_MAPPINGS:
-      if (Record.size() % 2 != 0) {
-        Error("invalid DIAG_USER_MAPPINGS block in AST file");
-        return true;
-      }
-        
       if (F.PragmaDiagMappings.empty())
         F.PragmaDiagMappings.swap(Record);
       else
@@ -4143,14 +4138,31 @@ HeaderFileInfo ASTReader::GetHeaderFileInfo(const FileEntry *FE) {
 }
 
 void ASTReader::ReadPragmaDiagnosticMappings(DiagnosticsEngine &Diag) {
+  // FIXME: Make it work properly with modules.
+  llvm::SmallVector<DiagnosticsEngine::DiagState *, 32> DiagStates;
   for (ModuleIterator I = ModuleMgr.begin(), E = ModuleMgr.end(); I != E; ++I) {
     ModuleFile &F = *(*I);
     unsigned Idx = 0;
+    DiagStates.clear();
+    assert(!Diag.DiagStates.empty());
+    DiagStates.push_back(&Diag.DiagStates.front()); // the command-line one.
     while (Idx < F.PragmaDiagMappings.size()) {
       SourceLocation Loc = ReadSourceLocation(F, F.PragmaDiagMappings[Idx++]);
+      unsigned DiagStateID = F.PragmaDiagMappings[Idx++];
+      if (DiagStateID != 0) {
+        Diag.DiagStatePoints.push_back(
+                    DiagnosticsEngine::DiagStatePoint(DiagStates[DiagStateID-1],
+                    FullSourceLoc(Loc, SourceMgr)));
+        continue;
+      }
+      
+      assert(DiagStateID == 0);
+      // A new DiagState was created here.
       Diag.DiagStates.push_back(*Diag.GetCurDiagState());
+      DiagnosticsEngine::DiagState *NewState = &Diag.DiagStates.back();
+      DiagStates.push_back(NewState);
       Diag.DiagStatePoints.push_back(
-          DiagnosticsEngine::DiagStatePoint(&Diag.DiagStates.back(),
+          DiagnosticsEngine::DiagStatePoint(NewState,
                                             FullSourceLoc(Loc, SourceMgr)));
       while (1) {
         assert(Idx < F.PragmaDiagMappings.size() &&
diff --git a/lib/Serialization/ASTWriter.cpp b/lib/Serialization/ASTWriter.cpp
index ea100864e3..500351f55f 100644
--- a/lib/Serialization/ASTWriter.cpp
+++ b/lib/Serialization/ASTWriter.cpp
@@ -2272,24 +2272,35 @@ ASTWriter::inferSubmoduleIDFromLocation(SourceLocation Loc) {
 }
 
 void ASTWriter::WritePragmaDiagnosticMappings(const DiagnosticsEngine &Diag) {
+  // FIXME: Make it work properly with modules.
+  llvm::SmallDenseMap<const DiagnosticsEngine::DiagState *, unsigned, 64>
+      DiagStateIDMap;
+  unsigned CurrID = 0;
+  DiagStateIDMap[&Diag.DiagStates.front()] = ++CurrID; // the command-line one.
   RecordData Record;
   for (DiagnosticsEngine::DiagStatePointsTy::const_iterator
          I = Diag.DiagStatePoints.begin(), E = Diag.DiagStatePoints.end();
          I != E; ++I) {
-    const DiagnosticsEngine::DiagStatePoint &point = *I; 
+    const DiagnosticsEngine::DiagStatePoint &point = *I;
     if (point.Loc.isInvalid())
       continue;
 
     Record.push_back(point.Loc.getRawEncoding());
-    for (DiagnosticsEngine::DiagState::const_iterator
-           I = point.State->begin(), E = point.State->end(); I != E; ++I) {
-      if (I->second.isPragma()) {
-        Record.push_back(I->first);
-        Record.push_back(I->second.getMapping());
+    unsigned &DiagStateID = DiagStateIDMap[point.State];
+    Record.push_back(DiagStateID);
+    
+    if (DiagStateID == 0) {
+      DiagStateID = ++CurrID;
+      for (DiagnosticsEngine::DiagState::const_iterator
+             I = point.State->begin(), E = point.State->end(); I != E; ++I) {
+        if (I->second.isPragma()) {
+          Record.push_back(I->first);
+          Record.push_back(I->second.getMapping());
+        }
       }
+      Record.push_back(-1); // mark the end of the diag/map pairs for this
+                            // location.
     }
-    Record.push_back(-1); // mark the end of the diag/map pairs for this
-                          // location.
   }
 
   if (!Record.empty())
diff --git a/test/PCH/pragma-diag-section.cpp b/test/PCH/pragma-diag-section.cpp
index 103b252416..627156f153 100644
--- a/test/PCH/pragma-diag-section.cpp
+++ b/test/PCH/pragma-diag-section.cpp
@@ -5,15 +5,13 @@
 // RUN: %clang_cc1 %s -emit-pch -o %t
 // RUN: %clang_cc1 %s -include-pch %t -verify -fsyntax-only
 
-// expected-no-diagnostics
-
 #ifndef HEADER
 #define HEADER
 
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wtautological-compare"
 template <typename T>
-struct TS {
+struct TS1 {
     void m() {
       T a = 0;
       T b = a==a;
@@ -23,9 +21,22 @@ struct TS {
 
 #else
 
+
+template <typename T>
+struct TS2 {
+    void m() {
+      T a = 0;
+      T b = a==a; // expected-warning {{self-comparison always evaluates to true}} expected-note@39 {{in instantiation of member function}}
+    }
+};
+
 void f() {
-    TS<int> ts;
-    ts.m();
+    TS1<int> ts1;
+    ts1.m();
+
+
+    TS2<int> ts2;
+    ts2.m();
 }
 
 #endif