From aea31f62d44adf9d6e51d2ee323692b6efadcede Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Sat, 12 Mar 2022 07:38:10 +0530
Subject: [PATCH] [MLIR] Fix block label parsing bug

Fix bug in `Block` label parsing:
https://github.com/llvm/llvm-project/issues/54343

The `parseOptionalBlockArgList` method was doing the wrong thing
(contrary to its doc comment) and its calling context was also
incorrect. This led to a parse failure for something like "^bb0()".

Fixes #54343

Differential Revision: https://reviews.llvm.org/D121503
---
 mlir/lib/Parser/Parser.cpp | 14 +++++++-------
 mlir/test/IR/invalid.mlir  |  8 +++++++-
 mlir/test/IR/parser.mlir   |  6 ++++++
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index affaa83..c750275 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -1958,11 +1958,9 @@ ParseResult OperationParser::parseBlock(Block *&block) {
     return emitError(nameLoc, "redefinition of block '") << name << "'";
 
   // If an argument list is present, parse it.
-  if (consumeIf(Token::l_paren)) {
-    if (parseOptionalBlockArgList(block) ||
-        parseToken(Token::r_paren, "expected ')' to end argument list"))
+  if (getToken().is(Token::l_paren))
+    if (parseOptionalBlockArgList(block))
       return failure();
-  }
 
   if (parseToken(Token::colon, "expected ':' after block name"))
     return failure();
@@ -2025,9 +2023,11 @@ Block *OperationParser::defineBlockNamed(StringRef name, SMLoc loc,
   return blockAndLoc.block;
 }
 
-/// Parse a (possibly empty) list of SSA operands with types as block arguments.
+/// Parse a (possibly empty) list of SSA operands with types as block arguments
+/// enclosed in parentheses.
 ///
-///   ssa-id-and-type-list ::= ssa-id-and-type (`,` ssa-id-and-type)*
+///   value-id-and-type-list ::= value-id-and-type (`,` ssa-id-and-type)*
+///   block-arg-list ::= `(` value-id-and-type-list? `)`
 ///
 ParseResult OperationParser::parseOptionalBlockArgList(Block *owner) {
   if (getToken().is(Token::r_brace))
@@ -2038,7 +2038,7 @@ ParseResult OperationParser::parseOptionalBlockArgList(Block *owner) {
   bool definingExistingArgs = owner->getNumArguments() != 0;
   unsigned nextArgument = 0;
 
-  return parseCommaSeparatedList([&]() -> ParseResult {
+  return parseCommaSeparatedList(Delimiter::Paren, [&]() -> ParseResult {
     return parseSSADefOrUseAndType(
         [&](SSAUseInfo useInfo, Type type) -> ParseResult {
           BlockArgument arg;
diff --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index a88bfda..26d67c1 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -136,7 +136,7 @@ func @no_terminator() {   // expected-error {{empty block: expect at least a ter
 // -----
 
 func @block_no_rparen() {
-^bb42 (%bb42 : i32: // expected-error {{expected ')' to end argument list}}
+^bb42 (%bb42 : i32: // expected-error {{expected ')'}}
   return
 }
 
@@ -188,6 +188,12 @@ func @no_terminator() {
   %y = arith.constant 1 : i32  // expected-error {{block with no terminator}}
 }
 
+// -----
+
+func @no_block_arg_enclosing_parens() {
+^bb %x: i32 : // expected-error {{expected ':' after block name}}
+  return
+}
 
 // -----
 
diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index f8c95a3..4984d85 100644
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -181,6 +181,12 @@ func @simpleCFGUsingBBArgs(i32, i64) {
 // CHECK: }
 }
 
+// CHECK-LABEL: func @block_label_empty_list
+func @block_label_empty_list() {
+^bb0():
+  return
+}
+
 // CHECK-LABEL: func @multiblock() {
 func @multiblock() {
   return     // CHECK:   return
-- 
2.7.4