[mlir] Add support for regex within `expected-*` diagnostics

[rriddle]

Summary

This can be enabled by using a -re suffix when defining the expected line,
e.g. expected-error-re. This support is similar to what clang provides in its "expected"
diagnostic framework(e.g. the -re is also the same). The regex definitions themselves are
similar to FileCheck in that regex blocks are specified within {{ }} blocks.

[jpienaar]

mlir/docs/Diagnostics.md

306

Any reason to not follow filecheck? E.g., re is default and these are special tokens and you have literal matches with special tags. Is this due to outer delimiters being {{?

[rriddle]

mlir/docs/Diagnostics.md

306

My thoughts there were:

a) This follows the same precedent as clang/other tools
b) We'd need to rewrite many checks

Seemed better to follow the standard set by clang/other tools in llvm for the expected-* things, and also avoid unnecessary churn.

[jpienaar]

Agreed, following clang SGTM

mlir/lib/IR/Diagnostics.cpp

612

Nit: prefer positive conditional

625

Do you want to allow escaping too? E.g., if string has literal {{ and actual regex's ? (I see clang's VerifyDiagnosticConsumer doesn't either, and can be worked around with same trick as we do FileCheck side)

[rriddle]

mlir/lib/IR/Diagnostics.cpp

625

We could, but I can't find any places right now where we would use it. I think if it becomes common enough, we could add something. It just doesn't look like it pays for itself right now (given that it has a workaround).

mlir/docs/Diagnostics.md

 the stack will still be shown.
 
 ### SourceMgr Diagnostic Verifier Handler
 
 This handler is a wrapper around a llvm::SourceMgr that is used to verify that
 certain diagnostics have been emitted to the context. To use this handler,
 annotate your source file with expected diagnostics in the form of:
 
-*   `expected-(error|note|remark|warning) {{ message }}`
+*   `expected-(error|note|remark|warning)(-re)? {{ message }}`
 
-A few examples are shown below:
+The provided `message` is a string expected to be contained within the generated
+diagnostic. The `-re` suffix may be used to enable regex matching within the
+`message`. When present, the `message` may define regex match sequences within
+`{{` `}}` blocks. The regular expression matcher supports Extended POSIX regular
+expressions (ERE). A few examples are shown below:
 
 ```mlir
 // Expect an error on the same line.
 func.func @bad_branch() {
   cf.br ^missing  // expected-error {{reference to an undefined block}}
 }
 
 // Expect an error on an adjacent line.
 func.func @foo(%a : f32) {
   // expected-error@+1 {{unknown comparison predicate "foo"}}
   %result = arith.cmpf "foo", %a, %a : f32
   return
 }
 
 // Expect an error on the next line that does not contain a designator.
 // expected-remark@below {{remark on function below}}
 // expected-remark@below {{another remark on function below}}
 func.func @bar(%a : f32)
 
 // Expect an error on the previous line that does not contain a designator.
 func.func @baz(%a : f32)
 // expected-remark@above {{remark on function above}}
 // expected-remark@above {{another remark on function above}}
 
+// Expect an error mentioning the parent function, but use regex to avoid
+// hardcoding the name.
+func.func @foo() -> i32 {
+  // expected-error-re@+1 {{'func.return' op has 0 operands, but enclosing function (@{{.*}}) returns 1}}
+  return
+}
 ```
 
 The handler will report an error if any unexpected diagnostics were seen, or if
 any expected diagnostics weren't.
 
 ```shell
 $ mlir-opt foo.mlir
 

mlir/lib/IR/Diagnostics.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // SourceMgrDiagnosticVerifierHandler
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
-// Record the expected diagnostic's position, substring and whether it was
-// seen.
+/// This class represents an expected output diagnostic.
 struct ExpectedDiag {
+  ExpectedDiag(DiagnosticSeverity kind, unsigned lineNo, SMLoc fileLoc,
+               StringRef substring)
+      : kind(kind), lineNo(lineNo), fileLoc(fileLoc), substring(substring) {}
+
+  /// Emit an error at the location referenced by this diagnostic.
+  LogicalResult emitError(raw_ostream &os, llvm::SourceMgr &mgr,
+                          const Twine &msg) {
+    SMRange range(fileLoc, SMLoc::getFromPointer(fileLoc.getPointer() +
+                                                 substring.size()));
+    mgr.PrintMessage(os, fileLoc, llvm::SourceMgr::DK_Error, msg, range);
+    return failure();
+  }
+
+  /// Returns true if this diagnostic matches the given string.
+  bool match(StringRef str) const {
+    // If this isn't a regex diagnostic, we simply check if the string was
+    // contained.
+    if (substringRegex)
+      return substringRegex->match(str);
+    return str.contains(substring);
+  }
+
+  /// Compute the regex matcher for this diagnostic, using the provided stream
+  /// and manager to emit diagnostics as necessary.
+  LogicalResult computeRegex(raw_ostream &os, llvm::SourceMgr &mgr) {
+    std::string regexStr;
+    llvm::raw_string_ostream regexOS(regexStr);
+    StringRef strToProcess = substring;
+    while (!strToProcess.empty()) {
+      // Find the next regex block.
+      size_t regexIt = strToProcess.find("{{");
+      if (regexIt == StringRef::npos) {
+        regexOS << llvm::Regex::escape(strToProcess);
+        break;
+      }
+      regexOS << llvm::Regex::escape(strToProcess.take_front(regexIt));
+      strToProcess = strToProcess.drop_front(regexIt + 2);
+
+      // Find the end of the regex block.
+      size_t regexEndIt = strToProcess.find("}}");
+      if (regexEndIt == StringRef::npos)
+        return emitError(os, mgr, "found start of regex with no end '}}'");
+      StringRef regexStr = strToProcess.take_front(regexEndIt);
+
+      // Validate that the regex is actually valid.
+      std::string regexError;
+      if (!llvm::Regex(regexStr).isValid(regexError))
+        return emitError(os, mgr, "invalid regex: " + regexError);
+
+      regexOS << '(' << regexStr << ')';
+      strToProcess = strToProcess.drop_front(regexEndIt + 2);
+    }
+    substringRegex = llvm::Regex(regexOS.str());
+    return success();
+  }
+
+  /// The severity of the diagnosic expected.
   DiagnosticSeverity kind;
+  /// The line number the expected diagnostic should be on.
   unsigned lineNo;
-  StringRef substring;
+  /// The location of the expected diagnostic within the input file.
   SMLoc fileLoc;
-  bool matched;
+  /// A flag indicating if the expected diagnostic has been matched yet.
+  bool matched = false;
+  /// The substring that is expected to be within the diagnostic.
+  StringRef substring;
+  /// An optional regex matcher, if the expected diagnostic sub-string was a
+  /// regex string.
+  Optional<llvm::Regex> substringRegex;
 };
 
 struct SourceMgrDiagnosticVerifierHandlerImpl {
   SourceMgrDiagnosticVerifierHandlerImpl() : status(success()) {}
 
   /// Returns the expected diagnostics for the given source file.
   Optional<MutableArrayRef<ExpectedDiag>> getExpectedDiags(StringRef bufName);
 
   /// Computes the expected diagnostics for the given source buffer.
   MutableArrayRef<ExpectedDiag>
-  computeExpectedDiags(const llvm::MemoryBuffer *buf);
+  computeExpectedDiags(raw_ostream &os, llvm::SourceMgr &mgr,
+                       const llvm::MemoryBuffer *buf);
 
   /// The current status of the verifier.
   LogicalResult status;
 
   /// A list of expected diagnostics for each buffer of the source manager.
   llvm::StringMap<SmallVector<ExpectedDiag, 2>> expectedDiagsPerFile;
 
   /// Regex to match the expected diagnostics format.
-  llvm::Regex expected = llvm::Regex("expected-(error|note|remark|warning) "
-                                     "*(@([+-][0-9]+|above|below))? *{{(.*)}}");
+  llvm::Regex expected =
+      llvm::Regex("expected-(error|note|remark|warning)(-re)? "
+                  "*(@([+-][0-9]+|above|below))? *{{(.*)}}$");
 };
 } // namespace detail
 } // namespace mlir
 
 /// Given a diagnostic kind, return a human readable string for it.
 static StringRef getDiagKindStr(DiagnosticSeverity kind) {
   switch (kind) {
   case DiagnosticSeverity::Note:
     return "note";
   case DiagnosticSeverity::Warning:
     return "warning";
   case DiagnosticSeverity::Error:
     return "error";
   case DiagnosticSeverity::Remark:
     return "remark";
   }
   llvm_unreachable("Unknown DiagnosticSeverity");
 }
 
-/// Returns the expected diagnostics for the given source file.
 Optional<MutableArrayRef<ExpectedDiag>>
 SourceMgrDiagnosticVerifierHandlerImpl::getExpectedDiags(StringRef bufName) {
   auto expectedDiags = expectedDiagsPerFile.find(bufName);
   if (expectedDiags != expectedDiagsPerFile.end())
     return MutableArrayRef<ExpectedDiag>(expectedDiags->second);
   return llvm::None;
 }
 
-/// Computes the expected diagnostics for the given source buffer.
 MutableArrayRef<ExpectedDiag>
 SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
-    const llvm::MemoryBuffer *buf) {
+    raw_ostream &os, llvm::SourceMgr &mgr, const llvm::MemoryBuffer *buf) {
   // If the buffer is invalid, return an empty list.
   if (!buf)
     return llvm::None;
   auto &expectedDiags = expectedDiagsPerFile[buf->getBufferIdentifier()];
 
   // The number of the last line that did not correlate to a designator.
   unsigned lastNonDesignatorLine = 0;
 
   // The indices of designators that apply to the next non designator line.
   SmallVector<unsigned, 1> designatorsForNextLine;
 
   // Scan the file for expected-* designators.
   SmallVector<StringRef, 100> lines;
   buf->getBuffer().split(lines, '\n');
   for (unsigned lineNo = 0, e = lines.size(); lineNo < e; ++lineNo) {
     SmallVector<StringRef, 4> matches;
-    if (!expected.match(lines[lineNo], &matches)) {
+    if (!expected.match(lines[lineNo].rtrim(), &matches)) {
       // Check for designators that apply to this line.
       if (!designatorsForNextLine.empty()) {
         for (unsigned diagIndex : designatorsForNextLine)
           expectedDiags[diagIndex].lineNo = lineNo + 1;
         designatorsForNextLine.clear();
       }
       lastNonDesignatorLine = lineNo;
       continue;
     }
 
     // Point to the start of expected-*.
-    auto expectedStart = SMLoc::getFromPointer(matches[0].data());
+    SMLoc expectedStart = SMLoc::getFromPointer(matches[0].data());
 
     DiagnosticSeverity kind;
     if (matches[1] == "error")
       kind = DiagnosticSeverity::Error;
     else if (matches[1] == "warning")
       kind = DiagnosticSeverity::Warning;
     else if (matches[1] == "remark")
       kind = DiagnosticSeverity::Remark;
     else {
       assert(matches[1] == "note");
       kind = DiagnosticSeverity::Note;
     }
+    ExpectedDiag record(kind, lineNo + 1, expectedStart, matches[5]);
 
-    ExpectedDiag record{kind, lineNo + 1, matches[4], expectedStart, false};
-    auto offsetMatch = matches[2];
+    // Check to see if this is a regex match, i.e. it includes the `-re`.
+    if (!matches[2].empty() && failed(record.computeRegex(os, mgr))) {
+      status = failure();
+      continue;
+    }
+
+    StringRef offsetMatch = matches[3];
     if (!offsetMatch.empty()) {
       offsetMatch = offsetMatch.drop_front(1);
 
       // Get the integer value without the @ and +/- prefix.
       if (offsetMatch[0] == '+' || offsetMatch[0] == '-') {
         int offset;
         offsetMatch.drop_front().getAsInteger(0, offset);
 
if (!offsetMatch.empty()) {
         assert(offsetMatch.consume_front("below"));
         designatorsForNextLine.push_back(expectedDiags.size());
 
         // Set the line number to the last in the case that this designator ends
         // up dangling.
         record.lineNo = e;
       }
     }
-    expectedDiags.push_back(record);
+    expectedDiags.emplace_back(std::move(record));
   }
   return expectedDiags;
 }
 
 SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
     llvm::SourceMgr &srcMgr, MLIRContext *ctx, raw_ostream &out)
     : SourceMgrDiagnosticHandler(srcMgr, ctx, out),
       impl(new SourceMgrDiagnosticVerifierHandlerImpl()) {
   // Compute the expected diagnostics for each of the current files in the
   // source manager.
   for (unsigned i = 0, e = mgr.getNumBuffers(); i != e; ++i)
-    (void)impl->computeExpectedDiags(mgr.getMemoryBuffer(i + 1));
+    (void)impl->computeExpectedDiags(out, mgr, mgr.getMemoryBuffer(i + 1));
 
   // Register a handler to verify the diagnostics.
   setHandler([&](Diagnostic &diag) {
     // Process the main diagnostics.
     process(diag);
 
     // Process each of the notes.
     for (auto &note : diag.getNotes())

 /// diagnostics were emitted. This return success if all diagnostics were
 /// verified correctly, failure otherwise.
 LogicalResult SourceMgrDiagnosticVerifierHandler::verify() {
   // Verify that all expected errors were seen.
   for (auto &expectedDiagsPair : impl->expectedDiagsPerFile) {
     for (auto &err : expectedDiagsPair.second) {
       if (err.matched)
         continue;
-      SMRange range(err.fileLoc,
-                          SMLoc::getFromPointer(err.fileLoc.getPointer() +
-                                                      err.substring.size()));
-      mgr.PrintMessage(os, err.fileLoc, llvm::SourceMgr::DK_Error,
+      impl->status =
+          err.emitError(os, mgr,
                         "expected " + getDiagKindStr(err.kind) + " \"" +
-                           err.substring + "\" was not produced",
-                       range);
-      impl->status = failure();
+                            err.substring + "\" was not produced");
     }
   }
   impl->expectedDiagsPerFile.clear();
   return impl->status;
 }
 
 /// Process a single diagnostic.
 void SourceMgrDiagnosticVerifierHandler::process(Diagnostic &diag) {

 }
 
 /// Process a FileLineColLoc diagnostic.
 void SourceMgrDiagnosticVerifierHandler::process(FileLineColLoc loc,
                                                  StringRef msg,
                                                  DiagnosticSeverity kind) {
   // Get the expected diagnostics for this file.
   auto diags = impl->getExpectedDiags(loc.getFilename());
-  if (!diags)
-    diags = impl->computeExpectedDiags(getBufferForFile(loc.getFilename()));
+  if (!diags) {
+    diags = impl->computeExpectedDiags(os, mgr,
+                                       getBufferForFile(loc.getFilename()));
+  }
 
   // Search for a matching expected diagnostic.
   // If we find something that is close then emit a more specific error.
   ExpectedDiag *nearMiss = nullptr;
 
   // If this was an expected error, remember that we saw it and return.
   unsigned line = loc.getLine();
   for (auto &e : *diags) {
-    if (line == e.lineNo && msg.contains(e.substring)) {
+    if (line == e.lineNo && e.match(msg)) {
       if (e.kind == kind) {
         e.matched = true;
         return;
       }
 
       // If this only differs based on the diagnostic kind, then consider it
       // to be a near miss.
       nearMiss = &e;

mlir/test/Dialect/SPIRV/IR/memory-ops.mlir

func.func @simple_store_missing_operand(%arg0 : f32) -> () {
   spv.Store  "Function" , %arg0 : f32
   return
 }
 
 // -----
 
 func.func @simple_store_missing_operand(%arg0 : f32) -> () {
   %0 = spv.Variable : !spv.ptr<f32, Function>
-  // expected-error @+1 {{custom op 'spv.Store' expected 2 operands}} : f32
+  // expected-error @+1 {{custom op 'spv.Store' expected 2 operands}}
   spv.Store  "Function" %0 : f32
   return
 }
 
 // -----
 
 func.func @volatile_store_missing_lbrace(%arg0 : f32) -> () {
   %0 = spv.Variable : !spv.ptr<f32, Function>

mlir/test/IR/diagnostic-handler-verify-regex.mlir

+// RUN: not mlir-opt -allow-unregistered-dialect %s -split-input-file -verify-diagnostics 2>&1 | FileCheck %s
+
+// CHECK: found start of regex with no end '}}'
+// expected-error-re {{{{}}
+
+// -----
+
+// CHECK: invalid regex: parentheses not balanced
+// expected-error-re {{ {{(}} }}
+
+// -----
+
+func.func @foo() -> i32 {
+  // expected-error-re@+1 {{'func.return' op has 0 operands, but enclosing function (@{{.*}}) returns 1}}
+  return
+}