7 files changed, 155 insertions, 69 deletions

diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 0311fa16420..b266183bb8d 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -17,6 +17,7 @@
 
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Location.h"
+#include "mlir/IR/Types.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Mutex.h"
 #include "llvm/Support/raw_ostream.h"
@@ -24,9 +25,86 @@
 using namespace mlir;
 using namespace mlir::detail;
 
+//===----------------------------------------------------------------------===//
+// DiagnosticArgument
+//===----------------------------------------------------------------------===//
+
+// Construct from a Type.
+DiagnosticArgument::DiagnosticArgument(Type val)
+    : kind(DiagnosticArgumentKind::Type),
+      opaqueVal(reinterpret_cast<intptr_t>(val.getAsOpaquePointer())) {}
+
+/// Returns this argument as a Type.
+Type DiagnosticArgument::getAsType() const {
+  assert(getKind() == DiagnosticArgumentKind::Type);
+  return Type::getFromOpaquePointer(reinterpret_cast<const void *>(opaqueVal));
+}
+
+/// Outputs this argument to a stream.
+void DiagnosticArgument::print(raw_ostream &os) const {
+  switch (kind) {
+  case DiagnosticArgumentKind::Integer:
+    os << getAsInteger();
+    break;
+  case DiagnosticArgumentKind::String:
+    os << getAsString();
+    break;
+  case DiagnosticArgumentKind::Type:
+    os << getAsType();
+    break;
+  case DiagnosticArgumentKind::Unsigned:
+    os << getAsUnsigned();
+    break;
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Diagnostic
+//===----------------------------------------------------------------------===//
+
+/// Outputs this diagnostic to a stream.
+void Diagnostic::print(raw_ostream &os) const {
+  for (auto &arg : getArguments())
+    arg.print(os);
+}
+
+/// Convert the diagnostic to a string.
+std::string Diagnostic::str() const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  print(os);
+  return os.str();
+}
+
+//===----------------------------------------------------------------------===//
+// InFlightDiagnostic
+//===----------------------------------------------------------------------===//
+
+/// Allow an inflight diagnostic to be converted to 'failure', otherwise
+/// 'success' if this is an empty diagnostic.
+InFlightDiagnostic::operator LogicalResult() const {
+  return failure(isInFlight());
+}
+
+/// Reports the diagnostic to the engine.
+void InFlightDiagnostic::report() {
+  if (isInFlight()) {
+    owner->emit(*impl);
+    impl.reset();
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// DiagnosticEngineImpl
+//===----------------------------------------------------------------------===//
+
 namespace mlir {
 namespace detail {
 struct DiagnosticEngineImpl {
+  /// Emit a diagnostic using the registered issue handle if present, or with
+  /// the default behavior if not.
+  void emit(Location loc, StringRef msg, DiagnosticSeverity severity);
+
   /// A mutex to ensure that diagnostics emission is thread-safe.
   llvm::sys::SmartMutex<true> mutex;
 
@@ -37,42 +115,18 @@ struct DiagnosticEngineImpl {
 } // namespace detail
 } // namespace mlir
 
-//===----------------------------------------------------------------------===//
-// DiagnosticEngine
-//===----------------------------------------------------------------------===//
-
-DiagnosticEngine::DiagnosticEngine() : impl(new DiagnosticEngineImpl()) {}
-DiagnosticEngine::~DiagnosticEngine() {}
-
-/// Register a diagnostic handler with this engine.  The handler is
-/// passed location information if present (nullptr if not) along with a
-/// message and a severity that indicates whether this is an error, warning,
-/// etc.
-void DiagnosticEngine::setHandler(const HandlerTy &handler) {
-  llvm::sys::SmartScopedLock<true> lock(impl->mutex);
-  impl->handler = handler;
-}
-
-/// Return the current diagnostic handler, or null if none is present.
-auto DiagnosticEngine::getHandler() -> HandlerTy {
-  llvm::sys::SmartScopedLock<true> lock(impl->mutex);
-  return impl->handler;
-}
-
 /// Emit a diagnostic using the registered issue handle if present, or with
-/// the default behavior if not.  The MLIR compiler should not generally
-/// interact with this, it should use methods on Operation instead.
-void DiagnosticEngine::emit(Location loc, const Twine &msg,
-                            DiagnosticSeverity severity) {
-  /// Lock access to the diagnostic engine.
-  llvm::sys::SmartScopedLock<true> lock(impl->mutex);
+/// the default behavior if not.
+void DiagnosticEngineImpl::emit(Location loc, StringRef msg,
+                                DiagnosticSeverity severity) {
+  // Lock access to the handler.
+  llvm::sys::SmartScopedLock<true> lock(mutex);
 
   // If we had a handler registered, emit the diagnostic using it.
-  if (impl->handler) {
+  if (handler) {
     // TODO(b/131756158) FusedLoc should be handled by the diagnostic handler
     // instead of here.
-    // Check to see if we are emitting a diagnostic on a fused
-    // location.
+    // Check to see if we are emitting a diagnostic on a fused location.
     if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
       auto fusedLocs = fusedLoc->getLocations();
 
@@ -85,7 +139,7 @@ void DiagnosticEngine::emit(Location loc, const Twine &msg,
       return;
     }
 
-    return impl->handler(loc, msg.str(), severity);
+    return handler(loc, msg, severity);
   }
 
   // Otherwise, if this is an error we emit it to stderr.
@@ -101,3 +155,31 @@ void DiagnosticEngine::emit(Location loc, const Twine &msg,
   os << msg << '\n';
   os.flush();
 }
+
+//===----------------------------------------------------------------------===//
+// DiagnosticEngine
+//===----------------------------------------------------------------------===//
+
+DiagnosticEngine::DiagnosticEngine() : impl(new DiagnosticEngineImpl()) {}
+DiagnosticEngine::~DiagnosticEngine() {}
+
+/// Set the diagnostic handler for this engine.  The handler is passed
+/// location information if present (nullptr if not) along with a message and
+/// a severity that indicates whether this is an error, warning, etc. Note
+/// that this replaces any existing handler.
+void DiagnosticEngine::setHandler(const HandlerTy &handler) {
+  llvm::sys::SmartScopedLock<true> lock(impl->mutex);
+  impl->handler = handler;
+}
+
+/// Return the current diagnostic handler, or null if none is present.
+auto DiagnosticEngine::getHandler() -> HandlerTy {
+  llvm::sys::SmartScopedLock<true> lock(impl->mutex);
+  return impl->handler;
+}
+
+/// Emit a diagnostic using the registered issue handler if present, or with
+/// the default behavior if not.
+void DiagnosticEngine::emit(const Diagnostic &diag) {
+  impl->emit(diag.getLocation(), diag.str(), diag.getSeverity());
+}
diff --git a/mlir/lib/IR/Dialect.cpp b/mlir/lib/IR/Dialect.cpp
index e87358c4bd4..89ff1b50ed4 100644
--- a/mlir/lib/IR/Dialect.cpp
+++ b/mlir/lib/IR/Dialect.cpp
@@ -16,6 +16,7 @@
 // =============================================================================
 
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectHooks.h"
 #include "mlir/IR/MLIRContext.h"
 #include "llvm/ADT/Twine.h"
diff --git a/mlir/lib/IR/Function.cpp b/mlir/lib/IR/Function.cpp
index e655e2ed774..66b296f2454 100644
--- a/mlir/lib/IR/Function.cpp
+++ b/mlir/lib/IR/Function.cpp
@@ -118,23 +118,24 @@ void Function::erase() {
 
 /// Emit a remark about this function, reporting up to any diagnostic
 /// handlers that may be listening.
-void Function::emitRemark(const Twine &message) {
-  getContext()->emitRemark(getLoc(), message);
+InFlightDiagnostic Function::emitRemark(const Twine &message) {
+  return getContext()->emitRemark(getLoc(), message);
 }
 
 /// Emit a warning about this function, reporting up to any diagnostic
 /// handlers that may be listening.
-void Function::emitWarning(const Twine &message) {
-  getContext()->getDiagEngine().emit(getLoc(), message,
-                                     DiagnosticSeverity::Warning);
+InFlightDiagnostic Function::emitWarning(const Twine &message) {
+  return getContext()->getDiagEngine().emit(getLoc(),
+                                            DiagnosticSeverity::Warning)
+         << message;
 }
 
 /// Emit an error about fatal conditions with this function, reporting up to
 /// any diagnostic handlers that may be listening.  This function always
 /// returns failure.  NOTE: This may terminate the containing application, only
 /// use when the IR is in an inconsistent state.
-LogicalResult Function::emitError(const Twine &message) {
-  return getContext()->emitError(getLoc(), message), failure();
+InFlightDiagnostic Function::emitError(const Twine &message) {
+  return getContext()->emitError(getLoc(), message);
 }
 
 /// Clone the internal blocks from this function into dest and all attributes
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 5161c27368c..6e3c7c9da5a 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -422,14 +422,17 @@ static ArrayRef<T> copyArrayRefInto(llvm::BumpPtrAllocator &allocator,
 // Diagnostic Handlers
 //===----------------------------------------------------------------------===//
 
-bool MLIRContext::emitError(Location location, const llvm::Twine &message) {
-  getImpl().diagEngine.emit(location, message, DiagnosticSeverity::Error);
-  return true;
+InFlightDiagnostic MLIRContext::emitError(Location location,
+                                          const llvm::Twine &message) {
+  return getImpl().diagEngine.emit(location, DiagnosticSeverity::Error)
+         << message;
 }
 
 /// Emit a remark message using the diagnostic engine.
-void MLIRContext::emitRemark(Location location, const Twine &message) {
-  getImpl().diagEngine.emit(location, message, DiagnosticSeverity::Remark);
+InFlightDiagnostic MLIRContext::emitRemark(Location location,
+                                           const Twine &message) {
+  return getImpl().diagEngine.emit(location, DiagnosticSeverity::Remark)
+         << message;
 }
 
 /// Returns the diagnostic engine for this context.
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 1a0aaa526c1..da02c58d366 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -307,30 +307,29 @@ void Operation::walk(const std::function<void(Operation *)> &callback) {
 
 /// Emit a remark about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void Operation::emitRemark(const Twine &message) {
-  getContext()->emitRemark(getLoc(), message);
+InFlightDiagnostic Operation::emitRemark(const Twine &message) {
+  return getContext()->emitRemark(getLoc(), message);
 }
 
 /// Emit a note about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void Operation::emitNote(const Twine &message) {
-  getContext()->getDiagEngine().emit(getLoc(), message,
-                                     DiagnosticSeverity::Note);
+InFlightDiagnostic Operation::emitNote(const Twine &message) {
+  return getContext()->getDiagEngine().emit(getLoc(), DiagnosticSeverity::Note)
+         << message;
 }
 
 /// Emit a warning about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void Operation::emitWarning(const Twine &message) {
-  getContext()->getDiagEngine().emit(getLoc(), message,
-                                     DiagnosticSeverity::Warning);
+InFlightDiagnostic Operation::emitWarning(const Twine &message) {
+  return getContext()->getDiagEngine().emit(getLoc(),
+                                            DiagnosticSeverity::Warning)
+         << message;
 }
 
 /// Emit an error about fatal conditions with this operation, reporting up to
-/// any diagnostic handlers that may be listening.  This function always
-/// returns failure.  NOTE: This may terminate the containing application, only
-/// use when the IR is in an inconsistent state.
-LogicalResult Operation::emitError(const Twine &message) {
-  return getContext()->emitError(getLoc(), message), failure();
+/// any diagnostic handlers that may be listening.
+InFlightDiagnostic Operation::emitError(const Twine &message) {
+  return getContext()->emitError(getLoc(), message);
 }
 
 /// Given an operation 'other' that is within the same parent block, return
@@ -551,7 +550,7 @@ LogicalResult Operation::fold(SmallVectorImpl<Value *> &results) {
 
 /// Emit an error with the op name prefixed, like "'dim' op " which is
 /// convenient for verifiers.
-LogicalResult Operation::emitOpError(const Twine &message) {
+InFlightDiagnostic Operation::emitOpError(const Twine &message) {
   return emitError(Twine('\'') + getName().getStringRef() + "' op " + message);
 }
 
@@ -652,35 +651,33 @@ bool OpState::parse(OpAsmParser *parser, OperationState *result) {
 void OpState::print(OpAsmPrinter *p) { p->printGenericOp(getOperation()); }
 
 /// Emit an error about fatal conditions with this operation, reporting up to
-/// any diagnostic handlers that may be listening.  NOTE: This may terminate
-/// the containing application, only use when the IR is in an inconsistent
-/// state.
-LogicalResult OpState::emitError(const Twine &message) {
+/// any diagnostic handlers that may be listening.
+InFlightDiagnostic OpState::emitError(const Twine &message) {
   return getOperation()->emitError(message);
 }
 
 /// Emit an error with the op name prefixed, like "'dim' op " which is
 /// convenient for verifiers.
-LogicalResult OpState::emitOpError(const Twine &message) {
+InFlightDiagnostic OpState::emitOpError(const Twine &message) {
   return getOperation()->emitOpError(message);
 }
 
 /// Emit a warning about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void OpState::emitWarning(const Twine &message) {
-  getOperation()->emitWarning(message);
+InFlightDiagnostic OpState::emitWarning(const Twine &message) {
+  return getOperation()->emitWarning(message);
 }
 
 /// Emit a note about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void OpState::emitNote(const Twine &message) {
-  getOperation()->emitNote(message);
+InFlightDiagnostic OpState::emitNote(const Twine &message) {
+  return getOperation()->emitNote(message);
 }
 
 /// Emit a remark about this operation, reporting up to any diagnostic
 /// handlers that may be listening.
-void OpState::emitRemark(const Twine &message) {
-  getOperation()->emitRemark(message);
+InFlightDiagnostic OpState::emitRemark(const Twine &message) {
+  return getOperation()->emitRemark(message);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/StandardTypes.cpp b/mlir/lib/IR/StandardTypes.cpp
index 1427f1f03c0..25a02332516 100644
--- a/mlir/lib/IR/StandardTypes.cpp
+++ b/mlir/lib/IR/StandardTypes.cpp
@@ -18,6 +18,7 @@
 #include "mlir/IR/StandardTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
diff --git a/mlir/lib/IR/Types.cpp b/mlir/lib/IR/Types.cpp
index 701935439f6..dba33a97886 100644
--- a/mlir/lib/IR/Types.cpp
+++ b/mlir/lib/IR/Types.cpp
@@ -17,6 +17,7 @@
 
 #include "mlir/IR/Types.h"
 #include "TypeDetail.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "llvm/ADT/Twine.h"