[mlir] Enable types to us custom assembly formats involving optional attributes

[laszlokindrat]

Summary

This patch adds a partial specialization of mlir::FieldParser that enables types to specify custom assembly formats that involve optional attributes.

[rriddle]

Nice!

mlir/include/mlir/IR/DialectImplementation.h

struct FieldParser<std::string> {
   static FailureOr<std::string> parse(AsmParser &parser) {
     std::string value;
     if (parser.parseString(&value))
       return failure();
     return value;
   }
 };
 
+/// Parse an Optional attribute.
+template <typename AttributeT>
+struct FieldParser<
+    std::optional<AttributeT>,
+    std::enable_if_t<std::is_base_of<Attribute, AttributeT>::value,
+                     std::optional<AttributeT>>> {
+  static FailureOr<std::optional<AttributeT>> parse(AsmParser &parser) {
+    AttributeT attr;
+    OptionalParseResult result = parser.parseOptionalAttribute(attr);
+    if (result.has_value()) {
+      if (succeeded(*result))
+        return {std::optional<AttributeT>(attr)};
+      return failure();
+    }
+    return {std::nullopt};
+  }
+};
+
 /// Parse an Optional integer.
 template <typename IntT>
 struct FieldParser<
     llvm::Optional<IntT>,
     std::enable_if_t<std::is_integral<IntT>::value, Optional<IntT>>> {
   static FailureOr<Optional<IntT>> parse(AsmParser &parser) {
     IntT value;
     OptionalParseResult result = parser.parseOptionalInteger(value);

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

let parameters = (
     "int":$v3
   );
 
   let mnemonic = "struct_capture_all";
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 def TestTypeOptionalParam : Test_Type<"TestTypeOptionalParam"> {
-  let parameters = (ins OptionalParameter<"mlir::Optional<int>">:$a, "int":$b);
+  let parameters = (ins
+    OptionalParameter<"mlir::Optional<int>">:$a,
+    "int":$b,
+    DefaultValuedParameter<"std::optional<::mlir::Attribute>",
+                           "std::nullopt">:$c
+  );
   let mnemonic = "optional_param";
-  let assemblyFormat = "`<` $a `,` $b `>`";
+  let assemblyFormat = "`<` $a `,` $b ( `,` $c^)? `>`";
 }
 
 def TestTypeOptionalParams : Test_Type<"TestTypeOptionalParams"> {
   let parameters = (ins OptionalParameter<"mlir::Optional<int>">:$a,
                         StringRefParameter<>:$b);
   let mnemonic = "optional_params";
   let assemblyFormat = "`<` params `>`";
 }

mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir

attributes {
   attr_10 = #test.attr_with_optional_unsigned<22>
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
 // CHECK: !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 // CHECK: !test.struct_capture_all<v0 = 0, v1 = 1, v2 = 2, v3 = 3>
 // CHECK: !test.optional_param<, 6>
 // CHECK: !test.optional_param<5, 6>
+// CHECK: !test.optional_param<5, 6, "foo">
+// CHECK: !test.optional_param<5, 6, {foo = "bar"}>
 // CHECK: !test.optional_params<"a">
 // CHECK: !test.optional_params<5, "a">
 // CHECK: !test.optional_struct<b = "a">
 // CHECK: !test.optional_struct<a = 5, b = "a">
 // CHECK: !test.optional_params_after<"a">
 // CHECK: !test.optional_params_after<"a", 5>
 // CHECK: !test.all_optional_params<>
 // CHECK: !test.all_optional_params<5>

 // CHECK: !test.custom_type_string<"bar" bar>
 
 func.func private @test_roundtrip_default_parsers_struct(
   !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 ) -> (
   !test.struct_capture_all<v3 = 3, v1 = 1, v2 = 2, v0 = 0>,
   !test.optional_param<, 6>,
   !test.optional_param<5, 6>,
+  !test.optional_param<5, 6, "foo">,
+  !test.optional_param<5, 6, {foo = "bar"}>,
   !test.optional_params<"a">,
   !test.optional_params<5, "a">,
   !test.optional_struct<b = "a">,
   !test.optional_struct<b = "a", a = 5>,
   !test.optional_params_after<"a">,
   !test.optional_params_after<"a", 5>,
   !test.all_optional_params<>,
   !test.all_optional_params<5>,