From 49af2a62758a9a526ce446111acaa2cbd5fd2a6e Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Fri, 3 Sep 2021 08:16:52 +0200
Subject: [PATCH] [mlir][flang] Do not prevent integer types from being parsed
 as MLIR keywords

DialectAsmParser::parseKeyword is rejecting `'i' digit+` while it is
a valid identifier according to mlir/docs/LangRef.md.

Integer types actually used to be TOK_KEYWORD a while back before the
change: https://github.com/llvm/llvm-project/commit/6af866c58d21813fb243906611d02bb2a8ffa43a.

This patch Modifies `isCurrentTokenAKeyword` to return true for tokens that
match integer types too.

The motivation for this change is the parsing of `!fir.type<{` `component-name: component-type,`+ `}>`
type in FIR that represent Fortran derived types. The component-names are
parsed as keywords, and can very well be i32 or any ixxx (which are
valid Fortran derived type component names).

The Quant dialect type parser had to be modified since it relied on `iw` not
being parsed as keywords.

Differential Revision: https://reviews.llvm.org/D108913
---
 flang/test/Fir/fir-types.fir             |  2 ++
 mlir/lib/Dialect/Quant/IR/TypeParser.cpp | 15 +++++++++------
 mlir/lib/Parser/DialectSymbolParser.cpp  |  2 +-
 mlir/test/mlir-tblgen/types.mlir         | 24 ++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/flang/test/Fir/fir-types.fir b/flang/test/Fir/fir-types.fir
index d9f46e9f0a03..82345d621331 100644
--- a/flang/test/Fir/fir-types.fir
+++ b/flang/test/Fir/fir-types.fir
@@ -25,12 +25,14 @@ func private @it7() -> !fir.char<4,?>
 // CHECK-LABEL: func private @dvd4() -> !fir.type<derived4(p:i8){f:f32}>
 // CHECK-LABEL: func private @dvd5() -> !fir.type<derived5(p1:i8,p2:i8,p3:i8,p4:i8,p5:i8){f1:f32,f2:f32,f3:f32,f4:f32,f5:f32,f6:f32,f7:f32,f8:f32}>
 // CHECK-LABEL: func private @dvd6() -> !fir.type<derived6{f:!fir.ptr<!fir.type<derived6>>}>
+// CHECK-LABEL: func private @dvd7() -> !fir.type<derived_with_field_name_same_as_integer_type{i32:f32}>
 func private @dvd1() -> !fir.type<derived1>
 func private @dvd2() -> !fir.type<derived2(p:i32)>
 func private @dvd3() -> !fir.type<derived3{f:f32}>
 func private @dvd4() -> !fir.type<derived4(p:i8){f:f32}>
 func private @dvd5() -> !fir.type<derived5(p1:i8,p2:i8,p3:i8,p4:i8,p5:i8){f1:f32,f2:f32,f3:f32,f4:f32,f5:f32,f6:f32,f7:f32,f8:f32}>
 func private @dvd6() -> !fir.type<derived6{f:!fir.ptr<!fir.type<derived6>>}>
+func private @dvd7() -> !fir.type<derived_with_field_name_same_as_integer_type{i32:f32}>
 
 // FIR array types
 // CHECK-LABEL: func private @arr1() -> !fir.array<10xf32>
diff --git a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
index 16fe1f0ebdee..72d679054308 100644
--- a/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
+++ b/mlir/lib/Dialect/Quant/IR/TypeParser.cpp
@@ -29,15 +29,16 @@ static IntegerType parseStorageType(DialectAsmParser &parser, bool &isSigned) {
   // Parse storage type (alpha_ident, integer_literal).
   StringRef identifier;
   unsigned storageTypeWidth = 0;
-  if (failed(parser.parseOptionalKeyword(&identifier))) {
-    // If we didn't parse a keyword, this must be a signed type.
-    if (parser.parseType(type))
+  OptionalParseResult result = parser.parseOptionalType(type);
+  if (result.hasValue()) {
+    if (!succeeded(*result)) {
+      parser.parseType(type);
       return nullptr;
-    isSigned = true;
+    }
+    isSigned = !type.isUnsigned();
     storageTypeWidth = type.getWidth();
-
+  } else if (succeeded(parser.parseKeyword(&identifier))) {
     // Otherwise, this must be an unsigned integer (`u` integer-literal).
-  } else {
     if (!identifier.consume_front("u")) {
       parser.emitError(typeLoc, "illegal storage type prefix");
       return nullptr;
@@ -48,6 +49,8 @@ static IntegerType parseStorageType(DialectAsmParser &parser, bool &isSigned) {
     }
     isSigned = false;
     type = parser.getBuilder().getIntegerType(storageTypeWidth);
+  } else {
+    return nullptr;
   }
 
   if (storageTypeWidth == 0 ||
diff --git a/mlir/lib/Parser/DialectSymbolParser.cpp b/mlir/lib/Parser/DialectSymbolParser.cpp
index 80b743675f44..5fea11309400 100644
--- a/mlir/lib/Parser/DialectSymbolParser.cpp
+++ b/mlir/lib/Parser/DialectSymbolParser.cpp
@@ -250,7 +250,7 @@ public:
 
   /// Returns true if the current token corresponds to a keyword.
   bool isCurrentTokenAKeyword() const {
-    return parser.getToken().is(Token::bare_identifier) ||
+    return parser.getToken().isAny(Token::bare_identifier, Token::inttype) ||
            parser.getToken().isKeyword();
   }
 
diff --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index 33ea3f9d044d..33ad65e780fa 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -504,3 +504,27 @@ func @elements_attr_not_index() {
   "test.indexElementsAttr"() {attr = dense<[1, 2]>:tensor<2xi32>} : () -> ()
   return
 }
+
+// -----
+
+// CHECK-LABEL: @struct_success
+func @struct_success() {
+  "test.simple_struct"() : () -> (!test.struct<{a, i32}, {b, f64}>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @struct_with_field_names_like_types
+func @struct_with_field_names_like_types() {
+  "test.struct_with_field_names_like_types"() : () -> (!test.struct<{i32, i32}, {f64, f64}>)
+  return
+}
+
+// -----
+
+func @struct_bad_keywords() {
+  // expected-error@+1 {{expected valid keyword}}
+  "test.struct_bad_keywords"() : () -> (!test.struct<{42, i32}>)
+  return
+}