From f4085c57ddfc1490560213e4611e434d25f84675 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Mon, 18 Apr 2022 09:30:35 -0700
Subject: [PATCH] [mlir] Fix two AttributeParser aborts

Reproducers that resulted in triggering the following asserts

mlir::NamedAttribute::NamedAttribute(mlir::StringAttr, mlir::Attribute)
mlir/lib/IR/Attributes.cpp:29:3
consumeToken
mlir/lib/Parser/Parser.h:126

Differential Revision: https://reviews.llvm.org/D122240
---
 mlir/lib/Parser/AttributeParser.cpp | 10 ++++++++--
 mlir/test/IR/invalid.mlir           | 11 +++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Parser/AttributeParser.cpp b/mlir/lib/Parser/AttributeParser.cpp
index 15d8b17fbdf3..562b4fe7e511 100644
--- a/mlir/lib/Parser/AttributeParser.cpp
+++ b/mlir/lib/Parser/AttributeParser.cpp
@@ -169,8 +169,10 @@ Attribute Parser::parseAttribute(Type type) {
       const char *curPointer = getToken().getLoc().getPointer();
       consumeToken(Token::colon);
       if (!consumeIf(Token::colon)) {
-        state.lex.resetPointer(curPointer);
-        consumeToken();
+        if (getToken().isNot(Token::eof, Token::error)) {
+          state.lex.resetPointer(curPointer);
+          consumeToken();
+        }
         break;
       }
       // Parse the reference itself.
@@ -271,6 +273,10 @@ ParseResult Parser::parseAttributeDict(NamedAttrList &attributes) {
       nameId = builder.getStringAttr(getTokenSpelling());
     else
       return emitError("expected attribute name");
+
+    if (nameId->size() == 0)
+      return emitError("expected valid attribute name");
+
     if (!seenKeys.insert(*nameId).second)
       return emitError("duplicate key '")
              << nameId->getValue() << "' in dictionary attribute";
diff --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index da3278b09b23..7c60f12fe6ff 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -1643,3 +1643,14 @@ func @invalid_region_dominance_with_dominance_free_regions() {
 // -----
 
 func @foo() {} // expected-error {{expected non-empty function body}}
+
+// -----
+
+// expected-error@+1 {{expected valid attribute name}}
+"t"(){""}
+
+// -----
+
+// expected-error@+2 {{expected ']'}}
+"f"() { b = [@m:
+