From 295d032762ad284068c72cc1904680a4db5e80d3 Mon Sep 17 00:00:00 2001
From: Benjamin Kramer <benny.kra@googlemail.com>
Date: Sat, 21 May 2022 13:25:24 +0200
Subject: [PATCH] [mlir] Move diagnostic handlers instead of copying

This also allows using unique_ptr instead of shared_ptr for the CAPI
user data. NFCI.
---
 mlir/include/mlir/IR/Diagnostics.h | 4 ++--
 mlir/lib/CAPI/IR/Diagnostics.cpp   | 7 ++++---
 mlir/lib/IR/Diagnostics.cpp        | 4 ++--
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index dd75aeea4842..82bd79390e61 100644
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -417,7 +417,7 @@ public:
   /// The handler type for MLIR diagnostics. This function takes a diagnostic as
   /// input, and returns success if the handler has fully processed this
   /// diagnostic. Returns failure otherwise.
-  using HandlerTy = std::function<LogicalResult(Diagnostic &)>;
+  using HandlerTy = llvm::unique_function<LogicalResult(Diagnostic &)>;
 
   /// A handle to a specific registered handler object.
   using HandlerID = uint64_t;
@@ -427,7 +427,7 @@ public:
   /// handlers will process diagnostics first. This function returns a unique
   /// identifier for the registered handler, which can be used to unregister
   /// this handler at a later time.
-  HandlerID registerHandler(const HandlerTy &handler);
+  HandlerID registerHandler(HandlerTy handler);
 
   /// Set the diagnostic handler with a function that returns void. This is a
   /// convenient wrapper for handlers that always completely process the given
diff --git a/mlir/lib/CAPI/IR/Diagnostics.cpp b/mlir/lib/CAPI/IR/Diagnostics.cpp
index 40639c7ba31b..4a13ae57611d 100644
--- a/mlir/lib/CAPI/IR/Diagnostics.cpp
+++ b/mlir/lib/CAPI/IR/Diagnostics.cpp
@@ -59,11 +59,12 @@ MlirDiagnosticHandlerID mlirContextAttachDiagnosticHandler(
   assert(handler && "unexpected null diagnostic handler");
   if (deleteUserData == nullptr)
     deleteUserData = deleteUserDataNoop;
-  std::shared_ptr<void> sharedUserData(userData, deleteUserData);
   DiagnosticEngine::HandlerID id =
       unwrap(context)->getDiagEngine().registerHandler(
-          [handler, sharedUserData](Diagnostic &diagnostic) {
-            return unwrap(handler(wrap(diagnostic), sharedUserData.get()));
+          [handler,
+           ownedUserData = std::unique_ptr<void, decltype(deleteUserData)>(
+               userData, deleteUserData)](Diagnostic &diagnostic) {
+            return unwrap(handler(wrap(diagnostic), ownedUserData.get()));
           });
   return static_cast<MlirDiagnosticHandlerID>(id);
 }
diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 975f6943fc67..98e4d40461c3 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -272,10 +272,10 @@ DiagnosticEngine::~DiagnosticEngine() = default;
 /// Register a new handler for diagnostics to the engine. This function returns
 /// a unique identifier for the registered handler, which can be used to
 /// unregister this handler at a later time.
-auto DiagnosticEngine::registerHandler(const HandlerTy &handler) -> HandlerID {
+auto DiagnosticEngine::registerHandler(HandlerTy handler) -> HandlerID {
   llvm::sys::SmartScopedLock<true> lock(impl->mutex);
   auto uniqueID = impl->uniqueHandlerId++;
-  impl->handlers.insert({uniqueID, handler});
+  impl->handlers.insert({uniqueID, std::move(handler)});
   return uniqueID;
 }