[mlir][NFC] Cleanup builtin dialect bytecode encoding

[rriddle]

Summary

Group the readers and writers for individual attributes/types together,
which makes the encoding more readable.

[mehdi_amini]

you know you don't review for this right? :)

[rriddle]

In D132583#3747260, @mehdi_amini wrote:

you know you don't review for this right? :)

Indeed, but it was in my branch stack and I struggle to selectively send reviews sometimes ;)

mlir/lib/IR/BuiltinDialectBytecode.cpp

 };
 } // namespace
 
 void builtin_dialect_detail::addBytecodeInterface(BuiltinDialect *dialect) {
   dialect->addInterfaces<BuiltinDialectBytecodeInterface>();
 }
 
 //===----------------------------------------------------------------------===//
-// Attributes: Reader
+// Attributes
+//===----------------------------------------------------------------------===//
 
 Attribute BuiltinDialectBytecodeInterface::readAttribute(
     DialectBytecodeReader &reader) const {
   uint64_t code;
   if (failed(reader.readVarInt(code)))
     return Attribute();
   switch (code) {
   case builtin_encoding::kArrayAttr:
Attribute BuiltinDialectBytecodeInterface::readAttribute(
   case builtin_encoding::kUnknownLoc:
     return UnknownLoc::get(getContext());
   default:
     reader.emitError() << "unknown builtin attribute code: " << code;
     return Attribute();
   }
 }
 
+LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
+    Attribute attr, DialectBytecodeWriter &writer) const {
+  return TypeSwitch<Attribute, LogicalResult>(attr)
+      .Case<ArrayAttr, DictionaryAttr, FloatAttr, IntegerAttr, StringAttr,
+            SymbolRefAttr, TypeAttr, CallSiteLoc, FileLineColLoc, FusedLoc,
+            NameLoc>([&](auto attr) {
+        write(attr, writer);
+        return success();
+      })
+      .Case([&](OpaqueLoc attr) { return write(attr, writer); })
+      .Case([&](UnitAttr) {
+        writer.writeVarInt(builtin_encoding::kUnitAttr);
+        return success();
+      })
+      .Case([&](UnknownLoc) {
+        writer.writeVarInt(builtin_encoding::kUnknownLoc);
+        return success();
+      })
+      .Default([&](Attribute) { return failure(); });
+}
+
+//===----------------------------------------------------------------------===//
+// ArrayAttr
+
 ArrayAttr BuiltinDialectBytecodeInterface::readArrayAttr(
     DialectBytecodeReader &reader) const {
   SmallVector<Attribute> elements;
   if (failed(reader.readAttributes(elements)))
     return ArrayAttr();
   return ArrayAttr::get(getContext(), elements);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    ArrayAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kArrayAttr);
+  writer.writeAttributes(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// DictionaryAttr
+
 DictionaryAttr BuiltinDialectBytecodeInterface::readDictionaryAttr(
     DialectBytecodeReader &reader) const {
   auto readNamedAttr = [&]() -> FailureOr<NamedAttribute> {
     StringAttr name;
     Attribute value;
     if (failed(reader.readAttribute(name)) ||
         failed(reader.readAttribute(value)))
       return failure();
     return NamedAttribute(name, value);
   };
   SmallVector<NamedAttribute> attrs;
   if (failed(reader.readList(attrs, readNamedAttr)))
     return DictionaryAttr();
   return DictionaryAttr::get(getContext(), attrs);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    DictionaryAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kDictionaryAttr);
+  writer.writeList(attr.getValue(), [&](NamedAttribute attr) {
+    writer.writeAttribute(attr.getName());
+    writer.writeAttribute(attr.getValue());
+  });
+}
+
+//===----------------------------------------------------------------------===//
+// FloatAttr
+
 FloatAttr BuiltinDialectBytecodeInterface::readFloatAttr(
     DialectBytecodeReader &reader) const {
   FloatType type;
   if (failed(reader.readType(type)))
     return FloatAttr();
   FailureOr<APFloat> value =
       reader.readAPFloatWithKnownSemantics(type.getFloatSemantics());
   if (failed(value))
     return FloatAttr();
   return FloatAttr::get(type, *value);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    FloatAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kFloatAttr);
+  writer.writeType(attr.getType());
+  writer.writeAPFloatWithKnownSemantics(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// IntegerAttr
+
 IntegerAttr BuiltinDialectBytecodeInterface::readIntegerAttr(
     DialectBytecodeReader &reader) const {
   Type type;
   if (failed(reader.readType(type)))
     return IntegerAttr();
 
   // Extract the value storage width from the type.
   unsigned bitWidth;
   if (auto intType = type.dyn_cast<IntegerType>()) {
     bitWidth = intType.getWidth();
   } else if (type.isa<IndexType>()) {
     bitWidth = IndexType::kInternalStorageBitWidth;
   } else {
     reader.emitError()
         << "expected integer or index type for IntegerAttr, but got: " << type;
     return IntegerAttr();
   }
 
   FailureOr<APInt> value = reader.readAPIntWithKnownWidth(bitWidth);
   if (failed(value))
     return IntegerAttr();
   return IntegerAttr::get(type, *value);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    IntegerAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kIntegerAttr);
+  writer.writeType(attr.getType());
+  writer.writeAPIntWithKnownWidth(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// StringAttr
+
 StringAttr
 BuiltinDialectBytecodeInterface::readStringAttr(DialectBytecodeReader &reader,
                                                 bool hasType) const {
   StringRef string;
   if (failed(reader.readString(string)))
     return StringAttr();
 
   // Read the type if present.
   Type type;
   if (!hasType)
     type = NoneType::get(getContext());
   else if (failed(reader.readType(type)))
     return StringAttr();
   return StringAttr::get(string, type);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    StringAttr attr, DialectBytecodeWriter &writer) const {
+  // We only encode the type if it isn't NoneType, which is significantly less
+  // common.
+  Type type = attr.getType();
+  if (!type.isa<NoneType>()) {
+    writer.writeVarInt(builtin_encoding::kStringAttrWithType);
+    writer.writeOwnedString(attr.getValue());
+    writer.writeType(type);
+    return;
+  }
+  writer.writeVarInt(builtin_encoding::kStringAttr);
+  writer.writeOwnedString(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// SymbolRefAttr
+
 SymbolRefAttr BuiltinDialectBytecodeInterface::readSymbolRefAttr(
     DialectBytecodeReader &reader, bool hasNestedRefs) const {
   StringAttr rootReference;
   if (failed(reader.readAttribute(rootReference)))
     return SymbolRefAttr();
   SmallVector<FlatSymbolRefAttr> nestedReferences;
   if (hasNestedRefs && failed(reader.readAttributes(nestedReferences)))
     return SymbolRefAttr();
   return SymbolRefAttr::get(rootReference, nestedReferences);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    SymbolRefAttr attr, DialectBytecodeWriter &writer) const {
+  ArrayRef<FlatSymbolRefAttr> nestedRefs = attr.getNestedReferences();
+  writer.writeVarInt(nestedRefs.empty() ? builtin_encoding::kFlatSymbolRefAttr
+                                        : builtin_encoding::kSymbolRefAttr);
+
+  writer.writeAttribute(attr.getRootReference());
+  if (!nestedRefs.empty())
+    writer.writeAttributes(nestedRefs);
+}
+
+//===----------------------------------------------------------------------===//
+// TypeAttr
+
 TypeAttr BuiltinDialectBytecodeInterface::readTypeAttr(
     DialectBytecodeReader &reader) const {
   Type type;
   if (failed(reader.readType(type)))
     return TypeAttr();
   return TypeAttr::get(type);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    TypeAttr attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kTypeAttr);
+  writer.writeType(attr.getValue());
+}
+
+//===----------------------------------------------------------------------===//
+// CallSiteLoc
+
 LocationAttr BuiltinDialectBytecodeInterface::readCallSiteLoc(
     DialectBytecodeReader &reader) const {
   LocationAttr callee, caller;
   if (failed(reader.readAttribute(callee)) ||
       failed(reader.readAttribute(caller)))
     return LocationAttr();
   return CallSiteLoc::get(callee, caller);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    CallSiteLoc attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kCallSiteLoc);
+  writer.writeAttribute(attr.getCallee());
+  writer.writeAttribute(attr.getCaller());
+}
+
+//===----------------------------------------------------------------------===//
+// FileLineColLoc
+
 LocationAttr BuiltinDialectBytecodeInterface::readFileLineColLoc(
     DialectBytecodeReader &reader) const {
   StringAttr filename;
   uint64_t line, column;
   if (failed(reader.readAttribute(filename)) ||
       failed(reader.readVarInt(line)) || failed(reader.readVarInt(column)))
     return LocationAttr();
   return FileLineColLoc::get(filename, line, column);
 }
 
+void BuiltinDialectBytecodeInterface::write(
+    FileLineColLoc attr, DialectBytecodeWriter &writer) const {
+  writer.writeVarInt(builtin_encoding::kFileLineColLoc);
+  writer.writeAttribute(attr.getFilename());
+  writer.writeVarInt(attr.getLine());
+  writer.writeVarInt(attr.getColumn());
+}
+
+//===----------------------------------------------------------------------===//
+// FusedLoc
+
 LocationAttr
 BuiltinDialectBytecodeInterface::readFusedLoc(DialectBytecodeReader &reader,
                                               bool hasMetadata) const {
   // Parse the child locations.
   auto readLoc = [&]() -> FailureOr<Location> {
     LocationAttr locAttr;
     if (failed(reader.readAttribute(locAttr)))
       return failure();
     return Location(locAttr);
   };
   SmallVector<Location> locations;
   if (failed(reader.readList(locations, readLoc)))
     return LocationAttr();
 
   // Parse the metadata if present.
   Attribute metadata;
   if (hasMetadata && failed(reader.readAttribute(metadata)))
     return LocationAttr();
 
   return FusedLoc::get(locations, metadata, getContext());
 }
 
-LocationAttr BuiltinDialectBytecodeInterface::readNameLoc(
-    DialectBytecodeReader &reader) const {
-  StringAttr name;
-  LocationAttr childLoc;
-  if (failed(reader.readAttribute(name)) ||
-      failed(reader.readAttribute(childLoc)))
-    return LocationAttr();
-  return NameLoc::get(name, childLoc);
-}
-
-//===----------------------------------------------------------------------===//
-// Attributes: Writer
-
-LogicalResult BuiltinDialectBytecodeInterface::writeAttribute(
-    Attribute attr, DialectBytecodeWriter &writer) const {
-  return TypeSwitch<Attribute, LogicalResult>(attr)
-      .Case<ArrayAttr, DictionaryAttr, FloatAttr, IntegerAttr, StringAttr,
-            SymbolRefAttr, TypeAttr, CallSiteLoc, FileLineColLoc, FusedLoc,
-            NameLoc>([&](auto attr) {
-        write(attr, writer);
-        return success();
-      })
-      .Case([&](OpaqueLoc attr) { return write(attr, writer); })
-      .Case([&](UnitAttr) {
-        writer.writeVarInt(builtin_encoding::kUnitAttr);
-        return success();
-      })
-      .Case([&](UnknownLoc) {
-        writer.writeVarInt(builtin_encoding::kUnknownLoc);
-        return success();
-      })
-      .Default([&](Attribute) { return failure(); });
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    ArrayAttr attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kArrayAttr);
-  writer.writeAttributes(attr.getValue());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    DictionaryAttr attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kDictionaryAttr);
-  writer.writeList(attr.getValue(), [&](NamedAttribute attr) {
-    writer.writeAttribute(attr.getName());
-    writer.writeAttribute(attr.getValue());
-  });
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    FloatAttr attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kFloatAttr);
-  writer.writeType(attr.getType());
-  writer.writeAPFloatWithKnownSemantics(attr.getValue());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    IntegerAttr attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kIntegerAttr);
-  writer.writeType(attr.getType());
-  writer.writeAPIntWithKnownWidth(attr.getValue());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    StringAttr attr, DialectBytecodeWriter &writer) const {
-  // We only encode the type if it isn't NoneType, which is significantly less
-  // common.
-  Type type = attr.getType();
-  if (!type.isa<NoneType>()) {
-    writer.writeVarInt(builtin_encoding::kStringAttrWithType);
-    writer.writeOwnedString(attr.getValue());
-    writer.writeType(type);
-    return;
-  }
-  writer.writeVarInt(builtin_encoding::kStringAttr);
-  writer.writeOwnedString(attr.getValue());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    SymbolRefAttr attr, DialectBytecodeWriter &writer) const {
-  ArrayRef<FlatSymbolRefAttr> nestedRefs = attr.getNestedReferences();
-  writer.writeVarInt(nestedRefs.empty() ? builtin_encoding::kFlatSymbolRefAttr
-                                        : builtin_encoding::kSymbolRefAttr);
-
-  writer.writeAttribute(attr.getRootReference());
-  if (!nestedRefs.empty())
-    writer.writeAttributes(nestedRefs);
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    TypeAttr attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kTypeAttr);
-  writer.writeType(attr.getValue());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    CallSiteLoc attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kCallSiteLoc);
-  writer.writeAttribute(attr.getCallee());
-  writer.writeAttribute(attr.getCaller());
-}
-
-void BuiltinDialectBytecodeInterface::write(
-    FileLineColLoc attr, DialectBytecodeWriter &writer) const {
-  writer.writeVarInt(builtin_encoding::kFileLineColLoc);
-  writer.writeAttribute(attr.getFilename());
-  writer.writeVarInt(attr.getLine());
-  writer.writeVarInt(attr.getColumn());
-}
-
 void BuiltinDialectBytecodeInterface::write(
     FusedLoc attr, DialectBytecodeWriter &writer) const {
   if (Attribute metadata = attr.getMetadata()) {
     writer.writeVarInt(builtin_encoding::kFusedLocWithMetadata);
     writer.writeAttributes(attr.getLocations());
     writer.writeAttribute(metadata);
   } else {
     writer.writeVarInt(builtin_encoding::kFusedLoc);
     writer.writeAttributes(attr.getLocations());
   }
 }
 
+//===----------------------------------------------------------------------===//
+// NameLoc
+
+LocationAttr BuiltinDialectBytecodeInterface::readNameLoc(
+    DialectBytecodeReader &reader) const {
+  StringAttr name;
+  LocationAttr childLoc;
+  if (failed(reader.readAttribute(name)) ||
+      failed(reader.readAttribute(childLoc)))
+    return LocationAttr();
+  return NameLoc::get(name, childLoc);
+}
+
 void BuiltinDialectBytecodeInterface::write(
     NameLoc attr, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kNameLoc);
   writer.writeAttribute(attr.getName());
   writer.writeAttribute(attr.getChildLoc());
 }
 
+//===----------------------------------------------------------------------===//
+// OpaqueLoc
+
 LogicalResult
 BuiltinDialectBytecodeInterface::write(OpaqueLoc attr,
                                        DialectBytecodeWriter &writer) const {
   // We can't encode an OpaqueLoc directly given that it is in-memory only, so
   // encode the fallback instead.
   return writeAttribute(attr.getFallbackLocation(), writer);
 }
 
 //===----------------------------------------------------------------------===//
-// Types: Reader
+// Types
+//===----------------------------------------------------------------------===//
 
 Type BuiltinDialectBytecodeInterface::readType(
     DialectBytecodeReader &reader) const {
   uint64_t code;
   if (failed(reader.readVarInt(code)))
     return Type();
   switch (code) {
   case builtin_encoding::kIntegerType:
     return readIntegerType(reader);
   case builtin_encoding::kIndexType:
     return IndexType::get(getContext());
 
   case builtin_encoding::kFunctionType:
     return readFunctionType(reader);
   default:
     reader.emitError() << "unknown builtin type code: " << code;
     return Type();
   }
 }
 
-IntegerType BuiltinDialectBytecodeInterface::readIntegerType(
-    DialectBytecodeReader &reader) const {
-  uint64_t encoding;
-  if (failed(reader.readVarInt(encoding)))
-    return IntegerType();
-  return IntegerType::get(
-      getContext(), encoding >> 2,
-      static_cast<IntegerType::SignednessSemantics>(encoding & 0x3));
-}
-
-FunctionType BuiltinDialectBytecodeInterface::readFunctionType(
-    DialectBytecodeReader &reader) const {
-  SmallVector<Type> inputs, results;
-  if (failed(reader.readTypes(inputs)) || failed(reader.readTypes(results)))
-    return FunctionType();
-  return FunctionType::get(getContext(), inputs, results);
-}
-
-//===----------------------------------------------------------------------===//
-// Types: Writer
-
 LogicalResult BuiltinDialectBytecodeInterface::writeType(
     Type type, DialectBytecodeWriter &writer) const {
   return TypeSwitch<Type, LogicalResult>(type)
       .Case<IntegerType, FunctionType>([&](auto type) {
         write(type, writer);
         return success();
       })
       .Case([&](IndexType) {
         return writer.writeVarInt(builtin_encoding::kIndexType), success();
       })
       .Default([&](Type) { return failure(); });
 }
 
+//===----------------------------------------------------------------------===//
+// IntegerType
+
+IntegerType BuiltinDialectBytecodeInterface::readIntegerType(
+    DialectBytecodeReader &reader) const {
+  uint64_t encoding;
+  if (failed(reader.readVarInt(encoding)))
+    return IntegerType();
+  return IntegerType::get(
+      getContext(), encoding >> 2,
+      static_cast<IntegerType::SignednessSemantics>(encoding & 0x3));
+}
+
 void BuiltinDialectBytecodeInterface::write(
     IntegerType type, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kIntegerType);
   writer.writeVarInt((type.getWidth() << 2) | type.getSignedness());
 }
 
+//===----------------------------------------------------------------------===//
+// FunctionType
+
+FunctionType BuiltinDialectBytecodeInterface::readFunctionType(
+    DialectBytecodeReader &reader) const {
+  SmallVector<Type> inputs, results;
+  if (failed(reader.readTypes(inputs)) || failed(reader.readTypes(results)))
+    return FunctionType();
+  return FunctionType::get(getContext(), inputs, results);
+}
+
 void BuiltinDialectBytecodeInterface::write(
     FunctionType type, DialectBytecodeWriter &writer) const {
   writer.writeVarInt(builtin_encoding::kFunctionType);
   writer.writeTypes(type.getInputs());
   writer.writeTypes(type.getResults());
 }