mlir-translate: support -verify-diagnostics

MLIR translation tools can emit diagnostics and we want to be able to check if
it is indeed the case in tests. Reuse the source manager error handlers
provided for mlir-opt to support the verification in mlir-translate. This
requires us to change the signature of the functions that are registered to
translate sources to MLIR: it now takes a source manager instead of a memory
buffer.

PiperOrigin-RevId: 279132972

1 files changed, 27 insertions, 6 deletions

diff --git a/mlir/lib/Translation/Translation.cpp b/mlir/lib/Translation/Translation.cpp
index 0f78d145710..8b5f98714e4 100644
--- a/mlir/lib/Translation/Translation.cpp
+++ b/mlir/lib/Translation/Translation.cpp
@@ -23,15 +23,16 @@
 #include "mlir/IR/Module.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
-#include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 // Get the mutable static map between registered "to MLIR" translations and the
 // TranslateToMLIRFunctions that perform those translations.
-static llvm::StringMap<TranslateToMLIRFunction> &
+static llvm::StringMap<TranslateSourceMgrToMLIRFunction> &
 getMutableTranslationToMLIRRegistry() {
-  static llvm::StringMap<TranslateToMLIRFunction> translationToMLIRRegistry;
+  static llvm::StringMap<TranslateSourceMgrToMLIRFunction>
+      translationToMLIRRegistry;
   return translationToMLIRRegistry;
 }
 // Get the mutable static map between registered "from MLIR" translations and
@@ -49,8 +50,10 @@ static llvm::StringMap<TranslateFunction> &getMutableTranslationRegistry() {
   return translationRegistry;
 }
 
-TranslateToMLIRRegistration::TranslateToMLIRRegistration(
-    StringRef name, const TranslateToMLIRFunction &function) {
+// Puts `function` into the to-MLIR translation registry unless there is already
+// a function registered for the same name.
+static void registerTranslateToMLIRFunction(
+    StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
   auto &translationToMLIRRegistry = getMutableTranslationToMLIRRegistry();
   if (translationToMLIRRegistry.find(name) != translationToMLIRRegistry.end())
     llvm::report_fatal_error(
@@ -59,6 +62,24 @@ TranslateToMLIRRegistration::TranslateToMLIRRegistration(
   translationToMLIRRegistry[name] = function;
 }
 
+TranslateToMLIRRegistration::TranslateToMLIRRegistration(
+    StringRef name, const TranslateSourceMgrToMLIRFunction &function) {
+  registerTranslateToMLIRFunction(name, function);
+}
+
+// Wraps `function` with a lambda that extracts a StringRef from a source
+// manager and registers the wrapper lambda as a to-MLIR conversion.
+TranslateToMLIRRegistration::TranslateToMLIRRegistration(
+    StringRef name, const TranslateStringRefToMLIRFunction &function) {
+  auto translationFunction = [function](llvm::SourceMgr &sourceMgr,
+                                        MLIRContext *ctx) {
+    const llvm::MemoryBuffer *buffer =
+        sourceMgr.getMemoryBuffer(sourceMgr.getMainFileID());
+    return function(buffer->getBuffer(), ctx);
+  };
+  registerTranslateToMLIRFunction(name, translationFunction);
+}
+
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, const TranslateFromMLIRFunction &function) {
   auto &translationFromMLIRRegistry = getMutableTranslationFromMLIRRegistry();
@@ -84,7 +105,7 @@ TranslateRegistration::TranslateRegistration(
 
 // Merely add the const qualifier to the mutable registry so that external users
 // cannot modify it.
-const llvm::StringMap<TranslateToMLIRFunction> &
+const llvm::StringMap<TranslateSourceMgrToMLIRFunction> &
 mlir::getTranslationToMLIRRegistry() {
   return getMutableTranslationToMLIRRegistry();
 }