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

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.

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

Author: River707

mlir/docs/Diagnostics.md

@@ -300,9 +300,13 @@ 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.
@@ -327,6 +331,12 @@ 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

mlir/lib/IR/Diagnostics.cpp

@@ -590,14 +590,77 @@ SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(FileLineColLoc loc) {
 
 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 {
@@ -608,7 +671,8 @@ struct SourceMgrDiagnosticVerifierHandlerImpl {
 
   /// 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;
@@ -617,8 +681,9 @@ struct SourceMgrDiagnosticVerifierHandlerImpl {
   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
@@ -638,7 +703,6 @@ static StringRef getDiagKindStr(DiagnosticSeverity kind) {
   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);
@@ -647,10 +711,9 @@ SourceMgrDiagnosticVerifierHandlerImpl::getExpectedDiags(StringRef bufName) {
   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;
@@ -667,7 +730,7 @@ SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
   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)
@@ -679,7 +742,7 @@ SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
     }
 
     // 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")
@@ -692,9 +755,15 @@ SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
       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);
 
@@ -722,7 +791,7 @@ SourceMgrDiagnosticVerifierHandlerImpl::computeExpectedDiags(
         record.lineNo = e;
       }
     }
-    expectedDiags.push_back(record);
+    expectedDiags.emplace_back(std::move(record));
   }
   return expectedDiags;
 }
@@ -734,7 +803,7 @@ SourceMgrDiagnosticVerifierHandler::SourceMgrDiagnosticVerifierHandler(
   // 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) {
@@ -765,14 +834,10 @@ LogicalResult SourceMgrDiagnosticVerifierHandler::verify() {
     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,
-                       "expected " + getDiagKindStr(err.kind) + " \"" +
-                           err.substring + "\" was not produced",
-                       range);
-      impl->status = failure();
+      impl->status =
+          err.emitError(os, mgr,
+                        "expected " + getDiagKindStr(err.kind) + " \"" +
+                            err.substring + "\" was not produced");
     }
   }
   impl->expectedDiagsPerFile.clear();
@@ -799,8 +864,10 @@ void SourceMgrDiagnosticVerifierHandler::process(FileLineColLoc loc,
                                                  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.
@@ -809,7 +876,7 @@ void SourceMgrDiagnosticVerifierHandler::process(FileLineColLoc loc,
   // 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;

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

@@ -405,7 +405,7 @@ func.func @simple_store_missing_operand(%arg0 : f32) -> () {
 
 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
 }

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

@@ -0,0 +1,16 @@
+// 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
+}