[MLIR] Add `KeywordOrString` handling to AsmParser

[youngar]

Summary

This adds a new parser and printer for text which may be a keyword or a
string. When printing, it will attempt to print the text as a keyword,
but if it has any special or non-printable characters, it will be
printed as an escaped string. When parsing, it will parse either a
valid keyword or a potentially escaped string. The printer allows for an
empty string, in which case it prints "".

This new function is used for printing the name in NamedAttributes, and
for printing the symbol name after the @. In CIRCT we are using this
to print module port names, which are conceptually similar to named
function arguments.

[youngar]

Fix description comment

[youngar]

fix description comment again

[youngar]

This was tested using CIRCT with this patch.
Any suggestions on how to add tests for this?

[mehdi_amini]

I think this is something that was intentionally excluded from the parser/printer to not have arbitrary form (I remember trying to add this a while back...).

Can you elaborate a bit on your use case? The CIRCT patch you're linking does not have tests or examples I believe.

Also you wrote that "In CIRCT we are using this to print module port names, which are conceptually similar to named function arguments", but if you want to print named SSA value you can already do it without this patch.

In D111683#3059895, @youngar wrote:

This was tested using CIRCT with this patch.
Any suggestions on how to add tests for this?

Yes, we have a TestDialect where you can implement custom printer/parser for ops.

[youngar]

Hi Mehdi

In D111683#3060024, @mehdi_amini wrote:

I think this is something that was intentionally excluded from the parser/printer to not have arbitrary form (I remember trying to add this a while back...).

Can you elaborate a bit on your use case?

Yeah. It makes more sense when you see an external module, where the port names are a part of the signature of the module:

firrtl.extmodule @Foo(in arg0: !firrtl.uint<2>)

In this example we are printing arg0 using the new printIdentifier function. I don't think it makes much sense to print it as %arg0 when there is no block or value. We could print it as "arg0" but I thought the extra quotes were not helpful. For regular modules we are printing port names as SSA values:

firrtl.module @Foo(in %arg0: !firrtl.uint<2>) {
  ...
}

The CIRCT patch you're linking does not have tests or examples I believe.

Since i'm upstreaming some code from CIRCT to MLIR, the examples are actually in an previous commit which I should have probably linked. You can see some examples in this open PR.

[lattner]

This patch LGTM, it is extending an existing pattern in the AsmParser and it causes boilerplate in other places. I'd love River's thoughts.

think this is something that was intentionally excluded from the parser/printer to not have arbitrary form (I remember trying to add this a while back...).

I don't think this exact one was excluded because of that, but I agree we decided we didn't want to have arbitrary form. That said, I think we should re-evaluate the premise there: LOTS of things (e.g. the eliding of dialect prefixes etc) have pushed back against this, and I think it makes sense to just concede to reality that dialect authors want to do weird things, because they spend LOTS of time looking at debug dumps etc. I think that we could have a much better design, e.g. have an enum for symbols instead of parseQuestion/parseOptionalQuestion pairs for everything etc. In any case, that is a different discussion.

Also you wrote that "In CIRCT we are using this to print module port names, which are conceptually similar to named function arguments", but if you want to print named SSA value you can already do it without this patch.

Yeah, as Andrew says, these aren't SSA names. The similar thing in MLIR is in printNamedAttribute which escapes the attribute name with quotes in the (really rare) case when it isn't a valid identifier according to the MLIR lexer.

mlir/include/mlir/IR/OpImplementation.h

697

typo "dentifier"

mlir/lib/IR/AsmPrinter.cpp

1566

Please add a doxygen comment, and please change printSymbolReference and AsmPrinter::Impl::printNamedAttribute to call this.

[youngar]

printSymbolReference and printNamedAttribute now call printIdentifier. isBareIdentifier no longer asserts on empty string and will return false.

[rriddle]

I think the concept is fine (as it is effectively just a combination of components we already support; i.e. keywords and strings), but that being said the name is very confusing. This isn't an identifier (in the parsing sense), as it could be a string (when not an identifier?). Also, how does this interact with the Keyword methods?

[lattner]

I agree, and I don't think that terseness is important here. How about "PotentiallyQuotedKeyword" or something explicit like that?

[youngar]

parsePotentiallyQuotedKeyword has the same sort of problem, where a keyword can't (by the tokenizer's definition) be quoted, but I like still like it better than the first suggestion. The most explicit thing I can think of is parseKeywordOrString, but the symmetrical API printKeywordOrString is a little weird (e.g. is it obvious that it prints a string when the text is not a valid keyword?).

[lattner]

I think "KeywordOrString" is close enough, despite the ambiguity on the print side of things.

[youngar]

identifier -> keywordOrString

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

469–478

I would have the default methods take a std::string * parameter, and potentially add additional overloads that construct a StringAttr. This is the base AsmParser class, which is also used by attributes/types and those generally don't want to go through StringAttr.

[youngar]

Parse std::strings instead of StringAttrs. No overloads for parsing
attributes were added.

[lattner]

nice, thanks Andrew

mlir/include/mlir/IR/OpImplementation.h

public:
 
   virtual void printType(Type type);
   virtual void printAttribute(Attribute attr);
 
   /// Print the given attribute without its type. The corresponding parser must
   /// provide a valid type for the attribute.
   virtual void printAttributeWithoutType(Attribute attr);
 
+  /// Print the given string as an identifier.  The identifer will be
+  /// surrounded with ""s and escaped if it has any special or non-printable
+  /// characters in it.
+  virtual void printIdentifier(StringRef identifier);
+
   /// Print the given string as a symbol reference, i.e. a form representable by
   /// a SymbolRefAttr. A symbol reference is represented as a string prefixed
   /// with '@'. The reference is surrounded with ""'s and escaped if it has any
   /// special or non-printable characters in it.
   virtual void printSymbolName(StringRef symbolRef);
 
   /// Print an optional arrow followed by a type list.
   template <typename TypeRange>
public:
   virtual ParseResult parseOptionalKeyword(StringRef *keyword) = 0;
 
   /// Parse a keyword, if present, and if one of the 'allowedValues',
   /// into 'keyword'
   virtual ParseResult
   parseOptionalKeyword(StringRef *keyword,
                        ArrayRef<StringRef> allowedValues) = 0;
 
   /// Parse a `(` token.
   virtual ParseResult parseLParen() = 0;
 
   /// Parse a `(` token if present.
   virtual ParseResult parseOptionalLParen() = 0;
 
   /// Parse a `)` token.
   virtual ParseResult parseRParen() = 0;
 
   /// Parse a `)` token if present.
   virtual ParseResult parseOptionalRParen() = 0;
 
   /// Parse a `[` token.
   virtual ParseResult parseLSquare() = 0;
 
   /// Parse a `[` token if present.
   virtual ParseResult parseOptionalLSquare() = 0;
 
public:
 
   /// Parse an integer set instance into 'set'.
   virtual ParseResult printIntegerSet(IntegerSet &set) = 0;
 
   //===--------------------------------------------------------------------===//
   // Identifier Parsing
   //===--------------------------------------------------------------------===//
 
+  /// Parse and identifer and store it in a string attribute named 'attrName'.
+  ParseResult parseIdentifier(StringAttr &result) {
+    if (failed(parseOptionalIdentifier(result)))
+      return emitError(getCurrentLocation()) << "expected valid dentifier";
+    return success();
+  }
+
+  /// Parse an optional identifier and store it in a string attribute named
+  /// 'attrName'.
+  virtual ParseResult parseOptionalIdentifier(StringAttr &result) = 0;
+
   /// Parse an @-identifier and store it (without the '@' symbol) in a string
   /// attribute named 'attrName'.
   ParseResult parseSymbolName(StringAttr &result, StringRef attrName,
                               NamedAttrList &attrs) {
     if (failed(parseOptionalSymbolName(result, attrName, attrs)))
       return emitError(getCurrentLocation())
              << "expected valid '@'-identifier for symbol name";
     return success();

mlir/lib/IR/AsmPrinter.cpp

private:
   void printNewline() override {}
   void printOperand(Value) override {}
   void printOperand(Value, raw_ostream &os) override {
     // Users expect the output string to have at least the prefixed % to signal
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
+  void printIdentifier(StringRef) override {}
   void printSymbolName(StringRef) override {}
   void printSuccessor(Block *) override {}
   void printSuccessorAndUseList(Block *, ValueRange) override {}
   void shadowRegionArgs(Region &, ValueRange) override {}
 
   /// The printer flags to use when determining potential aliases.
   const OpPrintingFlags &printerFlags;
 
static bool isBareIdentifier(StringRef name) {
   unsigned char firstChar = static_cast<unsigned char>(name[0]);
   if (!isalpha(firstChar) && firstChar != '_')
     return false;
   return llvm::all_of(name.drop_front(), [](unsigned char c) {
     return isalnum(c) || c == '_' || c == '$' || c == '.';
   });
 }
 
+static void printIdentifier(StringRef identifier, raw_ostream &os) {
+  // If it can be represented as a bare identifier, write it directly.
+  if (!identifier.empty() && isBareIdentifier(identifier)) {
+    os << identifier;
+    return;
+  }
+
+  // Otherwise, output the identifier wrapped in quotes with proper escaping.
+  os << "\"";
+  printEscapedString(identifier, os);
+  os << '"';
+}
+
 /// Print the given string as a symbol reference. A symbol reference is
 /// represented as a string prefixed with '@'. The reference is surrounded with
 /// ""'s and escaped if it has any special or non-printable characters in it.
 static void printSymbolReference(StringRef symbolRef, raw_ostream &os) {
   assert(!symbolRef.empty() && "expected valid symbol reference");
 
   // If the symbol can be represented as a bare identifier, write it directly.
   if (isBareIdentifier(symbolRef)) {

 }
 
 void AsmPrinter::printAttributeWithoutType(Attribute attr) {
   assert(impl &&
          "expected AsmPrinter::printAttributeWithoutType to be overriden");
   impl->printAttribute(attr, Impl::AttrTypeElision::Must);
 }
 
+void AsmPrinter::printIdentifier(StringRef identifier) {
+  assert(impl && "expected AsmPrinter::printIdentifier to be overriden");
+  ::printIdentifier(identifier, impl->getStream());
+}
+
 void AsmPrinter::printSymbolName(StringRef symbolRef) {
   assert(impl && "expected AsmPrinter::printSymbolName to be overriden");
   ::printSymbolReference(symbolRef, impl->getStream());
 }
 
 //===----------------------------------------------------------------------===//
 // Affine expressions and maps
 //===----------------------------------------------------------------------===//

mlir/lib/Parser/AsmParserImpl.h

public:
   ParseResult printIntegerSet(IntegerSet &set) override {
     return parser.parseIntegerSetReference(set);
   }
 
   //===--------------------------------------------------------------------===//
   // Identifier Parsing
   //===--------------------------------------------------------------------===//
 
+  /// Parse an optional identifier and store it in a string attribute named
+  /// 'attrName'.
+  ParseResult parseOptionalIdentifier(StringAttr &result) override {
+    StringRef keyword;
+    if (succeeded(parseOptionalKeyword(&keyword))) {
+      result = getBuilder().getStringAttr(keyword);
+      return success();
+    }
+
+    std::string string;
+    if (succeeded(parseOptionalString(&string))) {
+      result = getBuilder().getStringAttr(string);
+      return success();
+    }
+
+    return failure();
+  }
+
   /// Parse an optional @-identifier and store it (without the '@' symbol) in a
   /// string attribute named 'attrName'.
   ParseResult parseOptionalSymbolName(StringAttr &result, StringRef attrName,
                                       NamedAttrList &attrs) override {
     Token atToken = parser.getToken();
     if (atToken.isNot(Token::at_identifier))
       return failure();