[FileCheck] Fix -dump-input per-pattern diagnostic indexing

In input dump annotations, `check:2'1` indicates diagnostic 1 for the
`CHECK` directive on check file line 2.  Without this patch,
`-dump-input` computes the diagnostic index with the assumption that
FileCheck *consecutively* produces all diagnostics for the same
pattern.  Already, that can be a false assumption, as in the examples
below.  Moreover, it seems like a brittle assumption as FileCheck
evolves.  Finally, it actually complicates the implementation even if
it makes it slightly more efficient.

This patch avoids that assumption.  Examples below show results after
applying this patch.  Before applying this patch, `'N` is omitted
throughout these examples because the implementation doesn't notice
there's more than one diagnostic per pattern.

First, `CHECK-LABEL` violates the assumption because `CHECK-LABEL`
tries to match twice, and other directives can match in between:

```
$ cat check
CHECK: foobar
CHECK-LABEL: foobar

$ FileCheck -vv check < input |& tail -8
<<<<<<
           1: text
           2: foobar
label:2'0     ^~~~~~
check:1       ^~~~~~
label:2'1           X error: no match found
           3: text
>>>>>>
```

Second, `--implicit-check-not` is obviously processed many times among
other directives:

```
$ cat check
CHECK: foo
CHECK: foo

$ FileCheck -vv -dump-input=always -implicit-check-not=foo \
            check < input |& tail -16
<<<<<<
            1: text
not:imp1'0     X~~~~
            2: foo
check:1        ^~~
not:imp1'1        X
            3: text
not:imp1'1     ~~~~~
            4: foo
check:2        ^~~
not:imp1'2        X
            5: text
not:imp1'2     ~~~~~
            6:
eof:2          ^
>>>>>>
```

Reviewed By: thopre, jhenderson

Differential Revision: https://reviews.llvm.org/D97813

1 files changed, 17 insertions, 16 deletions

diff --git a/llvm/utils/FileCheck/FileCheck.cpp b/llvm/utils/FileCheck/FileCheck.cpp
index c1bb97f..0e97c711 100644
--- a/llvm/utils/FileCheck/FileCheck.cpp
+++ b/llvm/utils/FileCheck/FileCheck.cpp
@@ -22,6 +22,7 @@
 #include "llvm/Support/WithColor.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cmath>
+#include <map>
 using namespace llvm;
 
 static cl::extrahelp FileCheckOptsEnv(
@@ -378,16 +379,25 @@ BuildInputAnnotations(const SourceMgr &SM, unsigned CheckFileBufferID,
                       const std::vector<FileCheckDiag> &Diags,
                       std::vector<InputAnnotation> &Annotations,
                       unsigned &LabelWidth) {
-  // How many diagnostics have we seen so far?
-  unsigned DiagCount = 0;
-  // How many diagnostics has the current check seen so far?
-  unsigned CheckDiagCount = 0;
+  struct CompareSMLoc {
+    bool operator()(const SMLoc &LHS, const SMLoc &RHS) {
+      return LHS.getPointer() < RHS.getPointer();
+    }
+  };
+  // How many diagnostics does each pattern have?
+  std::map<SMLoc, unsigned, CompareSMLoc> DiagCountPerPattern;
+  for (auto Diag : Diags)
+    ++DiagCountPerPattern[Diag.CheckLoc];
+  // How many diagnostics have we seen so far per pattern?
+  std::map<SMLoc, unsigned, CompareSMLoc> DiagIndexPerPattern;
+  // How many total diagnostics have we seen so far?
+  unsigned DiagIndex = 0;
   // What's the widest label?
   LabelWidth = 0;
   for (auto DiagItr = Diags.begin(), DiagEnd = Diags.end(); DiagItr != DiagEnd;
        ++DiagItr) {
     InputAnnotation A;
-    A.DiagIndex = DiagCount++;
+    A.DiagIndex = DiagIndex++;
 
     // Build label, which uniquely identifies this check result.
     unsigned CheckBufferID = SM.FindBufferContainingLoc(DiagItr->CheckLoc);
@@ -403,17 +413,8 @@ BuildInputAnnotations(const SourceMgr &SM, unsigned CheckFileBufferID,
     else
       llvm_unreachable("expected diagnostic's check location to be either in "
                        "the check file or for an implicit pattern");
-    unsigned CheckDiagIndex = UINT_MAX;
-    auto DiagNext = std::next(DiagItr);
-    if (DiagNext != DiagEnd && DiagItr->CheckTy == DiagNext->CheckTy &&
-        DiagItr->CheckLoc == DiagNext->CheckLoc)
-      CheckDiagIndex = CheckDiagCount++;
-    else if (CheckDiagCount) {
-      CheckDiagIndex = CheckDiagCount;
-      CheckDiagCount = 0;
-    }
-    if (CheckDiagIndex != UINT_MAX)
-      Label << "'" << CheckDiagIndex;
+    if (DiagCountPerPattern[DiagItr->CheckLoc] > 1)
+      Label << "'" << DiagIndexPerPattern[DiagItr->CheckLoc]++;
     Label.flush();
     LabelWidth = std::max((std::string::size_type)LabelWidth, A.Label.size());