From 0d21596c1da5fd6b1c2d523937a7fe8903ad8589 Mon Sep 17 00:00:00 2001
From: Jeff Niu <jeff@modular.com>
Date: Fri, 30 Sep 2022 16:08:01 -0700
Subject: [PATCH] [mlir][ods] Allow references to the self type

The self type always "bound" since it is provided to the attribute
parser hook. Allow custom directives to reference it.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D134997
---
 mlir/test/mlir-tblgen/attr-or-type-format.td   | 17 +++++++++++++++++
 mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp | 18 +++++++++---------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/mlir/test/mlir-tblgen/attr-or-type-format.td b/mlir/test/mlir-tblgen/attr-or-type-format.td
index ed537c1..675104a 100644
--- a/mlir/test/mlir-tblgen/attr-or-type-format.td
+++ b/mlir/test/mlir-tblgen/attr-or-type-format.td
@@ -177,6 +177,23 @@ def AttrD : TestAttr<"TestG"> {
   let assemblyFormat = "$a";
 }
 
+// Check that the self-type parameter can be referenced without being bound.
+
+// ATTR-LABEL: Attribute TestHAttr::parse
+// ATTR: _result_type = reqType
+// ATTR: parseUseType(odsParser,
+// ATTR-NEXT: *_result_type
+
+// ATTR-LABEL: void TestHAttr::print
+// ATTR: printUseType(odsPrinter,
+// ATTR-NEXT: getType()
+
+def AttrE : TestAttr<"TestH"> {
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type);
+  let mnemonic = "attr_e";
+  let assemblyFormat = "custom<UseType>(ref($type))";
+}
+
 /// Test type parser and printer that mix variables and struct are generated
 /// correctly.
 
diff --git a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
index 72fa3cb..39bbdf3 100644
--- a/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp
@@ -290,6 +290,12 @@ void DefFormat::genParser(MethodBody &os) {
   os.indent();
   os << "::mlir::Builder odsBuilder(odsParser.getContext());\n";
 
+  // Store the initial location of the parser.
+  ctx.addSubst("_loc", "odsLoc");
+  os << tgfmt("::llvm::SMLoc $_loc = $_parser.getCurrentLocation();\n"
+              "(void) $_loc;\n",
+              &ctx);
+
   // Declare variables to store all of the parameters. Allocated parameters
   // such as `ArrayRef` and `StringRef` must provide a `storageType`. Store
   // FailureOr<T> to defer type construction for parameters that are parsed in
@@ -298,14 +304,10 @@ void DefFormat::genParser(MethodBody &os) {
   for (const AttrOrTypeParameter &param : params) {
     os << formatv("::mlir::FailureOr<{0}> _result_{1};\n",
                   param.getCppStorageType(), param.getName());
+    if (auto *selfTypeParam = dyn_cast<AttributeSelfTypeParameter>(&param))
+      genAttrSelfTypeParser(os, ctx, *selfTypeParam);
   }
 
-  // Store the initial location of the parser.
-  ctx.addSubst("_loc", "odsLoc");
-  os << tgfmt("::llvm::SMLoc $_loc = $_parser.getCurrentLocation();\n"
-              "(void) $_loc;\n",
-              &ctx);
-
   // Generate call to each parameter parser.
   for (FormatElement *el : elements)
     genElementParser(el, ctx, os);
@@ -313,8 +315,6 @@ void DefFormat::genParser(MethodBody &os) {
   // Emit an assert for each mandatory parameter. Triggering an assert means
   // the generated parser is incorrect (i.e. there is a bug in this code).
   for (const AttrOrTypeParameter &param : params) {
-    if (auto *selfTypeParam = dyn_cast<AttributeSelfTypeParameter>(&param))
-      genAttrSelfTypeParser(os, ctx, *selfTypeParam);
     if (param.isOptional())
       continue;
     os << formatv("assert(::mlir::succeeded(_result_{0}));\n", param.getName());
@@ -1033,7 +1033,7 @@ DefFormatParser::parseVariableImpl(SMLoc loc, StringRef name, Context ctx) {
     seenParams.set(idx);
 
     // Otherwise, to be referenced, a variable must have been bound.
-  } else if (!seenParams.test(idx)) {
+  } else if (!seenParams.test(idx) && !isa<AttributeSelfTypeParameter>(*it)) {
     return emitError(loc, "parameter '" + name +
                               "' must be bound before it is referenced");
   }
-- 
2.7.4