Add (parse|print)OptionalAttrDictWithKeyword hooks to simplify parsing attribute dictionaries with regions.

Many operations with regions add an additional 'attributes' prefix when printing the attribute dictionary to differentiate it from the region body. This leads to duplicated logic for detecting when to actually print the attribute dictionary.

PiperOrigin-RevId: 278747681

5 files changed, 40 insertions, 45 deletions

diff --git a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
index e8176191a8d..778e7cb3b24 100644
--- a/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/SPIRVOps.cpp
@@ -1628,10 +1628,8 @@ static ParseResult parseModuleOp(OpAsmParser &parser, OperationState &state) {
   if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
     return failure();
 
-  if (succeeded(parser.parseOptionalKeyword("attributes"))) {
-    if (parser.parseOptionalAttrDict(state.attributes))
-      return failure();
-  }
+  if (parser.parseOptionalAttrDictWithKeyword(state.attributes))
+    return failure();
 
   spirv::ModuleOp::ensureTerminator(*body, parser.getBuilder(), state.location);
   return success();
@@ -1657,19 +1655,7 @@ static void print(spirv::ModuleOp moduleOp, OpAsmPrinter &printer) {
 
   printer.printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
                       /*printBlockTerminators=*/false);
-
-  bool printAttrDict =
-      elidedAttrs.size() != 2 ||
-      llvm::any_of(op->getAttrs(), [&addressingModelAttrName,
-                                    &memoryModelAttrName](NamedAttribute attr) {
-        return attr.first != addressingModelAttrName &&
-               attr.first != memoryModelAttrName;
-      });
-
-  if (printAttrDict) {
-    printer << " attributes";
-    printer.printOptionalAttrDict(op->getAttrs(), elidedAttrs);
-  }
+  printer.printOptionalAttrDictWithKeyword(op->getAttrs(), elidedAttrs);
 }
 
 static LogicalResult verify(spirv::ModuleOp moduleOp) {
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 0e6b7882e14..af958c8e61f 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -421,7 +421,8 @@ public:
 
 protected:
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
-                             ArrayRef<StringRef> elidedAttrs = {});
+                             ArrayRef<StringRef> elidedAttrs = {},
+                             bool withKeyword = false);
   void printTrailingLocation(Location loc);
   void printLocationInternal(LocationAttr loc, bool pretty = false);
   void printDenseElementsAttr(DenseElementsAttr attr);
@@ -1327,27 +1328,26 @@ void ModulePrinter::printIntegerSet(IntegerSet set) {
 //===----------------------------------------------------------------------===//
 
 void ModulePrinter::printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
-                                          ArrayRef<StringRef> elidedAttrs) {
+                                          ArrayRef<StringRef> elidedAttrs,
+                                          bool withKeyword) {
   // If there are no attributes, then there is nothing to be done.
   if (attrs.empty())
     return;
 
   // Filter out any attributes that shouldn't be included.
-  SmallVector<NamedAttribute, 8> filteredAttrs;
-  for (auto attr : attrs) {
-    // If the caller has requested that this attribute be ignored, then drop it.
-    if (llvm::any_of(elidedAttrs,
-                     [&](StringRef elided) { return attr.first.is(elided); }))
-      continue;
-
-    // Otherwise add it to our filteredAttrs list.
-    filteredAttrs.push_back(attr);
-  }
+  SmallVector<NamedAttribute, 8> filteredAttrs(
+      llvm::make_filter_range(attrs, [&](NamedAttribute attr) {
+        return !llvm::is_contained(elidedAttrs, attr.first.strref());
+      }));
 
   // If there are no attributes left to print after filtering, then we're done.
   if (filteredAttrs.empty())
     return;
 
+  // Print the 'attributes' keyword if necessary.
+  if (withKeyword)
+    os << " attributes ";
+
   // Otherwise, print them all out in braces.
   os << " {";
   interleaveComma(filteredAttrs, [&](NamedAttribute attr) {
@@ -1389,8 +1389,14 @@ public:
 
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                              ArrayRef<StringRef> elidedAttrs = {}) override {
-    return ModulePrinter::printOptionalAttrDict(attrs, elidedAttrs);
-  };
+    ModulePrinter::printOptionalAttrDict(attrs, elidedAttrs);
+  }
+  void printOptionalAttrDictWithKeyword(
+      ArrayRef<NamedAttribute> attrs,
+      ArrayRef<StringRef> elidedAttrs = {}) override {
+    ModulePrinter::printOptionalAttrDict(attrs, elidedAttrs,
+                                         /*withKeyword=*/true);
+  }
 
   enum { nameSentinel = ~0U };
 
diff --git a/mlir/lib/IR/FunctionSupport.cpp b/mlir/lib/IR/FunctionSupport.cpp
index aa9965e89d1..d1ba2d30fa1 100644
--- a/mlir/lib/IR/FunctionSupport.cpp
+++ b/mlir/lib/IR/FunctionSupport.cpp
@@ -183,9 +183,8 @@ mlir::impl::parseFunctionLikeOp(OpAsmParser &parser, OperationState &result,
            << (errorMessage.empty() ? "" : ": ") << errorMessage;
 
   // If function attributes are present, parse them.
-  if (succeeded(parser.parseOptionalKeyword("attributes")))
-    if (parser.parseOptionalAttrDict(result.attributes))
-      return failure();
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
+    return failure();
 
   // Add the attributes to the function arguments.
   SmallString<8> attrNameBuf;
diff --git a/mlir/lib/IR/Module.cpp b/mlir/lib/IR/Module.cpp
index 3a08ee3bf27..f5cc98e39cf 100644
--- a/mlir/lib/IR/Module.cpp
+++ b/mlir/lib/IR/Module.cpp
@@ -48,9 +48,8 @@ ParseResult ModuleOp::parse(OpAsmParser &parser, OperationState &result) {
                                result.attributes);
 
   // If module attributes are present, parse them.
-  if (succeeded(parser.parseOptionalKeyword("attributes")))
-    if (parser.parseOptionalAttrDict(result.attributes))
-      return failure();
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
+    return failure();
 
   // Parse the module body.
   auto *body = result.addRegion();
@@ -65,18 +64,14 @@ ParseResult ModuleOp::parse(OpAsmParser &parser, OperationState &result) {
 void ModuleOp::print(OpAsmPrinter &p) {
   p << "module";
 
-  Optional<StringRef> name = getName();
-  if (name) {
+  if (Optional<StringRef> name = getName()) {
     p << ' ';
     p.printSymbolName(*name);
   }
 
   // Print the module attributes.
-  auto attrs = getAttrs();
-  if (!attrs.empty() && !(attrs.size() == 1 && name)) {
-    p << " attributes";
-    p.printOptionalAttrDict(attrs, {mlir::SymbolTable::getSymbolAttrName()});
-  }
+  p.printOptionalAttrDictWithKeyword(getAttrs(),
+                                     {mlir::SymbolTable::getSymbolAttrName()});
 
   // Print the region.
   p.printRegion(getOperation()->getRegion(0), /*printEntryBlockArgs=*/false,
diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 3a45933db87..35c694b6a43 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -1533,7 +1533,7 @@ Attribute Parser::parseAttribute(Type type) {
 ///
 ParseResult
 Parser::parseAttributeDict(SmallVectorImpl<NamedAttribute> &attributes) {
-  if (!consumeIf(Token::l_brace))
+  if (parseToken(Token::l_brace, "expected '{' in attribute dictionary"))
     return failure();
 
   auto parseElt = [&]() -> ParseResult {
@@ -3874,6 +3874,15 @@ public:
     return parser.parseAttributeDict(result);
   }
 
+  /// Parse a named dictionary into 'result' if the `attributes` keyword is
+  /// present.
+  ParseResult parseOptionalAttrDictWithKeyword(
+      SmallVectorImpl<NamedAttribute> &result) override {
+    if (failed(parseOptionalKeyword("attributes")))
+      return success();
+    return parser.parseAttributeDict(result);
+  }
+
   //===--------------------------------------------------------------------===//
   // Identifier Parsing
   //===--------------------------------------------------------------------===//