[mlir] Better Python diagnostics (#128581)

Updated the Python diagnostics handler to emit notes (in addition to
errors) into the output stream so that users have more context as to
where in the IR the error is occurring.

Author: nikalra

mlir/include/mlir/Bindings/Python/Diagnostics.h

@@ -9,12 +9,13 @@
 #ifndef MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
 #define MLIR_BINDINGS_PYTHON_DIAGNOSTICS_H
 
-#include <cassert>
-#include <string>
-
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
-#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
 
 namespace mlir {
 namespace python {
@@ -24,33 +25,45 @@ namespace python {
 class CollectDiagnosticsToStringScope {
 public:
   explicit CollectDiagnosticsToStringScope(MlirContext ctx) : context(ctx) {
-    handlerID = mlirContextAttachDiagnosticHandler(ctx, &handler, &errorMessage,
-                                                   /*deleteUserData=*/nullptr);
+    handlerID =
+        mlirContextAttachDiagnosticHandler(ctx, &handler, &messageStream,
+                                           /*deleteUserData=*/nullptr);
   }
   ~CollectDiagnosticsToStringScope() {
-    assert(errorMessage.empty() && "unchecked error message");
+    assert(message.empty() && "unchecked error message");
     mlirContextDetachDiagnosticHandler(context, handlerID);
   }
 
-  [[nodiscard]] std::string takeMessage() { return std::move(errorMessage); }
+  [[nodiscard]] std::string takeMessage() {
+    std::string newMessage;
+    std::swap(message, newMessage);
+    return newMessage;
+  }
 
 private:
   static MlirLogicalResult handler(MlirDiagnostic diag, void *data) {
     auto printer = +[](MlirStringRef message, void *data) {
-      *static_cast<std::string *>(data) +=
-          llvm::StringRef(message.data, message.length);
+      *static_cast<llvm::raw_string_ostream *>(data)
+          << std::string_view(message.data, message.length);
     };
     MlirLocation loc = mlirDiagnosticGetLocation(diag);
-    *static_cast<std::string *>(data) += "at ";
+    *static_cast<llvm::raw_string_ostream *>(data) << "at ";
     mlirLocationPrint(loc, printer, data);
-    *static_cast<std::string *>(data) += ": ";
+    *static_cast<llvm::raw_string_ostream *>(data) << ": ";
     mlirDiagnosticPrint(diag, printer, data);
+    for (intptr_t i = 0; i < mlirDiagnosticGetNumNotes(diag); i++) {
+      *static_cast<llvm::raw_string_ostream *>(data) << "\n";
+      MlirDiagnostic note = mlirDiagnosticGetNote(diag, i);
+      handler(note, data);
+    }
     return mlirLogicalResultSuccess();
   }
 
   MlirContext context;
   MlirDiagnosticHandlerID handlerID;
-  std::string errorMessage = "";
+
+  std::string message;
+  llvm::raw_string_ostream messageStream{message};
 };
 
 } // namespace python