65 changes: 50 additions & 15 deletions mlir/tools/mlir-tblgen/OpFormatGen.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,8 @@ class Element {
/// This element is a literal.
Literal,

/// This element prints or omits a space. It is ignored by the parser.
/// This element is a whitespace.
Newline,
Space,

/// This element is an variable value.
Expand Down Expand Up @@ -296,14 +297,35 @@ bool LiteralElement::isValidLiteral(StringRef value) {
}

//===----------------------------------------------------------------------===//
// SpaceElement
// 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 Element {
class SpaceElement : public WhitespaceElement {
public:
SpaceElement(bool value) : Element{Kind::Space}, value(value) {}
SpaceElement(bool value) : WhitespaceElement(Kind::Space), value(value) {}
static bool classof(const Element *element) {
return element->getKind() == Kind::Space;
}
Expand Down Expand Up @@ -347,7 +369,8 @@ class OptionalElement : public Element {
std::vector<std::unique_ptr<Element>> elements;
/// 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 SpaceElement).
/// The index of the first element that is parsed (is not a
/// WhitespaceElement).
unsigned parseStart;
};
} // end anonymous namespace
Expand Down Expand Up @@ -1098,8 +1121,8 @@ void OperationFormat::genElementParser(Element *element, OpMethodBody &body,
genLiteralParser(literal->getLiteral(), body);
body << ")\n return ::mlir::failure();\n";

/// Spaces.
} else if (isa<SpaceElement>(element)) {
/// Whitespaces.
} else if (isa<WhitespaceElement>(element)) {
// Nothing to parse.

/// Arguments.
Expand Down Expand Up @@ -1620,6 +1643,11 @@ void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
return genLiteralPrinter(literal->getLiteral(), body, shouldEmitSpace,
lastWasPunctuation);

// Emit a whitespace element.
if (NewlineElement *newline = dyn_cast<NewlineElement>(element)) {
body << " p.printNewline();\n";
return;
}
if (SpaceElement *space = dyn_cast<SpaceElement>(element))
return genSpacePrinter(space->getValue(), body, shouldEmitSpace,
lastWasPunctuation);
Expand Down Expand Up @@ -2272,9 +2300,10 @@ LogicalResult FormatParser::verifyAttributes(
for (auto &nextItPair : iteratorStack) {
ElementsIterT nextIt = nextItPair.first, nextE = nextItPair.second;
for (; nextIt != nextE; ++nextIt) {
// Skip any trailing spaces, attribute dictionaries, or optional groups.
if (isa<SpaceElement>(*nextIt) || isa<AttrDictDirective>(*nextIt) ||
isa<OptionalElement>(*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.
Expand Down Expand Up @@ -2600,6 +2629,11 @@ LogicalResult FormatParser::parseLiteral(std::unique_ptr<Element> &element) {
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 (!LiteralElement::isValidLiteral(value))
Expand Down Expand Up @@ -2635,8 +2669,9 @@ LogicalResult FormatParser::parseOptional(std::unique_ptr<Element> &element,

// The first parsable element of the group must be able to be parsed in an
// optional fashion.
auto parseBegin = llvm::find_if_not(
elements, [](auto &element) { return isa<SpaceElement>(element.get()); });
auto parseBegin = llvm::find_if_not(elements, [](auto &element) {
return isa<WhitespaceElement>(element.get());
});
Element *firstElement = parseBegin->get();
if (!isa<AttributeVariable>(firstElement) &&
!isa<LiteralElement>(firstElement) &&
Expand Down Expand Up @@ -2718,9 +2753,9 @@ LogicalResult FormatParser::parseOptionalChildElement(
// a check here.
return ::mlir::success();
})
// Literals, spaces, custom directives, and type directives may be used,
// but they can't anchor the group.
.Case<LiteralElement, SpaceElement, CustomDirective,
// Literals, whitespace, custom directives, and type directives may be
// used, but they can't anchor the group.
.Case<LiteralElement, WhitespaceElement, CustomDirective,
FunctionalTypeDirective, OptionalElement, TypeRefDirective,
TypeDirective>([&](Element *) {
if (isAnchor)
Expand Down