Add support for result type iteration in Operation/Instruction/OperationStmt.

PiperOrigin-RevId: 223264992

mlir/include/mlir/IR/Instructions.h

@@ -150,7 +150,7 @@ public:
     return {result_begin(), result_end()};
   }
 
-  // Support const operand iteration.
+  // Support const result iteration.
   using const_result_iterator =
       ResultIterator<const Instruction, const CFGValue>;
   const_result_iterator result_begin() const {
@@ -179,6 +179,21 @@ public:
     return getInstResults()[idx];
   }
 
+  // Support result type iteration.
+  using result_type_iterator =
+      ResultTypeIterator<const Instruction, const CFGValue>;
+  result_type_iterator result_type_begin() const {
+    return result_type_iterator(this, 0);
+  }
+
+  result_type_iterator result_type_end() const {
+    return result_type_iterator(this, getNumResults());
+  }
+
+  llvm::iterator_range<result_type_iterator> getResultTypes() const {
+    return {result_type_begin(), result_type_end()};
+  }
+
   //===--------------------------------------------------------------------===//
   // Terminators
   //===--------------------------------------------------------------------===//

mlir/include/mlir/IR/Operation.h

@@ -29,6 +29,7 @@ template <typename OpType> class ConstOpPointer;
 template <typename OpType> class OpPointer;
 template <typename ObjectType, typename ElementType> class OperandIterator;
 template <typename ObjectType, typename ElementType> class ResultIterator;
+template <typename ObjectType, typename ElementType> class ResultTypeIterator;
 class Function;
 class IROperandOwner;
 class Instruction;
@@ -97,12 +98,19 @@ public:
   result_iterator result_end();
   llvm::iterator_range<result_iterator> getResults();
 
-  // Support const operand iteration.
+  // Support const result iteration.
   using const_result_iterator = ResultIterator<const Operation, const SSAValue>;
   const_result_iterator result_begin() const;
   const_result_iterator result_end() const;
   llvm::iterator_range<const_result_iterator> getResults() const;
 
+  // Support for result type iteration.
+  using result_type_iterator =
+      ResultTypeIterator<const Operation, const SSAValue>;
+  result_type_iterator result_type_begin() const;
+  result_type_iterator result_type_end() const;
+  llvm::iterator_range<result_type_iterator> getResultTypes() const;
+
   // Support for successor querying.
   unsigned getNumSuccessors() const;
   unsigned getNumSuccessorOperands(unsigned index) const;
@@ -399,6 +407,31 @@ public:
   }
 };
 
+/// This template implements the result type iterators for the various IR
+/// classes in terms of getResult(idx)->getType().
+template <typename ObjectType, typename ElementType>
+class ResultTypeIterator final
+    : public IndexedAccessorIterator<
+          ResultTypeIterator<ObjectType, ElementType>, ObjectType,
+          ElementType> {
+public:
+  /// Initializes the result type iterator to the specified index.
+  ResultTypeIterator(ObjectType *object, unsigned index)
+      : IndexedAccessorIterator<ResultTypeIterator<ObjectType, ElementType>,
+                                ObjectType, ElementType>(object, index) {}
+
+  /// Support converting to the const variant. This will be a no-op for const
+  /// variant.
+  operator ResultTypeIterator<const ObjectType, const ElementType>() const {
+    return ResultTypeIterator<const ObjectType, const ElementType>(this->object,
+                                                                   this->index);
+  }
+
+  Type operator*() const {
+    return this->object->getResult(this->index)->getType();
+  }
+};
+
 // Implement the inline operand iterator methods.
 inline auto Operation::operand_begin() -> operand_iterator {
   return operand_iterator(this, 0);
@@ -450,6 +483,19 @@ inline auto Operation::getResults() const
     -> llvm::iterator_range<const_result_iterator> {
   return {result_begin(), result_end()};
 }
+
+inline auto Operation::result_type_begin() const -> result_type_iterator {
+  return result_type_iterator(this, 0);
+}
+
+inline auto Operation::result_type_end() const -> result_type_iterator {
+  return result_type_iterator(this, getNumResults());
+}
+
+inline auto Operation::getResultTypes() const
+    -> llvm::iterator_range<result_type_iterator> {
+  return {result_type_begin(), result_type_end()};
+}
 } // end namespace mlir
 
 /// We need to teach the LLVM cast/dyn_cast etc logic how to cast from an

mlir/include/mlir/IR/Statements.h

@@ -139,7 +139,7 @@ public:
     return {result_begin(), result_end()};
   }
 
-  // Support const operand iteration.
+  // Support const result iteration.
   using const_result_iterator =
       ResultIterator<const OperationStmt, const MLValue>;
   const_result_iterator result_begin() const {
@@ -168,6 +168,21 @@ public:
     return getStmtResults()[idx];
   }
 
+  // Support result type iteration.
+  using result_type_iterator =
+      ResultTypeIterator<const OperationStmt, const MLValue>;
+  result_type_iterator result_type_begin() const {
+    return result_type_iterator(this, 0);
+  }
+
+  result_type_iterator result_type_end() const {
+    return result_type_iterator(this, getNumResults());
+  }
+
+  llvm::iterator_range<result_type_iterator> getResultTypes() const {
+    return {result_type_begin(), result_type_end()};
+  }
+
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//