[llvm] [GOFF] Refactor writing GOFF records (PR #93855)

[James Henderson via llvm-commits]

================
@@ -62,119 +63,90 @@ ZerosImpl zeros(const size_t NumBytes) { return ZerosImpl{NumBytes}; }
 
 // 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 payload. 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.
-class GOFFOstream : public raw_ostream {
+// logical record, and writes the full logical block. The physical records are
+// created while the content is written to the underlying stream. Possible fill
+// bytes at the end of a physical record are written automatically.
+// The implementation aims at simplicity, not speed.
+class GOFFOStream {
 public:
-  explicit GOFFOstream(raw_ostream &OS)
-      : OS(OS), LogicalRecords(0), RemainingSize(0), NewLogicalRecord(false) {
-    SetBufferSize(GOFF::PayloadLength);
-  }
-
-  ~GOFFOstream() { finalize(); }
+  explicit GOFFOStream(raw_ostream &OS)
+      : OS(OS), CurrentType(GOFF::RecordType(-1)) {}
 
-  void makeNewRecord(GOFF::RecordType Type, size_t Size) {
-    fillRecord();
-    CurrentType = Type;
-    RemainingSize = Size;
-    if (size_t Gap = (RemainingSize % GOFF::PayloadLength))
-      RemainingSize += GOFF::PayloadLength - Gap;
-    NewLogicalRecord = true;
-    ++LogicalRecords;
+  GOFFOStream &operator<<(StringRef Str) {
+    write(Str);
+    return *this;
   }
 
-  void finalize() { fillRecord(); }
-
-  uint32_t logicalRecords() { return LogicalRecords; }
+  void newRecord(GOFF::RecordType Type) { CurrentType = Type; }
 
 private:
   // The underlying raw_ostream.
   raw_ostream &OS;
 
-  // The number of logical records emitted so far.
-  uint32_t LogicalRecords;
-
-  // The remaining size of this logical record, including fill bytes.
-  size_t RemainingSize;
-
   // The type of the current (logical) record.
   GOFF::RecordType CurrentType;
 
-  // Signals start of new record.
-  bool NewLogicalRecord;
-
-  // Return the number of bytes left to write until next physical record.
-  // Please note that we maintain the total number of bytes left, not the
-  // written size.
-  size_t bytesToNextPhysicalRecord() {
-    size_t Bytes = RemainingSize % GOFF::PayloadLength;
-    return Bytes ? Bytes : GOFF::PayloadLength;
-  }
-
   // Write the record prefix of a physical record, using the current record
   // type.
-  static void writeRecordPrefix(raw_ostream &OS, GOFF::RecordType Type,
-                                size_t RemainingSize,
-                                uint8_t Flags = Rec_Continuation) {
-    uint8_t TypeAndFlags = Flags | (Type << 4);
-    if (RemainingSize > GOFF::RecordLength)
-      TypeAndFlags |= Rec_Continued;
-    OS << binaryBe(static_cast<unsigned char>(GOFF::PTVPrefix))
-       << binaryBe(static_cast<unsigned char>(TypeAndFlags))
-       << binaryBe(static_cast<unsigned char>(0));
-  }
+  void writeRecordPrefix(uint8_t Flags);
 
-  // Fill the last physical record of a logical record with zero bytes.
-  void fillRecord() {
-    assert((GetNumBytesInBuffer() <= RemainingSize) &&
-           "More bytes in buffer than expected");
-    size_t Remains = RemainingSize - GetNumBytesInBuffer();
-    if (Remains) {
-      assert((Remains < GOFF::RecordLength) &&
-             "Attempting to fill more than one physical record");
-      raw_ostream::write_zeros(Remains);
-    }
-    flush();
-    assert(RemainingSize == 0 && "Not fully flushed");
-    assert(GetNumBytesInBuffer() == 0 && "Buffer not fully empty");
-  }
+  // Write a logical record.
+  void write(StringRef Str);
+};
 
-  // See raw_ostream::write_impl.
-  void write_impl(const char *Ptr, size_t Size) override {
-    assert((RemainingSize >= Size) && "Attempt to write too much data");
-    assert(RemainingSize && "Logical record overflow");
-    if (!(RemainingSize % GOFF::PayloadLength)) {
-      writeRecordPrefix(OS, CurrentType, RemainingSize,
-                        NewLogicalRecord ? 0 : Rec_Continuation);
-      NewLogicalRecord = false;
-    }
-    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);
-      }
-    }
+void GOFFOStream::writeRecordPrefix(uint8_t Flags) {
+  uint8_t TypeAndFlags = Flags | (CurrentType << 4);
+  OS << binaryBe(static_cast<unsigned char>(GOFF::PTVPrefix))
+     << binaryBe(static_cast<unsigned char>(TypeAndFlags))
+     << binaryBe(static_cast<unsigned char>(0));
+}
+
+void GOFFOStream::write(StringRef Str) {
+  // The flags are determined by the flags of the prvious record, and by the
+  // remaining size of data.
+  uint8_t Flags = 0;
+  size_t Ptr = 0;
+  size_t Size = Str.size();
+  while (Size >= GOFF::RecordContentLength) {
+    if (Flags) {
+      Flags |= Rec_Continuation;
+      if (Size == GOFF::RecordContentLength)
+        Flags &= ~Rec_Continued;
+    } else
+      Flags |= (Size == GOFF::RecordContentLength) ? 0 : Rec_Continued;
+    writeRecordPrefix(Flags);
+    OS.write(&Str.data()[Ptr], GOFF::RecordContentLength);
+    Size -= GOFF::RecordContentLength;
+    Ptr += GOFF::RecordContentLength;
   }
+  if (Size) {
+    Flags &= ~Rec_Continued;
+    writeRecordPrefix(Flags);
+    OS.write(&Str.data()[Ptr], Size);
+    OS.write_zeros(GOFF::RecordContentLength - Size);
+  }
----------------
jh7370 wrote:

This all feels a bit tricky to follow, and I wonder if an alternative implementation would be better? As I understand it, there are two main complexities that need resolving: calculating the flags and determining what to write. How about the following:

```
size_t Size = Str.size();
bool Continuation = false;
while (Size > 0) {
  uint8_t Flags = 0;
  if (Continuation)
    Flags |= Rec_Continuation;
  if (Size > GOFF::RecordContentLength) {
    Flags |= Rec_Continued;
    Continuation = true;
  }
  writeRecordPrefix(Flags);
  OS.write(&Str.data()[Ptr], Size);
  if (Size < GOFF::RecordContentLength) {
    OS.write_zeros(GOFF::RecordContentLength - Size);
    Size = 0;
  } else {
    Size -= GOFF::RecordContentLength;
  }
}
```

https://github.com/llvm/llvm-project/pull/93855