From a77cd55dea058395f787effff8411f1173bf69d8 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Mon, 18 Oct 2021 15:07:23 +0000
Subject: [PATCH] [mlir] Add support for specifying printing flags when adding
 an op to a Diagnostic

This removes edge cases where the default flags we want to use
during printing (e.g. local scope, eliding attributes, etc.)
get missed/dropped.

Differential Revision: https://reviews.llvm.org/D111761
---
 mlir/include/mlir/IR/Diagnostics.h |  3 +++
 mlir/lib/IR/Diagnostics.cpp        |  6 +++++-
 mlir/lib/IR/Operation.cpp          | 14 +++-----------
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/mlir/include/mlir/IR/Diagnostics.h b/mlir/include/mlir/IR/Diagnostics.h
index 7bf35c7b5de5..5f6d870fc805 100644
--- a/mlir/include/mlir/IR/Diagnostics.h
+++ b/mlir/include/mlir/IR/Diagnostics.h
@@ -29,6 +29,7 @@ struct LogicalResult;
 class MLIRContext;
 class Operation;
 class OperationName;
+class OpPrintingFlags;
 class Type;
 class Value;
 
@@ -218,6 +219,8 @@ public:
   Diagnostic &operator<<(Operation *val) {
     return *this << *val;
   }
+  /// Append an operation with the given printing flags.
+  Diagnostic &appendOp(Operation &val, const OpPrintingFlags &flags);
 
   /// Stream in a Value.
   Diagnostic &operator<<(Value val);
diff --git a/mlir/lib/IR/Diagnostics.cpp b/mlir/lib/IR/Diagnostics.cpp
index 339f5607abdb..31391b2b9405 100644
--- a/mlir/lib/IR/Diagnostics.cpp
+++ b/mlir/lib/IR/Diagnostics.cpp
@@ -127,9 +127,13 @@ Diagnostic &Diagnostic::operator<<(OperationName val) {
 
 /// Stream in an Operation.
 Diagnostic &Diagnostic::operator<<(Operation &val) {
+  return appendOp(val, OpPrintingFlags());
+}
+Diagnostic &Diagnostic::appendOp(Operation &val, const OpPrintingFlags &flags) {
   std::string str;
   llvm::raw_string_ostream os(str);
-  val.print(os, OpPrintingFlags().useLocalScope().elideLargeElementsAttrs());
+  val.print(os,
+            OpPrintingFlags(flags).useLocalScope().elideLargeElementsAttrs());
   return *this << os.str();
 }
 
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 8eb4c612dd61..e877449fa83f 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -276,17 +276,9 @@ void Operation::insertOperands(unsigned index, ValueRange operands) {
 InFlightDiagnostic Operation::emitError(const Twine &message) {
   InFlightDiagnostic diag = mlir::emitError(getLoc(), message);
   if (getContext()->shouldPrintOpOnDiagnostic()) {
-    // Print out the operation explicitly here so that we can print the generic
-    // form.
-    // TODO: It would be nice if we could instead provide the
-    // specific printing flags when adding the operation as an argument to the
-    // diagnostic.
-    std::string printedOp;
-    {
-      llvm::raw_string_ostream os(printedOp);
-      print(os, OpPrintingFlags().printGenericOpForm().useLocalScope());
-    }
-    diag.attachNote(getLoc()) << "see current operation: " << printedOp;
+    diag.attachNote(getLoc())
+        .append("see current operation: ")
+        .appendOp(*this, OpPrintingFlags().printGenericOpForm());
   }
   return diag;
 }