Add support to OpAsmParser for parsing unknown keywords.

This is useful in several cases, for example a user may want to sugar the syntax of a string(as we do with custom operation syntax), or avoid many nested ifs for  parsing a set of known keywords.

PiperOrigin-RevId: 269695451

third_party/mlir/include/mlir/IR/OpImplementation.h

@@ -246,15 +246,27 @@ public:
   /// Parse a `=` token.
   virtual ParseResult parseEqual() = 0;
 
-  /// Parse a keyword.
-  ParseResult parseKeyword(const char *keyword, const Twine &msg = "") {
+  /// Parse a given keyword.
+  ParseResult parseKeyword(StringRef keyword, const Twine &msg = "") {
+    auto loc = getCurrentLocation();
     if (parseOptionalKeyword(keyword))
-      return emitError(getNameLoc(), "expected '") << keyword << "'" << msg;
+      return emitError(loc, "expected '") << keyword << "'" << msg;
     return success();
   }
 
-  /// Parse a keyword if present.
-  virtual ParseResult parseOptionalKeyword(const char *keyword) = 0;
+  /// Parse a keyword into 'keyword'.
+  ParseResult parseKeyword(StringRef *keyword) {
+    auto loc = getCurrentLocation();
+    if (parseOptionalKeyword(keyword))
+      return emitError(loc, "expected valid keyword");
+    return success();
+  }
+
+  /// Parse the given keyword if present.
+  virtual ParseResult parseOptionalKeyword(StringRef keyword) = 0;
+
+  /// Parse a keyword, if present, into 'keyword'.
+  virtual ParseResult parseOptionalKeyword(StringRef *keyword) = 0;
 
   /// Parse a `(` token.
   virtual ParseResult parseLParen() = 0;

third_party/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -289,7 +289,7 @@ ParseResult LaunchOp::parse(OpAsmParser *parser, OperationState *result) {
   // so is the trailing type list.  Parse it as well and use the parsed types
   // to resolve the operands passed to the kernel arguments.
   SmallVector<Type, 4> dataTypes;
-  if (!parser->parseOptionalKeyword(getArgsKeyword().data())) {
+  if (!parser->parseOptionalKeyword(getArgsKeyword())) {
     llvm::SMLoc argsLoc = parser->getCurrentLocation();
 
     regionArgs.push_back({});

third_party/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp

@@ -269,11 +269,10 @@ static ParseResult parseVariableDecorations(OpAsmParser *parser,
         parser->parseRParen()) {
       return failure();
     }
-  } else if (succeeded(parser->parseOptionalKeyword(builtInName.c_str()))) {
+  } else if (succeeded(parser->parseOptionalKeyword(builtInName))) {
     StringAttr builtIn;
     if (parser->parseLParen() ||
-        parser->parseAttribute(builtIn, Type(), builtInName,
-                               state->attributes) ||
+        parser->parseAttribute(builtIn, builtInName, state->attributes) ||
         parser->parseRParen()) {
       return failure();
     }

third_party/mlir/lib/Parser/Parser.cpp

@@ -3353,20 +3353,6 @@ public:
     return parser.parseToken(Token::equal, "expected '='");
   }
 
-  /// Parse a keyword if present.
-  ParseResult parseOptionalKeyword(const char *keyword) override {
-    // Check that the current token is a bare identifier or keyword.
-    if (parser.getToken().isNot(Token::bare_identifier) &&
-        !parser.getToken().isKeyword())
-      return failure();
-
-    if (parser.getTokenSpelling() == keyword) {
-      parser.consumeToken();
-      return success();
-    }
-    return failure();
-  }
-
   /// Parse a `(` token.
   ParseResult parseLParen() override {
     return parser.parseToken(Token::l_paren, "expected '('");
@@ -3436,6 +3422,32 @@ public:
   // Identifier Parsing
   //===--------------------------------------------------------------------===//
 
+  /// Returns if the current token corresponds to a keyword.
+  bool isCurrentTokenAKeyword() const {
+    return parser.getToken().is(Token::bare_identifier) ||
+           parser.getToken().isKeyword();
+  }
+
+  /// Parse the given keyword if present.
+  ParseResult parseOptionalKeyword(StringRef keyword) override {
+    // Check that the current token has the same spelling.
+    if (!isCurrentTokenAKeyword() || parser.getTokenSpelling() != keyword)
+      return failure();
+    parser.consumeToken();
+    return success();
+  }
+
+  /// Parse a keyword, if present, into 'keyword'.
+  ParseResult parseOptionalKeyword(StringRef *keyword) override {
+    // Check that the current token is a keyword.
+    if (!isCurrentTokenAKeyword())
+      return failure();
+
+    *keyword = parser.getTokenSpelling();
+    parser.consumeToken();
+    return success();
+  }
+
   /// Parse an @-identifier and store it (without the '@' symbol) in a string
   /// attribute named 'attrName'.
   ParseResult parseSymbolName(StringAttr &result, StringRef attrName,

third_party/mlir/test/lib/TestDialect/TestDialect.cpp

@@ -123,12 +123,28 @@ static void print(OpAsmPrinter *p, IsolatedRegionOp op) {
 }
 
 //===----------------------------------------------------------------------===//
-// Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
+// Test parser.
 //===----------------------------------------------------------------------===//
 
+static ParseResult parseWrappedKeywordOp(OpAsmParser *parser,
+                                         OperationState *result) {
+  StringRef keyword;
+  if (parser->parseKeyword(&keyword))
+    return failure();
+  result->addAttribute("keyword", parser->getBuilder().getStringAttr(keyword));
+  return success();
+}
+
+static void print(OpAsmPrinter *p, WrappedKeywordOp op) {
+  *p << WrappedKeywordOp::getOperationName() << " " << op.keyword();
+}
+
+//===----------------------------------------------------------------------===//
+// Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
+
 static ParseResult parseWrappingRegionOp(OpAsmParser *parser,
                                          OperationState *result) {
-  if (parser->parseOptionalKeyword("wraps"))
+  if (parser->parseKeyword("wraps"))
     return failure();
 
   // Parse the wrapped op in a region

third_party/mlir/test/lib/TestDialect/TestOps.td

@@ -768,9 +768,18 @@ def TestValidOp : TEST_Op<"valid", [Terminator]>,
   Arguments<(ins Variadic<AnyType>:$inputs)>;
 
 //===----------------------------------------------------------------------===//
-// Test region argument list parsing.
+// Test parser.
 //===----------------------------------------------------------------------===//
 
+def WrappedKeywordOp : TEST_Op<"wrapped_keyword"> {
+  let arguments = (ins StrAttr:$keyword);
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+}
+
+//===----------------------------------------------------------------------===//
+// Test region argument list parsing.
+
 def IsolatedRegionOp : TEST_Op<"isolated_region", [IsolatedFromAbove]> {
   let summary =  "isolated region operation";
   let description = [{