[mlir][Diagnostic] Don't store Operation arguments as a DiagnosticArgument

[rriddle]

Summary

Diagnostics may be cached in the parallel diagnostic handler to preserve proper ordering. Storing the Operation as a DiagnosticArgument is problematic as the operation may be erased or changed before it finally gets printed.

[stellaraccident]

Thanks!

mlir/include/mlir/IR/Diagnostics.h

 //===- Diagnostics.h - MLIR Diagnostics -------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 class DiagnosticArgument {
 public:
   /// Enum that represents the different kinds of diagnostic arguments
   /// supported.
   enum class DiagnosticArgumentKind {
     Attribute,
     Double,
     Integer,
-    Operation,
     String,
     Type,
     Unsigned,
   };
 
   /// Outputs this argument to a stream.
   void print(raw_ostream &os) const;
 
public:
   }
 
   /// Returns this argument as a signed integer.
   int64_t getAsInteger() const {
     assert(getKind() == DiagnosticArgumentKind::Integer);
     return static_cast<int64_t>(opaqueVal);
   }
 
-  /// Returns this argument as an operation.
-  Operation &getAsOperation() const {
-    assert(getKind() == DiagnosticArgumentKind::Operation);
-    return *reinterpret_cast<Operation *>(opaqueVal);
-  }
-
   /// Returns this argument as a string.
   StringRef getAsString() const {
     assert(getKind() == DiagnosticArgumentKind::String);
     return stringVal;
   }
 
   /// Returns this argument as a Type.
   Type getAsType() const;
private:
   // Construct from an unsigned integer.
   template <typename T>
   explicit DiagnosticArgument(
       T val, typename std::enable_if<std::is_unsigned<T>::value &&
                                      std::numeric_limits<T>::is_integer &&
                                      sizeof(T) <= sizeof(uint64_t)>::type * = 0)
       : kind(DiagnosticArgumentKind::Unsigned), opaqueVal(uint64_t(val)) {}
 
-  // Construct from an operation reference.
-  explicit DiagnosticArgument(Operation &val) : DiagnosticArgument(&val) {}
-  explicit DiagnosticArgument(Operation *val)
-      : kind(DiagnosticArgumentKind::Operation),
-        opaqueVal(reinterpret_cast<intptr_t>(val)) {
-    assert(val && "expected valid operation");
-  }
-
   // Construct from a string reference.
   explicit DiagnosticArgument(StringRef val)
       : kind(DiagnosticArgumentKind::String), stringVal(val) {}
 
   // Construct from a Type.
   explicit DiagnosticArgument(Type val);
 
   /// The kind of this argument.
public:
   Diagnostic &operator<<(Twine &&val);
 
   /// Stream in an Identifier.
   Diagnostic &operator<<(Identifier val);
 
   /// Stream in an OperationName.
   Diagnostic &operator<<(OperationName val);
 
+  /// Stream in an Operation.
+  Diagnostic &operator<<(Operation &val);
+  Diagnostic &operator<<(Operation *val) {
+    return *this << *val;
+  }
+
   /// Stream in a range.
   template <typename T> Diagnostic &operator<<(iterator_range<T> range) {
     return appendRange(range);
   }
   template <typename T> Diagnostic &operator<<(ArrayRef<T> range) {
     return appendRange(range);
   }
 

mlir/lib/IR/Diagnostics.cpp

 //===- Diagnostics.cpp - MLIR Diagnostics ---------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
static llvm::cl::opt<bool> printStackTraceOnDiagnostic(
     "mlir-print-stacktrace-on-diagnostic",
     llvm::cl::desc("When a diagnostic is emitted, also print the stack trace "
                    "as an attached note"));
 
 //===----------------------------------------------------------------------===//
 // DiagnosticArgument
 //===----------------------------------------------------------------------===//
 
-// Construct from an Attribute.
+/// Construct from an Attribute.
 DiagnosticArgument::DiagnosticArgument(Attribute attr)
     : kind(DiagnosticArgumentKind::Attribute),
       opaqueVal(reinterpret_cast<intptr_t>(attr.getAsOpaquePointer())) {}
 
-// Construct from a Type.
+/// Construct from a Type.
 DiagnosticArgument::DiagnosticArgument(Type val)
     : kind(DiagnosticArgumentKind::Type),
       opaqueVal(reinterpret_cast<intptr_t>(val.getAsOpaquePointer())) {}
 
 /// Returns this argument as an Attribute.
 Attribute DiagnosticArgument::getAsAttribute() const {
   assert(getKind() == DiagnosticArgumentKind::Attribute);
   return Attribute::getFromOpaquePointer(
case DiagnosticArgumentKind::Attribute:
     os << getAsAttribute();
     break;
   case DiagnosticArgumentKind::Double:
     os << getAsDouble();
     break;
   case DiagnosticArgumentKind::Integer:
     os << getAsInteger();
     break;
-  case DiagnosticArgumentKind::Operation:
-    getAsOperation().print(os, OpPrintingFlags().useLocalScope());
-    break;
   case DiagnosticArgumentKind::String:
     os << getAsString();
     break;
   case DiagnosticArgumentKind::Type:
     os << '\'' << getAsType() << '\'';
     break;
   case DiagnosticArgumentKind::Unsigned:
     os << getAsUnsigned();

 /// Stream in an OperationName.
 Diagnostic &Diagnostic::operator<<(OperationName val) {
   // An OperationName is stored in the context, so we don't need to worry about
   // the lifetime of its data.
   arguments.push_back(DiagnosticArgument(val.getStringRef()));
   return *this;
 }
 
+/// Stream in an Operation.
+Diagnostic &Diagnostic::operator<<(Operation &val) {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  os << val;
+  return *this << os.str();
+}
+
 /// 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 {