From 74df13fdda9b0af514e7ff88d0c7e7cf0fb50da8 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <aminim@google.com>
Date: Mon, 17 Jun 2019 15:31:53 -0700
Subject: [PATCH] Refactor generic op printing: extract a public
 printFunctionalType() on OpAsmPrinter (NFC)

PiperOrigin-RevId: 253674584
---
 mlir/include/mlir/IR/OpImplementation.h | 17 +++++++++++++++++
 mlir/include/mlir/IR/Operation.h        |  1 +
 mlir/lib/IR/AsmPrinter.cpp              | 16 ++--------------
 mlir/lib/IR/Operation.cpp               |  3 ++-
 4 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index bd6ce8479cd1..c16a9c40e661 100644
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/mlir/include/mlir/IR/OpImplementation.h
@@ -98,6 +98,23 @@ public:
       os << ')';
   }
 
+  /// Print the complete type of an operation in functional form.
+  void printFunctionalType(Operation *op) {
+    auto &os = getStream();
+    os << "(";
+    interleaveComma(op->getNonSuccessorOperands(), os,
+                    [&](Value *operand) { printType(operand->getType()); });
+    os << ") -> ";
+    if (op->getNumResults() == 1 &&
+        !op->getResult(0)->getType().isa<FunctionType>()) {
+      printType(op->getResult(0)->getType());
+    } else {
+      os << '(';
+      interleaveComma(op->getResultTypes(), os);
+      os << ')';
+    }
+  }
+
 private:
   OpAsmPrinter(const OpAsmPrinter &) = delete;
   void operator=(const OpAsmPrinter &) = delete;
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index ec029b1c3e20..38eeac4960d6 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -320,6 +320,7 @@ public:
     return getOperand(getSuccessorOperandIndex(succIndex) + opIndex);
   }
 
+  bool hasSuccessors() { return numSuccs != 0; }
   unsigned getNumSuccessors() { return numSuccs; }
   unsigned getNumSuccessorOperands(unsigned index) {
     assert(!isKnownNonTerminator() && "only terminators may have successors");
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 37f8375b5092..545fd63a1752 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -1575,20 +1575,8 @@ void FunctionPrinter::printGenericOp(Operation *op) {
   printOptionalAttrDict(attrs);
 
   // Print the type signature of the operation.
-  os << " : (";
-  interleaveComma(properOperands,
-                  [&](Value *value) { printType(value->getType()); });
-  os << ") -> ";
-
-  if (op->getNumResults() == 1 &&
-      !op->getResult(0)->getType().isa<FunctionType>()) {
-    printType(op->getResult(0)->getType());
-  } else {
-    os << '(';
-    interleaveComma(op->getResults(),
-                    [&](Value *result) { printType(result->getType()); });
-    os << ')';
-  }
+  os << " : ";
+  printFunctionalType(op);
 }
 
 void FunctionPrinter::printSuccessorAndUseList(Operation *term,
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index 5b602d996795..2d4765eacc09 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -495,7 +495,8 @@ void Operation::setSuccessor(Block *block, unsigned index) {
 
 auto Operation::getNonSuccessorOperands() -> operand_range {
   return {operand_iterator(this, 0),
-          operand_iterator(this, getSuccessorOperandIndex(0))};
+          operand_iterator(this, hasSuccessors() ? getSuccessorOperandIndex(0)
+                                                 : getNumOperands())};
 }
 
 /// Get the index of the first operand of the successor at the provided