From d6296c3b0088915c6175587dda207bca81598293 Mon Sep 17 00:00:00 2001
From: Vladislav Vinogradov <vlad.vinogradov@intel.com>
Date: Mon, 11 Oct 2021 17:02:03 +0300
Subject: [PATCH] [mlir] Allow to use vformat utility with MLIR classes

Make `raw_ostream operator<<` follow const correctness semantic,
since it is a requirement of FormatVariadic implementation.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D111547
---
 mlir/include/mlir/Dialect/MemRef/IR/MemRef.h |  2 +-
 mlir/include/mlir/IR/OpDefinition.h          |  2 +-
 mlir/include/mlir/IR/Operation.h             |  4 ++--
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp     |  2 +-
 mlir/unittests/IR/OperationSupportTest.cpp   | 11 +++++++++++
 5 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h
index a022ffc..afb5a89 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h
@@ -24,7 +24,7 @@ namespace mlir {
 class Location;
 class OpBuilder;
 
-raw_ostream &operator<<(raw_ostream &os, Range &range);
+raw_ostream &operator<<(raw_ostream &os, const Range &range);
 
 /// Return the list of Range (i.e. offset, size, stride). Each Range
 /// entry contains either the dynamic value or a ConstantIndexOp constructed
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index e7f644f..b64eb7d 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -222,7 +222,7 @@ inline raw_ostream &operator<<(raw_ostream &os, OpFoldResult ofr) {
 }
 
 /// Allow printing to a stream.
-inline raw_ostream &operator<<(raw_ostream &os, OpState &op) {
+inline raw_ostream &operator<<(raw_ostream &os, OpState op) {
   op.print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 05f8e40..ac55edd 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -717,8 +717,8 @@ private:
   size_t numTrailingObjects(OverloadToken<Region>) const { return numRegions; }
 };
 
-inline raw_ostream &operator<<(raw_ostream &os, Operation &op) {
-  op.print(os, OpPrintingFlags().useLocalScope());
+inline raw_ostream &operator<<(raw_ostream &os, const Operation &op) {
+  const_cast<Operation &>(op).print(os, OpPrintingFlags().useLocalScope());
   return os;
 }
 
diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 32bd0a2..67041ab 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -2066,7 +2066,7 @@ static LogicalResult verify(SubViewOp op) {
   return produceSubViewErrorMsg(result, op, expectedType, errMsg);
 }
 
-raw_ostream &mlir::operator<<(raw_ostream &os, Range &range) {
+raw_ostream &mlir::operator<<(raw_ostream &os, const Range &range) {
   return os << "range " << range.offset << ":" << range.size << ":"
             << range.stride;
 }
diff --git a/mlir/unittests/IR/OperationSupportTest.cpp b/mlir/unittests/IR/OperationSupportTest.cpp
index b4fbc4be..3d03329 100644
--- a/mlir/unittests/IR/OperationSupportTest.cpp
+++ b/mlir/unittests/IR/OperationSupportTest.cpp
@@ -10,6 +10,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
@@ -214,4 +215,14 @@ TEST(OperationOrderTest, OrderIsAlwaysValid) {
   containerOp->destroy();
 }
 
+TEST(OperationFormatPrintTest, CanUseVariadicFormat) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  Operation *op = createOp(&context);
+
+  std::string str = formatv("{0}", *op).str();
+  ASSERT_STREQ(str.c_str(), "\"foo.bar\"() : () -> ()");
+}
+
 } // end namespace
-- 
2.7.4