From d259594be968b9a07bb937c27ba041ed14276bab Mon Sep 17 00:00:00 2001
From: Mogball <jeffniu22@gmail.com>
Date: Sun, 14 Nov 2021 23:01:59 +0000
Subject: [PATCH] [mlir][ods] AttrOrTypeDef format: parse types

Add template specialization to `FieldParser` for parsing types.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D113867
---
 mlir/include/mlir/IR/DialectImplementation.h             | 12 ++++++++++++
 mlir/test/lib/Dialect/Test/TestAttrDefs.td               |  8 ++++++++
 mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir |  4 +++-
 mlir/test/mlir-tblgen/attr-or-type-format.mlir           |  8 ++++++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/IR/DialectImplementation.h b/mlir/include/mlir/IR/DialectImplementation.h
index b841135c..1d0c488 100644
--- a/mlir/include/mlir/IR/DialectImplementation.h
+++ b/mlir/include/mlir/IR/DialectImplementation.h
@@ -70,6 +70,18 @@ struct FieldParser<
   }
 };
 
+/// Parse a type.
+template <typename TypeT>
+struct FieldParser<
+    TypeT, std::enable_if_t<std::is_base_of<Type, TypeT>::value, TypeT>> {
+  static FailureOr<TypeT> parse(AsmParser &parser) {
+    TypeT value;
+    if (parser.parseType(value))
+      return failure();
+    return value;
+  }
+};
+
 /// Parse any integer.
 template <typename IntT>
 struct FieldParser<IntT,
diff --git a/mlir/test/lib/Dialect/Test/TestAttrDefs.td b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
index af56d3c..a9721cf 100644
--- a/mlir/test/lib/Dialect/Test/TestAttrDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestAttrDefs.td
@@ -161,4 +161,12 @@ def TestAttrParams: Test_Attr<"TestAttrParams"> {
   let assemblyFormat = "`<` params `>`";
 }
 
+// Test types can be parsed/printed.
+def TestAttrWithTypeParam : Test_Attr<"TestAttrWithTypeParam"> {
+  let parameters = (ins "::mlir::IntegerType":$int_type,
+                        "::mlir::Type":$any_type);
+  let mnemonic = "attr_with_type";
+  let assemblyFormat = "`<` $int_type `,` $any_type `>`";
+}
+
 #endif // TEST_ATTRDEFS
diff --git a/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir b/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
index f403f6f..20aef66 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
+++ b/mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir
@@ -12,7 +12,9 @@ attributes {
   // CHECK: #test<"attr_ugly begin 5 : index end">
   attr2 = #test<"attr_ugly begin 5 : index end">,
   // CHECK: #test.attr_params<42, 24>
-  attr3 = #test.attr_params<42, 24>
+  attr3 = #test.attr_params<42, 24>,
+  // CHECK: #test.attr_with_type<i32, vector<4xi32>>
+  attr4 = #test.attr_with_type<i32, vector<4xi32>>
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
diff --git a/mlir/test/mlir-tblgen/attr-or-type-format.mlir b/mlir/test/mlir-tblgen/attr-or-type-format.mlir
index 3ff638c..c5b02eb 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format.mlir
+++ b/mlir/test/mlir-tblgen/attr-or-type-format.mlir
@@ -125,3 +125,11 @@ func private @test_verifier_fails() -> () attributes {
   // expected-error@+1 {{expected 'one' to equal 'four.size()'}}
   attr = #test.attr_with_format<42 : two = "hello", four = [1, 2, 3] : 42 : i64>
 }
+
+// -----
+
+func private @test_attr_with_type_failed_to_parse_type() -> () attributes {
+  // expected-error@+2 {{invalid kind of type specified}}
+  // expected-error@+1 {{failed to parse TestAttrWithTypeParam parameter 'int_type'}}
+  attr = #test.attr_with_type<vector<4xi32>, vector<4xi32>>
+}
-- 
2.7.4