3 files changed, 137 insertions, 22 deletions

diff --git a/clang/lib/Basic/Diagnostic.cpp b/clang/lib/Basic/Diagnostic.cpp
index 5098c2e4f8f..919ba8187ec 100644
--- a/clang/lib/Basic/Diagnostic.cpp
+++ b/clang/lib/Basic/Diagnostic.cpp
@@ -49,9 +49,10 @@ Diagnostic::Diagnostic(const llvm::IntrusiveRefCntPtr<DiagnosticIDs> &diags,
   ErrorLimit = 0;
   TemplateBacktraceLimit = 0;
 
-  // Set all mappings to 'unset'.
-  DiagMappingsStack.clear();
-  DiagMappingsStack.push_back(DiagMappings());
+  // Create a DiagState and DiagStatePoint representing diagnostic changes
+  // through command-line.
+  DiagStates.push_back(DiagState());
+  PushDiagStatePoint(&DiagStates.back(), SourceLocation());
 
   Reset();
 }
@@ -62,17 +63,19 @@ Diagnostic::~Diagnostic() {
 }
 
 
-void Diagnostic::pushMappings() {
-  // Avoids undefined behavior when the stack has to resize.
-  DiagMappingsStack.reserve(DiagMappingsStack.size() + 1);
-  DiagMappingsStack.push_back(DiagMappingsStack.back());
+void Diagnostic::pushMappings(SourceLocation Loc) {
+  DiagStateOnPushStack.push_back(GetCurDiagState());
 }
 
-bool Diagnostic::popMappings() {
-  if (DiagMappingsStack.size() == 1)
+bool Diagnostic::popMappings(SourceLocation Loc) {
+  if (DiagStateOnPushStack.empty())
     return false;
 
-  DiagMappingsStack.pop_back();
+  if (DiagStateOnPushStack.back() != GetCurDiagState()) {
+    // State changed at some point between push/pop.
+    PushDiagStatePoint(DiagStateOnPushStack.back(), Loc);
+  }
+  DiagStateOnPushStack.pop_back();
   return true;
 }
 
@@ -109,6 +112,91 @@ void Diagnostic::ReportDelayed() {
   DelayedDiagArg2.clear();
 }
 
+Diagnostic::DiagStatePointsTy::iterator
+Diagnostic::GetDiagStatePointForLoc(SourceLocation L) const {
+  assert(!DiagStatePoints.empty());
+  assert(DiagStatePoints.front().Loc.isInvalid() &&
+         "Should have created a DiagStatePoint for command-line");
+
+  FullSourceLoc Loc(L, *SourceMgr);
+  if (Loc.isInvalid())
+    return DiagStatePoints.end() - 1;
+
+  DiagStatePointsTy::iterator Pos = DiagStatePoints.end();
+  FullSourceLoc LastStateChangePos = DiagStatePoints.back().Loc;
+  if (LastStateChangePos.isValid() &&
+      Loc.isBeforeInTranslationUnitThan(LastStateChangePos))
+    Pos = std::upper_bound(DiagStatePoints.begin(), DiagStatePoints.end(),
+                           DiagStatePoint(0, Loc));
+  --Pos;
+  return Pos;
+}
+
+/// \brief This allows the client to specify that certain
+/// warnings are ignored.  Notes can never be mapped, errors can only be
+/// mapped to fatal, and WARNINGs and EXTENSIONs can be mapped arbitrarily.
+///
+/// \param The source location that this change of diagnostic state should
+/// take affect. It can be null if we are setting the latest state.
+void Diagnostic::setDiagnosticMapping(diag::kind Diag, diag::Mapping Map,
+                                      SourceLocation L) {
+  assert(Diag < diag::DIAG_UPPER_LIMIT &&
+         "Can only map builtin diagnostics");
+  assert((Diags->isBuiltinWarningOrExtension(Diag) ||
+          (Map == diag::MAP_FATAL || Map == diag::MAP_ERROR)) &&
+         "Cannot map errors into warnings!");
+  assert(!DiagStatePoints.empty());
+
+  FullSourceLoc Loc(L, *SourceMgr);
+  FullSourceLoc LastStateChangePos = DiagStatePoints.back().Loc;
+
+  // Common case; setting all the diagnostics of a group in one place.
+  if (Loc.isInvalid() || Loc == LastStateChangePos) {
+    setDiagnosticMappingInternal(Diag, Map, GetCurDiagState(), true);
+    return;
+  }
+
+  // Another common case; modifying diagnostic state in a source location
+  // after the previous one.
+  if ((Loc.isValid() && LastStateChangePos.isInvalid()) ||
+      LastStateChangePos.isBeforeInTranslationUnitThan(Loc)) {
+    // A diagnostic pragma occured, create a new DiagState initialized with
+    // the current one and a new DiagStatePoint to record at which location
+    // the new state became active.
+    DiagStates.push_back(*GetCurDiagState());
+    PushDiagStatePoint(&DiagStates.back(), Loc);
+    setDiagnosticMappingInternal(Diag, Map, GetCurDiagState(), true);
+    return;
+  }
+
+  // We allow setting the diagnostic state in random source order for
+  // completeness but it should not be actually happening in normal practice.
+
+  DiagStatePointsTy::iterator Pos = GetDiagStatePointForLoc(Loc);
+  assert(Pos != DiagStatePoints.end());
+
+  // Update all diagnostic states that are active after the given location.
+  for (DiagStatePointsTy::iterator
+         I = Pos+1, E = DiagStatePoints.end(); I != E; ++I) {
+    setDiagnosticMappingInternal(Diag, Map, I->State, true);
+  }
+
+  // If the location corresponds to an existing point, just update its state.
+  if (Pos->Loc == Loc) {
+    setDiagnosticMappingInternal(Diag, Map, Pos->State, true);
+    return;
+  }
+
+  // Create a new state/point and fit it into the vector of DiagStatePoints
+  // so that the vector is always ordered according to location.
+  Pos->Loc.isBeforeInTranslationUnitThan(Loc);
+  DiagStates.push_back(*Pos->State);
+  DiagState *NewState = &DiagStates.back();
+  setDiagnosticMappingInternal(Diag, Map, NewState, true);
+  DiagStatePoints.insert(Pos+1, DiagStatePoint(NewState,
+                                               FullSourceLoc(Loc, *SourceMgr)));
+}
+
 void DiagnosticBuilder::FlushCounts() {
   DiagObj->NumDiagArgs = NumArgs;
   DiagObj->NumDiagRanges = NumRanges;
diff --git a/clang/lib/Basic/DiagnosticIDs.cpp b/clang/lib/Basic/DiagnosticIDs.cpp
index 5d5bf7b72a2..29a8d9270a9 100644
--- a/clang/lib/Basic/DiagnosticIDs.cpp
+++ b/clang/lib/Basic/DiagnosticIDs.cpp
@@ -282,32 +282,42 @@ const char *DiagnosticIDs::getDescription(unsigned DiagID) const {
 /// getDiagnosticLevel - Based on the way the client configured the Diagnostic
 /// object, classify the specified diagnostic ID into a Level, consumable by
 /// the DiagnosticClient.
-DiagnosticIDs::Level DiagnosticIDs::getDiagnosticLevel(unsigned DiagID,
-                                        const Diagnostic &Diag) const {
+DiagnosticIDs::Level
+DiagnosticIDs::getDiagnosticLevel(unsigned DiagID, SourceLocation Loc,
+                                  const Diagnostic &Diag) const {
   // Handle custom diagnostics, which cannot be mapped.
   if (DiagID >= diag::DIAG_UPPER_LIMIT)
     return CustomDiagInfo->getLevel(DiagID);
 
   unsigned DiagClass = getBuiltinDiagClass(DiagID);
   assert(DiagClass != CLASS_NOTE && "Cannot get diagnostic level of a note!");
-  return getDiagnosticLevel(DiagID, DiagClass, Diag);
+  return getDiagnosticLevel(DiagID, DiagClass, Loc, Diag);
 }
 
-/// getDiagnosticLevel - Based on the way the client configured the Diagnostic
+/// \brief Based on the way the client configured the Diagnostic
 /// object, classify the specified diagnostic ID into a Level, consumable by
 /// the DiagnosticClient.
+///
+/// \param Loc The source location we are interested in finding out the
+/// diagnostic state. Can be null in order to query the latest state.
 DiagnosticIDs::Level
 DiagnosticIDs::getDiagnosticLevel(unsigned DiagID, unsigned DiagClass,
-                               const Diagnostic &Diag) const {
+                                  SourceLocation Loc,
+                                  const Diagnostic &Diag) const {
   // Specific non-error diagnostics may be mapped to various levels from ignored
   // to error.  Errors can only be mapped to fatal.
   DiagnosticIDs::Level Result = DiagnosticIDs::Fatal;
 
+  Diagnostic::DiagStatePointsTy::iterator
+    Pos = Diag.GetDiagStatePointForLoc(Loc);
+  Diagnostic::DiagState *State = Pos->State;
+
   // Get the mapping information, if unset, compute it lazily.
-  unsigned MappingInfo = Diag.getDiagnosticMappingInfo((diag::kind)DiagID);
+  unsigned MappingInfo = Diag.getDiagnosticMappingInfo((diag::kind)DiagID,
+                                                       State);
   if (MappingInfo == 0) {
     MappingInfo = GetDefaultDiagMapping(DiagID);
-    Diag.setDiagnosticMappingInternal(DiagID, MappingInfo, false);
+    Diag.setDiagnosticMappingInternal(DiagID, MappingInfo, State, false);
   }
 
   switch (MappingInfo & 7) {
@@ -404,17 +414,17 @@ static bool WarningOptionCompare(const WarningOption &LHS,
 }
 
 static void MapGroupMembers(const WarningOption *Group, diag::Mapping Mapping,
-                            Diagnostic &Diag) {
+                            SourceLocation Loc, Diagnostic &Diag) {
   // Option exists, poke all the members of its diagnostic set.
   if (const short *Member = Group->Members) {
     for (; *Member != -1; ++Member)
-      Diag.setDiagnosticMapping(*Member, Mapping);
+      Diag.setDiagnosticMapping(*Member, Mapping, Loc);
   }
 
   // Enable/disable all subgroups along with this one.
   if (const short *SubGroups = Group->SubGroups) {
     for (; *SubGroups != (short)-1; ++SubGroups)
-      MapGroupMembers(&OptionTable[(short)*SubGroups], Mapping, Diag);
+      MapGroupMembers(&OptionTable[(short)*SubGroups], Mapping, Loc, Diag);
   }
 }
 
@@ -423,7 +433,18 @@ static void MapGroupMembers(const WarningOption *Group, diag::Mapping Mapping,
 /// ignores the request if "Group" was unknown, false otherwise.
 bool DiagnosticIDs::setDiagnosticGroupMapping(const char *Group,
                                            diag::Mapping Map,
+                                           SourceLocation Loc,
                                            Diagnostic &Diag) const {
+  assert((Loc.isValid() ||
+          Diag.DiagStatePoints.empty() ||
+          Diag.DiagStatePoints.back().Loc.isInvalid()) &&
+         "Loc should be invalid only when the mapping comes from command-line");
+  assert((Loc.isInvalid() || Diag.DiagStatePoints.empty() ||
+          Diag.DiagStatePoints.back().Loc.isInvalid() ||
+          !Diag.SourceMgr->isBeforeInTranslationUnit(Loc,
+                                            Diag.DiagStatePoints.back().Loc)) &&
+         "Source location of new mapping is before the previous one!");
+
   WarningOption Key = { Group, 0, 0 };
   const WarningOption *Found =
   std::lower_bound(OptionTable, OptionTable + OptionTableSize, Key,
@@ -432,7 +453,7 @@ bool DiagnosticIDs::setDiagnosticGroupMapping(const char *Group,
       strcmp(Found->Name, Group) != 0)
     return true;  // Option not found.
 
-  MapGroupMembers(Found, Map, Diag);
+  MapGroupMembers(Found, Map, Loc, Diag);
   return false;
 }
 
@@ -475,7 +496,8 @@ bool DiagnosticIDs::ProcessDiag(Diagnostic &Diag) const {
       // *map* warnings/extensions to errors.
       ShouldEmitInSystemHeader = DiagClass == CLASS_ERROR;
 
-      DiagLevel = getDiagnosticLevel(DiagID, DiagClass, Diag);
+      DiagLevel = getDiagnosticLevel(DiagID, DiagClass, Info.getLocation(),
+                                     Diag);
     }
   }
 
diff --git a/clang/lib/Basic/SourceLocation.cpp b/clang/lib/Basic/SourceLocation.cpp
index 1571e2205f8..5062d43f58f 100644
--- a/clang/lib/Basic/SourceLocation.cpp
+++ b/clang/lib/Basic/SourceLocation.cpp
@@ -110,6 +110,11 @@ bool FullSourceLoc::isInSystemHeader() const {
   return SrcMgr->isInSystemHeader(*this);
 }
 
+bool FullSourceLoc::isBeforeInTranslationUnitThan(SourceLocation Loc) const {
+  assert(isValid());
+  return SrcMgr->isBeforeInTranslationUnit(*this, Loc);
+}
+
 const char *FullSourceLoc::getCharacterData(bool *Invalid) const {
   assert(isValid());
   return SrcMgr->getCharacterData(*this, Invalid);