PDLL: Parse integer literals into AttributeExpr

Author: mgehre-amd

mlir/include/mlir/Tools/PDLL/AST/Nodes.h

@@ -28,6 +28,8 @@ class NamedAttributeDecl;
 class OpNameDecl;
 class VariableDecl;
 
+StringRef copyStringWithNull(Context &ctx, StringRef str);
+
 //===----------------------------------------------------------------------===//
 // Name
 //===----------------------------------------------------------------------===//

mlir/lib/Tools/PDLL/AST/Nodes.cpp

@@ -16,7 +16,7 @@ using namespace mlir;
 using namespace mlir::pdll::ast;
 
 /// Copy a string reference into the context with a null terminator.
-static StringRef copyStringWithNull(Context &ctx, StringRef str) {
+StringRef mlir::pdll::ast::copyStringWithNull(Context &ctx, StringRef str) {
   if (str.empty())
     return str;
 

mlir/lib/Tools/PDLL/Parser/Parser.cpp

@@ -325,6 +325,7 @@ class Parser {
   FailureOr<ast::Expr *> parseInlineRewriteLambdaExpr();
   FailureOr<ast::Expr *> parseMemberAccessExpr(ast::Expr *parentExpr);
   FailureOr<ast::Expr *> parseNegatedExpr();
+  FailureOr<ast::Expr *> parseIntegerExpr();
   FailureOr<ast::OpNameDecl *> parseOperationName(bool allowEmptyName = false);
   FailureOr<ast::OpNameDecl *> parseWrappedOperationName(bool allowEmptyName);
   FailureOr<ast::Expr *>
@@ -1834,6 +1835,9 @@ FailureOr<ast::Expr *> Parser::parseExpr() {
   case Token::exclam:
     lhsExpr = parseNegatedExpr();
     break;
+  case Token::integer:
+    lhsExpr = parseIntegerExpr();
+    break;
   case Token::string_block:
     return emitError("expected expression. If you are trying to create an "
                      "ArrayAttr, use a space between `[` and `{`.");
@@ -2077,6 +2081,25 @@ FailureOr<ast::Expr *> Parser::parseNegatedExpr() {
   return parseCallExpr(*identifierExpr, /*isNegated = */ true);
 }
 
+/// Parse
+///   integer : identifier
+/// into an AttributeExpr.
+/// Examples: '4 : i32', '0 : si1'
+FailureOr<ast::Expr *> Parser::parseIntegerExpr() {
+  SMRange loc = curToken.getLoc();
+  StringRef value = curToken.getSpelling();
+  consumeToken();
+  if (!consumeIf(Token::colon))
+    return emitError("expected colon after integer literal");
+  if (!curToken.is(Token::identifier))
+    return emitError("expected integer type");
+  StringRef type = curToken.getSpelling();
+  consumeToken();
+
+  auto allocated = copyStringWithNull(ctx, (Twine(value) + ":" + type).str());
+  return ast::AttributeExpr::create(ctx, loc, allocated);
+}
+
 FailureOr<ast::OpNameDecl *> Parser::parseOperationName(bool allowEmptyName) {
   SMRange loc = curToken.getLoc();
 

mlir/test/mlir-pdll/Parser/expr-failure.pdll

@@ -196,7 +196,7 @@ Pattern {
 // -----
 
 Pattern {
-  // CHECK: expected expression
+  // CHECK: expected colon after integer literal
   let tuple = (10 = _: Value);
   erase op<>;
 }
@@ -239,6 +239,23 @@ Pattern {
   };;
 }
 
+// -----
+
+Pattern {
+  let root = op<func.func> -> ();
+  3;
+  // CHECK: expected colon after integer literal
+  replace root with root;
+}
+
+// -----
+
+Pattern {
+  let root = op<func.func> -> ();
+  3 :;
+  // CHECK: expected integer type
+  replace root with root;
+}
 
 // -----
 

mlir/test/mlir-pdll/Parser/expr.pdll

@@ -14,6 +14,15 @@ Pattern {
 
 // -----
 
+// CHECK: Module
+// CHECK: `-AttributeExpr {{.*}} Value<"10:i32">
+Pattern {
+  let attr = 10 : i32;
+  erase _: Op;
+}
+
+// -----
+
 // CHECK: |-NamedAttributeDecl {{.*}}  Name<some_array>
 // CHECK: `-UserRewriteDecl {{.*}} Name<addElemToArrayAttr> ResultType<Attr>
 // CHECK:   `Arguments`

mlir/test/mlir-pdll/Parser/stmt-failure.pdll

@@ -147,7 +147,7 @@ Pattern {
 // -----
 
 Pattern {
-  // CHECK: expected expression
+  // CHECK: expected colon after integer literal
   let foo: ValueRange<10>;
 }