[mlir:Bytecode][NFC] Refactor string section writing and reading

[rriddle]

Summary

This extracts the string section writer and reader into dedicated
classes, which better separates the logic and will also simplify future
patches that want to interact with the string section.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

295

Shall we check on load that there is no null char in the middle of the string?

302

Do you need totalStringDataSize? Isn't the same check if stringDataEndOffset == 0?

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

232

Still an opportunity to sort the string based on frequency, right? :)

[rriddle]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

295

StringAttr supports strings with null characters, so I've just been allowing null holding strings in the string section (to simplify things). It doesn't affect loading at all given we don't use the null character for size computation. Do you think we should disallow that here?

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

232

Yeah. I've got a TODO for that in the followup Attr/Type patch as well, it should just require a bit of preprocessing (with no format changes). It wasn't clear in the beginning how annoying it would be to support, but after attr/type support I don't think it'll be too bad.

[mehdi_amini]

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

295

Not necessarily, but I remember some doc on the bytecode that was explicit about null-terminated string which excluded null char?

mlir/lib/Bytecode/Reader/BytecodeReader.cpp

static LogicalResult parseEntry(EncodingReader &reader, RangeT &entries,
                                 T &entry, StringRef entryStr) {
   uint64_t entryIdx;
   if (failed(reader.parseVarInt(entryIdx)))
     return failure();
   return resolveEntry(reader, entries, entryIdx, entry, entryStr);
 }
 
 //===----------------------------------------------------------------------===//
+// StringSectionReader
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is used to read references to the string section from the
+/// bytecode.
+class StringSectionReader {
+public:
+  /// Initialize the string section reader with the given section data.
+  LogicalResult initialize(Location fileLoc, ArrayRef<uint8_t> sectionData);
+
+  /// Parse a shared string from the string section. The shared string is
+  /// encoded using an index to a corresponding string in the string section.
+  LogicalResult parseString(EncodingReader &reader, StringRef &result) {
+    return parseEntry(reader, strings, result, "string");
+  }
+
+private:
+  /// The table of strings referenced within the bytecode file.
+  SmallVector<StringRef> strings;
+};
+} // namespace
+
+LogicalResult StringSectionReader::initialize(Location fileLoc,
+                                              ArrayRef<uint8_t> sectionData) {
+  EncodingReader stringReader(sectionData, fileLoc);
+
+  // Parse the number of strings in the section.
+  uint64_t numStrings;
+  if (failed(stringReader.parseVarInt(numStrings)))
+    return failure();
+  strings.resize(numStrings);
+
+  // Parse each of the strings. The sizes of the strings are encoded in reverse
+  // order, so that's the order we populate the table.
+  size_t stringDataEndOffset = sectionData.size();
+  for (StringRef &string : llvm::reverse(strings)) {
+    uint64_t stringSize;
+    if (failed(stringReader.parseVarInt(stringSize)))
+      return failure();
+    if (stringDataEndOffset < stringSize) {
+      return stringReader.emitError(
+          "string size exceeds the available data size");
+    }
+
+    // Extract the string from the data, dropping the null character.
+    size_t stringOffset = stringDataEndOffset - stringSize;
+    string = StringRef(
+        reinterpret_cast<const char *>(sectionData.data() + stringOffset),
+        stringSize - 1);
+    stringDataEndOffset = stringOffset;
+  }
+
+  // Check that the only remaining data was for the strings, i.e. the reader
+  // should be at the same offset as the first string.
+  if ((sectionData.size() - stringReader.size()) != stringDataEndOffset) {
+    return stringReader.emitError("unexpected trailing data between the "
+                                  "offsets for strings and their data");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // BytecodeDialect
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This struct represents a dialect entry within the bytecode.
 struct BytecodeDialect {
   /// Load the dialect into the provided context if it hasn't been loaded yet.
   /// Returns failure if the dialect couldn't be loaded *and* the provided
FailureOr<Operation *> parseOpWithoutRegions(EncodingReader &reader,
                                                RegionReadState &readState,
                                                bool &isIsolatedFromAbove);
 
   LogicalResult parseRegion(EncodingReader &reader, RegionReadState &readState);
   LogicalResult parseBlock(EncodingReader &reader, RegionReadState &readState);
   LogicalResult parseBlockArguments(EncodingReader &reader, Block *block);
 
   //===--------------------------------------------------------------------===//
-  // String Section
-
-  LogicalResult parseStringSection(ArrayRef<uint8_t> sectionData);
-
-  /// Parse a shared string from the string section. The shared string is
-  /// encoded using an index to a corresponding string in the string section.
-  LogicalResult parseSharedString(EncodingReader &reader, StringRef &result) {
-    return parseEntry(reader, strings, result, "string");
-  }
-
-  //===--------------------------------------------------------------------===//
   // Value Processing
 
   /// Parse an operand reference using the given reader. Returns nullptr in the
   /// case of failure.
   Value parseOperand(EncodingReader &reader);
 
   /// Sequentially define the given value range.
   LogicalResult defineValues(EncodingReader &reader, ValueRange values);
private:
   /// The producer of the bytecode being read.
   StringRef producer;
 
   /// The table of IR units referenced within the bytecode file.
   SmallVector<BytecodeDialect> dialects;
   SmallVector<BytecodeOperationName> opNames;
 
   /// The table of strings referenced within the bytecode file.
-  SmallVector<StringRef> strings;
+  StringSectionReader stringReader;
 
   /// The current set of available IR value scopes.
   std::vector<ValueScope> valueScopes;
   /// A block containing the set of operations defined to create forward
   /// references.
   Block forwardRefOps;
   /// A block containing previously created, and no longer used, forward
   /// reference operations.
LogicalResult BytecodeReader::read(llvm::MemoryBufferRef buffer, Block *block) {
   for (int i = 0; i < bytecode::Section::kNumSections; ++i) {
     if (!sectionDatas[i]) {
       return reader.emitError("missing data for top-level section: ",
                               toString(bytecode::Section::ID(i)));
     }
   }
 
   // Process the string section first.
-  if (failed(parseStringSection(*sectionDatas[bytecode::Section::kString])))
+  if (failed(stringReader.initialize(
+          fileLoc, *sectionDatas[bytecode::Section::kString])))
     return failure();
 
   // Process the dialect section.
   if (failed(parseDialectSection(*sectionDatas[bytecode::Section::kDialect])))
     return failure();
 
   // Process the attribute and type section.
   if (failed(attrTypeReader.initialize(
BytecodeReader::parseDialectSection(ArrayRef<uint8_t> sectionData) {
   // Parse the number of dialects in the section.
   uint64_t numDialects;
   if (failed(sectionReader.parseVarInt(numDialects)))
     return failure();
   dialects.resize(numDialects);
 
   // Parse each of the dialects.
   for (uint64_t i = 0; i < numDialects; ++i)
-    if (failed(parseSharedString(sectionReader, dialects[i].name)))
+    if (failed(stringReader.parseString(sectionReader, dialects[i].name)))
       return failure();
 
   // Parse the operation names, which are grouped by dialect.
   auto parseOpName = [&](BytecodeDialect *dialect) {
     StringRef opName;
-    if (failed(parseSharedString(sectionReader, opName)))
+    if (failed(stringReader.parseString(sectionReader, opName)))
       return failure();
     opNames.emplace_back(dialect, opName);
     return success();
   };
   while (!sectionReader.empty())
     if (failed(parseDialectGrouping(sectionReader, dialects, parseOpName)))
       return failure();
   return success();
while (numArgs--) {
     argTypes.push_back(argType);
     argLocs.push_back(argLoc);
   }
   block->addArguments(argTypes, argLocs);
   return defineValues(reader, block->getArguments());
 }
 
 //===----------------------------------------------------------------------===//
-// String Section
-
-LogicalResult
-BytecodeReader::parseStringSection(ArrayRef<uint8_t> sectionData) {
-  EncodingReader stringReader(sectionData, fileLoc);
-
-  // Parse the number of strings in the section.
-  uint64_t numStrings;
-  if (failed(stringReader.parseVarInt(numStrings)))
-    return failure();
-  strings.resize(numStrings);
-
-  // Parse each of the strings. The sizes of the strings are encoded in reverse
-  // order, so that's the order we populate the table.
-  size_t stringDataEndOffset = sectionData.size();
-  size_t totalStringDataSize = 0;
-  for (StringRef &string : llvm::reverse(strings)) {
-    uint64_t stringSize;
-    if (failed(stringReader.parseVarInt(stringSize)))
-      return failure();
-    if (stringDataEndOffset < stringSize) {
-      return stringReader.emitError(
-          "string size exceeds the available data size");
-    }
-
-    // Extract the string from the data, dropping the null character.
-    size_t stringOffset = stringDataEndOffset - stringSize;
-    string = StringRef(
-        reinterpret_cast<const char *>(sectionData.data() + stringOffset),
-        stringSize - 1);
-    stringDataEndOffset = stringOffset;
-
-    // Update the total string data size.
-    totalStringDataSize += stringSize;
-  }
-
-  // Check that the only remaining data was for the strings
-  if (stringReader.size() != totalStringDataSize) {
-    return stringReader.emitError("unexpected trailing data between the "
-                                  "offsets for strings and their data");
-  }
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // Value Processing
 
 Value BytecodeReader::parseOperand(EncodingReader &reader) {
   std::vector<Value> &values = valueScopes.back().values;
   Value *value = nullptr;
   if (failed(parseEntry(reader, values, value, "value")))
     return Value();
 

mlir/lib/Bytecode/Writer/BytecodeWriter.cpp

void EncodingEmitter::emitMultiByteVarInt(uint64_t value) {
 
   // If the value is too large to encode in a single byte, emit a special all
   // zero marker byte and splat the value directly.
   emitByte(0);
   emitBytes({reinterpret_cast<uint8_t *>(&value), sizeof(value)});
 }
 
 //===----------------------------------------------------------------------===//
+// StringSectionBuilder
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is used to simplify the process of emitting the string section.
+class StringSectionBuilder {
+public:
+  /// Add the given string to the string section, and return the index of the
+  /// string within the section.
+  size_t insert(StringRef str) {
+    auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
+    return it.first->second;
+  }
+
+  /// Write the current set of strings to the given emitter.
+  void write(EncodingEmitter &emitter) {
+    emitter.emitVarInt(strings.size());
+
+    // Emit the sizes in reverse order, so that we don't need to backpatch an
+    // offset to the string data or have a separate section.
+    for (const auto &it : llvm::reverse(strings))
+      emitter.emitVarInt(it.first.size() + 1);
+    // Emit the string data itself.
+    for (const auto &it : strings)
+      emitter.emitNulTerminatedString(it.first.val());
+  }
+
+private:
+  /// A set of strings referenced within the bytecode. The value of the map is
+  /// unused.
+  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
 // Bytecode Writer
 //===----------------------------------------------------------------------===//
 
 namespace {
 class BytecodeWriter {
 public:
   BytecodeWriter(Operation *op) : numberingState(op) {}
 
private:
   void writeRegion(EncodingEmitter &emitter, Region *region);
   void writeIRSection(EncodingEmitter &emitter, Operation *op);
 
   //===--------------------------------------------------------------------===//
   // Strings
 
   void writeStringSection(EncodingEmitter &emitter);
 
-  /// Get the number for the given shared string, that is contained within the
-  /// string section.
-  size_t getSharedStringNumber(StringRef str);
-
   //===--------------------------------------------------------------------===//
   // Fields
 
+  /// The builder used for the string section.
+  StringSectionBuilder stringSection;
+
   /// The IR numbering state generated for the root operation.
   IRNumberingState numberingState;
-
-  /// A set of strings referenced within the bytecode. The value of the map is
-  /// unused.
-  llvm::MapVector<llvm::CachedHashStringRef, size_t> strings;
 };
 } // namespace
 
 void BytecodeWriter::write(Operation *rootOp, raw_ostream &os,
                            StringRef producer) {
   EncodingEmitter emitter;
 
   // Emit the bytecode file header. This is how we identify the output as a

 
 void BytecodeWriter::writeDialectSection(EncodingEmitter &emitter) {
   EncodingEmitter dialectEmitter;
 
   // Emit the referenced dialects.
   auto dialects = numberingState.getDialects();
   dialectEmitter.emitVarInt(llvm::size(dialects));
   for (DialectNumbering &dialect : dialects)
-    dialectEmitter.emitVarInt(getSharedStringNumber(dialect.name));
+    dialectEmitter.emitVarInt(stringSection.insert(dialect.name));
 
   // Emit the referenced operation names grouped by dialect.
   auto emitOpName = [&](OpNameNumbering &name) {
-    dialectEmitter.emitVarInt(getSharedStringNumber(name.name.stripDialect()));
+    dialectEmitter.emitVarInt(stringSection.insert(name.name.stripDialect()));
   };
   writeDialectGrouping(dialectEmitter, numberingState.getOpNames(), emitOpName);
 
   emitter.emitSection(bytecode::Section::kDialect, std::move(dialectEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Attributes and Types
void BytecodeWriter::writeIRSection(EncodingEmitter &emitter, Operation *op) {
   emitter.emitSection(bytecode::Section::kIR, std::move(irEmitter));
 }
 
 //===----------------------------------------------------------------------===//
 // Strings
 
 void BytecodeWriter::writeStringSection(EncodingEmitter &emitter) {
   EncodingEmitter stringEmitter;
-  stringEmitter.emitVarInt(strings.size());
-
-  // Emit the sizes in reverse order, so that we don't need to backpatch an
-  // offset to the string data or have a separate section.
-  for (const auto &it : llvm::reverse(strings))
-    stringEmitter.emitVarInt(it.first.size() + 1);
-  // Emit the string data itself.
-  for (const auto &it : strings)
-    stringEmitter.emitNulTerminatedString(it.first.val());
-
+  stringSection.write(stringEmitter);
   emitter.emitSection(bytecode::Section::kString, std::move(stringEmitter));
 }
 
-size_t BytecodeWriter::getSharedStringNumber(StringRef str) {
-  auto it = strings.insert({llvm::CachedHashStringRef(str), strings.size()});
-  return it.first->second;
-}
-
 //===----------------------------------------------------------------------===//
 // Entry Points
 //===----------------------------------------------------------------------===//
 
 void mlir::writeBytecodeToFile(Operation *op, raw_ostream &os,
                                StringRef producer) {
   BytecodeWriter writer(op);
   writer.write(op, os, producer);
 }