[PDLL] Add support for single line lambda-like patterns

[rriddle]

Summary

This allows for defining simple patterns in a single line. The lambda
body of a Pattern expects a single operation rewrite statement:

Pattern => replace op<my_dialect.foo>(operands: ValueRange) with operands;

[nicolasvasilache]

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

560

I get that you don't pop in the failure mode, I am wondering still if an RAAI guard (llvm::scope_exit IIRC or DeclScopeGuard) would be a nice thing to have in general (in case this pattern ends up repeating).

[jpienaar]

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

560

Could you add a comment as to why you don't pop in error state here? (is there any error recovery here or is a failure fatal or is this not an error if this fails?)

607

WDYT about factoring these out into (say) parseLambdaBodyDecl and parsePatternBodyDecl? (would read a bit simpler and these are two independent code paths, the lambda side is pretty much already in own function)

mlir/test/mlir-pdll/Parser/pattern.pdll

33

So all the patterns here simply delete any op with testing of different parts utilizing the AST generated. Could these be executed on test input to show difference in behavior too? And each case a doc comment added? (that way this becomes a little cookbook)

[rriddle]

mlir/test/mlir-pdll/Parser/pattern.pdll

33

Conceptually, but I'm hesitant to apply an integration test mindset to these tests. They are currently intended to test things in isolation, and we can likely construct better and more focused integration tests when we get the PDL generation pipeline setup. I'm definitely +1 on using our integration tests as a set of examples on how to use the language though, that is a great idea.

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

private:
   /// This structure contains the set of pattern metadata that may be parsed.
   struct ParsedPatternMetadata {
     Optional<uint16_t> benefit;
     bool hasBoundedRecursion = false;
   };
 
   FailureOr<ast::Decl *> parseTopLevelDecl();
   FailureOr<ast::NamedAttributeDecl *> parseNamedAttributeDecl();
+  FailureOr<ast::CompoundStmt *>
+  parseLambdaBody(function_ref<LogicalResult(ast::Stmt *&)> processStatementFn,
+                  bool expectTerminalSemicolon = true);
+  FailureOr<ast::CompoundStmt *> parsePatternLambdaBody();
   FailureOr<ast::Decl *> parsePatternDecl();
   LogicalResult parsePatternDeclMetadata(ParsedPatternMetadata &metadata);
 
   /// Check to see if a decl has already been defined with the given name, if
   /// one has emit and error and return failure. Returns success otherwise.
   LogicalResult checkDefineNamedDecl(const ast::Name &name);
 
   /// Try to define a variable decl with the given components, returns the
if (consumeIf(Token::equal)) {
     // If there isn't a concrete value, create an expression representing a
     // UnitAttr.
     attrValue = ast::AttributeExpr::create(ctx, name.getLoc(), "unit");
   }
 
   return ast::NamedAttributeDecl::create(ctx, name, attrValue);
 }
 
+FailureOr<ast::CompoundStmt *> Parser::parseLambdaBody(
+    function_ref<LogicalResult(ast::Stmt *&)> processStatementFn,
+    bool expectTerminalSemicolon) {
+  consumeToken(Token::equal_arrow);
+
+  // Parse the single statement of the lambda body.
+  SMLoc bodyStartLoc = curToken.getStartLoc();
+  pushDeclScope();
+  FailureOr<ast::Stmt *> singleStatement = parseStmt(expectTerminalSemicolon);
+  bool failedToParse =
+      failed(singleStatement) || failed(processStatementFn(*singleStatement));
+  popDeclScope();
+  if (failedToParse)
+    return failure();
+
+  SMRange bodyLoc(bodyStartLoc, curToken.getStartLoc());
+  return ast::CompoundStmt::create(ctx, bodyLoc, *singleStatement);
+}
+
+FailureOr<ast::CompoundStmt *> Parser::parsePatternLambdaBody() {
+  return parseLambdaBody([&](ast::Stmt *&statement) -> LogicalResult {
+    if (isa<ast::OpRewriteStmt>(statement))
+      return success();
+    return emitError(
+        statement->getLoc(),
+        "expected Pattern lambda body to contain a single operation "
+        "rewrite statement, such as `erase`, `replace`, or `rewrite`");
+  });
+}
+
 FailureOr<ast::Decl *> Parser::parsePatternDecl() {
   SMRange loc = curToken.getLoc();
   consumeToken(Token::kw_Pattern);
   llvm::SaveAndRestore<ParserContext> saveCtx(parserContext,
                                               ParserContext::PatternMatch);
 
   // Check for an optional identifier for the pattern name.
   const ast::Name *name = nullptr;
   if (curToken.is(Token::identifier)) {
     name = &ast::Name::create(ctx, curToken.getSpelling(), curToken.getLoc());
     consumeToken(Token::identifier);
   }
 
   // Parse any pattern metadata.
   ParsedPatternMetadata metadata;
   if (consumeIf(Token::kw_with) && failed(parsePatternDeclMetadata(metadata)))
     return failure();
 
   // Parse the pattern body.
   ast::CompoundStmt *body;
 
+  // Handle a lambda body.
+  if (curToken.is(Token::equal_arrow)) {
+    FailureOr<ast::CompoundStmt *> bodyResult = parsePatternLambdaBody();
+    if (failed(bodyResult))
+      return failure();
+    body = *bodyResult;
+  } else {
     if (curToken.isNot(Token::l_brace))
-    return emitError("expected `{` to start pattern body");
+      return emitError("expected `{` or `=>` to start pattern body");
     FailureOr<ast::CompoundStmt *> bodyResult = parseCompoundStmt();
     if (failed(bodyResult))
       return failure();
     body = *bodyResult;
 
     // Verify the body of the pattern.
     auto bodyIt = body->begin(), bodyE = body->end();
     for (; bodyIt != bodyE; ++bodyIt) {
       // Break when we've found the rewrite statement.
       if (isa<ast::OpRewriteStmt>(*bodyIt))
         break;
     }
     if (bodyIt == bodyE) {
       return emitError(loc,
                        "expected Pattern body to terminate with an operation "
                        "rewrite statement, such as `erase`");
     }
     if (std::next(bodyIt) != bodyE) {
       return emitError((*std::next(bodyIt))->getLoc(),
                        "Pattern body was terminated by an operation "
                        "rewrite statement, but found trailing statements");
     }
+  }
 
   return createPatternDecl(loc, name, metadata, body);
 }
 
 LogicalResult
 Parser::parsePatternDeclMetadata(ParsedPatternMetadata &metadata) {
   Optional<SMRange> benefitLoc;
   Optional<SMRange> hasBoundedRecursionLoc;

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

 // RUN: not mlir-pdll %s -split-input-file 2>&1  | FileCheck %s
 
-// CHECK: expected `{` to start pattern body
+// CHECK: expected `{` or `=>` to start pattern body
 Pattern }
 
 // -----
 
 // CHECK: :6:9: error: `Foo` has already been defined
 // CHECK: :5:9: note: see previous definition here
 Pattern Foo { erase root: Op; }
 Pattern Foo { erase root: Op; }

 // CHECK: Pattern body was terminated by an operation rewrite statement, but found trailing statements
 Pattern {
   erase root: Op;
   let value: Value;
 }
 
 // -----
 
+// CHECK: expected Pattern lambda body to contain a single operation rewrite statement, such as `erase`, `replace`, or `rewrite`
+Pattern => op<>;
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // Metadata
 //===----------------------------------------------------------------------===//
 
 // CHECK: expected pattern metadata identifier
 Pattern with {}
 
 // -----

mlir/test/mlir-pdll/Parser/pattern.pdll

 
 // -----
 
 // CHECK: Module
 // CHECK: `-PatternDecl {{.*}} Name<NamedPattern> Benefit<10> Recursion
 Pattern NamedPattern with benefit(10), recursion {
   erase _: Op;
 }
+
+// -----
+
+// CHECK: Module
+// CHECK: `-PatternDecl {{.*}} Name<NamedPattern>
+// CHECK:   `-CompoundStmt
+// CHECK:     `-EraseStmt
+Pattern NamedPattern => erase _: Op;