Switch return instruction to take its operand list separated from its type

list, for consistency with the rest of the language.  Consolidate some parsing
logic, add operand iterators to BranchInst.

PiperOrigin-RevId: 205699457

mlir/include/mlir/IR/Instructions.h

@@ -318,8 +318,6 @@ public:
 
   unsigned getNumOperands() const { return operands.size(); }
 
-  // TODO: Add a getOperands() custom sequence that provides a value projection
-  // of the operand list.
   CFGValue *getOperand(unsigned idx) { return getInstOperand(idx).get(); }
   const CFGValue *getOperand(unsigned idx) const {
     return getInstOperand(idx).get();
@@ -328,6 +326,35 @@ public:
     return getInstOperand(idx).set(value);
   }
 
+  // Support non-const operand iteration.
+  using operand_iterator = OperandIterator<BranchInst, CFGValue>;
+
+  operand_iterator operand_begin() { return operand_iterator(this, 0); }
+
+  operand_iterator operand_end() {
+    return operand_iterator(this, getNumOperands());
+  }
+
+  llvm::iterator_range<operand_iterator> getOperands() {
+    return {operand_begin(), operand_end()};
+  }
+
+  // Support const operand iteration.
+  typedef OperandIterator<const BranchInst, const CFGValue>
+      const_operand_iterator;
+
+  const_operand_iterator operand_begin() const {
+    return const_operand_iterator(this, 0);
+  }
+
+  const_operand_iterator operand_end() const {
+    return const_operand_iterator(this, getNumOperands());
+  }
+
+  llvm::iterator_range<const_operand_iterator> getOperands() const {
+    return {operand_begin(), operand_end()};
+  }
+
   ArrayRef<InstOperand> getInstOperands() const { return operands; }
   MutableArrayRef<InstOperand> getInstOperands() { return operands; }
 

mlir/lib/Parser/Parser.cpp

@@ -1249,7 +1249,8 @@ public:
 
   template <typename ValueTy>
   ParseResult
-  parseOptionalSSAUseAndTypeList(SmallVectorImpl<ValueTy *> &results);
+  parseOptionalSSAUseAndTypeList(SmallVectorImpl<ValueTy *> &results,
+                                 bool isParenthesized);
 
   // Operations
   ParseResult parseOperation(const CreateOperationFunction &createOpFunc);
@@ -1440,30 +1441,60 @@ ResultType FunctionParser::parseSSADefOrUseAndType(
   return action(useInfo, type);
 }
 
-/// Parse a (possibly empty) list of SSA operands with types.
+/// Parse a (possibly empty) list of SSA operands, followed by a colon, then
+/// followed by a type list.  If hasParens is true, then the operands are
+/// surrounded by parens.
 ///
-///   ssa-use-and-type-list ::= ssa-use-and-type (`,` ssa-use-and-type)*
+///   ssa-use-and-type-list[parens]
+///     ::= `(` ssa-use-list `)` ':' type-list-no-parens
+///
+///   ssa-use-and-type-list[!parens]
+///     ::= ssa-use-list ':' type-list-no-parens
 ///
 template <typename ValueTy>
 ParseResult FunctionParser::parseOptionalSSAUseAndTypeList(
-    SmallVectorImpl<ValueTy *> &results) {
-  if (getToken().isNot(Token::percent_identifier))
+    SmallVectorImpl<ValueTy *> &results, bool isParenthesized) {
+
+  // If we are in the parenthesized form and no paren exists, then we succeed
+  // with an empty list.
+  if (isParenthesized && !consumeIf(Token::l_paren))
     return ParseSuccess;
 
-  return parseCommaSeparatedList([&]() -> ParseResult {
-    if (auto *value = parseSSAUseAndType()) {
-      results.push_back(cast<ValueTy>(value));
-      return ParseSuccess;
-    }
+  SmallVector<SSAUseInfo, 4> valueIDs;
+  if (parseOptionalSSAUseList(valueIDs))
     return ParseFailure;
-  });
+
+  if (isParenthesized && !consumeIf(Token::r_paren))
+    return emitError("expected ')' in operand list");
+
+  // If there were no operands, then there is no colon or type lists.
+  if (valueIDs.empty())
+    return ParseSuccess;
+
+  if (!consumeIf(Token::colon))
+    return emitError("expected ':' in operand list");
+
+  SmallVector<Type *, 4> types;
+  if (parseTypeListNoParens(types))
+    return ParseFailure;
+
+  if (valueIDs.size() != types.size())
+    return emitError("expected " + Twine(valueIDs.size()) +
+                     " types to match operand list");
+
+  results.reserve(valueIDs.size());
+  for (unsigned i = 0, e = valueIDs.size(); i != e; ++i) {
+    if (auto *value = resolveSSAUse(valueIDs[i], types[i]))
+      results.push_back(cast<ValueTy>(value));
+    else
+      return ParseFailure;
+  }
+
+  return ParseSuccess;
 }
 
 /// Parse the CFG or MLFunc operation.
 ///
-/// TODO(clattner): This is a change from the MLIR spec as written, it is an
-/// experiment that will eliminate "builtin" instructions as a thing.
-///
 ///  operation ::=
 ///    (ssa-id `=`)? string '(' ssa-use-list? ')' attribute-dict?
 ///    `:` function-type
@@ -1608,8 +1639,7 @@ private:
 } // end anonymous namespace
 
 /// Parse a (possibly empty) list of SSA operands with types as basic block
-/// arguments. Unlike parseOptionalSsaUseAndTypeList the SSA IDs are treated as
-/// defs, not uses.
+/// arguments.
 ///
 ///   ssa-id-and-type-list ::= ssa-id-and-type (`,` ssa-id-and-type)*
 ///
@@ -1734,11 +1764,12 @@ TerminatorInst *CFGFunctionParser::parseTerminator() {
 
   case Token::kw_return: {
     consumeToken(Token::kw_return);
-    SmallVector<CFGValue *, 8> results;
-    if (parseOptionalSSAUseAndTypeList(results))
-      return nullptr;
 
-    return builder.createReturnInst(results);
+    // Parse any operands.
+    SmallVector<CFGValue *, 8> operands;
+    if (parseOptionalSSAUseAndTypeList(operands, /*isParenthesized*/ false))
+      return nullptr;
+    return builder.createReturnInst(operands);
   }
 
   case Token::kw_br: {
@@ -1748,37 +1779,10 @@ TerminatorInst *CFGFunctionParser::parseTerminator() {
       return (emitError("expected basic block name"), nullptr);
     auto branch = builder.createBranchInst(destBB);
 
-    // Parse the use list.
-    if (!consumeIf(Token::l_paren))
-      return branch;
-
-    SmallVector<SSAUseInfo, 4> valueIDs;
-    if (parseOptionalSSAUseList(valueIDs))
+    SmallVector<CFGValue *, 8> operands;
+    if (parseOptionalSSAUseAndTypeList(operands, /*isParenthesized*/ true))
       return nullptr;
-    if (!consumeIf(Token::r_paren))
-      return (emitError("expected ')' in branch argument list"), nullptr);
-    if (!consumeIf(Token::colon))
-      return (emitError("expected ':' in branch argument list"), nullptr);
-
-    auto typeLoc = getToken().getLoc();
-    SmallVector<Type *, 4> types;
-    if (parseTypeListNoParens(types))
-      return nullptr;
-
-    if (types.size() != valueIDs.size())
-      return (emitError(typeLoc, "expected " + Twine(valueIDs.size()) +
-                                     " types to match operand list"),
-              nullptr);
-
-    SmallVector<CFGValue *, 4> values;
-    values.reserve(valueIDs.size());
-    for (unsigned i = 0, e = valueIDs.size(); i != e; ++i) {
-      if (auto *value = resolveSSAUse(valueIDs[i], types[i]))
-        values.push_back(cast<CFGValue>(value));
-      else
-        return nullptr;
-    }
-    branch->addOperands(values);
+    branch->addOperands(operands);
     return branch;
   }
     // TODO: cond_br.

mlir/test/IR/parser.mlir

@@ -202,7 +202,7 @@ bb1:       // CHECK: bb1:
   %1 = "baz"(%2#1, %2#0, %0#1) : (f32, i11, i17) -> (i16, i8)
   
   // CHECK: return %1#0 : i16, %1#1 : i8
-  return %1#0 : i16, %1#1 : i8
+  return %1#0, %1#1 : i16, i8
 
 bb2:       // CHECK: bb2:
   // CHECK: %2 = "bar"(%0#0, %0#1) : (i1, i17) -> (i11, f32)
@@ -220,5 +220,5 @@ bb0:       // CHECK: bb0:
 bb1(%x: i17, %y: i1):       // CHECK: bb1(%1: i17, %2: i1):
   // CHECK: %3 = "baz"(%1, %2, %0#1) : (i17, i1, i17) -> (i16, i8)
   %1 = "baz"(%x, %y, %0#1) : (i17, i1, i17) -> (i16, i8)
-  return %1#0 : i16, %1#1 : i8
+  return %1#0, %1#1 : i16, i8
 }