[mlir] Add plus, star and optional less/greater parsing

[Kayjukh]

Summary

The tokens are already handled by the lexer. This revision exposes them
through the parser interface.

[rriddle]

Please add tests, which should just be extending the following:
https://github.com/llvm/llvm-project/blob/d93287cac89fd50a105ac4a59c079884b8e53e49/mlir/test/mlir-tblgen/op-format-spec.td#L312
https://github.com/llvm/llvm-project/blob/d93287cac89fd50a105ac4a59c079884b8e53e49/mlir/test/lib/Dialect/Test/TestOps.td#L1377
https://github.com/llvm/llvm-project/blob/6022a8b7e8d1796284541bbf6bad40449c785fef/mlir/test/mlir-tblgen/op-format.mlir#L9

[Kayjukh]

Add tests and rebase

[rriddle]

Can you update https://mlir.llvm.org/docs/OpDefinitions/#literals with the current set of punctuation? (It might be missing a few others as well)

[Kayjukh]

Add the missing tokens to the documentation and add question mark parsing functions

When completing the token list in the docs I noticed there were missing functions and associated tests for question mark parsing. I added them to this patch so as to bundle everything together.

mlir/include/mlir/IR/OpImplementation.h

public:
   virtual ParseResult parseEqual() = 0;
 
   /// Parse a `=` token if present.
   virtual ParseResult parseOptionalEqual() = 0;
 
   /// Parse a '<' token.
   virtual ParseResult parseLess() = 0;
 
+  /// Parse a '<' token if present.
+  virtual ParseResult parseOptionalLess() = 0;
+
   /// Parse a '>' token.
   virtual ParseResult parseGreater() = 0;
 
+  /// Parse a '>' token if present.
+  virtual ParseResult parseOptionalGreater() = 0;
+
+  /// Parse a '+' token.
+  virtual ParseResult parsePlus() = 0;
+
+  /// Parse a '+' token if present.
+  virtual ParseResult parseOptionalPlus() = 0;
+
+  /// Parse a '*' token.
+  virtual ParseResult parseStar() = 0;
+
+  /// Parse a '*' token if present.
+  virtual ParseResult parseOptionalStar() = 0;
+
   /// Parse a given keyword.
   ParseResult parseKeyword(StringRef keyword, const Twine &msg = "") {
     auto loc = getCurrentLocation();
     if (parseOptionalKeyword(keyword))
       return emitError(loc, "expected '") << keyword << "'" << msg;
     return success();
   }
 

mlir/lib/Parser/Parser.cpp

ParseResult parseOptionalEqual() override {
     return success(parser.consumeIf(Token::equal));
   }
 
   /// Parse a '<' token.
   ParseResult parseLess() override {
     return parser.parseToken(Token::less, "expected '<'");
   }
 
+  /// Parse a '<' token if present.
+  ParseResult parseOptionalLess() override {
+    return success(parser.consumeIf(Token::less));
+  }
+
   /// Parse a '>' token.
   ParseResult parseGreater() override {
     return parser.parseToken(Token::greater, "expected '>'");
   }
 
+  /// Parse a '>' token if present.
+  ParseResult parseOptionalGreater() override {
+    return success(parser.consumeIf(Token::greater));
+  }
+
   /// Parse a `(` token.
   ParseResult parseLParen() override {
     return parser.parseToken(Token::l_paren, "expected '('");
   }
 
   /// Parses a '(' if present.
   ParseResult parseOptionalLParen() override {
     return success(parser.consumeIf(Token::l_paren));
ParseResult parseRSquare() override {
     return parser.parseToken(Token::r_square, "expected ']'");
   }
 
   /// Parses a ']' if present.
   ParseResult parseOptionalRSquare() override {
     return success(parser.consumeIf(Token::r_square));
   }
 
+  /// Parses a '+' token.
+  ParseResult parsePlus() override {
+    return parser.parseToken(Token::plus, "expected '+'");
+  }
+
+  /// Parses a '+' token if present.
+  ParseResult parseOptionalPlus() override {
+    return success(parser.consumeIf(Token::plus));
+  }
+
+  /// Parses a '*' token.
+  ParseResult parseStar() override {
+    return parser.parseToken(Token::star, "expected '*'");
+  }
+
+  /// Parses a '*' token if present.
+  ParseResult parseOptionalStar() override {
+    return success(parser.consumeIf(Token::star));
+  }
+
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse an arbitrary attribute of a given type and return it in result.
   ParseResult parseAttribute(Attribute &result, Type type) override {
     result = parser.parseAttribute(type);
     return success(static_cast<bool>(result));

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 //===----------------------------------------------------------------------===//
 // Test Op Asm Format
 //===----------------------------------------------------------------------===//
 
 def FormatLiteralOp : TEST_Op<"format_literal_op"> {
   let assemblyFormat = [{
-    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` `` `(` ` ` `)` attr-dict
+    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` `` `(` ` ` `)` `+` `*` attr-dict
   }];
 }
 
 // Test that we elide attributes that are within the syntax.
 def FormatAttrOp : TEST_Op<"format_attr_op"> {
   let arguments = (ins I64Attr:$attr);
   let assemblyFormat = "$attr attr-dict";
 }

mlir/test/mlir-tblgen/op-format-spec.td

 }]>;
 // CHECK: error: unexpected end of file in literal
 // CHECK: error: expected directive, literal, variable, or optional group
 def LiteralInvalidB : TestFormat_Op<"literal_invalid_b", [{
   `
 }]>;
 // CHECK-NOT: error
 def LiteralValid : TestFormat_Op<"literal_valid", [{
-  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` ` ` `` `->` `abc$._`
+  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` `+` `*` ` ` `` `->` `abc$._`
   attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: optional groups can only be used as top-level elements

mlir/test/mlir-tblgen/op-format.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -verify-diagnostics | FileCheck %s
 
 // CHECK: %[[I64:.*]] =
 %i64 = "foo.op"() : () -> (i64)
 // CHECK: %[[MEMREF:.*]] =
 %memref = "foo.op"() : () -> (memref<1xf64>)
 
-// CHECK: test.format_literal_op keyword_$. -> :, = <> () []( ) {foo.some_attr}
-test.format_literal_op keyword_$. -> :, = <> () []( ) {foo.some_attr}
+// CHECK: test.format_literal_op keyword_$. -> :, = <> () []( ) + * {foo.some_attr}
+test.format_literal_op keyword_$. -> :, = <> () []( ) + * {foo.some_attr}
 
 // CHECK: test.format_attr_op 10
 // CHECK-NOT: {attr
 test.format_attr_op 10
 
 // CHECK: test.format_opt_attr_op_a(10)
 // CHECK-NOT: {opt_attr
 test.format_opt_attr_op_a(10)

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 bool LiteralElement::isValidLiteral(StringRef value) {
   if (value.empty())
     return false;
   char front = value.front();
 
   // If there is only one character, this must either be punctuation or a
   // single character bare identifier.
   if (value.size() == 1)
-    return isalpha(front) || StringRef("_:,=<>()[]{}?").contains(front);
+    return isalpha(front) || StringRef("_:,=<>()[]{}?+*").contains(front);
 
   // Check the punctuation that are larger than a single character.
   if (value == "->")
     return true;
 
   // Otherwise, this must be an identifier.
   if (!isalpha(front) && front != '_')
     return false;
body << (StringRef)StringSwitch<StringRef>(value)
               .Case("<", "Less()")
               .Case(">", "Greater()")
               .Case("{", "LBrace()")
               .Case("}", "RBrace()")
               .Case("(", "LParen()")
               .Case(")", "RParen()")
               .Case("[", "LSquare()")
               .Case("]", "RSquare()")
-              .Case("?", "Question()");
+              .Case("?", "Question()")
+              .Case("+", "Plus()")
+              .Case("*", "Star()");
 }
 
 /// Generate the storage code required for parsing the given element.
 static void genElementParserStorage(Element *element, OpMethodBody &body) {
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
     auto elements = optional->getElements();
 
     // If the anchor is a unit attribute, it won't be parsed directly so elide