Introduce a new common diagnostic handler ScopedDiagnosticHandler to simplify saving and restoring the currently registered handler.

--

PiperOrigin-RevId: 249735912

mlir/g3doc/Diagnostics.md

@@ -178,6 +178,29 @@ diagnostic handler with the [`DiagnosticEngine`](#diagnostic-engine).
 Recognizing the many users will want the same handler functionality, MLIR
 provides several common diagnostic handlers for immediate use.
 
+### Scoped Diagnostic Handler
+
+This diagnostic handler is a simple RAII class that saves and restores the
+current diagnostic handler registered to a given context. This class can be
+either be used directly, or in conjunction with a derived diagnostic handler.
+
+```c++
+// Construct the handler directly.
+MLIRContext context;
+ScopedDiagnosticHandler scopedHandler(&context, [](Diagnostic diag) {
+  ...
+});
+
+// Use this handler in conjunction with another.
+class MyDerivedHandler : public ScopedDiagnosticHandler {
+  MyDerivedHandler(MLIRContext *ctx) : ScopedDiagnosticHandler(ctx) {
+    ctx->getDiagEngine().setHandler([&](Diagnostic diag) {
+      ...
+    });
+  }
+};
+```
+
 ### SourceMgr Diagnostic Handler
 
 This diagnostic handler is a wrapper around an llvm::SourceMgr instance. It
@@ -270,9 +293,6 @@ with the 'orderID'. The thread that is processing 'a' should set the orderID to
 This provides a way for the handler to deterministically order the diagnostics
 that it receives given the thread that it is receiving on.
 
-Note: This handler automatically saves and restores the current handler
-registered with the context.
-
 A simple example is shown below:
 
 ```c++

mlir/include/mlir/IR/Diagnostics.h

@@ -443,6 +443,36 @@ private:
   std::unique_ptr<detail::DiagnosticEngineImpl> impl;
 };
 
+//===----------------------------------------------------------------------===//
+// ScopedDiagnosticHandler
+//===----------------------------------------------------------------------===//
+
+/// This diagnostic handler is a simple RAII class that saves and restores the
+/// current diagnostic handler registered to a given context. This class can
+/// be either be used directly, or in conjunction with a derived diagnostic
+/// handler.
+class ScopedDiagnosticHandler {
+public:
+  ScopedDiagnosticHandler(MLIRContext *ctx);
+  ScopedDiagnosticHandler(MLIRContext *ctx,
+                          const DiagnosticEngine::HandlerTy &handler);
+  ~ScopedDiagnosticHandler();
+
+  /// Propagate a diagnostic to the existing diagnostic handler.
+  void propagateDiagnostic(Diagnostic diag) {
+    if (existingHandler)
+      existingHandler(std::move(diag));
+  }
+
+private:
+  /// The existing diagnostic handler registered with the context at the time of
+  /// construction.
+  DiagnosticEngine::HandlerTy existingHandler;
+
+  /// The context to register the handler back to.
+  MLIRContext *ctx;
+};
+
 //===----------------------------------------------------------------------===//
 // SourceMgrDiagnosticHandler
 //===----------------------------------------------------------------------===//
@@ -452,7 +482,7 @@ struct SourceMgrDiagnosticHandlerImpl;
 } // end namespace detail
 
 /// This class is a utility diagnostic handler for use with llvm::SourceMgr.
-class SourceMgrDiagnosticHandler {
+class SourceMgrDiagnosticHandler : public ScopedDiagnosticHandler {
 public:
   SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr, MLIRContext *ctx);
   ~SourceMgrDiagnosticHandler();

mlir/lib/IR/Diagnostics.cpp

@@ -261,6 +261,21 @@ void DiagnosticEngine::emit(Diagnostic diag) {
   impl->emit(std::move(diag));
 }
 
+//===----------------------------------------------------------------------===//
+// ScopedDiagnosticHandler
+//===----------------------------------------------------------------------===//
+
+ScopedDiagnosticHandler::ScopedDiagnosticHandler(MLIRContext *ctx)
+    : existingHandler(ctx->getDiagEngine().getHandler()), ctx(ctx) {}
+ScopedDiagnosticHandler::ScopedDiagnosticHandler(
+    MLIRContext *ctx, const DiagnosticEngine::HandlerTy &handler)
+    : ScopedDiagnosticHandler(ctx) {
+  ctx->getDiagEngine().setHandler(handler);
+}
+ScopedDiagnosticHandler::~ScopedDiagnosticHandler() {
+  ctx->getDiagEngine().setHandler(existingHandler);
+}
+
 //===----------------------------------------------------------------------===//
 // SourceMgrDiagnosticHandler
 //===----------------------------------------------------------------------===//
@@ -328,7 +343,8 @@ static llvm::SourceMgr::DiagKind getDiagKind(DiagnosticSeverity kind) {
 
 SourceMgrDiagnosticHandler::SourceMgrDiagnosticHandler(llvm::SourceMgr &mgr,
                                                        MLIRContext *ctx)
-    : mgr(mgr), impl(new SourceMgrDiagnosticHandlerImpl()) {
+    : ScopedDiagnosticHandler(ctx), mgr(mgr),
+      impl(new SourceMgrDiagnosticHandlerImpl()) {
   // Register a simple diagnostic handler.
   ctx->getDiagEngine().setHandler(
       [this](Diagnostic diag) { emitDiagnostic(diag); });