[mlir][OpFormatGen] Update "custom" directives for attributes.

This tweaks the generated code for parsing attributes with a custom
directive to call `addAttribute` on the `OperationState` directly,
and adds a newline after this call. Previously, the generated code
would call `addAttribute` on the `OperationState` field `attributes`,
which has no such method and fails to compile. Furthermore, the lack
of newline would generate code with incorrectly formatted single line
`if` statements. Added tests for parsing and printing attributes with
a custom directive.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D87860

2 files changed, 29 insertions, 0 deletions

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 7bdd0fd..c84a771 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -364,6 +364,17 @@ parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,
   varSuccessors.append(2, varSuccessor);
   return success();
 }
+static ParseResult parseCustomDirectiveAttributes(OpAsmParser &parser,
+                                                  IntegerAttr &attr,
+                                                  IntegerAttr &optAttr) {
+  if (parser.parseAttribute(attr))
+    return failure();
+  if (succeeded(parser.parseOptionalComma())) {
+    if (parser.parseAttribute(optAttr))
+      return failure();
+  }
+  return success();
+}
 
 //===----------------------------------------------------------------------===//
 // Printing
@@ -417,6 +428,13 @@ static void printCustomDirectiveSuccessors(OpAsmPrinter &printer,
   if (!varSuccessors.empty())
     printer << ", " << varSuccessors.front();
 }
+static void printCustomDirectiveAttributes(OpAsmPrinter &printer,
+                                           Attribute attribute,
+                                           Attribute optAttribute) {
+  printer << attribute;
+  if (optAttribute)
+    printer << ", " << optAttribute;
+}
 
 //===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index a2cf5b0..6d4e58c 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1547,6 +1547,17 @@ def FormatCustomDirectiveSuccessors
   }];
 }
 
+def FormatCustomDirectiveAttributes
+    : TEST_Op<"format_custom_directive_attributes"> {
+  let arguments = (ins I64Attr:$attr, OptionalAttr<I64Attr>:$optAttr);
+  let assemblyFormat = [{
+    custom<CustomDirectiveAttributes>(
+      $attr, $optAttr
+    )
+    attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // AllTypesMatch type inference