Add a utility diagnostic handler class, SourceMgrDiagnosticHandler, to interface with llvm::SourceMgr. This lowers the barrier of entry for tools to get rich diagnostic handling when using llvm::SourceMgr.

--

PiperOrigin-RevId: 247358610

mlir/include/mlir/IR/Diagnostics.h

@@ -29,10 +29,16 @@
 #include "llvm/ADT/Twine.h"
 #include <functional>
 
+namespace llvm {
+class SMLoc;
+class SourceMgr;
+} // end namespace llvm
+
 namespace mlir {
 class DiagnosticEngine;
 class Identifier;
 struct LogicalResult;
+class MLIRContext;
 class Type;
 
 namespace detail {
@@ -414,6 +420,26 @@ private:
   /// The internal implementation of the DiagnosticEngine.
   std::unique_ptr<detail::DiagnosticEngineImpl> impl;
 };
+
+//===----------------------------------------------------------------------===//
+// SourceMgrDiagnosticHandler
+//===----------------------------------------------------------------------===//
+
+/// This class is a utility diagnostic handler for use with llvm::SourceMgr.
+class SourceMgrDiagnosticHandler {
+public:
+  SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr, MLIRContext *ctx);
+
+  /// Emit the given diagnostic information with the held source manager.
+  void emitDiagnostic(Location loc, Twine message, DiagnosticSeverity kind);
+
+private:
+  /// Convert a location into the given memory buffer into an SMLoc.
+  llvm::SMLoc convertLocToSMLoc(Location loc);
+
+  /// The source manager that we are wrapping.
+  llvm::SourceMgr &mgr;
+};
 } // namespace mlir
 
 #endif

mlir/lib/IR/Diagnostics.cpp

@@ -18,9 +18,11 @@
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Types.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Mutex.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
@@ -223,3 +225,86 @@ void DiagnosticEngine::emit(const Diagnostic &diag) {
   for (auto &note : diag.getNotes())
     impl->emit(note.getLocation(), note.str(), note.getSeverity());
 }
+
+//===----------------------------------------------------------------------===//
+// SourceMgrDiagnosticHandler
+//===----------------------------------------------------------------------===//
+
+/// Given a diagnostic kind, returns the LLVM DiagKind.
+static llvm::SourceMgr::DiagKind getDiagKind(DiagnosticSeverity kind) {
+  switch (kind) {
+  case DiagnosticSeverity::Note:
+    return llvm::SourceMgr::DK_Note;
+  case DiagnosticSeverity::Warning:
+    return llvm::SourceMgr::DK_Warning;
+  case DiagnosticSeverity::Error:
+    return llvm::SourceMgr::DK_Error;
+  case DiagnosticSeverity::Remark:
+    return llvm::SourceMgr::DK_Remark;
+  }
+}
+
+SourceMgrDiagnosticHandler::SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr,
+                                                       MLIRContext *ctx)
+    : mgr(mgr) {
+  // Register a simple diagnostic handler.
+  ctx->getDiagEngine().setHandler(
+      [this](Location loc, StringRef message, DiagnosticSeverity kind) {
+        emitDiagnostic(loc, message, kind);
+      });
+}
+
+void SourceMgrDiagnosticHandler::emitDiagnostic(Location loc, Twine message,
+                                                DiagnosticSeverity kind) {
+  mgr.PrintMessage(convertLocToSMLoc(loc), getDiagKind(kind), message);
+}
+
+/// Convert a location into the given memory buffer into an SMLoc.
+llvm::SMLoc SourceMgrDiagnosticHandler::convertLocToSMLoc(Location loc) {
+  auto fileLoc = loc.dyn_cast<FileLineColLoc>();
+
+  // We currently only support FileLineColLoc.
+  if (!fileLoc)
+    return llvm::SMLoc();
+
+  auto *membuf = mgr.getMemoryBuffer(mgr.getMainFileID());
+
+  // TODO: This should really be upstreamed to be a method on llvm::SourceMgr.
+  // Doing so would allow it to use the offset cache that is already maintained
+  // by SrcBuffer, making this more efficient.
+  unsigned lineNo = fileLoc->getLine();
+  unsigned columnNo = fileLoc->getColumn();
+
+  // Scan for the correct line number.
+  const char *position = membuf->getBufferStart();
+  const char *end = membuf->getBufferEnd();
+
+  // We start counting line and column numbers from 1.
+  --lineNo;
+  --columnNo;
+
+  while (position < end && lineNo) {
+    auto curChar = *position++;
+
+    // Scan for newlines.  If this isn't one, ignore it.
+    if (curChar != '\r' && curChar != '\n')
+      continue;
+
+    // We saw a line break, decrement our counter.
+    --lineNo;
+
+    // Check for \r\n and \n\r and treat it as a single escape.  We know that
+    // looking past one character is safe because MemoryBuffer's are always nul
+    // terminated.
+    if (*position != curChar && (*position == '\r' || *position == '\n'))
+      ++position;
+  }
+
+  // If the line/column counter was invalid, return a pointer to the start of
+  // the buffer.
+  if (lineNo || position + columnNo > end)
+    return llvm::SMLoc::getFromPointer(membuf->getBufferStart());
+
+  // Otherwise return the right pointer.
+  return llvm::SMLoc::getFromPointer(position + columnNo);
+}

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -73,48 +73,6 @@ static std::vector<const mlir::PassRegistryEntry *> *passList;
 
 enum OptResult { OptSuccess, OptFailure };
 
-/// Given a MemoryBuffer along with a line and column within it, return the
-/// location being referenced.
-static SMLoc getLocFromLineAndCol(MemoryBuffer &membuf, unsigned lineNo,
-                                  unsigned columnNo) {
-  // TODO: This should really be upstreamed to be a method on llvm::SourceMgr.
-  // Doing so would allow it to use the offset cache that is already maintained
-  // by SrcBuffer, making this more efficient.
-
-  // Scan for the correct line number.
-  const char *position = membuf.getBufferStart();
-  const char *end = membuf.getBufferEnd();
-
-  // We start counting line and column numbers from 1.
-  --lineNo;
-  --columnNo;
-
-  while (position < end && lineNo) {
-    auto curChar = *position++;
-
-    // Scan for newlines.  If this isn't one, ignore it.
-    if (curChar != '\r' && curChar != '\n')
-      continue;
-
-    // We saw a line break, decrement our counter.
-    --lineNo;
-
-    // Check for \r\n and \n\r and treat it as a single escape.  We know that
-    // looking past one character is safe because MemoryBuffer's are always nul
-    // terminated.
-    if (*position != curChar && (*position == '\r' || *position == '\n'))
-      ++position;
-  }
-
-  // If the line/column counter was invalid, return a pointer to the start of
-  // the buffer.
-  if (lineNo || position + columnNo > end)
-    return SMLoc::getFromPointer(membuf.getBufferStart());
-
-  // Otherwise return the right pointer.
-  return SMLoc::getFromPointer(position + columnNo);
-}
-
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
@@ -165,20 +123,6 @@ static StringRef getDiagnosticKindString(DiagnosticSeverity kind) {
   }
 }
 
-/// Given a diagnostic kind, returns the LLVM DiagKind.
-static llvm::SourceMgr::DiagKind getDiagKind(DiagnosticSeverity kind) {
-  switch (kind) {
-  case DiagnosticSeverity::Note:
-    return llvm::SourceMgr::DK_Note;
-  case DiagnosticSeverity::Warning:
-    return llvm::SourceMgr::DK_Warning;
-  case DiagnosticSeverity::Error:
-    return llvm::SourceMgr::DK_Error;
-  case DiagnosticSeverity::Remark:
-    return llvm::SourceMgr::DK_Remark;
-  }
-}
-
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static OptResult processFile(std::unique_ptr<MemoryBuffer> ownedBuffer) {
@@ -189,24 +133,11 @@ static OptResult processFile(std::unique_ptr<MemoryBuffer> ownedBuffer) {
 
   // Parse the input file.
   MLIRContext context;
+  SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
 
   // If we are in verify mode then we have a lot of work to do, otherwise just
   // perform the actions without worrying about it.
   if (!verifyDiagnostics) {
-    // Register a simple diagnostic handler that prints out info with context.
-    context.getDiagEngine().setHandler(
-        [&](Location location, StringRef message, DiagnosticSeverity kind) {
-          unsigned line = 1, column = 1;
-          SMLoc loc;
-          if (auto fileLoc = location.dyn_cast<FileLineColLoc>()) {
-            line = fileLoc->getLine();
-            column = fileLoc->getColumn();
-            loc = getLocFromLineAndCol(buffer, line, column);
-          }
-
-          sourceMgr.PrintMessage(loc, getDiagKind(kind), message);
-        });
-
     // Run the test actions.
     return performActions(sourceMgr, &context);
   }
@@ -264,9 +195,8 @@ static OptResult processFile(std::unique_ptr<MemoryBuffer> ownedBuffer) {
 
     // If this error wasn't expected, produce an error out of mlir-opt saying
     // so.
-    auto unexpectedLoc = getLocFromLineAndCol(buffer, line, column);
+    sourceMgrHandler.emitDiagnostic(location, "unexpected error: " + message,
-    sourceMgr.PrintMessage(unexpectedLoc, SourceMgr::DK_Error,
+                                    DiagnosticSeverity::Error);
-                           "unexpected error: " + Twine(message));
     result = OptFailure;
   };