[mlir:Bytecode][NFC] Cleanup Attribute/Type reading

[rriddle]

Summary

This moves some parsing functionality from BytecodeReader to
AttrTypeReader, and removes some duplication between the attribute/type
code paths.

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

public:
 
   /// Resolve the attribute or type at the given index. Returns nullptr on
   /// failure.
   Attribute resolveAttribute(size_t index) {
     return resolveEntry(attributes, index, "Attribute");
   }
   Type resolveType(size_t index) { return resolveEntry(types, index, "Type"); }
 
+  /// Parse a reference to an attribute or type using the given reader.
+  LogicalResult parseAttribute(EncodingReader &reader, Attribute &result) {
+    uint64_t attrIdx;
+    if (failed(reader.parseVarInt(attrIdx)))
+      return failure();
+    result = resolveAttribute(attrIdx);
+    return success(!!result);
+  }
+  LogicalResult parseType(EncodingReader &reader, Type &result) {
+    uint64_t typeIdx;
+    if (failed(reader.parseVarInt(typeIdx)))
+      return failure();
+    result = resolveType(typeIdx);
+    return success(!!result);
+  }
+
+  template <typename T>
+  LogicalResult parseAttribute(EncodingReader &reader, T &result) {
+    Attribute baseResult;
+    if (failed(parseAttribute(reader, baseResult)))
+      return failure();
+    if ((result = baseResult.dyn_cast<T>()))
+      return success();
+    return reader.emitError("expected attribute of type: ",
+                            llvm::getTypeName<T>(), ", but got: ", baseResult);
+  }
+
 private:
   /// Resolve the given entry at `index`.
   template <typename T>
   T resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
                  StringRef entryType);
 
-  /// Parse the value defined within the given reader. `code` indicates how the
-  /// entry was encoded.
-  LogicalResult parseEntry(EncodingReader &reader, bool hasCustomEncoding,
-                           Attribute &result);
-  LogicalResult parseEntry(EncodingReader &reader, bool hasCustomEncoding,
-                           Type &result);
+  /// Parse an entry using the given reader that was encoded using the textual
+  /// assembly format.
+  template <typename T>
+  LogicalResult parseAsmEntry(T &result, EncodingReader &reader,
+                              StringRef entryType);
+
+  /// Parse an entry using the given reader that was encoded using a custom
+  /// bytecode format.
+  template <typename T>
+  LogicalResult parseCustomEntry(Entry<T> &entry, EncodingReader &reader,
+                                 StringRef entryType);
 
   /// The set of attribute and type entries.
   SmallVector<AttrEntry> attributes;
   SmallVector<TypeEntry> types;
 
   /// A location used for error emission.
   Location fileLoc;
 };
T AttrTypeReader::resolveEntry(SmallVectorImpl<Entry<T>> &entries, size_t index,
 
   // If the entry has already been resolved, there is nothing left to do.
   Entry<T> &entry = entries[index];
   if (entry.entry)
     return entry.entry;
 
   // Parse the entry.
   EncodingReader reader(entry.data, fileLoc);
-  if (failed(parseEntry(reader, entry.hasCustomEncoding, entry.entry)))
+
+  // Parse based on how the entry was encoded.
+  if (entry.hasCustomEncoding) {
+    if (failed(parseCustomEntry(entry, reader, entryType)))
+      return T();
+  } else if (failed(parseAsmEntry(entry.entry, reader, entryType))) {
     return T();
+  }
+
   if (!reader.empty()) {
     (void)reader.emitError("unexpected trailing bytes after " + entryType +
                            " entry");
     return T();
   }
   return entry.entry;
 }
 
-LogicalResult AttrTypeReader::parseEntry(EncodingReader &reader,
-                                         bool hasCustomEncoding,
-                                         Attribute &result) {
-  // Handle the fallback case, where the attribute was encoded using its
-  // assembly format.
-  if (!hasCustomEncoding) {
-    StringRef attrStr;
-    if (failed(reader.parseNullTerminatedString(attrStr)))
-      return failure();
+template <typename T>
+LogicalResult AttrTypeReader::parseAsmEntry(T &result, EncodingReader &reader,
+                                            StringRef entryType) {
+  StringRef asmStr;
+  if (failed(reader.parseNullTerminatedString(asmStr)))
 
-    size_t numRead = 0;
-    if (!(result = parseAttribute(attrStr, fileLoc->getContext(), numRead)))
-      return failure();
-    if (numRead != attrStr.size()) {
-      return reader.emitError(
-          "trailing characters found after Attribute assembly format: ",
-          attrStr.drop_front(numRead));
-    }
-    return success();
-  }
-
-  return reader.emitError("unexpected Attribute encoding");
-}
-
-LogicalResult AttrTypeReader::parseEntry(EncodingReader &reader,
-                                         bool hasCustomEncoding, Type &result) {
-  // Handle the fallback case, where the type was encoded using its
-  // assembly format.
-  if (!hasCustomEncoding) {
-    StringRef typeStr;
-    if (failed(reader.parseNullTerminatedString(typeStr)))
-      return failure();
+  // Invoke the MLIR assembly parser to parse the entry text.
+  MLIRContext *context = fileLoc->getContext();
+  if constexpr (std::is_same_v<T, Type>)
+    result = ::parseType(asmStr, context, numRead);
+  else
+    result = ::parseAttribute(asmStr, context, numRead);
+  if (!result)
 
-    size_t numRead = 0;
-    if (!(result = parseType(typeStr, fileLoc->getContext(), numRead)))
-      return failure();
-    if (numRead != typeStr.size()) {
-      return reader.emitError(
-          "trailing characters found after Type assembly format: " +
-          typeStr.drop_front(numRead));
-    }
-    return success();
-  }
+  // 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));
 
-  return reader.emitError("unexpected Type encoding");
+template <typename T>
+LogicalResult AttrTypeReader::parseCustomEntry(Entry<T> &entry,
+                                               EncodingReader &reader,
+                                               StringRef entryType) {
+  // FIXME: Add support for reading custom attribute/type encodings.
+  return reader.emitError("unexpected Attribute encoding");
 }
 
 //===----------------------------------------------------------------------===//
 // Bytecode Reader
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class is used to read a bytecode buffer and translate it into MLIR.
private:
   LogicalResult parseDialectSection(ArrayRef<uint8_t> sectionData);
 
   /// Parse an operation name reference using the given reader.
   FailureOr<OperationName> parseOpName(EncodingReader &reader);
 
   //===--------------------------------------------------------------------===//
   // Attribute/Type Section
 
-  /// Parse an attribute or type using the given reader. Returns nullptr in the
-  /// case of failure.
-  Attribute parseAttribute(EncodingReader &reader);
-  Type parseType(EncodingReader &reader);
-
+  /// Parse an attribute or type using the given reader.
   template <typename T>
-  T parseAttribute(EncodingReader &reader) {
-    if (Attribute attr = parseAttribute(reader)) {
-      if (auto derivedAttr = attr.dyn_cast<T>())
-        return derivedAttr;
-      (void)reader.emitError("expected attribute of type: ",
-                             llvm::getTypeName<T>(), ", but got: ", attr);
-    }
-    return T();
+  LogicalResult parseAttribute(EncodingReader &reader, T &result) {
+    return attrTypeReader.parseAttribute(reader, result);
+  LogicalResult parseType(EncodingReader &reader, Type &result) {
+    return attrTypeReader.parseType(reader, result);
   }
 
   //===--------------------------------------------------------------------===//
   // IR Section
 
   /// This struct represents the current read state of a range of regions. This
   /// struct is used to enable iterative parsing of regions.
   struct RegionReadState {
if (failed(opName->dialect->load(reader, getContext())))
       return failure();
     opName->opName.emplace((opName->dialect->name + "." + opName->name).str(),
                            getContext());
   }
   return *opName->opName;
 }
 
 //===----------------------------------------------------------------------===//
-// Attribute/Type Section
-
-Attribute BytecodeReader::parseAttribute(EncodingReader &reader) {
-  uint64_t attrIdx;
-  if (failed(reader.parseVarInt(attrIdx)))
-    return Attribute();
-  return attrTypeReader.resolveAttribute(attrIdx);
-}
-
-Type BytecodeReader::parseType(EncodingReader &reader) {
-  uint64_t typeIdx;
-  if (failed(reader.parseVarInt(typeIdx)))
-    return Type();
-  return attrTypeReader.resolveType(typeIdx);
-}
-
-//===----------------------------------------------------------------------===//
 // IR Section
 
 LogicalResult BytecodeReader::parseIRSection(ArrayRef<uint8_t> sectionData,
                                              Block *block) {
   EncodingReader reader(sectionData, fileLoc);
 
   // A stack of operation regions currently being read from the bytecode.
   std::vector<RegionReadState> regionStack;
BytecodeReader::parseOpWithoutRegions(EncodingReader &reader,
 
   // Parse the operation mask, which indicates which components of the operation
   // are present.
   uint8_t opMask;
   if (failed(reader.parseByte(opMask)))
     return failure();
 
   /// Parse the location.
-  LocationAttr opLoc = parseAttribute<LocationAttr>(reader);
-  if (!opLoc)
+  LocationAttr opLoc;
+  if (failed(parseAttribute(reader, opLoc)))
     return failure();
 
   // With the location and name resolved, we can start building the operation
   // state.
   OperationState opState(opLoc, *opName);
 
   // Parse the attributes of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasAttrs) {
-    DictionaryAttr dictAttr = parseAttribute<DictionaryAttr>(reader);
-    if (!dictAttr)
+    DictionaryAttr dictAttr;
+    if (failed(parseAttribute(reader, dictAttr)))
       return failure();
     opState.attributes = dictAttr;
   }
 
   /// Parse the results of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasResults) {
     uint64_t numResults;
     if (failed(reader.parseVarInt(numResults)))
       return failure();
     opState.types.resize(numResults);
     for (int i = 0, e = numResults; i < e; ++i)
-      if (!(opState.types[i] = parseType(reader)))
+      if (failed(parseType(reader, opState.types[i])))
         return failure();
   }
 
   /// Parse the operands of the operation.
   if (opMask & bytecode::OpEncodingMask::kHasOperands) {
     uint64_t numOperands;
     if (failed(reader.parseVarInt(numOperands)))
       return failure();
if (failed(reader.parseVarInt(numArgs)))
     return failure();
 
   SmallVector<Type> argTypes;
   SmallVector<Location> argLocs;
   argTypes.reserve(numArgs);
   argLocs.reserve(numArgs);
 
   while (numArgs--) {
-    Type argType = parseType(reader);
-    if (!argType)
-      return failure();
-    LocationAttr argLoc = parseAttribute<LocationAttr>(reader);
-    if (!argLoc)
+    Type argType;
+    LocationAttr argLoc;
+    if (failed(parseType(reader, argType)) ||
+        failed(parseAttribute(reader, argLoc)))
       return failure();
 
     argTypes.push_back(argType);
     argLocs.push_back(argLoc);
   }
   block->addArguments(argTypes, argLocs);
   return defineValues(reader, block->getArguments());
 }