[GOFF] Refactor GOFFOstream (#131143)

GOFFOstream writes the physical 80 byte records. The records are
connected by flags indicating if there is a successor or a predecessor.
Using the length of the logical record is prone to errors. The new
implementation buffers the last physical record, and writes it out when
new data is written. In this way, the flags can be easily determined.
No obversable change in functionality, therefore no tests.

Author: redstar

llvm/lib/MC/GOFFObjectWriter.cpp

@@ -63,158 +63,167 @@ constexpr uint8_t RecContinued = Flags(7, 1, 1);
 constexpr uint8_t RecContinuation = Flags(6, 1, 1);
 
 // The GOFFOstream is responsible to write the data into the fixed physical
-// records of the format. A user of this class announces the start of a new
-// logical record and the size of its content. While writing the content, the
-// physical records are created for the data. Possible fill bytes at the end of
-// a physical record are written automatically. In principle, the GOFFOstream
-// is agnostic of the endianness of the content. However, it also supports
-// writing data in big endian byte order.
-class GOFFOstream : public raw_ostream {
+// records of the format. A user of this class announces the begin of a new
+// logical record. While writing the payload, the physical records are created
+// for the data. Possible fill bytes at the end of a physical record are written
+// automatically. In principle, the GOFFOstream is agnostic of the endianness of
+// the payload. However, it also supports writing data in big endian byte order.
+//
+// The physical records use the flag field to indicate if the there is a
+// successor and predecessor record. To be able to set these flags while
+// writing, the basic implementation idea is to always buffer the last seen
+// physical record.
+class GOFFOstream {
   /// The underlying raw_pwrite_stream.
   raw_pwrite_stream &OS;
 
-  /// The remaining size of this logical record, including fill bytes.
-  size_t RemainingSize;
-
-#ifndef NDEBUG
-  /// The number of bytes needed to fill up the last physical record.
-  size_t Gap = 0;
-#endif
-
-  /// The number of logical records emitted to far.
-  uint32_t LogicalRecords;
-
-  /// The type of the current (logical) record.
-  GOFF::RecordType CurrentType;
-
-  /// Signals start of new record.
-  bool NewLogicalRecord;
+  /// The number of logical records emitted so far.
+  uint32_t LogicalRecords = 0;
 
-  /// Static allocated buffer for the stream, used by the raw_ostream class. The
-  /// buffer is sized to hold the content of a physical record.
-  char Buffer[GOFF::RecordContentLength];
+  /// The number of physical records emitted so far.
+  uint32_t PhysicalRecords = 0;
 
-  // Return the number of bytes left to write until next physical record.
-  // Please note that we maintain the total numbers of byte left, not the
-  // written size.
-  size_t bytesToNextPhysicalRecord() {
-    size_t Bytes = RemainingSize % GOFF::RecordContentLength;
-    return Bytes ? Bytes : GOFF::RecordContentLength;
-  }
-
-  /// Write the record prefix of a physical record, using the given record type.
-  static void writeRecordPrefix(raw_ostream &OS, GOFF::RecordType Type,
-                                size_t RemainingSize,
-                                uint8_t Flags = RecContinuation);
+  /// The size of the buffer. Same as the payload size of a physical record.
+  static constexpr uint8_t BufferSize = GOFF::PayloadLength;
 
-  /// Fill the last physical record of a logical record with zero bytes.
-  void fillRecord();
+  /// Current position in buffer.
+  char *BufferPtr = Buffer;
 
-  /// See raw_ostream::write_impl.
-  void write_impl(const char *Ptr, size_t Size) override;
+  /// Static allocated buffer for the stream.
+  char Buffer[BufferSize];
 
-  /// Return the current position within the stream, not counting the bytes
-  /// currently in the buffer.
-  uint64_t current_pos() const override { return OS.tell(); }
+  /// The type of the current logical record, and the flags (aka continued and
+  /// continuation indicators) for the previous (physical) record.
+  uint8_t TypeAndFlags = 0;
 
 public:
-  explicit GOFFOstream(raw_pwrite_stream &OS)
-      : OS(OS), RemainingSize(0), LogicalRecords(0), NewLogicalRecord(false) {
-    SetBuffer(Buffer, sizeof(Buffer));
-  }
-
-  ~GOFFOstream() { finalize(); }
+  GOFFOstream(raw_pwrite_stream &OS);
+  ~GOFFOstream();
 
   raw_pwrite_stream &getOS() { return OS; }
+  size_t getWrittenSize() const { return PhysicalRecords * GOFF::RecordLength; }
+  uint32_t getNumLogicalRecords() { return LogicalRecords; }
 
-  void newRecord(GOFF::RecordType Type, size_t Size);
-
-  void finalize() { fillRecord(); }
+  /// Write the specified bytes.
+  void write(const char *Ptr, size_t Size);
 
-  uint32_t logicalRecords() { return LogicalRecords; }
+  /// Write zeroes, up to a maximum of 16 bytes.
+  void write_zeros(unsigned NumZeros);
 
-  // Support for endian-specific data.
+  /// Support for endian-specific data.
   template <typename value_type> void writebe(value_type Value) {
     Value =
         support::endian::byte_swap<value_type>(Value, llvm::endianness::big);
-    write(reinterpret_cast<const char *>(&Value), sizeof(value_type));
+    write((const char *)&Value, sizeof(value_type));
   }
+
+  /// Begin a new logical record. Implies finalizing the previous record.
+  void newRecord(GOFF::RecordType Type);
+
+  /// Ends a logical record.
+  void finalizeRecord();
+
+private:
+  /// Updates the continued/continuation flags, and writes the record prefix of
+  /// a physical record.
+  void updateFlagsAndWritePrefix(bool IsContinued);
+
+  /// Returns the remaining size in the buffer.
+  size_t getRemainingSize();
 };
+} // namespace
+
+GOFFOstream::GOFFOstream(raw_pwrite_stream &OS) : OS(OS) {}
 
-void GOFFOstream::writeRecordPrefix(raw_ostream &OS, GOFF::RecordType Type,
-                                    size_t RemainingSize, uint8_t Flags) {
-  uint8_t TypeAndFlags = Flags | (Type << 4);
-  if (RemainingSize > GOFF::RecordLength)
+GOFFOstream::~GOFFOstream() { finalizeRecord(); }
+
+void GOFFOstream::updateFlagsAndWritePrefix(bool IsContinued) {
+  // Update the flags based on the previous state and the flag IsContinued.
+  if (TypeAndFlags & RecContinued)
+    TypeAndFlags |= RecContinuation;
+  if (IsContinued)
     TypeAndFlags |= RecContinued;
+  else
+    TypeAndFlags &= ~RecContinued;
+
   OS << static_cast<unsigned char>(GOFF::PTVPrefix) // Record Type
      << static_cast<unsigned char>(TypeAndFlags)    // Continuation
      << static_cast<unsigned char>(0);              // Version
+
+  ++PhysicalRecords;
 }
 
-void GOFFOstream::newRecord(GOFF::RecordType Type, size_t Size) {
-  fillRecord();
-  CurrentType = Type;
-  RemainingSize = Size;
-#ifdef NDEBUG
-  size_t Gap;
-#endif
-  Gap = (RemainingSize % GOFF::RecordContentLength);
-  if (Gap) {
-    Gap = GOFF::RecordContentLength - Gap;
-    RemainingSize += Gap;
-  }
-  NewLogicalRecord = true;
-  ++LogicalRecords;
+size_t GOFFOstream::getRemainingSize() {
+  return size_t(&Buffer[BufferSize] - BufferPtr);
 }
 
-void GOFFOstream::fillRecord() {
-  assert((GetNumBytesInBuffer() <= RemainingSize) &&
-         "More bytes in buffer than expected");
-  size_t Remains = RemainingSize - GetNumBytesInBuffer();
-  if (Remains) {
-    assert(Remains == Gap && "Wrong size of fill gap");
-    assert((Remains < GOFF::RecordLength) &&
-           "Attempt to fill more than one physical record");
-    raw_ostream::write_zeros(Remains);
+void GOFFOstream::write(const char *Ptr, size_t Size) {
+  size_t RemainingSize = getRemainingSize();
+
+  // Data fits into the buffer.
+  if (LLVM_LIKELY(Size <= RemainingSize)) {
+    memcpy(BufferPtr, Ptr, Size);
+    BufferPtr += Size;
+    return;
+  }
+
+  // Otherwise the buffer is partially filled or full, and data does not fit
+  // into it.
+  updateFlagsAndWritePrefix(/*IsContinued=*/true);
+  OS.write(Buffer, size_t(BufferPtr - Buffer));
+  if (RemainingSize > 0) {
+    OS.write(Ptr, RemainingSize);
+    Ptr += RemainingSize;
+    Size -= RemainingSize;
   }
-  flush();
-  assert(RemainingSize == 0 && "Not fully flushed");
-  assert(GetNumBytesInBuffer() == 0 && "Buffer not fully empty");
-}
 
-// This function is called from the raw_ostream implementation if:
-// - The internal buffer is full. Size is excactly the size of the buffer.
-// - Data larger than the internal buffer is written. Size is a multiple of the
-//   buffer size.
-// - flush() has been called. Size is at most the buffer size.
-// The GOFFOstream implementation ensures that flush() is called before a new
-// logical record begins. Therefore it is sufficient to check for a new block
-// only once.
-void GOFFOstream::write_impl(const char *Ptr, size_t Size) {
-  assert((RemainingSize >= Size) && "Attempt to write too much data");
-  assert(RemainingSize && "Logical record overflow");
-  if (!(RemainingSize % GOFF::RecordContentLength)) {
-    writeRecordPrefix(OS, CurrentType, RemainingSize,
-                      NewLogicalRecord ? 0 : RecContinuation);
-    NewLogicalRecord = false;
+  while (Size > BufferSize) {
+    updateFlagsAndWritePrefix(/*IsContinued=*/true);
+    OS.write(Ptr, BufferSize);
+    Ptr += BufferSize;
+    Size -= BufferSize;
   }
-  assert(!NewLogicalRecord &&
-         "New logical record not on physical record boundary");
-
-  size_t Idx = 0;
-  while (Size > 0) {
-    size_t BytesToWrite = bytesToNextPhysicalRecord();
-    if (BytesToWrite > Size)
-      BytesToWrite = Size;
-    OS.write(Ptr + Idx, BytesToWrite);
-    Idx += BytesToWrite;
-    Size -= BytesToWrite;
-    RemainingSize -= BytesToWrite;
-    if (Size)
-      writeRecordPrefix(OS, CurrentType, RemainingSize);
+
+  // The remaining bytes fit into the buffer.
+  memcpy(Buffer, Ptr, Size);
+  BufferPtr = &Buffer[Size];
+}
+
+void GOFFOstream::write_zeros(unsigned NumZeros) {
+  assert(NumZeros <= 16 && "Range for zeros too large");
+
+  // Handle the common case first: all fits in the buffer.
+  size_t RemainingSize = getRemainingSize();
+  if (LLVM_LIKELY(RemainingSize >= NumZeros)) {
+    memset(BufferPtr, 0, NumZeros);
+    BufferPtr += NumZeros;
+    return;
   }
+
+  // Otherwise some field value is cleared.
+  static char Zeros[16] = {
+      0,
+  };
+  write(Zeros, NumZeros);
 }
 
+void GOFFOstream::newRecord(GOFF::RecordType Type) {
+  finalizeRecord();
+  TypeAndFlags = Type << 4;
+  ++LogicalRecords;
+}
+
+void GOFFOstream::finalizeRecord() {
+  if (Buffer == BufferPtr)
+    return;
+  updateFlagsAndWritePrefix(/*IsContinued=*/false);
+  OS.write(Buffer, size_t(BufferPtr - Buffer));
+  OS.write_zeros(getRemainingSize());
+  BufferPtr = Buffer;
+}
+
+namespace {
+
 class GOFFObjectWriter : public MCObjectWriter {
   // The target specific GOFF writer instance.
   std::unique_ptr<MCGOFFObjectTargetWriter> TargetObjectWriter;
@@ -242,7 +251,7 @@ class GOFFObjectWriter : public MCObjectWriter {
 } // end anonymous namespace
 
 void GOFFObjectWriter::writeHeader() {
-  OS.newRecord(GOFF::RT_HDR, /*Size=*/57);
+  OS.newRecord(GOFF::RT_HDR);
   OS.write_zeros(1);       // Reserved
   OS.writebe<uint32_t>(0); // Target Hardware Environment
   OS.writebe<uint32_t>(0); // Target Operating System Environment
@@ -262,7 +271,7 @@ void GOFFObjectWriter::writeEnd() {
 
   // TODO Set Flags/AMODE/ESDID for entry point.
 
-  OS.newRecord(GOFF::RT_END, /*Size=*/13);
+  OS.newRecord(GOFF::RT_END);
   OS.writebe<uint8_t>(Flags(6, 2, F)); // Indicator flags
   OS.writebe<uint8_t>(AMODE);          // AMODE
   OS.write_zeros(3);                   // Reserved
@@ -271,18 +280,19 @@ void GOFFObjectWriter::writeEnd() {
   // being zero.
   OS.writebe<uint32_t>(0);     // Record Count
   OS.writebe<uint32_t>(ESDID); // ESDID (of entry point)
-  OS.finalize();
 }
 
 uint64_t GOFFObjectWriter::writeObject(MCAssembler &Asm) {
-  uint64_t StartOffset = OS.tell();
-
   writeHeader();
   writeEnd();
 
-  LLVM_DEBUG(dbgs() << "Wrote " << OS.logicalRecords() << " logical records.");
+  // Make sure all records are written.
+  OS.finalizeRecord();
+
+  LLVM_DEBUG(dbgs() << "Wrote " << OS.getNumLogicalRecords()
+                    << " logical records.");
 
-  return OS.tell() - StartOffset;
+  return OS.getWrittenSize();
 }
 
 std::unique_ptr<MCObjectWriter>