[mlir][Parser] Make parse{Attribute,Type} null-terminate input

[rkayaith]

Summary

parseAttribute and parseType require null-terminated strings as
input, but this isn't great considering the argument type is
StringRef. This changes them to copy to a null-terminated buffer by
default, with a isKnownNullTerminated flag added to disable the
copying.

closes #58964

[kuhar]

mlir/include/mlir/AsmParser/AsmParser.h

72

nit: typeStr.str()? Also everywhere else

[rkayaith]

std::string() -> .str()

[kuhar]

mlir/unittests/Parser/ParserTest.cpp

99

nit: StringRef attrAsm("999", 1);

[rkayaith]

address review comment

[kuhar]

LGTM but please wait for a second review before submitting

[lattner]

Cool. Suggestion: instead of duplicating the APIs, you could add a , bool isKnownNullTerminated = false flag. That keeps the API simpler and maintains a safe default while allowing clients to avoid the copy if they are sufficiently smart and already null terminated.

[rkayaith]

use flag instead of separate functions

[lattner]

Nice, thank you!

mlir/include/mlir/AsmParser/AsmParser.h

parseAsmSourceFile(const llvm::SourceMgr &sourceMgr, Block *block,
                    AsmParserCodeCompleteContext *codeCompleteContext = nullptr);
 
 /// This parses a single MLIR attribute to an MLIR context if it was valid. If
 /// not, an error diagnostic is emitted to the context and a null value is
 /// returned.
 /// If `numRead` is provided, it is set to the number of consumed characters on
 /// succesful parse. Otherwise, parsing fails if the entire string is not
 /// consumed.
-Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context,
-                         Type type = {}, size_t *numRead = nullptr);
+Attribute parseNullTerminatedAttribute(llvm::StringRef attrStr,
+                                       MLIRContext *context, Type type = {},
+                                       size_t *numRead = nullptr);
+inline Attribute parseAttribute(llvm::StringRef attrStr, MLIRContext *context,
+                                Type type = {}, size_t *numRead = nullptr) {
+  // Ensure source is null-terminated by copying to std::string
+  return parseNullTerminatedAttribute(std::string(attrStr), context, type,
+                                      numRead);
+}
 
 /// This parses a single MLIR type to an MLIR context if it was valid. If not,
 /// an error diagnostic is emitted to the context.
 /// If `numRead` is provided, it is set to the number of consumed characters on
 /// succesful parse. Otherwise, parsing fails if the entire string is not
 /// consumed.
-Type parseType(llvm::StringRef typeStr, MLIRContext *context,
+Type parseNullTerminatedType(llvm::StringRef typeStr, MLIRContext *context,
                              size_t *numRead = nullptr);
+inline Type parseType(llvm::StringRef typeStr, MLIRContext *context,
+                      size_t *numRead = nullptr) {
+  // Ensure source is null-terminated by copying to std::string
+  return parseNullTerminatedType(std::string(typeStr), context, numRead);
+}
 
 /// This parses a single IntegerSet/AffineMap to an MLIR context if it was
 /// valid. If not, an error message is emitted through a new
 /// SourceMgrDiagnosticHandler constructed from a new SourceMgr with a single
 /// MemoryBuffer wrapping `str`. If the passed `str` has additional tokens that
 /// were not part of the IntegerSet/AffineMap, a failure is returned.
 AffineMap parseAffineMap(llvm::StringRef str, MLIRContext *context);
 IntegerSet parseIntegerSet(llvm::StringRef str, MLIRContext *context);
 
 } // namespace mlir
 
 #endif // MLIR_ASMPARSER_ASMPARSER_H

mlir/lib/AsmParser/DialectSymbolParser.cpp

static T parseSymbol(StringRef inputStr, MLIRContext *context,
   } else if (numRead != inputStr.size()) {
     parser.emitError(endTok.getLoc()) << "found trailing characters: '"
                                       << inputStr.drop_front(numRead) << "'";
     return T();
   }
   return symbol;
 }
 
-Attribute mlir::parseAttribute(StringRef attrStr, MLIRContext *context,
-                               Type type, size_t *numRead) {
+Attribute mlir::parseNullTerminatedAttribute(StringRef attrStr,
+                                             MLIRContext *context, Type type,
+                                             size_t *numRead) {
   return parseSymbol<Attribute>(
       attrStr, context, numRead,
       [type](Parser &parser) { return parser.parseAttribute(type); });
 }
-Type mlir::parseType(StringRef typeStr, MLIRContext *context, size_t *numRead) {
+Type mlir::parseNullTerminatedType(StringRef typeStr, MLIRContext *context,
+                                   size_t *numRead) {
   return parseSymbol<Type>(typeStr, context, numRead,
                            [](Parser &parser) { return parser.parseType(); });
 }

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
   StringRef asmStr;
   if (failed(reader.parseNullTerminatedString(asmStr)))
     return failure();
 
   // Invoke the MLIR assembly parser to parse the entry text.
   size_t numRead = 0;
   MLIRContext *context = fileLoc->getContext();
   if constexpr (std::is_same_v<T, Type>)
-    result = ::parseType(asmStr, context, &numRead);
+    result = ::parseNullTerminatedType(asmStr, context, &numRead);
   else
-    result = ::parseAttribute(asmStr, context, Type(), &numRead);
+    result = ::parseNullTerminatedAttribute(asmStr, context, Type(), &numRead);
   if (!result)
     return failure();
 
   // Ensure there weren't dangling characters after the entry.
   if (numRead != asmStr.size()) {
     return reader.emitError("trailing characters found after ", entryType,
                             " assembly format: ", asmStr.drop_front(numRead));
   }

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

for (auto const &it :
     if (!attr) {
       emitOpError("expects padding values to be typed attributes");
       return DiagnosedSilenceableFailure::definiteFailure();
     }
     Type elementType = getElementTypeOrSelf(std::get<1>(it));
     // Try to parse string attributes to obtain an attribute of element type.
     if (auto stringAttr = attr.dyn_cast<StringAttr>()) {
       auto parsedAttr = dyn_cast_if_present<TypedAttr>(
-          parseAttribute(stringAttr, getContext(), elementType));
+          parseNullTerminatedAttribute(stringAttr, getContext(), elementType));
       if (!parsedAttr || parsedAttr.getType() != elementType) {
         auto diag = this->emitOpError("expects a padding that parses to ")
                     << elementType << ", got " << std::get<0>(it);
         diag.attachNote(target.getLoc()) << "when applied to this op";
         return DiagnosedSilenceableFailure::definiteFailure();
       }
       paddingValues.push_back(parsedAttr);
       continue;

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

static void output(const SerializedAffineMap &value, void *rawYamlContext,
                      raw_ostream &out) {
     assert(value.affineMapAttr);
     value.affineMapAttr.print(out);
   }
   static StringRef input(StringRef scalar, void *rawYamlContext,
                          SerializedAffineMap &value) {
     assert(rawYamlContext);
     auto *yamlContext = static_cast<LinalgYAMLContext *>(rawYamlContext);
-    std::string nullTerminatedScalar(scalar);
-    if (auto attr =
-            mlir::parseAttribute(nullTerminatedScalar, yamlContext->mlirContext)
+    if (auto attr = mlir::parseAttribute(scalar, yamlContext->mlirContext)
                         .dyn_cast_or_null<AffineMapAttr>())
       value.affineMapAttr = attr;
     else if (!value.affineMapAttr || !value.affineMapAttr.isa<AffineMapAttr>())
       return "could not parse as an affine map attribute";
     return StringRef();
   }
   static QuotingType mustQuote(StringRef) { return QuotingType::None; }
 };
 

mlir/unittests/Parser/ParserTest.cpp

EXPECT_THAT(
         diagnostics,
         ElementsAre("loc(\"10  foo\":1:5): found trailing characters: 'foo'"));
 
     size_t numRead = 0;
     EXPECT_EQ(parseAttribute("10  foo", &context, Type(), &numRead),
               b.getI64IntegerAttr(10));
     EXPECT_EQ(numRead, size_t(4)); // includes trailing whitespace
   }
+  { // Parse without null-terminator
+    StringRef attrAsm = StringRef("999", 1);
+    Attribute attr = parseAttribute(attrAsm, &context);
+    EXPECT_EQ(attr, b.getI64IntegerAttr(9));
+  }
 }
 } // namespace