Remove and replace DiagStatePoint tracking and lookup data structure.

Rather than storing a single flat list of SourceLocations where the diagnostic
state changes (in source order), we now store a separate list for each FileID
in which there is a diagnostic state transition. (State for other files is
built and cached lazily, on demand.) This has two consequences:

1) We can now sensibly support modules, and properly track the diagnostic state
for modular headers (this matters when, for instance, triggering instantiation
of a template defined within a module triggers diagnostics).

2) It's much faster than the old approach, since we can now just do a binary
search on the offsets within the FileID rather than needing to call
isBeforeInTranslationUnit to determine source order (which is surprisingly
slow). For some pathological (but real world) files, this reduces total
compilation time by more than 10%.

For now, the diagnostic state points for modules are loaded eagerly. It seems
feasible to defer this until diagnostic state information for one of the
module's files is needed, but that's not part of this patch.

llvm-svn: 293123

2 files changed, 101 insertions, 89 deletions

diff --git a/clang/lib/Basic/Diagnostic.cpp b/clang/lib/Basic/Diagnostic.cpp
index 89b704da011..0c27ae4638d 100644
--- a/clang/lib/Basic/Diagnostic.cpp
+++ b/clang/lib/Basic/Diagnostic.cpp
@@ -132,13 +132,13 @@ void DiagnosticsEngine::Reset() {
 
   // Clear state related to #pragma diagnostic.
   DiagStates.clear();
-  DiagStatePoints.clear();
+  DiagStatesByLoc.clear();
   DiagStateOnPushStack.clear();
 
   // Create a DiagState and DiagStatePoint representing diagnostic changes
   // through command-line.
   DiagStates.emplace_back();
-  DiagStatePoints.emplace_back(&DiagStates.back(), SourceLocation());
+  DiagStatesByLoc.appendFirst(&DiagStates.back());
 }
 
 void DiagnosticsEngine::SetDelayedDiagnostic(unsigned DiagID, StringRef Arg1,
@@ -158,49 +158,94 @@ void DiagnosticsEngine::ReportDelayed() {
   DelayedDiagArg2.clear();
 }
 
+void DiagnosticsEngine::DiagStateMap::appendFirst(
+                                             DiagState *State) {
+  assert(Files.empty() && "not first");
+  FirstDiagState = CurDiagState = State;
+  CurDiagStateLoc = SourceLocation();
+}
+
+void DiagnosticsEngine::DiagStateMap::append(SourceManager &SrcMgr,
+                                             SourceLocation Loc,
+                                             DiagState *State) {
+  CurDiagState = State;
+  CurDiagStateLoc = Loc;
+
+  std::pair<FileID, unsigned> Decomp = SrcMgr.getDecomposedLoc(Loc);
+  unsigned Offset = Decomp.second;
+  for (File *F = getFile(SrcMgr, Decomp.first); F;
+       Offset = F->ParentOffset, F = F->Parent) {
+    F->HasLocalTransitions = true;
+    auto &Last = F->StateTransitions.back();
+    assert(Last.Offset <= Offset && "state transitions added out of order");
+
+    if (Last.Offset == Offset) {
+      if (Last.State == State)
+        break;
+      Last.State = State;
+      continue;
+    }
+
+    F->StateTransitions.push_back({State, Offset});
+  }
+}
+
+DiagnosticsEngine::DiagState *
+DiagnosticsEngine::DiagStateMap::lookup(SourceManager &SrcMgr,
+                                        SourceLocation Loc) const {
+  // Common case: we have not seen any diagnostic pragmas.
+  if (Files.empty())
+    return FirstDiagState;
+
+  std::pair<FileID, unsigned> Decomp = SrcMgr.getDecomposedLoc(Loc);
+  const File *F = getFile(SrcMgr, Decomp.first);
+  return F->lookup(Decomp.second);
+}
+
+DiagnosticsEngine::DiagState *
+DiagnosticsEngine::DiagStateMap::File::lookup(unsigned Offset) const {
+  auto OnePastIt = std::upper_bound(
+      StateTransitions.begin(), StateTransitions.end(), Offset,
+      [](unsigned Offset, const DiagStatePoint &P) {
+        return Offset < P.Offset;
+      });
+  assert(OnePastIt != StateTransitions.begin() && "missing initial state");
+  return OnePastIt[-1].State;
+}
+
+DiagnosticsEngine::DiagStateMap::File *
+DiagnosticsEngine::DiagStateMap::getFile(SourceManager &SrcMgr,
+                                         FileID ID) const {
+  // Get or insert the File for this ID.
+  auto Range = Files.equal_range(ID);
+  if (Range.first != Range.second)
+    return &Range.first->second;
+  auto &F = Files.insert(Range.first, std::make_pair(ID, File()))->second;
+
+  // We created a new File; look up the diagnostic state at the start of it and
+  // initialize it.
+  if (ID.isValid()) {
+    std::pair<FileID, unsigned> Decomp = SrcMgr.getDecomposedIncludedLoc(ID);
+    F.Parent = getFile(SrcMgr, Decomp.first);
+    F.ParentOffset = Decomp.second;
+    F.StateTransitions.push_back({F.Parent->lookup(Decomp.second), 0});
+  } else {
+    // This is the (imaginary) root file into which we pretend all top-level
+    // files are included; it descends from the initial state.
+    //
+    // FIXME: This doesn't guarantee that we use the same ordering as
+    // isBeforeInTranslationUnit in the cases where someone invented another
+    // top-level file and added diagnostic pragmas to it. See the code at the
+    // end of isBeforeInTranslationUnit for the quirks it deals with.
+    F.StateTransitions.push_back({FirstDiagState, 0});
+  }
+  return &F;
+}
+
 void DiagnosticsEngine::PushDiagStatePoint(DiagState *State,
                                            SourceLocation Loc) {
-  // Make sure that DiagStatePoints is always sorted according to Loc.
   assert(Loc.isValid() && "Adding invalid loc point");
-  assert(!DiagStatePoints.empty() &&
-         (DiagStatePoints.back().Loc.isInvalid() ||
-          getSourceManager().isBeforeInTranslationUnit(
-              DiagStatePoints.back().Loc, Loc)) &&
-         "Previous point loc comes after or is the same as new one");
-  DiagStatePoints.push_back(DiagStatePoint(State, Loc));
-}
-
-DiagnosticsEngine::DiagStatePointsTy::iterator
-DiagnosticsEngine::GetDiagStatePointForLoc(SourceLocation L) const {
-  assert(!DiagStatePoints.empty());
-  assert(DiagStatePoints.front().Loc.isInvalid() &&
-         "Should have created a DiagStatePoint for command-line");
-
-  if (!SourceMgr)
-    return DiagStatePoints.end() - 1;
-
-  FullSourceLoc Loc(L, *SourceMgr);
-  if (Loc.isInvalid())
-    return DiagStatePoints.end() - 1;
-
-  DiagStatePointsTy::iterator Pos = DiagStatePoints.end();
-  SourceLocation LastStateChangePos = DiagStatePoints.back().Loc;
-  if (LastStateChangePos.isValid() &&
-      Loc.isBeforeInTranslationUnitThan(LastStateChangePos))
-    Pos = std::upper_bound(
-        DiagStatePoints.begin(), DiagStatePoints.end(),
-        DiagStatePoint(nullptr, Loc),
-        [this](const DiagStatePoint &LHS, const DiagStatePoint &RHS) {
-          // If Loc is invalid it means it came from <command-line>, in which
-          // case we regard it as coming before any valid source location.
-          if (RHS.Loc.isInvalid())
-            return false;
-          if (LHS.Loc.isInvalid())
-            return true;
-          return SourceMgr->isBeforeInTranslationUnit(LHS.Loc, RHS.Loc);
-        });
-  --Pos;
-  return Pos;
+  DiagStatesByLoc.append(*SourceMgr, Loc, State);
 }
 
 void DiagnosticsEngine::setSeverity(diag::kind Diag, diag::Severity Map,
@@ -210,11 +255,8 @@ void DiagnosticsEngine::setSeverity(diag::kind Diag, diag::Severity Map,
   assert((Diags->isBuiltinWarningOrExtension(Diag) ||
           (Map == diag::Severity::Fatal || Map == diag::Severity::Error)) &&
          "Cannot map errors into warnings!");
-  assert(!DiagStatePoints.empty());
   assert((L.isInvalid() || SourceMgr) && "No SourceMgr for valid location");
 
-  SourceLocation Loc = SourceMgr ? L : SourceLocation();
-  SourceLocation LastStateChangePos = DiagStatePoints.back().Loc;
   // Don't allow a mapping to a warning override an error/fatal mapping.
   if (Map == diag::Severity::Warning) {
     DiagnosticMapping &Info = GetCurDiagState()->getOrAddMapping(Diag);
@@ -225,49 +267,22 @@ void DiagnosticsEngine::setSeverity(diag::kind Diag, diag::Severity Map,
   DiagnosticMapping Mapping = makeUserMapping(Map, L);
 
   // Common case; setting all the diagnostics of a group in one place.
-  if (Loc.isInvalid() || Loc == LastStateChangePos) {
-    GetCurDiagState()->setMapping(Diag, Mapping);
-    return;
-  }
-
-  // Another common case; modifying diagnostic state in a source location
-  // after the previous one.
-  if ((Loc.isValid() && LastStateChangePos.isInvalid()) ||
-      SourceMgr->isBeforeInTranslationUnit(LastStateChangePos, Loc)) {
-    // A diagnostic pragma occurred, 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);
-    GetCurDiagState()->setMapping(Diag, Mapping);
-    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) {
-    I->State->setMapping(Diag, Mapping);
-  }
-
-  // If the location corresponds to an existing point, just update its state.
-  if (Pos->Loc == Loc) {
-    Pos->State->setMapping(Diag, Mapping);
+  if ((L.isInvalid() || L == DiagStatesByLoc.getCurDiagStateLoc()) &&
+      DiagStatesByLoc.getCurDiagState()) {
+    // FIXME: This is theoretically wrong: if the current state is shared with
+    // some other location (via push/pop) we will change the state for that
+    // other location as well. This cannot currently happen, as we can't update
+    // the diagnostic state at the same location at which we pop.
+    DiagStatesByLoc.getCurDiagState()->setMapping(Diag, Mapping);
     return;
   }
 
-  // Create a new state/point and fit it into the vector of DiagStatePoints
-  // so that the vector is always ordered according to location.
-  assert(SourceMgr->isBeforeInTranslationUnit(Pos->Loc, Loc));
-  DiagStates.push_back(*Pos->State);
-  DiagState *NewState = &DiagStates.back();
-  NewState->setMapping(Diag, Mapping);
-  DiagStatePoints.insert(Pos + 1, DiagStatePoint(NewState, Loc));
+  // A diagnostic pragma occurred, 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());
+  DiagStates.back().setMapping(Diag, Mapping);
+  PushDiagStatePoint(&DiagStates.back(), L);
 }
 
 bool DiagnosticsEngine::setSeverityForGroup(diag::Flavor Flavor,
diff --git a/clang/lib/Basic/DiagnosticIDs.cpp b/clang/lib/Basic/DiagnosticIDs.cpp
index 3c370f67fa3..e0580af45b5 100644
--- a/clang/lib/Basic/DiagnosticIDs.cpp
+++ b/clang/lib/Basic/DiagnosticIDs.cpp
@@ -411,11 +411,8 @@ DiagnosticIDs::getDiagnosticSeverity(unsigned DiagID, SourceLocation Loc,
   // to error.  Errors can only be mapped to fatal.
   diag::Severity Result = diag::Severity::Fatal;
 
-  DiagnosticsEngine::DiagStatePointsTy::iterator
-    Pos = Diag.GetDiagStatePointForLoc(Loc);
-  DiagnosticsEngine::DiagState *State = Pos->State;
-
   // Get the mapping information, or compute it lazily.
+  DiagnosticsEngine::DiagState *State = Diag.GetDiagStateForLoc(Loc);
   DiagnosticMapping &Mapping = State->getOrAddMapping((diag::kind)DiagID);
 
   // TODO: Can a null severity really get here?