[mlir][ods] Unify Attr/TypeDef and Operation Format Parsing

[Mogball]

Summary

Part 2 of 3 of unifying the assembly formats of attributes/types and operations.The last patch that introduced attribute/type formats (D111594) factored out the format lexer entirely. This patch factors out most of the format parsers such that the attribute/type and op parsers only need to implement handling for specific elements.

Certain things could be factored better (element verification, 'seen' variables) but the primary goal of factoring is so that features can be used across both assembly formats.

[Mogball]

fix MSVC build

[rriddle]

Letting you know I haven't missed this, will get to it in the next day or two.

[rriddle]

All of the llvm::SMLoc can be SMLoc now, but otherwise LGTM.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

65–66

Why do these need llvm:: now?

2169

Can we drop the mlir:: on all of these?

[Mogball]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

65–66

Conflicting with VariableKind::Optional

[Mogball]

review comments

[Mogball]

review comments

[nicolasvasilache]

gcc-5 did not seem to like this: https://buildkite.com/mlir/mlir-core/builds/20053#e10e18fb-21a3-4acf-8e22-37b6e8408cae

[jpienaar]

This seems to be hitting ASAN error:

Direct leak of 128 byte(s) in 8 object(s) allocated from:
...
#6 0x569acb in mlir::tblgen::FormatParser::parseCustomDirective(llvm::SMLoc, mlir::tblgen::FormatParser::Context) Compiles/llvm/mlir/tools/mlir-tblgen/FormatGen.cpp:360:15

mlir/test/mlir-tblgen/attr-or-type-format-invalid.td

def InvalidTypeB : InvalidType<"InvalidTypeB", "invalid_b"> {
   let parameters = (ins "int":$v0, "int":$v1);
   // CHECK: duplicate parameter 'v0'
   let assemblyFormat = "`<` $v0 `,` $v1 `,` $v0 `>`";
 }
 
 /// Test format has invalid syntax.
 def InvalidTypeC : InvalidType<"InvalidTypeC", "invalid_c"> {
   let parameters = (ins "int":$v0, "int":$v1);
-  // CHECK: expected literal, directive, or variable
+  // CHECK: expected literal, variable, directive, or optional group
   let assemblyFormat = "`<` $v0, $v1 `>`";
 }
 
 /// Test struct directive has invalid syntax.
 def InvalidTypeD : InvalidType<"InvalidTypeD", "invalid_d"> {
   let parameters = (ins "int":$v0);
   // CHECK: literals may only be used in the top-level section of the format
   // CHECK: expected a variable in `struct` argument list

mlir/test/mlir-tblgen/op-format-spec.td

 // CHECK: error: 'functional-type' is only valid as a top-level directive
 def DirectiveFunctionalTypeInvalidA : TestFormat_Op<[{
   functional-type(functional-type)
 }]>;
 // CHECK: error: expected '(' before argument list
 def DirectiveFunctionalTypeInvalidB : TestFormat_Op<[{
   functional-type
 }]>;
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def DirectiveFunctionalTypeInvalidC : TestFormat_Op<[{
   functional-type(
 }]>;
 // CHECK: error: expected ',' after inputs argument
 def DirectiveFunctionalTypeInvalidD : TestFormat_Op<[{
   functional-type(operands
 }]>;
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def DirectiveFunctionalTypeInvalidE : TestFormat_Op<[{
   functional-type(operands,
 }]>;
 // CHECK: error: expected ')' after argument list
 def DirectiveFunctionalTypeInvalidF : TestFormat_Op<[{
   functional-type(operands, results
 }]>;
 // CHECK-NOT: error

 
 //===----------------------------------------------------------------------===//
 // type
 
 // CHECK: error: expected '(' before argument list
 def DirectiveTypeInvalidA : TestFormat_Op<[{
   type
 }]>;
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def DirectiveTypeInvalidB : TestFormat_Op<[{
   type(
 }]>;
 // CHECK: error: expected ')' after argument list
 def DirectiveTypeInvalidC : TestFormat_Op<[{
   type(operands
 }]>;
 // CHECK-NOT: error:
 def DirectiveTypeValid : TestFormat_Op<[{
   type(operands) attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
 // functional-type/type operands
 
-// CHECK: error: literals may only be used in a top-level section of the format
+// CHECK: error: literals may only be used in the top-level section of the format
 def DirectiveTypeZOperandInvalidA : TestFormat_Op<[{
   type(`literal`)
 }]>;
 // CHECK: error: 'operands' 'type' is already bound
 def DirectiveTypeZOperandInvalidB : TestFormat_Op<[{
   type(operands) type(operands)
 }]>;
 // CHECK: error: 'operands' 'type' is already bound
def LiteralInvalidB : TestFormat_Op<[{
   `1`
 }]>;
 // CHECK: error: expected valid literal but got ':abc': valid keyword starts with a letter or '_'
 def LiteralInvalidC : TestFormat_Op<[{
   `:abc`
 }]>;
 
 // CHECK: error: unexpected end of file in literal
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def LiteralInvalidD : TestFormat_Op<[{
   `
 }]>;
 // CHECK-NOT: error
 def LiteralValid : TestFormat_Op<[{
   `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` `?` `+` `*` ` ` `` `->` `\n` `abc$._`
   attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: optional groups can only be used as top-level elements
 def OptionalInvalidA : TestFormat_Op<[{
   type(($attr^)?) attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def OptionalInvalidB : TestFormat_Op<[{
   () attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: optional group specified no anchor element
+// CHECK: error: optional group has no anchor element
 def OptionalInvalidC : TestFormat_Op<[{
   ($attr)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: first parsable element of an operand group must be an attribute, literal, operand, or region
+// CHECK: error: first parsable element of an optional group must be a literal or variable
 def OptionalInvalidD : TestFormat_Op<[{
   (type($operand) $operand^)? attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidE : TestFormat_Op<[{
   (`,` $attr^ type(operands))? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: only one element can be marked as the anchor of an optional group
 def OptionalInvalidF : TestFormat_Op<[{
-  ($attr^ $attr2^) attr-dict
+  ($attr^ $attr2^)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr, OptionalAttr<I64Attr>:$attr2)>;
 // CHECK: error: only optional attributes can be used to anchor an optional group
 def OptionalInvalidG : TestFormat_Op<[{
-  ($attr^) attr-dict
+  ($attr^)? attr-dict
 }]>, Arguments<(ins I64Attr:$attr)>;
 // CHECK: error: only variable length operands can be used within an optional group
 def OptionalInvalidH : TestFormat_Op<[{
-  ($arg^) attr-dict
+  ($arg^)? attr-dict
 }]>, Arguments<(ins I64:$arg)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidI : TestFormat_Op<[{
   (functional-type($arg, results)^)? attr-dict
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidJ : TestFormat_Op<[{
-  (attr-dict)
+  (attr-dict^)?
 }]>;
 // CHECK: error: expected '?' after optional group
 def OptionalInvalidK : TestFormat_Op<[{
   ($arg^)
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
 // CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidL : TestFormat_Op<[{
   (custom<MyDirective>($arg)^)?
 }]>, Arguments<(ins I64:$arg)>;
 // CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidM : TestFormat_Op<[{
   (` `^)?
 }]>, Arguments<(ins)>;
 // CHECK: error: expected '(' to start else branch of optional group
 def OptionalInvalidN : TestFormat_Op<[{
   ($arg^):
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
-// CHECK: error: expected directive, literal, variable, or optional group
+// CHECK: error: expected literal, variable, directive, or optional group
 def OptionalInvalidO : TestFormat_Op<[{
   ($arg^):(`test`
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
 // CHECK: error: expected '?' after optional group
 def OptionalInvalidP : TestFormat_Op<[{
   ($arg^):(`test`)
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
 

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

 
 using llvm::formatv;
 
 //===----------------------------------------------------------------------===//
 // Element
 //===----------------------------------------------------------------------===//
 
 namespace {
-
-/// This class represents a single format element.
-class Element {
-public:
-  /// LLVM-style RTTI.
-  enum class Kind {
-    /// This element is a directive.
-    ParamsDirective,
-    StructDirective,
-
-    /// This element is a literal.
-    Literal,
-
-    /// This element is a variable.
-    Variable,
-  };
-  Element(Kind kind) : kind(kind) {}
-  virtual ~Element() = default;
-
-  /// Return the kind of this element.
-  Kind getKind() const { return kind; }
-
-private:
-  /// The kind of this element.
-  Kind kind;
-};
-
-/// This class represents an instance of a literal element.
-class LiteralElement : public Element {
-public:
-  LiteralElement(StringRef literal)
-      : Element(Kind::Literal), literal(literal) {}
-
-  static bool classof(const Element *el) {
-    return el->getKind() == Kind::Literal;
-  }
-
-  /// Get the literal spelling.
-  StringRef getSpelling() const { return literal; }
-
-private:
-  /// The spelling of the literal for this element.
-  StringRef literal;
-};
-
 /// This class represents an instance of a variable element. A variable refers
 /// to an attribute or type parameter.
-class VariableElement : public Element {
+class ParameterElement
+    : public VariableElementBase<VariableElement::Parameter> {
 public:
-  VariableElement(AttrOrTypeParameter param)
-      : Element(Kind::Variable), param(param) {}
-
-  static bool classof(const Element *el) {
-    return el->getKind() == Kind::Variable;
-  }
+  ParameterElement(AttrOrTypeParameter param) : param(param) {}
 
   /// Get the parameter in the element.
   const AttrOrTypeParameter &getParam() const { return param; }
 
   /// Indicate if this variable is printed "qualified" (that is it is
   /// prefixed with the `#dialect.mnemonic`).
   bool shouldBeQualified() { return shouldBeQualifiedFlag; }
   void setShouldBeQualified(bool qualified = true) {
     shouldBeQualifiedFlag = qualified;
   }
 
 private:
   bool shouldBeQualifiedFlag = false;
   AttrOrTypeParameter param;
 };
 
 /// Base class for a directive that contains references to multiple variables.
-template <Element::Kind ElementKind>
-class ParamsDirectiveBase : public Element {
+template <DirectiveElement::Kind DirectiveKind>
+class ParamsDirectiveBase : public DirectiveElementBase<DirectiveKind> {
 public:
-  using Base = ParamsDirectiveBase<ElementKind>;
-
-  ParamsDirectiveBase(SmallVector<std::unique_ptr<Element>> &&params)
-      : Element(ElementKind), params(std::move(params)) {}
+  using Base = ParamsDirectiveBase<DirectiveKind>;
 
-  static bool classof(const Element *el) {
-    return el->getKind() == ElementKind;
-  }
+  ParamsDirectiveBase(std::vector<FormatElement *> &&params)
+      : params(std::move(params)) {}
 
   /// Get the parameters contained in this directive.
   auto getParams() const {
-    return llvm::map_range(params, [](auto &el) {
-      return cast<VariableElement>(el.get())->getParam();
+    return llvm::map_range(params, [](FormatElement *el) {
+      return cast<ParameterElement>(el)->getParam();
     });
   }
 
   /// Get the number of parameters.
   unsigned getNumParams() const { return params.size(); }
 
   /// Take all of the parameters from this directive.
-  SmallVector<std::unique_ptr<Element>> takeParams() {
-    return std::move(params);
-  }
+  std::vector<FormatElement *> takeParams() { return std::move(params); }
 
 private:
   /// The parameters captured by this directive.
-  SmallVector<std::unique_ptr<Element>> params;
+  std::vector<FormatElement *> params;
 };
 
 /// This class represents a `params` directive that refers to all parameters
 /// of an attribute or type. When used as a top-level directive, it generates
 /// a format of the form:
 ///
 ///   (param-value (`,` param-value)*)?
 ///
 /// When used as an argument to another directive that accepts variables,
 /// `params` can be used in place of manually listing all parameters of an
 /// attribute or type.
-class ParamsDirective
-    : public ParamsDirectiveBase<Element::Kind::ParamsDirective> {
+class ParamsDirective : public ParamsDirectiveBase<DirectiveElement::Params> {
 public:
   using Base::Base;
 };
 
 /// This class represents a `struct` directive that generates a struct format
 /// of the form:
 ///
 ///   `{` param-name `=` param-value (`,` param-name `=` param-value)* `}`
 ///
-class StructDirective
-    : public ParamsDirectiveBase<Element::Kind::StructDirective> {
+class StructDirective : public ParamsDirectiveBase<DirectiveElement::Struct> {
 public:
   using Base::Base;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Format Strings

 //===----------------------------------------------------------------------===//
 // AttrOrTypeFormat
 //===----------------------------------------------------------------------===//
 
 namespace {
 class AttrOrTypeFormat {
 public:
   AttrOrTypeFormat(const AttrOrTypeDef &def,
-                   std::vector<std::unique_ptr<Element>> &&elements)
+                   std::vector<FormatElement *> &&elements)
       : def(def), elements(std::move(elements)) {}
 
   /// Generate the attribute or type parser.
   void genParser(MethodBody &os);
   /// Generate the attribute or type printer.
   void genPrinter(MethodBody &os);
 
 private:
   /// Generate the parser code for a specific format element.
-  void genElementParser(Element *el, FmtContext &ctx, MethodBody &os);
+  void genElementParser(FormatElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a literal.
   void genLiteralParser(StringRef value, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a variable.
   void genVariableParser(const AttrOrTypeParameter &param, FmtContext &ctx,
                          MethodBody &os);
   /// Generate the parser code for a `params` directive.
   void genParamsParser(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `struct` directive.
   void genStructParser(StructDirective *el, FmtContext &ctx, MethodBody &os);
 
   /// Generate the printer code for a specific format element.
-  void genElementPrinter(Element *el, FmtContext &ctx, MethodBody &os);
+  void genElementPrinter(FormatElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a literal.
   void genLiteralPrinter(StringRef value, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a variable.
   void genVariablePrinter(const AttrOrTypeParameter &param, FmtContext &ctx,
                           MethodBody &os, bool printQualified = false);
   /// Generate the printer code for a `params` directive.
   void genParamsPrinter(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a `struct` directive.
   void genStructPrinter(StructDirective *el, FmtContext &ctx, MethodBody &os);
 
   /// The ODS definition of the attribute or type whose format is being used to
   /// generate a parser and printer.
   const AttrOrTypeDef &def;
   /// The list of top-level format elements returned by the assembly format
   /// parser.
-  std::vector<std::unique_ptr<Element>> elements;
+  std::vector<FormatElement *> elements;
 
   /// Flags for printing spaces.
   bool shouldEmitSpace = false;
   bool lastWasPunctuation = false;
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
void AttrOrTypeFormat::genParser(MethodBody &os) {
 
   /// Store the initial location of the parser.
   ctx.addSubst("_loc", "loc");
   os << tgfmt("::llvm::SMLoc $_loc = $_parser.getCurrentLocation();\n"
               "(void) $_loc;\n",
               &ctx);
 
   /// Generate call to each parameter parser.
-  for (auto &el : elements)
-    genElementParser(el.get(), ctx, os);
+  for (FormatElement *el : elements)
+    genElementParser(el, ctx, os);
 
   /// Generate call to the attribute or type builder. Use the checked getter
   /// if one was generated.
   if (def.genVerifyDecl()) {
     os << tgfmt("return $_parser.getChecked<$0>($_loc, $_parser.getContext()",
                 &ctx, def.getCppClassName());
   } else {
     os << tgfmt("return $0::get($_parser.getContext()", &ctx,
                 def.getCppClassName());
   }
   for (const AttrOrTypeParameter &param : params)
     os << formatv(",\n    _result_{0}.getValue()", param.getName());
   os << ");";
 }
 
-void AttrOrTypeFormat::genElementParser(Element *el, FmtContext &ctx,
+void AttrOrTypeFormat::genElementParser(FormatElement *el, FmtContext &ctx,
                                         MethodBody &os) {
   if (auto *literal = dyn_cast<LiteralElement>(el))
     return genLiteralParser(literal->getSpelling(), ctx, os);
-  if (auto *var = dyn_cast<VariableElement>(el))
+  if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariableParser(var->getParam(), ctx, os);
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsParser(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructParser(strct, ctx, os);
 
   llvm_unreachable("unknown format element");
 }

 
 void AttrOrTypeFormat::genPrinter(MethodBody &os) {
   FmtContext ctx;
   ctx.addSubst("_printer", "printer");
 
   /// Generate printers.
   shouldEmitSpace = true;
   lastWasPunctuation = false;
-  for (auto &el : elements)
-    genElementPrinter(el.get(), ctx, os);
+  for (FormatElement *el : elements)
+    genElementPrinter(el, ctx, os);
 }
 
-void AttrOrTypeFormat::genElementPrinter(Element *el, FmtContext &ctx,
+void AttrOrTypeFormat::genElementPrinter(FormatElement *el, FmtContext &ctx,
                                          MethodBody &os) {
   if (auto *literal = dyn_cast<LiteralElement>(el))
     return genLiteralPrinter(literal->getSpelling(), ctx, os);
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsPrinter(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructPrinter(strct, ctx, os);
-  if (auto *var = dyn_cast<VariableElement>(el))
+  if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariablePrinter(var->getParam(), ctx, os,
                               var->shouldBeQualified());
 
   llvm_unreachable("unknown format element");
 }
 
 void AttrOrTypeFormat::genLiteralPrinter(StringRef value, FmtContext &ctx,
                                          MethodBody &os) {
else
     os << tgfmt(defaultParameterPrinter, &ctx) << ";\n";
 }
 
 void AttrOrTypeFormat::genParamsPrinter(ParamsDirective *el, FmtContext &ctx,
                                         MethodBody &os) {
   llvm::interleave(
       el->getParams(),
       [&](auto param) { this->genVariablePrinter(param, ctx, os); },
-      [&]() { this->genLiteralPrinter(",", ctx, os); });
+      [&] { this->genLiteralPrinter(",", ctx, os); });
 }
 
 void AttrOrTypeFormat::genStructPrinter(StructDirective *el, FmtContext &ctx,
                                         MethodBody &os) {
   llvm::interleave(
       el->getParams(),
       [&](auto param) {
         this->genLiteralPrinter(param.getName(), ctx, os);
         this->genLiteralPrinter("=", ctx, os);
         this->genVariablePrinter(param, ctx, os);
       },
-      [&]() { this->genLiteralPrinter(",", ctx, os); });
+      [&] { this->genLiteralPrinter(",", ctx, os); });
 }
 
 //===----------------------------------------------------------------------===//
-// FormatParser
+// DefFormatParser
 //===----------------------------------------------------------------------===//
 
 namespace {
-class FormatParser {
+class DefFormatParser : public FormatParser {
 public:
-  FormatParser(llvm::SourceMgr &mgr, const AttrOrTypeDef &def)
-      : lexer(mgr, def.getLoc()[0]), curToken(lexer.lexToken()), def(def),
+  DefFormatParser(llvm::SourceMgr &mgr, const AttrOrTypeDef &def)
+      : FormatParser(mgr, def.getLoc()[0]), def(def),
         seenParams(def.getNumParameters()) {}
 
   /// Parse the attribute or type format and create the format elements.
   FailureOr<AttrOrTypeFormat> parse();
 
-private:
-  /// The current context of the parser when parsing an element.
-  enum ParserContext {
-    /// The element is being parsed in the default context - at the top of the
-    /// format
-    TopLevelContext,
-    /// The element is being parsed as a child to a `struct` directive.
-    StructDirective,
-  };
-
-  /// Emit an error.
-  LogicalResult emitError(const Twine &msg) {
-    lexer.emitError(curToken.getLoc(), msg);
-    return failure();
-  }
-
-  /// Parse an expected token.
-  LogicalResult parseToken(FormatToken::Kind kind, const Twine &msg) {
-    if (curToken.getKind() != kind)
-      return emitError(msg);
-    consumeToken();
-    return success();
-  }
+protected:
+  /// Verify the parsed elements.
+  LogicalResult verify(llvm::SMLoc loc,
+                       ArrayRef<FormatElement *> elements) override;
+  /// Verify the elements of a custom directive.
+  LogicalResult
+  verifyCustomDirectiveArguments(llvm::SMLoc loc,
+                                 ArrayRef<FormatElement *> arguments) override {
+    return emitError(loc, "'custom' not supported (yet)");
+  }
+  /// Verify the elements of an optional group.
+  LogicalResult
+  verifyOptionalGroupElements(llvm::SMLoc loc,
+                              ArrayRef<FormatElement *> elements,
+                              Optional<unsigned> anchorIndex) override {
+    return emitError(loc, "optional groups not (yet) supported");
+  }
+
+  /// Parse an attribute or type variable.
+  FailureOr<FormatElement *> parseVariableImpl(llvm::SMLoc loc, StringRef name,
+                                               Context ctx) override;
+  /// Parse an attribute or type format directive.
+  FailureOr<FormatElement *> parseDirectiveImpl(llvm::SMLoc loc,
+                                                FormatToken::Kind kind,
+                                                Context ctx) override;
 
-  /// Advance the lexer to the next token.
-  void consumeToken() {
-    assert(curToken.getKind() != FormatToken::eof &&
-           curToken.getKind() != FormatToken::error &&
-           "shouldn't advance past EOF or errors");
-    curToken = lexer.lexToken();
-  }
-
-  /// Parse any element.
-  FailureOr<std::unique_ptr<Element>> parseElement(ParserContext ctx);
-  /// Parse a literal element.
-  FailureOr<std::unique_ptr<Element>> parseLiteral(ParserContext ctx);
-  /// Parse a variable element.
-  FailureOr<std::unique_ptr<Element>> parseVariable(ParserContext ctx);
-  /// Parse a directive.
-  FailureOr<std::unique_ptr<Element>> parseDirective(ParserContext ctx);
+private:
   /// Parse a `params` directive.
-  FailureOr<std::unique_ptr<Element>> parseParamsDirective();
+  FailureOr<FormatElement *> parseParamsDirective(llvm::SMLoc loc);
   /// Parse a `qualified` directive.
-  FailureOr<std::unique_ptr<Element>>
-  parseQualifiedDirective(ParserContext ctx);
+  FailureOr<FormatElement *> parseQualifiedDirective(llvm::SMLoc loc,
+                                                     Context ctx);
   /// Parse a `struct` directive.
-  FailureOr<std::unique_ptr<Element>> parseStructDirective();
+  FailureOr<FormatElement *> parseStructDirective(llvm::SMLoc loc);
 
-  /// The current format lexer.
-  FormatLexer lexer;
-  /// The current token in the stream.
-  FormatToken curToken;
   /// Attribute or type tablegen def.
   const AttrOrTypeDef &def;
 
   /// Seen attribute or type parameters.
   llvm::BitVector seenParams;
 };
 } // namespace
 
-FailureOr<AttrOrTypeFormat> FormatParser::parse() {
-  std::vector<std::unique_ptr<Element>> elements;
-  elements.reserve(16);
-
-  /// Parse the format elements.
-  while (curToken.getKind() != FormatToken::eof) {
-    auto element = parseElement(TopLevelContext);
-    if (failed(element))
-      return failure();
-
-    /// Add the format element and continue.
-    elements.push_back(std::move(*element));
-  }
-
-  /// Check that all parameters have been seen.
+LogicalResult DefFormatParser::verify(llvm::SMLoc loc,
+                                      ArrayRef<FormatElement *> elements) {
   for (auto &it : llvm::enumerate(def.getParameters())) {
     if (!seenParams.test(it.index())) {
-      return emitError("format is missing reference to parameter: " +
+      return emitError(loc, "format is missing reference to parameter: " +
                                 it.value().getName());
     }
   }
-
-  return AttrOrTypeFormat(def, std::move(elements));
-}
-
-FailureOr<std::unique_ptr<Element>>
-FormatParser::parseElement(ParserContext ctx) {
-  if (curToken.getKind() == FormatToken::literal)
-    return parseLiteral(ctx);
-  if (curToken.getKind() == FormatToken::variable)
-    return parseVariable(ctx);
-  if (curToken.isKeyword())
-    return parseDirective(ctx);
-
-  return emitError("expected literal, directive, or variable");
-}
-
-FailureOr<std::unique_ptr<Element>>
-FormatParser::parseLiteral(ParserContext ctx) {
-  if (ctx != TopLevelContext) {
-    return emitError(
-        "literals may only be used in the top-level section of the format");
-  }
+  return success();
 
-  /// Get the literal spelling without the surrounding "`".
-  auto value = curToken.getSpelling().drop_front().drop_back();
-  if (!isValidLiteral(value, [&](Twine diag) {
-        (void)emitError("expected valid literal but got '" + value +
-                        "': " + diag);
-      }))
+FailureOr<AttrOrTypeFormat> DefFormatParser::parse() {
+  FailureOr<std::vector<FormatElement *>> elements = FormatParser::parse();
+  if (failed(elements))
     return failure();
-
-  consumeToken();
-  return {std::make_unique<LiteralElement>(value)};
+  return AttrOrTypeFormat(def, std::move(*elements));
 }
 
-FailureOr<std::unique_ptr<Element>>
-FormatParser::parseVariable(ParserContext ctx) {
-  /// Get the parameter name without the preceding "$".
-  auto name = curToken.getSpelling().drop_front();
-
+FailureOr<FormatElement *> DefFormatParser::parseVariableImpl(llvm::SMLoc loc,
+                                                              StringRef name,
+                                                              Context ctx) {
   /// Lookup the parameter.
   ArrayRef<AttrOrTypeParameter> params = def.getParameters();
   auto *it = llvm::find_if(
       params, [&](auto &param) { return param.getName() == name; });
 
   /// Check that the parameter reference is valid.
-  if (it == params.end())
-    return emitError(def.getName() + " has no parameter named '" + name + "'");
+  if (it == params.end()) {
+    return emitError(loc,
+                     def.getName() + " has no parameter named '" + name + "'");
+  }
   auto idx = std::distance(params.begin(), it);
   if (seenParams.test(idx))
-    return emitError("duplicate parameter '" + name + "'");
+    return emitError(loc, "duplicate parameter '" + name + "'");
   seenParams.set(idx);
 
-  consumeToken();
-  return {std::make_unique<VariableElement>(*it)};
+  return create<ParameterElement>(*it);
 }
 
-FailureOr<std::unique_ptr<Element>>
-FormatParser::parseDirective(ParserContext ctx) {
+FailureOr<FormatElement *>
+DefFormatParser::parseDirectiveImpl(llvm::SMLoc loc, FormatToken::Kind kind,
+                                    Context ctx) {
 
-  switch (curToken.getKind()) {
+  switch (kind) {
   case FormatToken::kw_qualified:
-    return parseQualifiedDirective(ctx);
+    return parseQualifiedDirective(loc, ctx);
   case FormatToken::kw_params:
-    return parseParamsDirective();
+    return parseParamsDirective(loc);
   case FormatToken::kw_struct:
     if (ctx != TopLevelContext) {
       return emitError(
+          loc,
           "`struct` may only be used in the top-level section of the format");
     }
-    return parseStructDirective();
+    return parseStructDirective(loc);
+
   default:
-    return emitError("unknown directive in format: " + curToken.getSpelling());
+    return emitError(loc, "unsupported directive kind");
   }
 }
 
-FailureOr<std::unique_ptr<Element>>
-FormatParser::parseQualifiedDirective(ParserContext ctx) {
-  consumeToken();
+FailureOr<FormatElement *>
+DefFormatParser::parseQualifiedDirective(llvm::SMLoc loc, Context ctx) {
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")))
     return failure();
-  FailureOr<std::unique_ptr<Element>> var = parseElement(ctx);
+  FailureOr<FormatElement *> var = parseElement(ctx);
   if (failed(var))
     return var;
-  if (!isa<VariableElement>(*var))
-    return emitError("`qualified` argument list expected a variable");
-  cast<VariableElement>(var->get())->setShouldBeQualified();
+  if (!isa<ParameterElement>(*var))
+    return emitError(loc, "`qualified` argument list expected a variable");
+  cast<ParameterElement>(*var)->setShouldBeQualified();
   if (failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
     return failure();
   return var;
 }
 
-FailureOr<std::unique_ptr<Element>> FormatParser::parseParamsDirective() {
-  consumeToken();
+FailureOr<FormatElement *>
+DefFormatParser::parseParamsDirective(llvm::SMLoc loc) {
   /// Collect all of the attribute's or type's parameters.
-  SmallVector<std::unique_ptr<Element>> vars;
+  std::vector<FormatElement *> vars;
   /// Ensure that none of the parameters have already been captured.
   for (const auto &it : llvm::enumerate(def.getParameters())) {
     if (seenParams.test(it.index())) {
-      return emitError("`params` captures duplicate parameter: " +
+      return emitError(loc, "`params` captures duplicate parameter: " +
                                 it.value().getName());
     }
     seenParams.set(it.index());
-    vars.push_back(std::make_unique<VariableElement>(it.value()));
+    vars.push_back(create<ParameterElement>(it.value()));
   }
-  return {std::make_unique<ParamsDirective>(std::move(vars))};
+  return create<ParamsDirective>(std::move(vars));
 }
 
-FailureOr<std::unique_ptr<Element>> FormatParser::parseStructDirective() {
-  consumeToken();
+FailureOr<FormatElement *>
+DefFormatParser::parseStructDirective(llvm::SMLoc loc) {
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before `struct` argument list")))
     return failure();
 
   /// Parse variables captured by `struct`.
-  SmallVector<std::unique_ptr<Element>> vars;
+  std::vector<FormatElement *> vars;
 
   /// Parse first captured parameter or a `params` directive.
-  FailureOr<std::unique_ptr<Element>> var = parseElement(StructDirective);
-  if (failed(var) || !isa<VariableElement, ParamsDirective>(*var))
-    return emitError("`struct` argument list expected a variable or directive");
+  FailureOr<FormatElement *> var = parseElement(StructDirectiveContext);
+  if (failed(var) || !isa<VariableElement, ParamsDirective>(*var)) {
+    return emitError(loc,
+                     "`struct` argument list expected a variable or directive");
+  }
   if (isa<VariableElement>(*var)) {
     /// Parse any other parameters.
     vars.push_back(std::move(*var));
-    while (curToken.getKind() == FormatToken::comma) {
+    while (peekToken().is(FormatToken::comma)) {
       consumeToken();
-      var = parseElement(StructDirective);
+      var = parseElement(StructDirectiveContext);
       if (failed(var) || !isa<VariableElement>(*var))
-        return emitError("expected a variable in `struct` argument list");
+        return emitError(loc, "expected a variable in `struct` argument list");
       vars.push_back(std::move(*var));
     }
   } else {
     /// `struct(params)` captures all parameters in the attribute or type.
-    vars = cast<ParamsDirective>(var->get())->takeParams();
+    vars = cast<ParamsDirective>(*var)->takeParams();
   }
 
-  if (curToken.getKind() != FormatToken::r_paren)
-    return emitError("expected ')' at the end of an argument list");
+  if (failed(parseToken(FormatToken::r_paren,
+                        "expected ')' at the end of an argument list")))
+    return failure();
 
-  consumeToken();
-  return {std::make_unique<::StructDirective>(std::move(vars))};
+  return create<StructDirective>(std::move(vars));
 }
 
 //===----------------------------------------------------------------------===//
 // Interface
 //===----------------------------------------------------------------------===//
 
 void mlir::tblgen::generateAttrOrTypeFormat(const AttrOrTypeDef &def,
                                             MethodBody &parser,
                                             MethodBody &printer) {
   llvm::SourceMgr mgr;
   mgr.AddNewSourceBuffer(
       llvm::MemoryBuffer::getMemBuffer(*def.getAssemblyFormat()),
       llvm::SMLoc());
 
   /// Parse the custom assembly format>
-  FormatParser fmtParser(mgr, def);
+  DefFormatParser fmtParser(mgr, def);
   FailureOr<AttrOrTypeFormat> format = fmtParser.parse();
   if (failed(format)) {
     if (formatErrorIsFatal)
       PrintFatalError(def.getLoc(), "failed to parse assembly format");
     return;
   }
 
   /// Generate the parser and printer.
   format->genParser(parser);
   format->genPrinter(printer);
 }

mlir/tools/mlir-tblgen/FormatGen.h

 // and generators.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TOOLS_MLIRTBLGEN_FORMATGEN_H_
 #define MLIR_TOOLS_MLIRTBLGEN_FORMATGEN_H_
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/SMLoc.h"
+#include <vector>
 
 namespace llvm {
 class SourceMgr;
 } // namespace llvm
 
 namespace mlir {
 namespace tblgen {
 
public:
   StringRef getSpelling() const { return spelling; }
 
   /// Return the kind of this token.
   Kind getKind() const { return kind; }
 
   /// Return a location for this token.
   llvm::SMLoc getLoc() const;
 
+  /// Returns true if the token is of the given kind.
+  bool is(Kind kind) { return getKind() == kind; }
+
   /// Return if this token is a keyword.
   bool isKeyword() const {
     return getKind() > Kind::keyword_start && getKind() < Kind::keyword_end;
   }
 
 private:
   /// Discriminator that indicates the kind of token this is.
   Kind kind;
private:
   /// Location of the format string.
   llvm::SMLoc loc;
   /// Buffer containing the format string.
   StringRef curBuffer;
   /// Current pointer in the buffer.
   const char *curPtr;
 };
 
+//===----------------------------------------------------------------------===//
+// FormatElement
+//===----------------------------------------------------------------------===//
+
+/// This class represents a single format element.
+///
+/// If you squint and take a close look, you can see the outline of a `Format`
+/// dialect.
+class FormatElement {
+public:
+  /// The top-level kinds of format elements.
+  enum Kind { Literal, Variable, Whitespace, Directive, Optional };
+
+  /// Support LLVM-style RTTI.
+  static bool classof(const FormatElement *el) { return true; }
+
+  /// Get the element kind.
+  Kind getKind() const { return kind; }
+
+protected:
+  /// Create a format element with the given kind.
+  FormatElement(Kind kind) : kind(kind) {}
+
+private:
+  /// The kind of the element.
+  Kind kind;
+};
+
+/// The base class for all format elements. This class implements common methods
+/// for LLVM-style RTTI.
+template <FormatElement::Kind ElementKind>
+class FormatElementBase : public FormatElement {
+public:
+  /// Support LLVM-style RTTI.
+  static bool classof(const FormatElement *el) {
+    return ElementKind == el->getKind();
+  }
+
+protected:
+  /// Create a format element with the given kind.
+  FormatElementBase() : FormatElement(ElementKind) {}
+};
+
+/// This class represents a literal element. A literal is either one of the
+/// supported punctuation characters (e.g. `(` or `,`) or a string literal (e.g.
+/// `literal`).
+class LiteralElement : public FormatElementBase<FormatElement::Literal> {
+public:
+  /// Create a literal element with the given spelling.
+  explicit LiteralElement(StringRef spelling) : spelling(spelling) {}
+
+  /// Get the spelling of the literal.
+  StringRef getSpelling() const { return spelling; }
+
+private:
+  /// The spelling of the variable, i.e. the string contained within the
+  /// backticks.
+  StringRef spelling;
+};
+
+/// This class represents a variable element. A variable refers to some part of
+/// the object being parsed, e.g. an attribute or operand on an operation or a
+/// parameter on an attribute.
+class VariableElement : public FormatElementBase<FormatElement::Variable> {
+public:
+  /// These are the kinds of variables.
+  enum Kind { Attribute, Operand, Region, Result, Successor, Parameter };
+
+  /// Get the kind of variable.
+  Kind getKind() const { return kind; }
+
+protected:
+  /// Create a variable with a kind.
+  VariableElement(Kind kind) : kind(kind) {}
+
+private:
+  /// The kind of variable.
+  Kind kind;
+};
+
+/// Base class for variable elements. This class implements common methods for
+/// LLVM-style RTTI.
+template <VariableElement::Kind VariableKind>
+class VariableElementBase : public VariableElement {
+public:
+  /// An element is of this class if it is a variable and has the same variable
+  /// type.
+  static bool classof(const FormatElement *el) {
+    if (auto *varEl = dyn_cast<VariableElement>(el))
+      return VariableKind == varEl->getKind();
+    return false;
+  }
+
+protected:
+  /// Create a variable element with the given variable kind.
+  VariableElementBase() : VariableElement(VariableKind) {}
+};
+
+/// This class represents a whitespace element, e.g. a newline or space. It is a
+/// literal that is printed but never parsed. When the value is empty, i.e. ``,
+/// a space is elided where one would have been printed automatically.
+class WhitespaceElement : public FormatElementBase<FormatElement::Whitespace> {
+public:
+  /// Create a whitespace element.
+  explicit WhitespaceElement(StringRef value) : value(value) {}
+
+  /// Get the whitespace value.
+  StringRef getValue() const { return value; }
+
+private:
+  /// The value of the whitespace element. Can be empty.
+  StringRef value;
+};
+
+class DirectiveElement : public FormatElementBase<FormatElement::Directive> {
+public:
+  /// These are the kinds of directives.
+  enum Kind {
+    AttrDict,
+    Custom,
+    FunctionalType,
+    Operands,
+    Ref,
+    Regions,
+    Results,
+    Successors,
+    Type,
+    Params,
+    Struct
+  };
+
+  /// Get the directive kind.
+  Kind getKind() const { return kind; }
+
+protected:
+  /// Create a directive element with a kind.
+  DirectiveElement(Kind kind) : kind(kind) {}
+
+private:
+  /// The directive kind.
+  Kind kind;
+};
+
+/// Base class for directive elements. This class implements common methods for
+/// LLVM-style RTTI.
+template <DirectiveElement::Kind DirectiveKind>
+class DirectiveElementBase : public DirectiveElement {
+public:
+  /// Create a directive element with the specified kind.
+  DirectiveElementBase() : DirectiveElement(DirectiveKind) {}
+
+  /// A format element is of this class if it is a directive element and has the
+  /// same kind.
+  static bool classof(const FormatElement *el) {
+    if (auto *directiveEl = dyn_cast<DirectiveElement>(el))
+      return DirectiveKind == directiveEl->getKind();
+    return false;
+  }
+};
+
+/// This class represents a custom format directive that is implemented by the
+/// user in C++. The directive accepts a list of arguments that is passed to the
+/// C++ function.
+class CustomDirective : public DirectiveElementBase<DirectiveElement::Custom> {
+public:
+  /// Create a custom directive with a name and list of arguments.
+  CustomDirective(StringRef name, std::vector<FormatElement *> &&arguments)
+      : name(name), arguments(std::move(arguments)) {}
+
+  /// Get the custom directive name.
+  StringRef getName() const { return name; }
+
+  /// Get the arguments to the custom directive.
+  ArrayRef<FormatElement *> getArguments() const { return arguments; }
+
+private:
+  /// The name of the custom directive. The name is used to call two C++
+  /// methods: `parse{name}` and `print{name}` with the given arguments.
+  StringRef name;
+  /// The arguments with which to call the custom functions. These are either
+  /// variables (for which the functions are responsible for populating) or
+  /// references to variables.
+  std::vector<FormatElement *> arguments;
+};
+
+/// This class represents a group of elements that are optionally emitted based
+/// on an optional variable "anchor" and a group of elements that are emitted
+/// when the anchor element is not present.
+class OptionalElement : public FormatElementBase<FormatElement::Optional> {
+public:
+  /// Create an optional group with the given child elements.
+  OptionalElement(std::vector<FormatElement *> &&thenElements,
+                  std::vector<FormatElement *> &&elseElements,
+                  unsigned anchorIndex, unsigned parseStart)
+      : thenElements(std::move(thenElements)),
+        elseElements(std::move(elseElements)), anchorIndex(anchorIndex),
+        parseStart(parseStart) {}
+
+  /// Return the `then` elements of the optional group.
+  ArrayRef<FormatElement *> getThenElements() const { return thenElements; }
+
+  /// Return the `else` elements of the optional group.
+  ArrayRef<FormatElement *> getElseElements() const { return elseElements; }
+
+  /// Return the anchor of the optional group.
+  FormatElement *getAnchor() const { return thenElements[anchorIndex]; }
+
+  /// Return the index of the first element to be parsed.
+  unsigned getParseStart() const { return parseStart; }
+
+private:
+  /// The child elements emitted when the anchor is present.
+  std::vector<FormatElement *> thenElements;
+  /// The child elements emitted when the anchor is not present.
+  std::vector<FormatElement *> elseElements;
+  /// The index of the anchor element of the optional group within
+  /// `thenElements`.
+  unsigned anchorIndex;
+  /// The index of the first element that is parsed in `thenElements`. That is,
+  /// the first non-whitespace element.
+  unsigned parseStart;
+};
+
+//===----------------------------------------------------------------------===//
+// FormatParserBase
+//===----------------------------------------------------------------------===//
+
+/// Base class for a parser that implements an assembly format. This class
+/// defines a common assembly format syntax and the creation of format elements.
+/// Subclasses will need to implement parsing for the format elements they
+/// support.
+class FormatParser {
+public:
+  /// Vtable anchor.
+  virtual ~FormatParser();
+
+  /// Parse the assembly format.
+  FailureOr<std::vector<FormatElement *>> parse();
+
+protected:
+  /// The current context of the parser when parsing an element.
+  enum Context {
+    /// The element is being parsed in a "top-level" context, i.e. at the top of
+    /// the format or in an optional group.
+    TopLevelContext,
+    /// The element is being parsed as a custom directive child.
+    CustomDirectiveContext,
+    /// The element is being parsed as a type directive child.
+    TypeDirectiveContext,
+    /// The element is being parsed as a reference directive child.
+    RefDirectiveContext,
+    /// The element is being parsed as a struct directive child.
+    StructDirectiveContext
+  };
+
+  /// Create a format parser with the given source manager and a location.
+  explicit FormatParser(llvm::SourceMgr &mgr, llvm::SMLoc loc)
+      : lexer(mgr, loc), curToken(lexer.lexToken()) {}
+
+  /// Allocate and construct a format element.
+  template <typename FormatElementT, typename... Args>
+  FormatElementT *create(Args &&...args) {
+    FormatElementT *ptr = allocator.Allocate<FormatElementT>();
+    ::new (ptr) FormatElementT(std::forward<Args>(args)...);
+    return ptr;
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Element Parsing
+
+  /// Parse a single element of any kind.
+  FailureOr<FormatElement *> parseElement(Context ctx);
+  /// Parse a literal.
+  FailureOr<FormatElement *> parseLiteral(Context ctx);
+  /// Parse a variable.
+  FailureOr<FormatElement *> parseVariable(Context ctx);
+  /// Parse a directive.
+  FailureOr<FormatElement *> parseDirective(Context ctx);
+  /// Parse an optional group.
+  FailureOr<FormatElement *> parseOptionalGroup(Context ctx);
+
+  /// Parse a custom directive.
+  FailureOr<FormatElement *> parseCustomDirective(llvm::SMLoc loc, Context ctx);
+
+  /// Parse a format-specific variable kind.
+  virtual FailureOr<FormatElement *>
+  parseVariableImpl(llvm::SMLoc loc, StringRef name, Context ctx) = 0;
+  /// Parse a format-specific directive kind.
+  virtual FailureOr<FormatElement *>
+  parseDirectiveImpl(llvm::SMLoc loc, FormatToken::Kind kind, Context ctx) = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Format Verification
+
+  /// Verify that the format is well-formed.
+  virtual LogicalResult verify(llvm::SMLoc loc,
+                               ArrayRef<FormatElement *> elements) = 0;
+  /// Verify the arguments to a custom directive.
+  virtual LogicalResult
+  verifyCustomDirectiveArguments(llvm::SMLoc loc,
+                                 ArrayRef<FormatElement *> arguments) = 0;
+  /// Verify the elements of an optional group.
+  virtual LogicalResult
+  verifyOptionalGroupElements(llvm::SMLoc loc,
+                              ArrayRef<FormatElement *> elements,
+                              Optional<unsigned> anchorIndex) = 0;
+
+  //===--------------------------------------------------------------------===//
+  // Lexer Utilities
+
+  /// Emit an error at the given location.
+  LogicalResult emitError(llvm::SMLoc loc, const Twine &msg) {
+    lexer.emitError(loc, msg);
+    return failure();
+  }
+
+  /// Emit an error and a note at the given notation.
+  LogicalResult emitErrorAndNote(llvm::SMLoc loc, const Twine &msg,
+                                 const Twine &note) {
+    lexer.emitErrorAndNote(loc, msg, note);
+    return failure();
+  }
+
+  /// Parse a single token of the expected kind.
+  FailureOr<FormatToken> parseToken(FormatToken::Kind kind, const Twine &msg) {
+    if (!curToken.is(kind))
+      return emitError(curToken.getLoc(), msg);
+    FormatToken tok = curToken;
+    consumeToken();
+    return tok;
+  }
+
+  /// Advance the lexer to the next token.
+  void consumeToken() {
+    assert(!curToken.is(FormatToken::eof) && !curToken.is(FormatToken::error) &&
+           "shouldn't advance past EOF or errors");
+    curToken = lexer.lexToken();
+  }
+
+  /// Get the current token.
+  FormatToken peekToken() { return curToken; }
+
+private:
+  /// The format parser retains ownership of the format elements in a bump
+  /// pointer allocator.
+  llvm::BumpPtrAllocator allocator;
+  /// The format lexer to use.
+  FormatLexer lexer;
+  /// The current token in the lexer.
+  FormatToken curToken;
+};
+
+//===----------------------------------------------------------------------===//
+// Utility Functions
+//===----------------------------------------------------------------------===//
+
 /// Whether a space needs to be emitted before a literal. E.g., two keywords
 /// back-to-back require a space separator, but a keyword followed by '<' does
 /// not require a space.
 bool shouldEmitSpaceBefore(StringRef value, bool lastWasPunctuation);
 
 /// Returns true if the given string can be formatted as a keyword.
 bool canFormatStringAsKeyword(StringRef value,
                               function_ref<void(Twine)> emitError = nullptr);

mlir/tools/mlir-tblgen/FormatGen.cpp

auto kind =
           .Case("successors", FormatToken::kw_successors)
           .Case("type", FormatToken::kw_type)
           .Case("qualified", FormatToken::kw_qualified)
           .Default(FormatToken::identifier);
   return FormatToken(kind, str);
 }
 
 //===----------------------------------------------------------------------===//
+// FormatParser
+//===----------------------------------------------------------------------===//
+
+FormatParser::~FormatParser() = default;
+
+FailureOr<std::vector<FormatElement *>> FormatParser::parse() {
+  llvm::SMLoc loc = curToken.getLoc();
+
+  // Parse each of the format elements into the main format.
+  std::vector<FormatElement *> elements;
+  while (curToken.getKind() != FormatToken::eof) {
+    FailureOr<FormatElement *> element = parseElement(TopLevelContext);
+    if (failed(element))
+      return failure();
+    elements.push_back(*element);
+  }
+
+  // Verify the format.
+  if (failed(verify(loc, elements)))
+    return failure();
+  return elements;
+}
+
+//===----------------------------------------------------------------------===//
+// Element Parsing
+
+FailureOr<FormatElement *> FormatParser::parseElement(Context ctx) {
+  if (curToken.is(FormatToken::literal))
+    return parseLiteral(ctx);
+  if (curToken.is(FormatToken::variable))
+    return parseVariable(ctx);
+  if (curToken.isKeyword())
+    return parseDirective(ctx);
+  if (curToken.is(FormatToken::l_paren))
+    return parseOptionalGroup(ctx);
+  return emitError(curToken.getLoc(),
+                   "expected literal, variable, directive, or optional group");
+}
+
+FailureOr<FormatElement *> FormatParser::parseLiteral(Context ctx) {
+  FormatToken tok = curToken;
+  llvm::SMLoc loc = tok.getLoc();
+  consumeToken();
+
+  if (ctx != TopLevelContext) {
+    return emitError(
+        loc,
+        "literals may only be used in the top-level section of the format");
+  }
+  // Get the spelling without the surrounding backticks.
+  StringRef value = tok.getSpelling().drop_front().drop_back();
+
+  // The parsed literal is a space element (`` or ` `) or a newline.
+  if (value.empty() || value == " " || value == "\\n")
+    return create<WhitespaceElement>(value);
+
+  // Check that the parsed literal is valid.
+  if (!isValidLiteral(value, [&](Twine msg) {
+        (void)emitError(loc, "expected valid literal but got '" + value +
+                                 "': " + msg);
+      }))
+    return failure();
+  return create<LiteralElement>(value);
+}
+
+FailureOr<FormatElement *> FormatParser::parseVariable(Context ctx) {
+  FormatToken tok = curToken;
+  llvm::SMLoc loc = tok.getLoc();
+  consumeToken();
+
+  // Get the name of the variable without the leading `$`.
+  StringRef name = tok.getSpelling().drop_front();
+  return parseVariableImpl(loc, name, ctx);
+}
+
+FailureOr<FormatElement *> FormatParser::parseDirective(Context ctx) {
+  FormatToken tok = curToken;
+  llvm::SMLoc loc = tok.getLoc();
+  consumeToken();
+
+  if (tok.is(FormatToken::kw_custom))
+    return parseCustomDirective(loc, ctx);
+  return parseDirectiveImpl(loc, tok.getKind(), ctx);
+}
+
+FailureOr<FormatElement *> FormatParser::parseOptionalGroup(Context ctx) {
+  llvm::SMLoc loc = curToken.getLoc();
+  consumeToken();
+  if (ctx != TopLevelContext) {
+    return emitError(loc,
+                     "optional groups can only be used as top-level elements");
+  }
+
+  // Parse the child elements for this optional group.
+  std::vector<FormatElement *> thenElements, elseElements;
+  Optional<unsigned> anchorIndex;
+  do {
+    FailureOr<FormatElement *> element = parseElement(TopLevelContext);
+    if (failed(element))
+      return failure();
+    // Check for an anchor.
+    if (curToken.is(FormatToken::caret)) {
+      if (anchorIndex)
+        return emitError(curToken.getLoc(), "only one element can be marked as "
+                                            "the anchor of an optional group");
+      anchorIndex = thenElements.size();
+      consumeToken();
+    }
+    thenElements.push_back(*element);
+  } while (!curToken.is(FormatToken::r_paren));
+  consumeToken();
+
+  // Parse the `else` elements of this optional group.
+  if (curToken.is(FormatToken::colon)) {
+    consumeToken();
+    if (failed(
+            parseToken(FormatToken::l_paren,
+                       "expected '(' to start else branch of optional group")))
+      return failure();
+    do {
+      FailureOr<FormatElement *> element = parseElement(TopLevelContext);
+      if (failed(element))
+        return failure();
+      elseElements.push_back(*element);
+    } while (!curToken.is(FormatToken::r_paren));
+    consumeToken();
+  }
+  if (failed(parseToken(FormatToken::question,
+                        "expected '?' after optional group")))
+    return failure();
+
+  // The optional group is required to have an anchor.
+  if (!anchorIndex)
+    return emitError(loc, "optional group has no anchor element");
+
+  // Verify the child elements.
+  if (failed(verifyOptionalGroupElements(loc, thenElements, anchorIndex)) ||
+      failed(verifyOptionalGroupElements(loc, elseElements, llvm::None)))
+    return failure();
+
+  // Get the first parsable element. It must be an element that can be
+  // optionally-parsed.
+  auto parseBegin = llvm::find_if_not(thenElements, [](FormatElement *element) {
+    return isa<WhitespaceElement>(element);
+  });
+  if (!isa<LiteralElement, VariableElement>(*parseBegin)) {
+    return emitError(loc, "first parsable element of an optional group must be "
+                          "a literal or variable");
+  }
+
+  unsigned parseStart = std::distance(thenElements.begin(), parseBegin);
+  return create<OptionalElement>(std::move(thenElements),
+                                 std::move(elseElements), *anchorIndex,
+                                 parseStart);
+}
+
+FailureOr<FormatElement *> FormatParser::parseCustomDirective(llvm::SMLoc loc,
+                                                              Context ctx) {
+  if (ctx != TopLevelContext)
+    return emitError(loc, "'custom' is only valid as a top-level directive");
+
+  FailureOr<FormatToken> nameTok;
+  if (failed(parseToken(FormatToken::less,
+                        "expected '<' before custom directive name")) ||
+      failed(nameTok =
+                 parseToken(FormatToken::identifier,
+                            "expected custom directive name identifier")) ||
+      failed(parseToken(FormatToken::greater,
+                        "expected '>' after custom directive name")) ||
+      failed(parseToken(FormatToken::l_paren,
+                        "expected '(' before custom directive parameters")))
+    return failure();
+
+  // Parse the arguments.
+  std::vector<FormatElement *> arguments;
+  while (true) {
+    FailureOr<FormatElement *> argument = parseElement(CustomDirectiveContext);
+    if (failed(argument))
+      return failure();
+    arguments.push_back(*argument);
+    if (!curToken.is(FormatToken::comma))
+      break;
+    consumeToken();
+  }
+
+  if (failed(parseToken(FormatToken::r_paren,
+                        "expected ')' after custom directive parameters")))
+    return failure();
+
+  if (failed(verifyCustomDirectiveArguments(loc, arguments)))
+    return failure();
+  return create<CustomDirective>(nameTok->getSpelling(), std::move(arguments));
+}
+
+//===----------------------------------------------------------------------===//
 // Utility Functions
 //===----------------------------------------------------------------------===//
 
 bool mlir::tblgen::shouldEmitSpaceBefore(StringRef value,
                                          bool lastWasPunctuation) {
   if (value.size() != 1 && value != "->")
     return true;
   if (lastWasPunctuation)

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 //===- OpFormatGen.cpp - MLIR operation asm format generator --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "OpFormatGen.h"
 #include "FormatGen.h"
 #include "OpClass.h"
+#include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/TableGen/Class.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Interfaces.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/Trait.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
 #define DEBUG_TYPE "mlir-tblgen-opformatgen"
 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 //===----------------------------------------------------------------------===//
-// Element
-//===----------------------------------------------------------------------===//
+// VariableElement
 
 namespace {
-/// This class represents a single format element.
-class Element {
+/// This class represents an instance of an op variable element. A variable
+/// refers to something registered on the operation itself, e.g. an operand,
+/// result, attribute, region, or successor.
+template <typename VarT, VariableElement::Kind Kind>
+class OpVariableElement : public VariableElementBase<Kind> {
 public:
-  enum class Kind {
-    /// This element is a directive.
-    AttrDictDirective,
-    CustomDirective,
-    FunctionalTypeDirective,
-    OperandsDirective,
-    RefDirective,
-    RegionsDirective,
-    ResultsDirective,
-    SuccessorsDirective,
-    TypeDirective,
-
-    /// This element is a literal.
-    Literal,
-
-    /// This element is a whitespace.
-    Newline,
-    Space,
-
-    /// This element is an variable value.
-    AttributeVariable,
-    OperandVariable,
-    RegionVariable,
-    ResultVariable,
-    SuccessorVariable,
-
-    /// This element is an optional element.
-    Optional,
-  };
-  Element(Kind kind) : kind(kind) {}
-  virtual ~Element() = default;
+  using Base = OpVariableElement<VarT, Kind>;
 
-  /// Return the kind of this element.
-  Kind getKind() const { return kind; }
+  /// Create an op variable element with the variable value.
+  OpVariableElement(const VarT *var) : var(var) {}
 
-private:
-  /// The kind of this element.
-  Kind kind;
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
-// VariableElement
-
-namespace {
-/// This class represents an instance of an variable element. A variable refers
-/// to something registered on the operation itself, e.g. an argument, result,
-/// etc.
-template <typename VarT, Element::Kind kindVal>
-class VariableElement : public Element {
-public:
-  VariableElement(const VarT *var) : Element(kindVal), var(var) {}
-  static bool classof(const Element *element) {
-    return element->getKind() == kindVal;
-  }
+  /// Get the variable.
   const VarT *getVar() { return var; }
 
 protected:
+  /// The op variable, e.g. a type or attribute constraint.
   const VarT *var;
 };
 
 /// This class represents a variable that refers to an attribute argument.
 struct AttributeVariable
-    : public VariableElement<NamedAttribute, Element::Kind::AttributeVariable> {
-  using VariableElement<NamedAttribute,
-                        Element::Kind::AttributeVariable>::VariableElement;
+    : public OpVariableElement<NamedAttribute, VariableElement::Attribute> {
+  using Base::Base;
 
   /// Return the constant builder call for the type of this attribute, or None
   /// if it doesn't have one.
-  Optional<StringRef> getTypeBuilder() const {
-    Optional<Type> attrType = var->attr.getValueType();
+  llvm::Optional<StringRef> getTypeBuilder() const {
+    llvm::Optional<Type> attrType = var->attr.getValueType();
     return attrType ? attrType->getBuilderCall() : llvm::None;
   }
 
   /// Return if this attribute refers to a UnitAttr.
   bool isUnitAttr() const {
     return var->attr.getBaseAttr().getAttrDefName() == "UnitAttr";
   }
 
   /// Indicate if this attribute is printed "qualified" (that is it is
   /// prefixed with the `#dialect.mnemonic`).
   bool shouldBeQualified() { return shouldBeQualifiedFlag; }
   void setShouldBeQualified(bool qualified = true) {
     shouldBeQualifiedFlag = qualified;
   }
 
 private:
   bool shouldBeQualifiedFlag = false;
 };
 
 /// This class represents a variable that refers to an operand argument.
 using OperandVariable =
-    VariableElement<NamedTypeConstraint, Element::Kind::OperandVariable>;
-
-/// This class represents a variable that refers to a region.
-using RegionVariable =
-    VariableElement<NamedRegion, Element::Kind::RegionVariable>;
+    OpVariableElement<NamedTypeConstraint, VariableElement::Operand>;
 
 /// This class represents a variable that refers to a result.
 using ResultVariable =
-    VariableElement<NamedTypeConstraint, Element::Kind::ResultVariable>;
+    OpVariableElement<NamedTypeConstraint, VariableElement::Result>;
+
+/// This class represents a variable that refers to a region.
+using RegionVariable = OpVariableElement<NamedRegion, VariableElement::Region>;
 
 /// This class represents a variable that refers to a successor.
 using SuccessorVariable =
-    VariableElement<NamedSuccessor, Element::Kind::SuccessorVariable>;
+    OpVariableElement<NamedSuccessor, VariableElement::Successor>;
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // DirectiveElement
 
 namespace {
-/// This class implements single kind directives.
-template <Element::Kind type> class DirectiveElement : public Element {
-public:
-  DirectiveElement() : Element(type){};
-  static bool classof(const Element *ele) { return ele->getKind() == type; }
-};
 /// This class represents the `operands` directive. This directive represents
 /// all of the operands of an operation.
-using OperandsDirective = DirectiveElement<Element::Kind::OperandsDirective>;
-
-/// This class represents the `regions` directive. This directive represents
-/// all of the regions of an operation.
-using RegionsDirective = DirectiveElement<Element::Kind::RegionsDirective>;
+using OperandsDirective = DirectiveElementBase<DirectiveElement::Operands>;
 
 /// This class represents the `results` directive. This directive represents
 /// all of the results of an operation.
-using ResultsDirective = DirectiveElement<Element::Kind::ResultsDirective>;
+using ResultsDirective = DirectiveElementBase<DirectiveElement::Results>;
+
+/// This class represents the `regions` directive. This directive represents
+/// all of the regions of an operation.
+using RegionsDirective = DirectiveElementBase<DirectiveElement::Regions>;
 
 /// This class represents the `successors` directive. This directive represents
 /// all of the successors of an operation.
-using SuccessorsDirective =
-    DirectiveElement<Element::Kind::SuccessorsDirective>;
+using SuccessorsDirective = DirectiveElementBase<DirectiveElement::Successors>;
 
 /// This class represents the `attr-dict` directive. This directive represents
 /// the attribute dictionary of the operation.
 class AttrDictDirective
-    : public DirectiveElement<Element::Kind::AttrDictDirective> {
+    : public DirectiveElementBase<DirectiveElement::AttrDict> {
 public:
   explicit AttrDictDirective(bool withKeyword) : withKeyword(withKeyword) {}
+
+  /// Return whether the dictionary should be printed with the 'attributes'
+  /// keyword.
   bool isWithKeyword() const { return withKeyword; }
 
 private:
   /// If the dictionary should be printed with the 'attributes' keyword.
   bool withKeyword;
 };
 
-/// This class represents a custom format directive that is implemented by the
-/// user in C++.
-class CustomDirective : public Element {
-public:
-  CustomDirective(StringRef name,
-                  std::vector<std::unique_ptr<Element>> &&arguments)
-      : Element{Kind::CustomDirective}, name(name),
-        arguments(std::move(arguments)) {}
-
-  static bool classof(const Element *element) {
-    return element->getKind() == Kind::CustomDirective;
-  }
-
-  /// Return the name of the custom directive.
-  StringRef getName() const { return name; }
-
-  /// Return the arguments to the custom directive.
-  auto getArguments() const { return llvm::make_pointee_range(arguments); }
-
-private:
-  /// The user provided name of the directive.
-  StringRef name;
-
-  /// The arguments to the custom directive.
-  std::vector<std::unique_ptr<Element>> arguments;
-};
-
 /// This class represents the `functional-type` directive. This directive takes
 /// two arguments and formats them, respectively, as the inputs and results of a
 /// FunctionType.
 class FunctionalTypeDirective
-    : public DirectiveElement<Element::Kind::FunctionalTypeDirective> {
+    : public DirectiveElementBase<DirectiveElement::FunctionalType> {
 public:
-  FunctionalTypeDirective(std::unique_ptr<Element> inputs,
-                          std::unique_ptr<Element> results)
-      : inputs(std::move(inputs)), results(std::move(results)) {}
-  Element *getInputs() const { return inputs.get(); }
-  Element *getResults() const { return results.get(); }
+  FunctionalTypeDirective(FormatElement *inputs, FormatElement *results)
+      : inputs(inputs), results(results) {}
+
+  FormatElement *getInputs() const { return inputs; }
+  FormatElement *getResults() const { return results; }
 
 private:
   /// The input and result arguments.
-  std::unique_ptr<Element> inputs, results;
+  FormatElement *inputs, *results;
 };
 
 /// This class represents the `ref` directive.
-class RefDirective : public DirectiveElement<Element::Kind::RefDirective> {
+class RefDirective : public DirectiveElementBase<DirectiveElement::Ref> {
 public:
-  RefDirective(std::unique_ptr<Element> arg) : operand(std::move(arg)) {}
-  Element *getOperand() const { return operand.get(); }
+  RefDirective(FormatElement *arg) : arg(arg) {}
+
+  FormatElement *getArg() const { return arg; }
 
 private:
-  /// The operand that is used to format the directive.
-  std::unique_ptr<Element> operand;
+  /// The argument that is used to format the directive.
+  FormatElement *arg;
 };
 
 /// This class represents the `type` directive.
-class TypeDirective : public DirectiveElement<Element::Kind::TypeDirective> {
+class TypeDirective : public DirectiveElementBase<DirectiveElement::Type> {
 public:
-  TypeDirective(std::unique_ptr<Element> arg) : operand(std::move(arg)) {}
-  Element *getOperand() const { return operand.get(); }
+  TypeDirective(FormatElement *arg) : arg(arg) {}
+
+  FormatElement *getArg() const { return arg; }
 
   /// Indicate if this type is printed "qualified" (that is it is
   /// prefixed with the `!dialect.mnemonic`).
   bool shouldBeQualified() { return shouldBeQualifiedFlag; }
   void setShouldBeQualified(bool qualified = true) {
     shouldBeQualifiedFlag = qualified;
   }
 
 private:
-  /// The operand that is used to format the directive.
-  std::unique_ptr<Element> operand;
+  /// The argument that is used to format the directive.
+  FormatElement *arg;
 
   bool shouldBeQualifiedFlag = false;
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// LiteralElement
-
-namespace {
-/// This class represents an instance of a literal element.
-class LiteralElement : public Element {
-public:
-  LiteralElement(StringRef literal)
-      : Element{Kind::Literal}, literal(literal) {}
-  static bool classof(const Element *element) {
-    return element->getKind() == Kind::Literal;
-  }
-
-  /// Return the literal for this element.
-  StringRef getLiteral() const { return literal; }
-
-private:
-  /// The spelling of the literal for this element.
-  StringRef literal;
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
-// WhitespaceElement
-
-namespace {
-/// This class represents a whitespace element, e.g. newline or space. It's a
-/// literal that is printed but never parsed.
-class WhitespaceElement : public Element {
-public:
-  WhitespaceElement(Kind kind) : Element{kind} {}
-  static bool classof(const Element *element) {
-    Kind kind = element->getKind();
-    return kind == Kind::Newline || kind == Kind::Space;
-  }
-};
-
-/// This class represents an instance of a newline element. It's a literal that
-/// prints a newline. It is ignored by the parser.
-class NewlineElement : public WhitespaceElement {
-public:
-  NewlineElement() : WhitespaceElement(Kind::Newline) {}
-  static bool classof(const Element *element) {
-    return element->getKind() == Kind::Newline;
-  }
-};
-
-/// This class represents an instance of a space element. It's a literal that
-/// prints or omits printing a space. It is ignored by the parser.
-class SpaceElement : public WhitespaceElement {
-public:
-  SpaceElement(bool value) : WhitespaceElement(Kind::Space), value(value) {}
-  static bool classof(const Element *element) {
-    return element->getKind() == Kind::Space;
-  }
-
-  /// Returns true if this element should print as a space. Otherwise, the
-  /// element should omit printing a space between the surrounding elements.
-  bool getValue() const { return value; }
-
-private:
-  bool value;
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
-// OptionalElement
-
-namespace {
-/// This class represents a group of elements that are optionally emitted based
-/// upon an optional variable of the operation, and a group of elements that are
-/// emotted when the anchor element is not present.
-class OptionalElement : public Element {
-public:
-  OptionalElement(std::vector<std::unique_ptr<Element>> &&thenElements,
-                  std::vector<std::unique_ptr<Element>> &&elseElements,
-                  unsigned anchor, unsigned parseStart)
-      : Element{Kind::Optional}, thenElements(std::move(thenElements)),
-        elseElements(std::move(elseElements)), anchor(anchor),
-        parseStart(parseStart) {}
-  static bool classof(const Element *element) {
-    return element->getKind() == Kind::Optional;
-  }
-
-  /// Return the `then` elements of this grouping.
-  auto getThenElements() const {
-    return llvm::make_pointee_range(thenElements);
-  }
-
-  /// Return the `else` elements of this grouping.
-  auto getElseElements() const {
-    return llvm::make_pointee_range(elseElements);
-  }
-
-  /// Return the anchor of this optional group.
-  Element *getAnchor() const { return thenElements[anchor].get(); }
-
-  /// Return the index of the first element that needs to be parsed.
-  unsigned getParseStart() const { return parseStart; }
-
-private:
-  /// The child elements of `then` branch of this optional.
-  std::vector<std::unique_ptr<Element>> thenElements;
-  /// The child elements of `else` branch of this optional.
-  std::vector<std::unique_ptr<Element>> elseElements;
-  /// The index of the element that acts as the anchor for the optional group.
-  unsigned anchor;
-  /// The index of the first element that is parsed (is not a
-  /// WhitespaceElement).
-  unsigned parseStart;
-};
-} // namespace
-
-//===----------------------------------------------------------------------===//
 // OperationFormat
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 using ConstArgument =
     llvm::PointerUnion<const NamedAttribute *, const NamedTypeConstraint *>;
 
OperationFormat(const Operator &op)
 
     hasSingleBlockTrait =
         hasImplicitTermTrait || op.getTrait("::mlir::OpTrait::SingleBlock");
   }
 
   /// Generate the operation parser from this format.
   void genParser(Operator &op, OpClass &opClass);
   /// Generate the parser code for a specific format element.
-  void genElementParser(Element *element, MethodBody &body,
+  void genElementParser(FormatElement *element, MethodBody &body,
                         FmtContext &attrTypeCtx,
                         GenContext genCtx = GenContext::Normal);
   /// Generate the C++ to resolve the types of operands and results during
   /// parsing.
   void genParserTypeResolution(Operator &op, MethodBody &body);
   /// Generate the C++ to resolve the types of the operands during parsing.
   void genParserOperandTypeResolution(
       Operator &op, MethodBody &body,
       function_ref<void(TypeResolution &, StringRef)> emitTypeResolver);
   /// Generate the C++ to resolve regions during parsing.
   void genParserRegionResolution(Operator &op, MethodBody &body);
   /// Generate the C++ to resolve successors during parsing.
   void genParserSuccessorResolution(Operator &op, MethodBody &body);
   /// Generate the C++ to handling variadic segment size traits.
   void genParserVariadicSegmentResolution(Operator &op, MethodBody &body);
 
   /// Generate the operation printer from this format.
   void genPrinter(Operator &op, OpClass &opClass);
 
   /// Generate the printer code for a specific format element.
-  void genElementPrinter(Element *element, MethodBody &body, Operator &op,
+  void genElementPrinter(FormatElement *element, MethodBody &body, Operator &op,
                          bool &shouldEmitSpace, bool &lastWasPunctuation);
 
   /// The various elements in this format.
-  std::vector<std::unique_ptr<Element>> elements;
+  std::vector<FormatElement *> elements;
 
   /// A flag indicating if all operand/result types were seen. If the format
   /// contains these, it can not contain individual type resolvers.
   bool allOperands, allOperandTypes, allResultTypes;
 
   /// A flag indicating if this operation infers its result types
   bool infersResultTypes;
 
if (var->isVariadicOfVariadic())
     return ArgumentLengthKind::VariadicOfVariadic;
   if (var->isVariadic())
     return ArgumentLengthKind::Variadic;
   return ArgumentLengthKind::Single;
 }
 
 /// Get the name used for the type list for the given type directive operand.
 /// 'lengthKind' to the corresponding kind for the given argument.
-static StringRef getTypeListName(Element *arg, ArgumentLengthKind &lengthKind) {
+static StringRef getTypeListName(FormatElement *arg,
+                                 ArgumentLengthKind &lengthKind) {
   if (auto *operand = dyn_cast<OperandVariable>(arg)) {
     lengthKind = getArgumentLengthKind(operand->getVar());
     return operand->getVar()->name;
   }
   if (auto *result = dyn_cast<ResultVariable>(arg)) {
     lengthKind = getArgumentLengthKind(result->getVar());
     return result->getVar()->name;
   }
body << (StringRef)StringSwitch<StringRef>(value)
               .Case("[", "LSquare()")
               .Case("]", "RSquare()")
               .Case("?", "Question()")
               .Case("+", "Plus()")
               .Case("*", "Star()");
 }
 
 /// Generate the storage code required for parsing the given element.
-static void genElementParserStorage(Element *element, const Operator &op,
+static void genElementParserStorage(FormatElement *element, const Operator &op,
                                     MethodBody &body) {
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    auto elements = optional->getThenElements();
+    ArrayRef<FormatElement *> elements = optional->getThenElements();
 
     // If the anchor is a unit attribute, it won't be parsed directly so elide
     // it.
     auto *anchor = dyn_cast<AttributeVariable>(optional->getAnchor());
-    Element *elidedAnchorElement = nullptr;
-    if (anchor && anchor != &*elements.begin() && anchor->isUnitAttr())
+    FormatElement *elidedAnchorElement = nullptr;
+    if (anchor && anchor != elements.front() && anchor->isUnitAttr())
       elidedAnchorElement = anchor;
-    for (auto &childElement : elements)
-      if (&childElement != elidedAnchorElement)
-        genElementParserStorage(&childElement, op, body);
-    for (auto &childElement : optional->getElseElements())
-      genElementParserStorage(&childElement, op, body);
+    for (FormatElement *childElement : elements)
+      if (childElement != elidedAnchorElement)
+        genElementParserStorage(childElement, op, body);
+    for (FormatElement *childElement : optional->getElseElements())
+      genElementParserStorage(childElement, op, body);
 
   } else if (auto *custom = dyn_cast<CustomDirective>(element)) {
-    for (auto &paramElement : custom->getArguments())
-      genElementParserStorage(&paramElement, op, body);
+    for (FormatElement *paramElement : custom->getArguments())
+      genElementParserStorage(paramElement, op, body);
 
   } else if (isa<OperandsDirective>(element)) {
     body << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
             "allOperands;\n";
 
   } else if (isa<RegionsDirective>(element)) {
     body << "  ::llvm::SmallVector<std::unique_ptr<::mlir::Region>, 2> "
             "fullRegions;\n";
if (successor->getVar()->isVariadic()) {
                             "{0}Successors;\n",
                             name);
     } else {
       body << llvm::formatv("  ::mlir::Block *{0}Successor = nullptr;\n", name);
     }
 
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     ArgumentLengthKind lengthKind;
-    StringRef name = getTypeListName(dir->getOperand(), lengthKind);
+    StringRef name = getTypeListName(dir->getArg(), lengthKind);
     if (lengthKind != ArgumentLengthKind::Single)
       body << "  ::mlir::SmallVector<::mlir::Type, 1> " << name << "Types;\n";
     else
       body << llvm::formatv("  ::mlir::Type {0}RawTypes[1];\n", name)
            << llvm::formatv(
                   "  ::llvm::ArrayRef<::mlir::Type> {0}Types({0}RawTypes);\n",
                   name);
   } else if (auto *dir = dyn_cast<FunctionalTypeDirective>(element)) {
     ArgumentLengthKind ignored;
     body << "  ::llvm::ArrayRef<::mlir::Type> "
          << getTypeListName(dir->getInputs(), ignored) << "Types;\n";
     body << "  ::llvm::ArrayRef<::mlir::Type> "
          << getTypeListName(dir->getResults(), ignored) << "Types;\n";
   }
 }
 
 /// Generate the parser for a parameter to a custom directive.
-static void genCustomParameterParser(Element &param, MethodBody &body) {
-  if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
+static void genCustomParameterParser(FormatElement *param, MethodBody &body) {
+  if (auto *attr = dyn_cast<AttributeVariable>(param)) {
     body << attr->getVar()->name << "Attr";
-  } else if (isa<AttrDictDirective>(&param)) {
+  } else if (isa<AttrDictDirective>(param)) {
     body << "result.attributes";
-  } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
+  } else if (auto *operand = dyn_cast<OperandVariable>(param)) {
     StringRef name = operand->getVar()->name;
     ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
     if (lengthKind == ArgumentLengthKind::VariadicOfVariadic)
       body << llvm::formatv("{0}OperandGroups", name);
     else if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv("{0}Operands", name);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv("{0}Operand", name);
     else
       body << formatv("{0}RawOperands[0]", name);
 
-  } else if (auto *region = dyn_cast<RegionVariable>(&param)) {
+  } else if (auto *region = dyn_cast<RegionVariable>(param)) {
     StringRef name = region->getVar()->name;
     if (region->getVar()->isVariadic())
       body << llvm::formatv("{0}Regions", name);
     else
       body << llvm::formatv("*{0}Region", name);
 
-  } else if (auto *successor = dyn_cast<SuccessorVariable>(&param)) {
+  } else if (auto *successor = dyn_cast<SuccessorVariable>(param)) {
     StringRef name = successor->getVar()->name;
     if (successor->getVar()->isVariadic())
       body << llvm::formatv("{0}Successors", name);
     else
       body << llvm::formatv("{0}Successor", name);
 
-  } else if (auto *dir = dyn_cast<RefDirective>(&param)) {
-    genCustomParameterParser(*dir->getOperand(), body);
+  } else if (auto *dir = dyn_cast<RefDirective>(param)) {
+    genCustomParameterParser(dir->getArg(), body);
 
-  } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
+  } else if (auto *dir = dyn_cast<TypeDirective>(param)) {
     ArgumentLengthKind lengthKind;
-    StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+    StringRef listName = getTypeListName(dir->getArg(), lengthKind);
     if (lengthKind == ArgumentLengthKind::VariadicOfVariadic)
       body << llvm::formatv("{0}TypeGroups", listName);
     else if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv("{0}Types", listName);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv("{0}Type", listName);
     else
       body << formatv("{0}RawTypes[0]", listName);
   } else {
     llvm_unreachable("unknown custom directive parameter");
   }
 }
 
 /// Generate the parser for a custom directive.
 static void genCustomDirectiveParser(CustomDirective *dir, MethodBody &body) {
   body << "  {\n";
 
   // Preprocess the directive variables.
   // * Add a local variable for optional operands and types. This provides a
   //   better API to the user defined parser methods.
   // * Set the location of operand variables.
-  for (Element &param : dir->getArguments()) {
-    if (auto *operand = dyn_cast<OperandVariable>(&param)) {
+  for (FormatElement *param : dir->getArguments()) {
+    if (auto *operand = dyn_cast<OperandVariable>(param)) {
       auto *var = operand->getVar();
       body << "    " << var->name
            << "OperandsLoc = parser.getCurrentLocation();\n";
       if (var->isOptional()) {
         body << llvm::formatv(
             "    llvm::Optional<::mlir::OpAsmParser::OperandType> "
             "{0}Operand;\n",
             var->name);
       } else if (var->isVariadicOfVariadic()) {
         body << llvm::formatv("    "
                               "llvm::SmallVector<llvm::SmallVector<::mlir::"
                               "OpAsmParser::OperandType>> "
                               "{0}OperandGroups;\n",
                               var->name);
       }
-    } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
+    } else if (auto *dir = dyn_cast<TypeDirective>(param)) {
       ArgumentLengthKind lengthKind;
-      StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+      StringRef listName = getTypeListName(dir->getArg(), lengthKind);
       if (lengthKind == ArgumentLengthKind::Optional) {
         body << llvm::formatv("    ::mlir::Type {0}Type;\n", listName);
       } else if (lengthKind == ArgumentLengthKind::VariadicOfVariadic) {
         body << llvm::formatv(
             "    llvm::SmallVector<llvm::SmallVector<::mlir::Type>> "
             "{0}TypeGroups;\n",
             listName);
       }
-    } else if (auto *dir = dyn_cast<RefDirective>(&param)) {
-      Element *input = dir->getOperand();
+    } else if (auto *dir = dyn_cast<RefDirective>(param)) {
+      FormatElement *input = dir->getArg();
       if (auto *operand = dyn_cast<OperandVariable>(input)) {
         if (!operand->getVar()->isOptional())
           continue;
         body << llvm::formatv(
             "    {0} {1}Operand = {1}Operands.empty() ? {0}() : "
             "{1}Operands[0];\n",
             "llvm::Optional<::mlir::OpAsmParser::OperandType>",
             operand->getVar()->name);
 
       } else if (auto *type = dyn_cast<TypeDirective>(input)) {
         ArgumentLengthKind lengthKind;
-        StringRef listName = getTypeListName(type->getOperand(), lengthKind);
+        StringRef listName = getTypeListName(type->getArg(), lengthKind);
         if (lengthKind == ArgumentLengthKind::Optional) {
           body << llvm::formatv("    ::mlir::Type {0}Type = {0}Types.empty() ? "
                                 "::mlir::Type() : {0}Types[0];\n",
                                 listName);
         }
       }
     }
   }
 
   body << "    if (parse" << dir->getName() << "(parser";
-  for (Element &param : dir->getArguments()) {
+  for (FormatElement *param : dir->getArguments()) {
     body << ", ";
     genCustomParameterParser(param, body);
   }
 
   body << "))\n"
        << "      return ::mlir::failure();\n";
 
   // After parsing, add handling for any of the optional constructs.
-  for (Element &param : dir->getArguments()) {
-    if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
+  for (FormatElement *param : dir->getArguments()) {
+    if (auto *attr = dyn_cast<AttributeVariable>(param)) {
       const NamedAttribute *var = attr->getVar();
       if (var->attr.isOptional())
         body << llvm::formatv("    if ({0}Attr)\n  ", var->name);
 
       body << llvm::formatv("    result.addAttribute(\"{0}\", {0}Attr);\n",
                             var->name);
-    } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
+    } else if (auto *operand = dyn_cast<OperandVariable>(param)) {
       const NamedTypeConstraint *var = operand->getVar();
       if (var->isOptional()) {
         body << llvm::formatv("    if ({0}Operand.hasValue())\n"
                               "      {0}Operands.push_back(*{0}Operand);\n",
                               var->name);
       } else if (var->isVariadicOfVariadic()) {
         body << llvm::formatv(
             "    for (const auto &subRange : {0}OperandGroups) {{\n"
             "      {0}Operands.append(subRange.begin(), subRange.end());\n"
             "      {0}OperandGroupSizes.push_back(subRange.size());\n"
             "    }\n",
             var->name, var->constraint.getVariadicOfVariadicSegmentSizeAttr());
       }
-    } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
+    } else if (auto *dir = dyn_cast<TypeDirective>(param)) {
       ArgumentLengthKind lengthKind;
-      StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+      StringRef listName = getTypeListName(dir->getArg(), lengthKind);
       if (lengthKind == ArgumentLengthKind::Optional) {
         body << llvm::formatv("    if ({0}Type)\n"
                               "      {0}Types.push_back({0}Type);\n",
                               listName);
       } else if (lengthKind == ArgumentLengthKind::VariadicOfVariadic) {
         body << llvm::formatv(
             "    for (const auto &subRange : {0}TypeGroups)\n"
             "      {0}Types.append(subRange.begin(), subRange.end());\n",
void OperationFormat::genParser(Operator &op, OpClass &opClass) {
 
   auto *method = opClass.addStaticMethod("::mlir::ParseResult", "parse",
                                          std::move(paramList));
   auto &body = method->body();
 
   // Generate variables to store the operands and type within the format. This
   // allows for referencing these variables in the presence of optional
   // groupings.
-  for (auto &element : elements)
-    genElementParserStorage(&*element, op, body);
+  for (FormatElement *element : elements)
+    genElementParserStorage(element, op, body);
 
   // A format context used when parsing attributes with buildable types.
   FmtContext attrTypeCtx;
   attrTypeCtx.withBuilder("parser.getBuilder()");
 
   // Generate parsers for each of the elements.
-  for (auto &element : elements)
-    genElementParser(element.get(), body, attrTypeCtx);
+  for (FormatElement *element : elements)
+    genElementParser(element, body, attrTypeCtx);
 
   // Generate the code to resolve the operand/result types and successors now
   // that they have been parsed.
   genParserRegionResolution(op, body);
   genParserSuccessorResolution(op, body);
   genParserVariadicSegmentResolution(op, body);
   genParserTypeResolution(op, body);
 
   body << "  return ::mlir::success();\n";
 }
 
-void OperationFormat::genElementParser(Element *element, MethodBody &body,
+void OperationFormat::genElementParser(FormatElement *element, MethodBody &body,
                                        FmtContext &attrTypeCtx,
                                        GenContext genCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    auto elements = llvm::drop_begin(optional->getThenElements(),
-                                     optional->getParseStart());
+    ArrayRef<FormatElement *> elements =
+        optional->getThenElements().drop_front(optional->getParseStart());
 
     // Generate a special optional parser for the first element to gate the
     // parsing of the rest of the elements.
-    Element *firstElement = &*elements.begin();
+    FormatElement *firstElement = elements.front();
     if (auto *attrVar = dyn_cast<AttributeVariable>(firstElement)) {
       genElementParser(attrVar, body, attrTypeCtx);
       body << "  if (" << attrVar->getVar()->name << "Attr) {\n";
     } else if (auto *literal = dyn_cast<LiteralElement>(firstElement)) {
       body << "  if (succeeded(parser.parseOptional";
-      genLiteralParser(literal->getLiteral(), body);
+      genLiteralParser(literal->getSpelling(), body);
       body << ")) {\n";
     } else if (auto *opVar = dyn_cast<OperandVariable>(firstElement)) {
       genElementParser(opVar, body, attrTypeCtx);
       body << "  if (!" << opVar->getVar()->name << "Operands.empty()) {\n";
     } else if (auto *regionVar = dyn_cast<RegionVariable>(firstElement)) {
       const NamedRegion *region = regionVar->getVar();
       if (region->isVariadic()) {
         genElementParser(regionVar, body, attrTypeCtx);
         body << "  if (!" << region->name << "Regions.empty()) {\n";
       } else {
         body << llvm::formatv(optionalRegionParserCode, region->name);
         body << "  if (!" << region->name << "Region->empty()) {\n  ";
         if (hasImplicitTermTrait)
           body << llvm::formatv(regionEnsureTerminatorParserCode, region->name);
         else if (hasSingleBlockTrait)
           body << llvm::formatv(regionEnsureSingleBlockParserCode,
                                 region->name);
       }
     }
 
     // If the anchor is a unit attribute, we don't need to print it. When
     // parsing, we will add this attribute if this group is present.
-    Element *elidedAnchorElement = nullptr;
+    FormatElement *elidedAnchorElement = nullptr;
     auto *anchorAttr = dyn_cast<AttributeVariable>(optional->getAnchor());
     if (anchorAttr && anchorAttr != firstElement && anchorAttr->isUnitAttr()) {
       elidedAnchorElement = anchorAttr;
 
       // Add the anchor unit attribute to the operation state.
       body << "    result.addAttribute(\"" << anchorAttr->getVar()->name
            << "\", parser.getBuilder().getUnitAttr());\n";
     }
 
     // Generate the rest of the elements inside an optional group. Elements in
     // an optional group after the guard are parsed as required.
-    for (Element &childElement : llvm::drop_begin(elements, 1)) {
-      if (&childElement != elidedAnchorElement) {
-        genElementParser(&childElement, body, attrTypeCtx,
-                         GenContext::Optional);
-      }
-    }
+    for (FormatElement *childElement : llvm::drop_begin(elements, 1))
+      if (childElement != elidedAnchorElement)
+        genElementParser(childElement, body, attrTypeCtx, GenContext::Optional);
     body << "  }";
 
     // Generate the else elements.
     auto elseElements = optional->getElseElements();
     if (!elseElements.empty()) {
       body << " else {\n";
-      for (Element &childElement : elseElements)
-        genElementParser(&childElement, body, attrTypeCtx);
+      for (FormatElement *childElement : elseElements)
+        genElementParser(childElement, body, attrTypeCtx);
       body << "  }";
     }
     body << "\n";
 
     /// Literals.
   } else if (LiteralElement *literal = dyn_cast<LiteralElement>(element)) {
     body << "  if (parser.parse";
-    genLiteralParser(literal->getLiteral(), body);
+    genLiteralParser(literal->getSpelling(), body);
     body << ")\n    return ::mlir::failure();\n";
 
     /// Whitespaces.
   } else if (isa<WhitespaceElement>(element)) {
     // Nothing to parse.
 
     /// Arguments.
   } else if (auto *attr = dyn_cast<AttributeVariable>(element)) {
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     else if (hasSingleBlockTrait)
       body << llvm::formatv(regionListEnsureSingleBlockParserCode, "full");
 
   } else if (isa<SuccessorsDirective>(element)) {
     body << llvm::formatv(successorListParserCode, "full");
 
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     ArgumentLengthKind lengthKind;
-    StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+    StringRef listName = getTypeListName(dir->getArg(), lengthKind);
     if (lengthKind == ArgumentLengthKind::VariadicOfVariadic) {
       body << llvm::formatv(variadicOfVariadicTypeParserCode, listName);
     } else if (lengthKind == ArgumentLengthKind::Variadic) {
       body << llvm::formatv(variadicTypeParserCode, listName);
     } else if (lengthKind == ArgumentLengthKind::Optional) {
       body << llvm::formatv(optionalTypeParserCode, listName);
     } else {
       const char *parserCode =
           dir->shouldBeQualified() ? qualifiedTypeParserCode : typeParserCode;
-      TypeSwitch<Element *>(dir->getOperand())
+      TypeSwitch<FormatElement *>(dir->getArg())
           .Case<OperandVariable, ResultVariable>([&](auto operand) {
             body << formatv(parserCode,
                             operand->getVar()->constraint.getCPPClassName(),
                             listName);
           })
           .Default([&](auto operand) {
             body << formatv(parserCode, "::mlir::Type", listName);
           });
for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     body << ", result.operands))\n    return ::mlir::failure();\n";
   }
 }
 
 void OperationFormat::genParserRegionResolution(Operator &op,
                                                 MethodBody &body) {
   // Check for the case where all regions were parsed.
   bool hasAllRegions = llvm::any_of(
-      elements, [](auto &elt) { return isa<RegionsDirective>(elt.get()); });
+      elements, [](FormatElement *elt) { return isa<RegionsDirective>(elt); });
   if (hasAllRegions) {
     body << "  result.addRegions(fullRegions);\n";
     return;
   }
 
   // Otherwise, handle each region individually.
   for (const NamedRegion &region : op.getRegions()) {
     if (region.isVariadic())
       body << "  result.addRegions(" << region.name << "Regions);\n";
     else
       body << "  result.addRegion(std::move(" << region.name << "Region));\n";
   }
 }
 
 void OperationFormat::genParserSuccessorResolution(Operator &op,
                                                    MethodBody &body) {
   // Check for the case where all successors were parsed.
-  bool hasAllSuccessors = llvm::any_of(
-      elements, [](auto &elt) { return isa<SuccessorsDirective>(elt.get()); });
+  bool hasAllSuccessors = llvm::any_of(elements, [](FormatElement *elt) {
+    return isa<SuccessorsDirective>(elt);
+  });
   if (hasAllSuccessors) {
     body << "  result.addSuccessors(fullSuccessors);\n";
     return;
   }
 
   // Otherwise, handle each successor individually.
   for (const NamedSuccessor &successor : op.getSuccessors()) {
     if (successor.isVariadic())
if (value) {
     lastWasPunctuation = false;
   } else {
     lastWasPunctuation = true;
   }
   shouldEmitSpace = false;
 }
 
 /// Generate the printer for a custom directive parameter.
-static void genCustomDirectiveParameterPrinter(Element *element,
+static void genCustomDirectiveParameterPrinter(FormatElement *element,
                                                const Operator &op,
                                                MethodBody &body) {
   if (auto *attr = dyn_cast<AttributeVariable>(element)) {
     body << op.getGetterName(attr->getVar()->name) << "Attr()";
 
   } else if (isa<AttrDictDirective>(element)) {
     body << "getOperation()->getAttrDictionary()";
 
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     body << op.getGetterName(operand->getVar()->name) << "()";
 
   } else if (auto *region = dyn_cast<RegionVariable>(element)) {
     body << op.getGetterName(region->getVar()->name) << "()";
 
   } else if (auto *successor = dyn_cast<SuccessorVariable>(element)) {
     body << op.getGetterName(successor->getVar()->name) << "()";
 
   } else if (auto *dir = dyn_cast<RefDirective>(element)) {
-    genCustomDirectiveParameterPrinter(dir->getOperand(), op, body);
+    genCustomDirectiveParameterPrinter(dir->getArg(), op, body);
 
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
-    auto *typeOperand = dir->getOperand();
+    auto *typeOperand = dir->getArg();
     auto *operand = dyn_cast<OperandVariable>(typeOperand);
     auto *var = operand ? operand->getVar()
                         : cast<ResultVariable>(typeOperand)->getVar();
     std::string name = op.getGetterName(var->name);
     if (var->isVariadic())
       body << name << "().getTypes()";
     else if (var->isOptional())
       body << llvm::formatv("({0}() ? {0}().getType() : Type())", name);
     else
       body << name << "().getType()";
   } else {
     llvm_unreachable("unknown custom directive parameter");
   }
 }
 
 /// Generate the printer for a custom directive.
 static void genCustomDirectivePrinter(CustomDirective *customDir,
                                       const Operator &op, MethodBody &body) {
   body << "  print" << customDir->getName() << "(_odsPrinter, *this";
-  for (Element &param : customDir->getArguments()) {
+  for (FormatElement *param : customDir->getArguments()) {
     body << ", ";
-    genCustomDirectiveParameterPrinter(&param, op, body);
+    genCustomDirectiveParameterPrinter(param, op, body);
   }
   body << ");\n";
 }
 
 /// Generate the printer for a region with the given variable name.
 static void genRegionPrinter(const Twine &regionName, MethodBody &body,
                              bool hasImplicitTermTrait) {
   if (hasImplicitTermTrait)
     body << llvm::formatv(regionSingleBlockImplicitTerminatorPrinterCode,
                           regionName);
   else
     body << "  _odsPrinter.printRegion(" << regionName << ");\n";
 }
 static void genVariadicRegionPrinter(const Twine &regionListName,
                                      MethodBody &body,
                                      bool hasImplicitTermTrait) {
   body << "    llvm::interleaveComma(" << regionListName
        << ", _odsPrinter, [&](::mlir::Region &region) {\n      ";
   genRegionPrinter("region", body, hasImplicitTermTrait);
   body << "    });\n";
 }
 
 /// Generate the C++ for an operand to a (*-)type directive.
-static MethodBody &genTypeOperandPrinter(Element *arg, const Operator &op,
+static MethodBody &genTypeOperandPrinter(FormatElement *arg, const Operator &op,
                                          MethodBody &body,
                                          bool useArrayRef = true) {
   if (isa<OperandsDirective>(arg))
     return body << "getOperation()->getOperandTypes()";
   if (isa<ResultsDirective>(arg))
     return body << "getOperation()->getResultTypes()";
   auto *operand = dyn_cast<OperandVariable>(arg);
   auto *var = operand ? operand->getVar() : cast<ResultVariable>(arg)->getVar();
if (nonKeywordCases.any()) {
     return;
   }
 
   body << "    _odsPrinter << caseValueStr;\n"
           "  }\n";
 }
 
 /// Generate the check for the anchor of an optional group.
-static void genOptionalGroupPrinterAnchor(Element *anchor, const Operator &op,
+static void genOptionalGroupPrinterAnchor(FormatElement *anchor,
+                                          const Operator &op,
                                           MethodBody &body) {
-  TypeSwitch<Element *>(anchor)
+  TypeSwitch<FormatElement *>(anchor)
       .Case<OperandVariable, ResultVariable>([&](auto *element) {
         const NamedTypeConstraint *var = element->getVar();
         std::string name = op.getGetterName(var->name);
         if (var->isOptional())
           body << "  if (" << name << "()) {\n";
         else if (var->isVariadic())
           body << "  if (!" << name << "().empty()) {\n";
       })
       .Case<RegionVariable>([&](RegionVariable *element) {
         const NamedRegion *var = element->getVar();
         std::string name = op.getGetterName(var->name);
         // TODO: Add a check for optional regions here when ODS supports it.
         body << "  if (!" << name << "().empty()) {\n";
       })
       .Case<TypeDirective>([&](TypeDirective *element) {
-        genOptionalGroupPrinterAnchor(element->getOperand(), op, body);
+        genOptionalGroupPrinterAnchor(element->getArg(), op, body);
       })
       .Case<FunctionalTypeDirective>([&](FunctionalTypeDirective *element) {
         genOptionalGroupPrinterAnchor(element->getInputs(), op, body);
       })
       .Case<AttributeVariable>([&](AttributeVariable *attr) {
         body << "  if ((*this)->getAttr(\"" << attr->getVar()->name
              << "\")) {\n";
       });
 }
 
-void OperationFormat::genElementPrinter(Element *element, MethodBody &body,
-                                        Operator &op, bool &shouldEmitSpace,
+void OperationFormat::genElementPrinter(FormatElement *element,
+                                        MethodBody &body, Operator &op,
+                                        bool &shouldEmitSpace,
                                         bool &lastWasPunctuation) {
   if (LiteralElement *literal = dyn_cast<LiteralElement>(element))
-    return genLiteralPrinter(literal->getLiteral(), body, shouldEmitSpace,
+    return genLiteralPrinter(literal->getSpelling(), body, shouldEmitSpace,
                              lastWasPunctuation);
 
   // Emit a whitespace element.
-  if (isa<NewlineElement>(element)) {
+  if (auto *space = dyn_cast<WhitespaceElement>(element)) {
+    if (space->getValue() == "\\n") {
       body << "  _odsPrinter.printNewline();\n";
+    } else {
+      genSpacePrinter(!space->getValue().empty(), body, shouldEmitSpace,
+                      lastWasPunctuation);
+    }
     return;
   }
-  if (SpaceElement *space = dyn_cast<SpaceElement>(element))
-    return genSpacePrinter(space->getValue(), body, shouldEmitSpace,
-                           lastWasPunctuation);
 
   // Emit an optional group.
   if (OptionalElement *optional = dyn_cast<OptionalElement>(element)) {
     // Emit the check for the presence of the anchor element.
-    Element *anchor = optional->getAnchor();
+    FormatElement *anchor = optional->getAnchor();
     genOptionalGroupPrinterAnchor(anchor, op, body);
 
     // If the anchor is a unit attribute, we don't need to print it. When
     // parsing, we will add this attribute if this group is present.
     auto elements = optional->getThenElements();
-    Element *elidedAnchorElement = nullptr;
+    FormatElement *elidedAnchorElement = nullptr;
     auto *anchorAttr = dyn_cast<AttributeVariable>(anchor);
-    if (anchorAttr && anchorAttr != &*elements.begin() &&
+    if (anchorAttr && anchorAttr != elements.front() &&
         anchorAttr->isUnitAttr()) {
       elidedAnchorElement = anchorAttr;
     }
 
     // Emit each of the elements.
-    for (Element &childElement : elements) {
-      if (&childElement != elidedAnchorElement) {
-        genElementPrinter(&childElement, body, op, shouldEmitSpace,
+    for (FormatElement *childElement : elements) {
+      if (childElement != elidedAnchorElement) {
+        genElementPrinter(childElement, body, op, shouldEmitSpace,
                           lastWasPunctuation);
       }
     }
     body << "  }";
 
     // Emit each of the else elements.
     auto elseElements = optional->getElseElements();
     if (!elseElements.empty()) {
       body << " else {\n";
-      for (Element &childElement : elseElements) {
-        genElementPrinter(&childElement, body, op, shouldEmitSpace,
+      for (FormatElement *childElement : elseElements) {
+        genElementPrinter(childElement, body, op, shouldEmitSpace,
                           lastWasPunctuation);
       }
       body << "  }";
     }
 
     body << "\n";
     return;
   }
if (auto *attr = dyn_cast<AttributeVariable>(element)) {
     body << "  _odsPrinter << getOperation()->getOperands();\n";
   } else if (isa<RegionsDirective>(element)) {
     genVariadicRegionPrinter("getOperation()->getRegions()", body,
                              hasImplicitTermTrait);
   } else if (isa<SuccessorsDirective>(element)) {
     body << "  ::llvm::interleaveComma(getOperation()->getSuccessors(), "
             "_odsPrinter);\n";
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
-    if (auto *operand = dyn_cast<OperandVariable>(dir->getOperand())) {
+    if (auto *operand = dyn_cast<OperandVariable>(dir->getArg())) {
       if (operand->getVar()->isVariadicOfVariadic()) {
         body << llvm::formatv(
             "  ::llvm::interleaveComma({0}().getTypes(), _odsPrinter, "
             "[&](::mlir::TypeRange types) {{ _odsPrinter << \"(\" << "
             "types << \")\"; });\n",
             op.getGetterName(operand->getVar()->name));
         return;
       }
     }
     const NamedTypeConstraint *var = nullptr;
     {
-      if (auto *operand = dyn_cast<OperandVariable>(dir->getOperand()))
+      if (auto *operand = dyn_cast<OperandVariable>(dir->getArg()))
         var = operand->getVar();
-      else if (auto *operand = dyn_cast<ResultVariable>(dir->getOperand()))
+      else if (auto *operand = dyn_cast<ResultVariable>(dir->getArg()))
         var = operand->getVar();
     }
     if (var && !var->isVariadicOfVariadic() && !var->isVariadic() &&
         !var->isOptional()) {
       std::string cppClass = var->constraint.getCPPClassName();
       if (dir->shouldBeQualified()) {
         body << "   _odsPrinter << " << op.getGetterName(var->name)
              << "().getType();\n";
         return;
       }
       body << "  {\n"
            << "    auto type = " << op.getGetterName(var->name)
            << "().getType();\n"
            << "    if (auto validType = type.dyn_cast<" << cppClass << ">())\n"
            << "      _odsPrinter.printStrippedAttrOrType(validType);\n"
            << "   else\n"
            << "     _odsPrinter << type;\n"
            << "  }\n";
       return;
     }
     body << "  _odsPrinter << ";
-    genTypeOperandPrinter(dir->getOperand(), op, body, /*useArrayRef=*/false)
+    genTypeOperandPrinter(dir->getArg(), op, body, /*useArrayRef=*/false)
         << ";\n";
   } else if (auto *dir = dyn_cast<FunctionalTypeDirective>(element)) {
     body << "  _odsPrinter.printFunctionalType(";
     genTypeOperandPrinter(dir->getInputs(), op, body) << ", ";
     genTypeOperandPrinter(dir->getResults(), op, body) << ");\n";
   } else {
     llvm_unreachable("unknown format element");
   }
 }
 
 void OperationFormat::genPrinter(Operator &op, OpClass &opClass) {
   auto *method = opClass.addMethod(
       "void", "print",
       MethodParameter("::mlir::OpAsmPrinter &", "_odsPrinter"));
   auto &body = method->body();
 
   // Flags for if we should emit a space, and if the last element was
   // punctuation.
   bool shouldEmitSpace = true, lastWasPunctuation = false;
-  for (auto &element : elements)
-    genElementPrinter(element.get(), body, op, shouldEmitSpace,
-                      lastWasPunctuation);
+  for (FormatElement *element : elements)
+    genElementPrinter(element, body, op, shouldEmitSpace, lastWasPunctuation);
 }
 
 //===----------------------------------------------------------------------===//
-// FormatParser
+// OpFormatParser
 //===----------------------------------------------------------------------===//
 
 /// Function to find an element within the given range that has the same name as
 /// 'name'.
 template <typename RangeT> static auto findArg(RangeT &&range, StringRef name) {
   auto it = llvm::find_if(range, [=](auto &arg) { return arg.name == name; });
   return it != range.end() ? &*it : nullptr;
 }
 
 namespace {
 /// This class implements a parser for an instance of an operation assembly
 /// format.
-class FormatParser {
+class OpFormatParser : public FormatParser {
 public:
-  FormatParser(llvm::SourceMgr &mgr, OperationFormat &format, Operator &op)
-      : lexer(mgr, op.getLoc()[0]), curToken(lexer.lexToken()), fmt(format),
-        op(op), seenOperandTypes(op.getNumOperands()),
+  OpFormatParser(llvm::SourceMgr &mgr, OperationFormat &format, Operator &op)
+      : FormatParser(mgr, op.getLoc()[0]), fmt(format), op(op),
+        seenOperandTypes(op.getNumOperands()),
         seenResultTypes(op.getNumResults()) {}
 
-  /// Parse the operation assembly format.
-  LogicalResult parse();
+protected:
+  /// Verify the format elements.
+  LogicalResult verify(llvm::SMLoc loc,
+                       ArrayRef<FormatElement *> elements) override;
+  /// Verify the arguments to a custom directive.
+  LogicalResult
+  verifyCustomDirectiveArguments(llvm::SMLoc loc,
+                                 ArrayRef<FormatElement *> arguments) override;
+  /// Verify the elements of an optional group.
+  LogicalResult
+  verifyOptionalGroupElements(llvm::SMLoc loc,
+                              ArrayRef<FormatElement *> elements,
+                              Optional<unsigned> anchorIndex) override;
+  LogicalResult verifyOptionalGroupElement(llvm::SMLoc loc,
+                                           FormatElement *element,
+                                           bool isAnchor);
+
+  /// Parse an operation variable.
+  FailureOr<FormatElement *> parseVariableImpl(llvm::SMLoc loc, StringRef name,
+                                               Context ctx) override;
+  /// Parse an operation format directive.
+  FailureOr<FormatElement *> parseDirectiveImpl(llvm::SMLoc loc,
+                                                FormatToken::Kind kind,
+                                                Context ctx) override;
 
 private:
-  /// The current context of the parser when parsing an element.
-  enum ParserContext {
-    /// The element is being parsed in a "top-level" context, i.e. at the top of
-    /// the format or in an optional group.
-    TopLevelContext,
-    /// The element is being parsed as a custom directive child.
-    CustomDirectiveContext,
-    /// The element is being parsed as a type directive child.
-    TypeDirectiveContext,
-    /// The element is being parsed as a reference directive child.
-    RefDirectiveContext
-  };
-
   /// This struct represents a type resolution instance. It includes a specific
   /// type as well as an optional transformer to apply to that type in order to
   /// properly resolve the type of a variable.
   struct TypeResolutionInstance {
     ConstArgument resolver;
     Optional<StringRef> transformer;
   };
 
-  /// An iterator over the elements of a format group.
-  using ElementsIterT = llvm::pointee_iterator<
-      std::vector<std::unique_ptr<Element>>::const_iterator>;
+  using ElementsItT = ArrayRef<FormatElement *>::iterator;
 
   /// Verify the state of operation attributes within the format.
-  LogicalResult verifyAttributes(llvm::SMLoc loc);
+  LogicalResult verifyAttributes(llvm::SMLoc loc,
+                                 ArrayRef<FormatElement *> elements);
   /// Verify the attribute elements at the back of the given stack of iterators.
   LogicalResult verifyAttributes(
       llvm::SMLoc loc,
-      SmallVectorImpl<std::pair<ElementsIterT, ElementsIterT>> &iteratorStack);
+      SmallVectorImpl<std::pair<ElementsItT, ElementsItT>> &iteratorStack);
 
   /// Verify the state of operation operands within the format.
   LogicalResult
   verifyOperands(llvm::SMLoc loc,
                  llvm::StringMap<TypeResolutionInstance> &variableTyResolver);
 
   /// Verify the state of operation regions within the format.
   LogicalResult verifyRegions(llvm::SMLoc loc);
private:
   void handleTypesMatchConstraint(
       llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
       const llvm::Record &def);
 
   /// Returns an argument or attribute with the given name that has been seen
   /// within the format.
   ConstArgument findSeenArg(StringRef name);
 
-  /// Parse a specific element.
-  LogicalResult parseElement(std::unique_ptr<Element> &element,
-                             ParserContext context);
-  LogicalResult parseVariable(std::unique_ptr<Element> &element,
-                              ParserContext context);
-  LogicalResult parseDirective(std::unique_ptr<Element> &element,
-                               ParserContext context);
-  LogicalResult parseLiteral(std::unique_ptr<Element> &element,
-                             ParserContext context);
-  LogicalResult parseOptional(std::unique_ptr<Element> &element,
-                              ParserContext context);
-  LogicalResult parseOptionalChildElement(
-      std::vector<std::unique_ptr<Element>> &childElements,
-      Optional<unsigned> &anchorIdx);
-  LogicalResult verifyOptionalChildElement(Element *element,
-                                           llvm::SMLoc childLoc, bool isAnchor);
-
   /// Parse the various different directives.
-  LogicalResult parseAttrDictDirective(std::unique_ptr<Element> &element,
-                                       llvm::SMLoc loc, ParserContext context,
-                                       bool withKeyword);
-  LogicalResult parseCustomDirective(std::unique_ptr<Element> &element,
-                                     llvm::SMLoc loc, ParserContext context);
-  LogicalResult parseCustomDirectiveParameter(
-      std::vector<std::unique_ptr<Element>> &parameters);
-  LogicalResult parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
-                                             FormatToken tok,
-                                             ParserContext context);
-  LogicalResult parseOperandsDirective(std::unique_ptr<Element> &element,
-                                       llvm::SMLoc loc, ParserContext context);
-  LogicalResult parseQualifiedDirective(std::unique_ptr<Element> &element,
-                                        FormatToken tok, ParserContext context);
-  LogicalResult parseReferenceDirective(std::unique_ptr<Element> &element,
-                                        llvm::SMLoc loc, ParserContext context);
-  LogicalResult parseRegionsDirective(std::unique_ptr<Element> &element,
-                                      llvm::SMLoc loc, ParserContext context);
-  LogicalResult parseResultsDirective(std::unique_ptr<Element> &element,
-                                      llvm::SMLoc loc, ParserContext context);
-  LogicalResult parseSuccessorsDirective(std::unique_ptr<Element> &element,
-                                         llvm::SMLoc loc,
-                                         ParserContext context);
-  LogicalResult parseTypeDirective(std::unique_ptr<Element> &element,
-                                   FormatToken tok, ParserContext context);
-  LogicalResult parseTypeDirectiveOperand(std::unique_ptr<Element> &element,
+  FailureOr<FormatElement *>
+  parseAttrDictDirective(llvm::SMLoc loc, Context context, bool withKeyword);
+  FailureOr<FormatElement *> parseFunctionalTypeDirective(llvm::SMLoc loc,
+                                                          Context context);
+  FailureOr<FormatElement *> parseOperandsDirective(llvm::SMLoc loc,
+                                                    Context context);
+  FailureOr<FormatElement *> parseQualifiedDirective(llvm::SMLoc loc,
+                                                     Context context);
+  FailureOr<FormatElement *> parseReferenceDirective(llvm::SMLoc loc,
+                                                     Context context);
+  FailureOr<FormatElement *> parseRegionsDirective(llvm::SMLoc loc,
+                                                   Context context);
+  FailureOr<FormatElement *> parseResultsDirective(llvm::SMLoc loc,
+                                                   Context context);
+  FailureOr<FormatElement *> parseSuccessorsDirective(llvm::SMLoc loc,
+                                                      Context context);
+  FailureOr<FormatElement *> parseTypeDirective(llvm::SMLoc loc,
+                                                Context context);
+  FailureOr<FormatElement *> parseTypeDirectiveOperand(llvm::SMLoc loc,
                                                        bool isRefChild = false);
 
   //===--------------------------------------------------------------------===//
-  // Lexer Utilities
-  //===--------------------------------------------------------------------===//
-
-  /// Advance the current lexer onto the next token.
-  void consumeToken() {
-    assert(curToken.getKind() != FormatToken::eof &&
-           curToken.getKind() != FormatToken::error &&
-           "shouldn't advance past EOF or errors");
-    curToken = lexer.lexToken();
-  }
-  LogicalResult parseToken(FormatToken::Kind kind, const Twine &msg) {
-    if (curToken.getKind() != kind)
-      return emitError(curToken.getLoc(), msg);
-    consumeToken();
-    return ::mlir::success();
-  }
-  LogicalResult emitError(llvm::SMLoc loc, const Twine &msg) {
-    lexer.emitError(loc, msg);
-    return ::mlir::failure();
-  }
-  LogicalResult emitErrorAndNote(llvm::SMLoc loc, const Twine &msg,
-                                 const Twine &note) {
-    lexer.emitErrorAndNote(loc, msg, note);
-    return ::mlir::failure();
-  }
-
-  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
-  FormatLexer lexer;
-  FormatToken curToken;
   OperationFormat &fmt;
   Operator &op;
 
   // The following are various bits of format state used for verification
   // during parsing.
   bool hasAttrDict = false;
   bool hasAllRegions = false, hasAllSuccessors = false;
   bool canInferResultTypes = false;
   llvm::SmallBitVector seenOperandTypes, seenResultTypes;
   llvm::SmallSetVector<const NamedAttribute *, 8> seenAttrs;
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
 };
 } // namespace
 
-LogicalResult FormatParser::parse() {
-  llvm::SMLoc loc = curToken.getLoc();
-
-  // Parse each of the format elements into the main format.
-  while (curToken.getKind() != FormatToken::eof) {
-    std::unique_ptr<Element> element;
-    if (failed(parseElement(element, TopLevelContext)))
-      return ::mlir::failure();
-    fmt.elements.push_back(std::move(element));
-  }
-
+LogicalResult OpFormatParser::verify(llvm::SMLoc loc,
+                                     ArrayRef<FormatElement *> elements) {
   // Check that the attribute dictionary is in the format.
   if (!hasAttrDict)
     return emitError(loc, "'attr-dict' directive not found in "
                           "custom assembly format");
 
   // Check for any type traits that we can use for inferring types.
   llvm::StringMap<TypeResolutionInstance> variableTyResolver;
   for (const Trait &trait : op.getTraits()) {
if (def.isSubClassOf("AllTypesMatch")) {
         if (*name == "InferTypeOpInterface" &&
             def.getValueAsString("cppNamespace") == "::mlir")
           canInferResultTypes = true;
       }
     }
   }
 
   // Verify the state of the various operation components.
-  if (failed(verifyAttributes(loc)) ||
+  if (failed(verifyAttributes(loc, elements)) ||
       failed(verifyResults(loc, variableTyResolver)) ||
       failed(verifyOperands(loc, variableTyResolver)) ||
       failed(verifyRegions(loc)) || failed(verifySuccessors(loc)))
-    return ::mlir::failure();
+    return mlir::failure();
 
   // Collect the set of used attributes in the format.
   fmt.usedAttributes = seenAttrs.takeVector();
-  return ::mlir::success();
+  return success();
 }
 
-LogicalResult FormatParser::verifyAttributes(llvm::SMLoc loc) {
+LogicalResult
+OpFormatParser::verifyAttributes(llvm::SMLoc loc,
+                                 ArrayRef<FormatElement *> elements) {
   // Check that there are no `:` literals after an attribute without a constant
   // type. The attribute grammar contains an optional trailing colon type, which
   // can lead to unexpected and generally unintended behavior. Given that, it is
   // better to just error out here instead.
-  using ElementsIterT = llvm::pointee_iterator<
-      std::vector<std::unique_ptr<Element>>::const_iterator>;
-  SmallVector<std::pair<ElementsIterT, ElementsIterT>, 1> iteratorStack;
-  iteratorStack.emplace_back(fmt.elements.begin(), fmt.elements.end());
+  SmallVector<std::pair<ElementsItT, ElementsItT>, 1> iteratorStack;
+  iteratorStack.emplace_back(elements.begin(), elements.end());
   while (!iteratorStack.empty())
     if (failed(verifyAttributes(loc, iteratorStack)))
-      return ::mlir::failure();
+      return ::failure();
 
   // Check for VariadicOfVariadic variables. The segment attribute of those
   // variables will be infered.
   for (const NamedTypeConstraint *var : seenOperands) {
     if (var->constraint.isVariadicOfVariadic()) {
       fmt.inferredAttributes.insert(
           var->constraint.getVariadicOfVariadicSegmentSizeAttr());
     }
   }
 
-  return ::mlir::success();
+  return success();
 }
 /// Verify the attribute elements at the back of the given stack of iterators.
-LogicalResult FormatParser::verifyAttributes(
+LogicalResult OpFormatParser::verifyAttributes(
     llvm::SMLoc loc,
-    SmallVectorImpl<std::pair<ElementsIterT, ElementsIterT>> &iteratorStack) {
+    SmallVectorImpl<std::pair<ElementsItT, ElementsItT>> &iteratorStack) {
   auto &stackIt = iteratorStack.back();
-  ElementsIterT &it = stackIt.first, e = stackIt.second;
+  ElementsItT &it = stackIt.first, e = stackIt.second;
   while (it != e) {
-    Element *element = &*(it++);
+    FormatElement *element = *(it++);
 
     // Traverse into optional groups.
     if (auto *optional = dyn_cast<OptionalElement>(element)) {
       auto thenElements = optional->getThenElements();
       iteratorStack.emplace_back(thenElements.begin(), thenElements.end());
 
       auto elseElements = optional->getElseElements();
       iteratorStack.emplace_back(elseElements.begin(), elseElements.end());
-      return ::mlir::success();
+      return success();
     }
 
     // We are checking for an attribute element followed by a `:`, so there is
     // no need to check the end.
     if (it == e && iteratorStack.size() == 1)
       break;
 
     // Check for an attribute with a constant type builder, followed by a `:`.
     auto *prevAttr = dyn_cast<AttributeVariable>(element);
     if (!prevAttr || prevAttr->getTypeBuilder())
       continue;
 
     // Check the next iterator within the stack for literal elements.
     for (auto &nextItPair : iteratorStack) {
-      ElementsIterT nextIt = nextItPair.first, nextE = nextItPair.second;
+      ElementsItT nextIt = nextItPair.first, nextE = nextItPair.second;
       for (; nextIt != nextE; ++nextIt) {
         // Skip any trailing whitespace, attribute dictionaries, or optional
         // groups.
         if (isa<WhitespaceElement>(*nextIt) ||
             isa<AttrDictDirective>(*nextIt) || isa<OptionalElement>(*nextIt))
           continue;
 
         // We are only interested in `:` literals.
-        auto *literal = dyn_cast<LiteralElement>(&*nextIt);
-        if (!literal || literal->getLiteral() != ":")
+        auto *literal = dyn_cast<LiteralElement>(*nextIt);
+        if (!literal || literal->getSpelling() != ":")
           break;
 
         // TODO: Use the location of the literal element itself.
         return emitError(
             loc, llvm::formatv("format ambiguity caused by `:` literal found "
                                "after attribute `{0}` which does not have "
                                "a buildable type",
                                prevAttr->getVar()->name));
       }
     }
   }
   iteratorStack.pop_back();
-  return ::mlir::success();
+  return success();
 }
 
-LogicalResult FormatParser::verifyOperands(
+LogicalResult OpFormatParser::verifyOperands(
     llvm::SMLoc loc,
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   // Check that all of the operands are within the format, and their types can
   // be inferred.
   auto &buildableTypes = fmt.buildableTypes;
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);
 
if (!builder || (fmt.allOperands && operand.isVariableLength())) {
               "', is not buildable and a buildable type cannot be inferred",
           "suggest adding a type constraint to the operation or adding a "
           "'type($" +
               operand.name + ")' directive to the " + "custom assembly format");
     }
     auto it = buildableTypes.insert({*builder, buildableTypes.size()});
     fmt.operandTypes[i].setBuilderIdx(it.first->second);
   }
-  return ::mlir::success();
+  return success();
 }
 
-LogicalResult FormatParser::verifyRegions(llvm::SMLoc loc) {
+LogicalResult OpFormatParser::verifyRegions(llvm::SMLoc loc) {
   // Check that all of the regions are within the format.
   if (hasAllRegions)
     return ::mlir::success();
 
   for (unsigned i = 0, e = op.getNumRegions(); i != e; ++i) {
     const NamedRegion &region = op.getRegion(i);
     if (!seenRegions.count(&region)) {
       return emitErrorAndNote(loc,
                               "region #" + Twine(i) + ", named '" +
                                   region.name + "', not found",
                               "suggest adding a '$" + region.name +
                                   "' directive to the custom assembly format");
     }
   }
   return ::mlir::success();
 }
 
-LogicalResult FormatParser::verifyResults(
+LogicalResult OpFormatParser::verifyResults(
     llvm::SMLoc loc,
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   // If we format all of the types together, there is nothing to check.
   if (fmt.allResultTypes)
     return ::mlir::success();
 
   // If no result types are specified and we can infer them, infer all result
   // types
for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
     }
     // Note in the format that this result uses the custom builder.
     auto it = buildableTypes.insert({*builder, buildableTypes.size()});
     fmt.resultTypes[i].setBuilderIdx(it.first->second);
   }
   return ::mlir::success();
 }
 
-LogicalResult FormatParser::verifySuccessors(llvm::SMLoc loc) {
+LogicalResult OpFormatParser::verifySuccessors(llvm::SMLoc loc) {
   // Check that all of the successors are within the format.
   if (hasAllSuccessors)
     return ::mlir::success();
 
   for (unsigned i = 0, e = op.getNumSuccessors(); i != e; ++i) {
     const NamedSuccessor &successor = op.getSuccessor(i);
     if (!seenSuccessors.count(&successor)) {
       return emitErrorAndNote(loc,
                               "successor #" + Twine(i) + ", named '" +
                                   successor.name + "', not found",
                               "suggest adding a '$" + successor.name +
                                   "' directive to the custom assembly format");
     }
   }
   return ::mlir::success();
 }
 
-void FormatParser::handleAllTypesMatchConstraint(
+void OpFormatParser::handleAllTypesMatchConstraint(
     ArrayRef<StringRef> values,
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   for (unsigned i = 0, e = values.size(); i != e; ++i) {
     // Check to see if this value matches a resolved operand or result type.
     ConstArgument arg = findSeenArg(values[i]);
     if (!arg)
       continue;
 
     // Mark this value as the type resolver for the other variables.
     for (unsigned j = 0; j != i; ++j)
       variableTyResolver[values[j]] = {arg, llvm::None};
     for (unsigned j = i + 1; j != e; ++j)
       variableTyResolver[values[j]] = {arg, llvm::None};
   }
 }
 
-void FormatParser::handleSameTypesConstraint(
+void OpFormatParser::handleSameTypesConstraint(
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
     bool includeResults) {
   const NamedTypeConstraint *resolver = nullptr;
   int resolvedIt = -1;
 
   // Check to see if there is an operand or result to use for the resolution.
   if ((resolvedIt = seenOperandTypes.find_first()) != -1)
     resolver = &op.getOperand(resolvedIt);
   else if (includeResults && (resolvedIt = seenResultTypes.find_first()) != -1)
     resolver = &op.getResult(resolvedIt);
   else
     return;
 
   // Set the resolvers for each operand and result.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i)
     if (!seenOperandTypes.test(i) && !op.getOperand(i).name.empty())
       variableTyResolver[op.getOperand(i).name] = {resolver, llvm::None};
   if (includeResults) {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i)
       if (!seenResultTypes.test(i) && !op.getResultName(i).empty())
         variableTyResolver[op.getResultName(i)] = {resolver, llvm::None};
   }
 }
 
-void FormatParser::handleTypesMatchConstraint(
+void OpFormatParser::handleTypesMatchConstraint(
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
     const llvm::Record &def) {
   StringRef lhsName = def.getValueAsString("lhs");
   StringRef rhsName = def.getValueAsString("rhs");
   StringRef transformer = def.getValueAsString("transformer");
   if (ConstArgument arg = findSeenArg(lhsName))
     variableTyResolver[rhsName] = {arg, transformer};
 }
 
-ConstArgument FormatParser::findSeenArg(StringRef name) {
+ConstArgument OpFormatParser::findSeenArg(StringRef name) {
   if (const NamedTypeConstraint *arg = findArg(op.getOperands(), name))
     return seenOperandTypes.test(arg - op.operand_begin()) ? arg : nullptr;
   if (const NamedTypeConstraint *arg = findArg(op.getResults(), name))
     return seenResultTypes.test(arg - op.result_begin()) ? arg : nullptr;
   if (const NamedAttribute *attr = findArg(op.getAttributes(), name))
     return seenAttrs.count(attr) ? attr : nullptr;
   return nullptr;
 }
 
-LogicalResult FormatParser::parseElement(std::unique_ptr<Element> &element,
-                                         ParserContext context) {
-  // Directives.
-  if (curToken.isKeyword())
-    return parseDirective(element, context);
-  // Literals.
-  if (curToken.getKind() == FormatToken::literal)
-    return parseLiteral(element, context);
-  // Optionals.
-  if (curToken.getKind() == FormatToken::l_paren)
-    return parseOptional(element, context);
-  // Variables.
-  if (curToken.getKind() == FormatToken::variable)
-    return parseVariable(element, context);
-  return emitError(curToken.getLoc(),
-                   "expected directive, literal, variable, or optional group");
-}
-
-LogicalResult FormatParser::parseVariable(std::unique_ptr<Element> &element,
-                                          ParserContext context) {
-  FormatToken varTok = curToken;
-  consumeToken();
-
-  StringRef name = varTok.getSpelling().drop_front();
-  llvm::SMLoc loc = varTok.getLoc();
-
-  // Check that the parsed argument is something actually registered on the
-  // op.
-  /// Attributes
+FailureOr<FormatElement *> OpFormatParser::parseVariableImpl(llvm::SMLoc loc,
+                                                             StringRef name,
+                                                             Context ctx) {
+  // Check that the parsed argument is something actually registered on the op.
+  // Attributes
   if (const NamedAttribute *attr = findArg(op.getAttributes(), name)) {
-    if (context == TypeDirectiveContext)
+    if (ctx == TypeDirectiveContext)
       return emitError(
           loc, "attributes cannot be used as children to a `type` directive");
-    if (context == RefDirectiveContext) {
+    if (ctx == RefDirectiveContext) {
       if (!seenAttrs.count(attr))
         return emitError(loc, "attribute '" + name +
                                   "' must be bound before it is referenced");
     } else if (!seenAttrs.insert(attr)) {
       return emitError(loc, "attribute '" + name + "' is already bound");
     }
 
-    element = std::make_unique<AttributeVariable>(attr);
-    return ::mlir::success();
+    return create<AttributeVariable>(attr);
   }
-  /// Operands
+  // Operands
   if (const NamedTypeConstraint *operand = findArg(op.getOperands(), name)) {
-    if (context == TopLevelContext || context == CustomDirectiveContext) {
+    if (ctx == TopLevelContext || ctx == CustomDirectiveContext) {
       if (fmt.allOperands || !seenOperands.insert(operand).second)
         return emitError(loc, "operand '" + name + "' is already bound");
-    } else if (context == RefDirectiveContext && !seenOperands.count(operand)) {
+    } else if (ctx == RefDirectiveContext && !seenOperands.count(operand)) {
       return emitError(loc, "operand '" + name +
                                 "' must be bound before it is referenced");
     }
-    element = std::make_unique<OperandVariable>(operand);
-    return ::mlir::success();
+    return create<OperandVariable>(operand);
   }
-  /// Regions
+  // Regions
   if (const NamedRegion *region = findArg(op.getRegions(), name)) {
-    if (context == TopLevelContext || context == CustomDirectiveContext) {
+    if (ctx == TopLevelContext || ctx == CustomDirectiveContext) {
       if (hasAllRegions || !seenRegions.insert(region).second)
         return emitError(loc, "region '" + name + "' is already bound");
-    } else if (context == RefDirectiveContext && !seenRegions.count(region)) {
+    } else if (ctx == RefDirectiveContext && !seenRegions.count(region)) {
       return emitError(loc, "region '" + name +
                                 "' must be bound before it is referenced");
     } else {
       return emitError(loc, "regions can only be used at the top level");
     }
-    element = std::make_unique<RegionVariable>(region);
-    return ::mlir::success();
+    return create<RegionVariable>(region);
   }
-  /// Results.
+  // Results.
   if (const auto *result = findArg(op.getResults(), name)) {
-    if (context != TypeDirectiveContext)
+    if (ctx != TypeDirectiveContext)
       return emitError(loc, "result variables can can only be used as a child "
                             "to a 'type' directive");
-    element = std::make_unique<ResultVariable>(result);
-    return ::mlir::success();
+    return create<ResultVariable>(result);
   }
-  /// Successors.
+  // Successors.
   if (const auto *successor = findArg(op.getSuccessors(), name)) {
-    if (context == TopLevelContext || context == CustomDirectiveContext) {
+    if (ctx == TopLevelContext || ctx == CustomDirectiveContext) {
       if (hasAllSuccessors || !seenSuccessors.insert(successor).second)
         return emitError(loc, "successor '" + name + "' is already bound");
-    } else if (context == RefDirectiveContext &&
-               !seenSuccessors.count(successor)) {
+    } else if (ctx == RefDirectiveContext && !seenSuccessors.count(successor)) {
       return emitError(loc, "successor '" + name +
                                 "' must be bound before it is referenced");
     } else {
       return emitError(loc, "successors can only be used at the top level");
     }
 
-    element = std::make_unique<SuccessorVariable>(successor);
-    return ::mlir::success();
+    return create<SuccessorVariable>(successor);
   }
   return emitError(loc, "expected variable to refer to an argument, region, "
                         "result, or successor");
 }
 
-LogicalResult FormatParser::parseDirective(std::unique_ptr<Element> &element,
-                                           ParserContext context) {
-  FormatToken dirTok = curToken;
-  consumeToken();
-
-  switch (dirTok.getKind()) {
+FailureOr<FormatElement *>
+OpFormatParser::parseDirectiveImpl(llvm::SMLoc loc, FormatToken::Kind kind,
+                                   Context ctx) {
+  switch (kind) {
   case FormatToken::kw_attr_dict:
-    return parseAttrDictDirective(element, dirTok.getLoc(), context,
+    return parseAttrDictDirective(loc, ctx,
                                   /*withKeyword=*/false);
   case FormatToken::kw_attr_dict_w_keyword:
-    return parseAttrDictDirective(element, dirTok.getLoc(), context,
+    return parseAttrDictDirective(loc, ctx,
                                   /*withKeyword=*/true);
-  case FormatToken::kw_custom:
-    return parseCustomDirective(element, dirTok.getLoc(), context);
   case FormatToken::kw_functional_type:
-    return parseFunctionalTypeDirective(element, dirTok, context);
+    return parseFunctionalTypeDirective(loc, ctx);
   case FormatToken::kw_operands:
-    return parseOperandsDirective(element, dirTok.getLoc(), context);
+    return parseOperandsDirective(loc, ctx);
   case FormatToken::kw_qualified:
-    return parseQualifiedDirective(element, dirTok, context);
+    return parseQualifiedDirective(loc, ctx);
   case FormatToken::kw_regions:
-    return parseRegionsDirective(element, dirTok.getLoc(), context);
+    return parseRegionsDirective(loc, ctx);
   case FormatToken::kw_results:
-    return parseResultsDirective(element, dirTok.getLoc(), context);
+    return parseResultsDirective(loc, ctx);
   case FormatToken::kw_successors:
-    return parseSuccessorsDirective(element, dirTok.getLoc(), context);
+    return parseSuccessorsDirective(loc, ctx);
   case FormatToken::kw_ref:
-    return parseReferenceDirective(element, dirTok.getLoc(), context);
+    return parseReferenceDirective(loc, ctx);
   case FormatToken::kw_type:
-    return parseTypeDirective(element, dirTok, context);
+    return parseTypeDirective(loc, ctx);
 
   default:
-    llvm_unreachable("unknown directive token");
+    return emitError(loc, "unsupported directive kind");
   }
 }
 
-LogicalResult FormatParser::parseLiteral(std::unique_ptr<Element> &element,
-                                         ParserContext context) {
-  FormatToken literalTok = curToken;
-  if (context != TopLevelContext) {
-    return emitError(
-        literalTok.getLoc(),
-        "literals may only be used in a top-level section of the format");
-  }
-  consumeToken();
-
-  StringRef value = literalTok.getSpelling().drop_front().drop_back();
-
-  // The parsed literal is a space element (`` or ` `).
-  if (value.empty() || (value.size() == 1 && value.front() == ' ')) {
-    element = std::make_unique<SpaceElement>(!value.empty());
-    return ::mlir::success();
-  }
-  // The parsed literal is a newline element.
-  if (value == "\\n") {
-    element = std::make_unique<NewlineElement>();
-    return ::mlir::success();
-  }
-
-  // Check that the parsed literal is valid.
-  if (!isValidLiteral(value, [&](Twine diag) {
-        (void)emitError(literalTok.getLoc(),
-                        "expected valid literal but got '" + value +
-                            "': " + diag);
-      }))
-    return failure();
-  element = std::make_unique<LiteralElement>(value);
-  return ::mlir::success();
-}
-
-LogicalResult FormatParser::parseOptional(std::unique_ptr<Element> &element,
-                                          ParserContext context) {
-  llvm::SMLoc curLoc = curToken.getLoc();
-  if (context != TopLevelContext)
-    return emitError(curLoc, "optional groups can only be used as top-level "
-                             "elements");
-  consumeToken();
-
-  // Parse the child elements for this optional group.
-  std::vector<std::unique_ptr<Element>> thenElements, elseElements;
-  Optional<unsigned> anchorIdx;
-  do {
-    if (failed(parseOptionalChildElement(thenElements, anchorIdx)))
-      return ::mlir::failure();
-  } while (curToken.getKind() != FormatToken::r_paren);
-  consumeToken();
-
-  // Parse the `else` elements of this optional group.
-  if (curToken.getKind() == FormatToken::colon) {
-    consumeToken();
-    if (failed(parseToken(FormatToken::l_paren,
-                          "expected '(' to start else branch "
-                          "of optional group")))
-      return failure();
-    do {
-      llvm::SMLoc childLoc = curToken.getLoc();
-      elseElements.push_back({});
-      if (failed(parseElement(elseElements.back(), TopLevelContext)) ||
-          failed(verifyOptionalChildElement(elseElements.back().get(), childLoc,
-                                            /*isAnchor=*/false)))
-        return failure();
-    } while (curToken.getKind() != FormatToken::r_paren);
-    consumeToken();
-  }
-
-  if (failed(parseToken(FormatToken::question,
-                        "expected '?' after optional group")))
-    return ::mlir::failure();
-
-  // The optional group is required to have an anchor.
-  if (!anchorIdx)
-    return emitError(curLoc, "optional group specified no anchor element");
-
-  // The first parsable element of the group must be able to be parsed in an
-  // optional fashion.
-  auto parseBegin = llvm::find_if_not(thenElements, [](auto &element) {
-    return isa<WhitespaceElement>(element.get());
-  });
-  Element *firstElement = parseBegin->get();
-  if (!isa<AttributeVariable>(firstElement) &&
-      !isa<LiteralElement>(firstElement) &&
-      !isa<OperandVariable>(firstElement) && !isa<RegionVariable>(firstElement))
-    return emitError(curLoc,
-                     "first parsable element of an operand group must be "
-                     "an attribute, literal, operand, or region");
-
-  auto parseStart = parseBegin - thenElements.begin();
-  element = std::make_unique<OptionalElement>(
-      std::move(thenElements), std::move(elseElements), *anchorIdx, parseStart);
-  return ::mlir::success();
-}
-
-LogicalResult FormatParser::parseOptionalChildElement(
-    std::vector<std::unique_ptr<Element>> &childElements,
-    Optional<unsigned> &anchorIdx) {
-  llvm::SMLoc childLoc = curToken.getLoc();
-  childElements.push_back({});
-  if (failed(parseElement(childElements.back(), TopLevelContext)))
-    return ::mlir::failure();
-
-  // Check to see if this element is the anchor of the optional group.
-  bool isAnchor = curToken.getKind() == FormatToken::caret;
-  if (isAnchor) {
-    if (anchorIdx)
-      return emitError(childLoc, "only one element can be marked as the anchor "
-                                 "of an optional group");
-    anchorIdx = childElements.size() - 1;
-    consumeToken();
-  }
-
-  return verifyOptionalChildElement(childElements.back().get(), childLoc,
-                                    isAnchor);
-}
-
-LogicalResult FormatParser::verifyOptionalChildElement(Element *element,
-                                                       llvm::SMLoc childLoc,
-                                                       bool isAnchor) {
-  return TypeSwitch<Element *, LogicalResult>(element)
-      // All attributes can be within the optional group, but only optional
-      // attributes can be the anchor.
-      .Case([&](AttributeVariable *attrEle) {
-        if (isAnchor && !attrEle->getVar()->attr.isOptional())
-          return emitError(childLoc, "only optional attributes can be used to "
-                                     "anchor an optional group");
-        return ::mlir::success();
-      })
-      // Only optional-like(i.e. variadic) operands can be within an optional
-      // group.
-      .Case([&](OperandVariable *ele) {
-        if (!ele->getVar()->isVariableLength())
-          return emitError(childLoc, "only variable length operands can be "
-                                     "used within an optional group");
-        return ::mlir::success();
-      })
-      // Only optional-like(i.e. variadic) results can be within an optional
-      // group.
-      .Case([&](ResultVariable *ele) {
-        if (!ele->getVar()->isVariableLength())
-          return emitError(childLoc, "only variable length results can be "
-                                     "used within an optional group");
-        return ::mlir::success();
-      })
-      .Case([&](RegionVariable *) {
-        // TODO: When ODS has proper support for marking "optional" regions, add
-        // a check here.
-        return ::mlir::success();
-      })
-      .Case([&](TypeDirective *ele) {
-        return verifyOptionalChildElement(ele->getOperand(), childLoc,
-                                          /*isAnchor=*/false);
-      })
-      .Case([&](FunctionalTypeDirective *ele) {
-        if (failed(verifyOptionalChildElement(ele->getInputs(), childLoc,
-                                              /*isAnchor=*/false)))
-          return failure();
-        return verifyOptionalChildElement(ele->getResults(), childLoc,
-                                          /*isAnchor=*/false);
-      })
-      // Literals, whitespace, and custom directives may be used, but they can't
-      // anchor the group.
-      .Case<LiteralElement, WhitespaceElement, CustomDirective,
-            FunctionalTypeDirective, OptionalElement>([&](Element *) {
-        if (isAnchor)
-          return emitError(childLoc, "only variables and types can be used "
-                                     "to anchor an optional group");
-        return ::mlir::success();
-      })
-      .Default([&](Element *) {
-        return emitError(childLoc, "only literals, types, and variables can be "
-                                   "used within an optional group");
-      });
-}
-
-LogicalResult
-FormatParser::parseAttrDictDirective(std::unique_ptr<Element> &element,
-                                     llvm::SMLoc loc, ParserContext context,
+FailureOr<FormatElement *>
+OpFormatParser::parseAttrDictDirective(llvm::SMLoc loc, Context context,
                                        bool withKeyword) {
   if (context == TypeDirectiveContext)
     return emitError(loc, "'attr-dict' directive can only be used as a "
                           "top-level directive");
 
   if (context == RefDirectiveContext) {
     if (!hasAttrDict)
       return emitError(loc, "'ref' of 'attr-dict' is not bound by a prior "
                             "'attr-dict' directive");
 
     // Otherwise, this is a top-level context.
   } else {
     if (hasAttrDict)
       return emitError(loc, "'attr-dict' directive has already been seen");
     hasAttrDict = true;
   }
 
-  element = std::make_unique<AttrDictDirective>(withKeyword);
-  return ::mlir::success();
+  return create<AttrDictDirective>(withKeyword);
 }
 
-LogicalResult
-FormatParser::parseCustomDirective(std::unique_ptr<Element> &element,
-                                   llvm::SMLoc loc, ParserContext context) {
-  llvm::SMLoc curLoc = curToken.getLoc();
-  if (context != TopLevelContext)
-    return emitError(loc, "'custom' is only valid as a top-level directive");
-
-  // Parse the custom directive name.
-  if (failed(parseToken(FormatToken::less,
-                        "expected '<' before custom directive name")))
-    return ::mlir::failure();
-
-  FormatToken nameTok = curToken;
-  if (failed(parseToken(FormatToken::identifier,
-                        "expected custom directive name identifier")) ||
-      failed(parseToken(FormatToken::greater,
-                        "expected '>' after custom directive name")) ||
-      failed(parseToken(FormatToken::l_paren,
-                        "expected '(' before custom directive parameters")))
-    return ::mlir::failure();
-
-  // Parse the child elements for this optional group.=
-  std::vector<std::unique_ptr<Element>> elements;
-  do {
-    if (failed(parseCustomDirectiveParameter(elements)))
-      return ::mlir::failure();
-    if (curToken.getKind() != FormatToken::comma)
-      break;
-    consumeToken();
-  } while (true);
-
-  if (failed(parseToken(FormatToken::r_paren,
-                        "expected ')' after custom directive parameters")))
-    return ::mlir::failure();
-
-  // After parsing all of the elements, ensure that all type directives refer
-  // only to variables.
-  for (auto &ele : elements) {
-    if (auto *typeEle = dyn_cast<TypeDirective>(ele.get())) {
-      if (!isa<OperandVariable, ResultVariable>(typeEle->getOperand())) {
-        return emitError(curLoc, "type directives within a custom directive "
-                                 "may only refer to variables");
-      }
+LogicalResult OpFormatParser::verifyCustomDirectiveArguments(
+    llvm::SMLoc loc, ArrayRef<FormatElement *> arguments) {
+  for (FormatElement *argument : arguments) {
+    if (!isa<RefDirective, TypeDirective, AttrDictDirective, AttributeVariable,
+             OperandVariable, RegionVariable, SuccessorVariable>(argument)) {
+      // TODO: FormatElement should have location info attached.
+      return emitError(loc, "only variables and types may be used as "
+                            "parameters to a custom directive");
     }
+    if (auto *type = dyn_cast<TypeDirective>(argument)) {
+      if (!isa<OperandVariable, ResultVariable>(type->getArg())) {
+        return emitError(loc, "type directives within a custom directive may "
+                              "only refer to variables");
       }
-
-  element = std::make_unique<CustomDirective>(nameTok.getSpelling(),
-                                              std::move(elements));
-  return ::mlir::success();
     }
-
-LogicalResult FormatParser::parseCustomDirectiveParameter(
-    std::vector<std::unique_ptr<Element>> &parameters) {
-  llvm::SMLoc childLoc = curToken.getLoc();
-  parameters.push_back({});
-  if (failed(parseElement(parameters.back(), CustomDirectiveContext)))
-    return ::mlir::failure();
-
-  // Verify that the element can be placed within a custom directive.
-  if (!isa<RefDirective, TypeDirective, AttrDictDirective, AttributeVariable,
-           OperandVariable, RegionVariable, SuccessorVariable>(
-          parameters.back().get())) {
-    return emitError(childLoc, "only variables and types may be used as "
-                               "parameters to a custom directive");
   }
-  return ::mlir::success();
+  return success();
 }
 
-LogicalResult FormatParser::parseFunctionalTypeDirective(
-    std::unique_ptr<Element> &element, FormatToken tok, ParserContext context) {
-  llvm::SMLoc loc = tok.getLoc();
+FailureOr<FormatElement *>
+OpFormatParser::parseFunctionalTypeDirective(llvm::SMLoc loc, Context context) {
   if (context != TopLevelContext)
     return emitError(
         loc, "'functional-type' is only valid as a top-level directive");
 
   // Parse the main operand.
-  std::unique_ptr<Element> inputs, results;
+  FailureOr<FormatElement *> inputs, results;
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")) ||
-      failed(parseTypeDirectiveOperand(inputs)) ||
+      failed(inputs = parseTypeDirectiveOperand(loc)) ||
       failed(parseToken(FormatToken::comma,
                         "expected ',' after inputs argument")) ||
-      failed(parseTypeDirectiveOperand(results)) ||
+      failed(results = parseTypeDirectiveOperand(loc)) ||
       failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
-    return ::mlir::failure();
-  element = std::make_unique<FunctionalTypeDirective>(std::move(inputs),
-                                                      std::move(results));
-  return ::mlir::success();
+    return failure();
+  return create<FunctionalTypeDirective>(*inputs, *results);
 }
 
-LogicalResult
-FormatParser::parseOperandsDirective(std::unique_ptr<Element> &element,
-                                     llvm::SMLoc loc, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseOperandsDirective(llvm::SMLoc loc, Context context) {
   if (context == RefDirectiveContext) {
     if (!fmt.allOperands)
       return emitError(loc, "'ref' of 'operands' is not bound by a prior "
                             "'operands' directive");
 
   } else if (context == TopLevelContext || context == CustomDirectiveContext) {
     if (fmt.allOperands || !seenOperands.empty())
       return emitError(loc, "'operands' directive creates overlap in format");
     fmt.allOperands = true;
   }
-  element = std::make_unique<OperandsDirective>();
-  return ::mlir::success();
+  return create<OperandsDirective>();
 }
 
-LogicalResult
-FormatParser::parseReferenceDirective(std::unique_ptr<Element> &element,
-                                      llvm::SMLoc loc, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseReferenceDirective(llvm::SMLoc loc, Context context) {
   if (context != CustomDirectiveContext)
     return emitError(loc, "'ref' is only valid within a `custom` directive");
 
-  std::unique_ptr<Element> operand;
+  FailureOr<FormatElement *> arg;
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")) ||
-      failed(parseElement(operand, RefDirectiveContext)) ||
+      failed(arg = parseElement(RefDirectiveContext)) ||
       failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
-    return ::mlir::failure();
+    return failure();
 
-  element = std::make_unique<RefDirective>(std::move(operand));
-  return ::mlir::success();
+  return create<RefDirective>(*arg);
 }
 
-LogicalResult
-FormatParser::parseRegionsDirective(std::unique_ptr<Element> &element,
-                                    llvm::SMLoc loc, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseRegionsDirective(llvm::SMLoc loc, Context context) {
   if (context == TypeDirectiveContext)
     return emitError(loc, "'regions' is only valid as a top-level directive");
   if (context == RefDirectiveContext) {
     if (!hasAllRegions)
       return emitError(loc, "'ref' of 'regions' is not bound by a prior "
                             "'regions' directive");
 
     // Otherwise, this is a TopLevel directive.
   } else {
     if (hasAllRegions || !seenRegions.empty())
       return emitError(loc, "'regions' directive creates overlap in format");
     hasAllRegions = true;
   }
-  element = std::make_unique<RegionsDirective>();
-  return ::mlir::success();
+  return create<RegionsDirective>();
 }
 
-LogicalResult
-FormatParser::parseResultsDirective(std::unique_ptr<Element> &element,
-                                    llvm::SMLoc loc, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseResultsDirective(llvm::SMLoc loc, Context context) {
   if (context != TypeDirectiveContext)
     return emitError(loc, "'results' directive can can only be used as a child "
                           "to a 'type' directive");
-  element = std::make_unique<ResultsDirective>();
-  return ::mlir::success();
+  return create<ResultsDirective>();
 }
 
-LogicalResult
-FormatParser::parseSuccessorsDirective(std::unique_ptr<Element> &element,
-                                       llvm::SMLoc loc, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseSuccessorsDirective(llvm::SMLoc loc, Context context) {
   if (context == TypeDirectiveContext)
     return emitError(loc,
                      "'successors' is only valid as a top-level directive");
   if (context == RefDirectiveContext) {
     if (!hasAllSuccessors)
       return emitError(loc, "'ref' of 'successors' is not bound by a prior "
                             "'successors' directive");
 
     // Otherwise, this is a TopLevel directive.
   } else {
     if (hasAllSuccessors || !seenSuccessors.empty())
       return emitError(loc, "'successors' directive creates overlap in format");
     hasAllSuccessors = true;
   }
-  element = std::make_unique<SuccessorsDirective>();
-  return ::mlir::success();
+  return create<SuccessorsDirective>();
 }
 
-LogicalResult
-FormatParser::parseTypeDirective(std::unique_ptr<Element> &element,
-                                 FormatToken tok, ParserContext context) {
-  llvm::SMLoc loc = tok.getLoc();
+FailureOr<FormatElement *> OpFormatParser::parseTypeDirective(llvm::SMLoc loc,
+                                                              Context context) {
   if (context == TypeDirectiveContext)
     return emitError(loc, "'type' cannot be used as a child of another `type`");
 
   bool isRefChild = context == RefDirectiveContext;
-  std::unique_ptr<Element> operand;
+  FailureOr<FormatElement *> operand;
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")) ||
-      failed(parseTypeDirectiveOperand(operand, isRefChild)) ||
+      failed(operand = parseTypeDirectiveOperand(loc, isRefChild)) ||
       failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
-    return ::mlir::failure();
+    return failure();
 
-  element = std::make_unique<TypeDirective>(std::move(operand));
-  return ::mlir::success();
+  return create<TypeDirective>(*operand);
 }
 
-LogicalResult
-FormatParser::parseQualifiedDirective(std::unique_ptr<Element> &element,
-                                      FormatToken tok, ParserContext context) {
+FailureOr<FormatElement *>
+OpFormatParser::parseQualifiedDirective(llvm::SMLoc loc, Context context) {
+  FailureOr<FormatElement *> element;
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")) ||
-      failed(parseElement(element, context)) ||
+      failed(element = parseElement(context)) ||
       failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
     return failure();
-  if (auto *attr = dyn_cast<AttributeVariable>(element.get())) {
-    attr->setShouldBeQualified();
-  } else if (auto *type = dyn_cast<TypeDirective>(element.get())) {
-    type->setShouldBeQualified();
-  } else {
+  return TypeSwitch<FormatElement *, FailureOr<FormatElement *>>(*element)
+      .Case<AttributeVariable, TypeDirective>([](auto *element) {
+        element->setShouldBeQualified();
+        return element;
+      })
+      .Default([&](auto *element) {
         return emitError(
-        tok.getLoc(),
+            loc,
             "'qualified' directive expects an attribute or a `type` directive");
-  }
-  return success();
+      });
 }
 
-LogicalResult
-FormatParser::parseTypeDirectiveOperand(std::unique_ptr<Element> &element,
-                                        bool isRefChild) {
-  llvm::SMLoc loc = curToken.getLoc();
-  if (failed(parseElement(element, TypeDirectiveContext)))
-    return ::mlir::failure();
-  if (isa<LiteralElement>(element.get()))
+FailureOr<FormatElement *>
+OpFormatParser::parseTypeDirectiveOperand(llvm::SMLoc loc, bool isRefChild) {
+  FailureOr<FormatElement *> result = parseElement(TypeDirectiveContext);
+  if (failed(result))
+    return failure();
+
+  FormatElement *element = *result;
+  if (isa<LiteralElement>(element))
     return emitError(
         loc, "'type' directive operand expects variable or directive operand");
 
-  if (auto *var = dyn_cast<OperandVariable>(element.get())) {
+  if (auto *var = dyn_cast<OperandVariable>(element)) {
     unsigned opIdx = var->getVar() - op.operand_begin();
     if (!isRefChild && (fmt.allOperandTypes || seenOperandTypes.test(opIdx)))
       return emitError(loc, "'type' of '" + var->getVar()->name +
                                 "' is already bound");
     if (isRefChild && !(fmt.allOperandTypes || seenOperandTypes.test(opIdx)))
       return emitError(loc, "'ref' of 'type($" + var->getVar()->name +
                                 ")' is not bound by a prior 'type' directive");
     seenOperandTypes.set(opIdx);
-  } else if (auto *var = dyn_cast<ResultVariable>(element.get())) {
+  } else if (auto *var = dyn_cast<ResultVariable>(element)) {
     unsigned resIdx = var->getVar() - op.result_begin();
     if (!isRefChild && (fmt.allResultTypes || seenResultTypes.test(resIdx)))
       return emitError(loc, "'type' of '" + var->getVar()->name +
                                 "' is already bound");
     if (isRefChild && !(fmt.allResultTypes || seenResultTypes.test(resIdx)))
       return emitError(loc, "'ref' of 'type($" + var->getVar()->name +
                                 ")' is not bound by a prior 'type' directive");
     seenResultTypes.set(resIdx);
if (!isRefChild && (fmt.allResultTypes || seenResultTypes.any()))
       return emitError(loc, "'results' 'type' is already bound");
     if (isRefChild && !fmt.allResultTypes)
       return emitError(loc, "'ref' of 'type(results)' is not bound by a prior "
                             "'type' directive");
     fmt.allResultTypes = true;
   } else {
     return emitError(loc, "invalid argument to 'type' directive");
   }
-  return ::mlir::success();
+  return element;
+}
+
+LogicalResult
+OpFormatParser::verifyOptionalGroupElements(llvm::SMLoc loc,
+                                            ArrayRef<FormatElement *> elements,
+                                            Optional<unsigned> anchorIndex) {
+  for (auto &it : llvm::enumerate(elements)) {
+    if (failed(verifyOptionalGroupElement(
+            loc, it.value(), anchorIndex && *anchorIndex == it.index())))
+      return failure();
+  }
+  return success();
+}
+
+LogicalResult OpFormatParser::verifyOptionalGroupElement(llvm::SMLoc loc,
+                                                         FormatElement *element,
+                                                         bool isAnchor) {
+  return TypeSwitch<FormatElement *, LogicalResult>(element)
+      // All attributes can be within the optional group, but only optional
+      // attributes can be the anchor.
+      .Case([&](AttributeVariable *attrEle) {
+        if (isAnchor && !attrEle->getVar()->attr.isOptional())
+          return emitError(loc, "only optional attributes can be used to "
+                                "anchor an optional group");
+        return success();
+      })
+      // Only optional-like(i.e. variadic) operands can be within an optional
+      // group.
+      .Case([&](OperandVariable *ele) {
+        if (!ele->getVar()->isVariableLength())
+          return emitError(loc, "only variable length operands can be used "
+                                "within an optional group");
+        return success();
+      })
+      // Only optional-like(i.e. variadic) results can be within an optional
+      // group.
+      .Case([&](ResultVariable *ele) {
+        if (!ele->getVar()->isVariableLength())
+          return emitError(loc, "only variable length results can be used "
+                                "within an optional group");
+        return success();
+      })
+      .Case([&](RegionVariable *) {
+        // TODO: When ODS has proper support for marking "optional" regions, add
+        // a check here.
+        return success();
+      })
+      .Case([&](TypeDirective *ele) {
+        return verifyOptionalGroupElement(loc, ele->getArg(),
+                                          /*isAnchor=*/false);
+      })
+      .Case([&](FunctionalTypeDirective *ele) {
+        if (failed(verifyOptionalGroupElement(loc, ele->getInputs(),
+                                              /*isAnchor=*/false)))
+          return failure();
+        return verifyOptionalGroupElement(loc, ele->getResults(),
+                                          /*isAnchor=*/false);
+      })
+      // Literals, whitespace, and custom directives may be used, but they can't
+      // anchor the group.
+      .Case<LiteralElement, WhitespaceElement, CustomDirective,
+            FunctionalTypeDirective, OptionalElement>([&](FormatElement *) {
+        if (isAnchor)
+          return emitError(loc, "only variables and types can be used "
+                                "to anchor an optional group");
+        return success();
+      })
+      .Default([&](FormatElement *) {
+        return emitError(loc, "only literals, types, and variables can be "
+                              "used within an optional group");
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // Interface
 //===----------------------------------------------------------------------===//
 
 void mlir::tblgen::generateOpFormat(const Operator &constOp, OpClass &opClass) {
   // TODO: Operator doesn't expose all necessary functionality via
   // the const interface.
   Operator &op = const_cast<Operator &>(constOp);
   if (!op.hasAssemblyFormat())
     return;
 
   // Parse the format description.
   llvm::SourceMgr mgr;
   mgr.AddNewSourceBuffer(
       llvm::MemoryBuffer::getMemBuffer(op.getAssemblyFormat()), llvm::SMLoc());
   OperationFormat format(op);
-  if (failed(FormatParser(mgr, format, op).parse())) {
+  OpFormatParser parser(mgr, format, op);
+  FailureOr<std::vector<FormatElement *>> elements = parser.parse();
+  if (failed(elements)) {
     // Exit the process if format errors are treated as fatal.
     if (formatErrorIsFatal) {
       // Invoke the interrupt handlers to run the file cleanup handlers.
       llvm::sys::RunInterruptHandlers();
       std::exit(1);
     }
     return;
   }
+  format.elements = std::move(*elements);
 
   // Generate the printer and parser based on the parsed format.
   format.genParser(op, opClass);
   format.genPrinter(op, opClass);
 }