Unittests and usability for BitstreamWriter incremental flushing (llvm#92983)

- added unittests for the raw_fd_stream output case.

- the `BitstreamWriter` ctor was confusing, the relationship between the buffer and the file stream wasn't clear and in fact there was a potential bug in `BitcodeWriter` in the mach-o case, because that code assumed in-buffer only serialization. The incremental flushing behavior of flushing at end of block boundaries was an implementation detail that meant serializers not using blocks (for example) would need to know to check the buffer and flush. The bug was latent - in the sense that, today, because the stream being passed was not a `raw_fd_stream`, incremental buffering never kicked in.

The new design moves the responsibility of flushing to the `BitstreamWriter`, and makes it work with any `raw_ostream` (but incrementally flush only in the `raw_fd_stream` case). If the `raw_ostream` is over a buffer - i.e. a `raw_svector_stream` - then it's equivalent to today's buffer case. For all other `raw_ostream` cases, buffering is an implementation detail. In all cases, the buffer is flushed (well, in the buffer case, that's a moot statement).

This simplifies the state and state transitions the user has to track: you have a raw_ostream -> BitstreamWrite in it -> destroy the writer => the bitstream is completely written in your raw_ostream. The "buffer" case and the "raw_fd_stream" case become optimizations rather than imposing state transition concerns to the user.

Author: mtrofin

llvm/include/llvm/Bitcode/BitcodeWriter.h

@@ -30,7 +30,6 @@ class Module;
 class raw_ostream;
 
 class BitcodeWriter {
-  SmallVectorImpl<char> &Buffer;
   std::unique_ptr<BitstreamWriter> Stream;
 
   StringTableBuilder StrtabBuilder{StringTableBuilder::RAW};
@@ -47,7 +46,8 @@ class BitcodeWriter {
 
 public:
   /// Create a BitcodeWriter that writes to Buffer.
-  BitcodeWriter(SmallVectorImpl<char> &Buffer, raw_fd_stream *FS = nullptr);
+  BitcodeWriter(SmallVectorImpl<char> &Buffer);
+  BitcodeWriter(raw_ostream &FS);
 
   ~BitcodeWriter();
 

llvm/include/llvm/Bitstream/BitstreamWriter.h

@@ -18,6 +18,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Bitstream/BitCodes.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
@@ -28,34 +29,44 @@
 namespace llvm {
 
 class BitstreamWriter {
-  /// Out - The buffer that keeps unflushed bytes.
-  SmallVectorImpl<char> &Out;
-
-  /// FS - The file stream that Out flushes to. If FS is nullptr, it does not
-  /// support read or seek, Out cannot be flushed until all data are written.
-  raw_fd_stream *FS;
-
-  /// FlushThreshold - If FS is valid, this is the threshold (unit B) to flush
-  /// FS.
+  /// Owned buffer, used to init Buffer if the provided stream doesn't happen to
+  /// be a buffer itself.
+  SmallVector<char, 0> OwnBuffer;
+  /// Internal buffer for unflushed bytes (unless there is no stream to flush
+  /// to, case in which these are "the bytes"). The writer backpatches, so it is
+  /// efficient to buffer.
+  SmallVectorImpl<char> &Buffer;
+
+  /// FS - The file stream that Buffer flushes to. If FS is a raw_fd_stream, the
+  /// writer will incrementally flush at subblock boundaries. Otherwise flushing
+  /// will happen at the end of BitstreamWriter's lifetime.
+  raw_ostream *const FS;
+
+  /// FlushThreshold - this is the threshold (unit B) to flush to FS, if FS is a
+  /// raw_fd_stream.
   const uint64_t FlushThreshold;
 
   /// CurBit - Always between 0 and 31 inclusive, specifies the next bit to use.
-  unsigned CurBit;
+  unsigned CurBit = 0;
 
   /// CurValue - The current value. Only bits < CurBit are valid.
-  uint32_t CurValue;
+  uint32_t CurValue = 0;
 
   /// CurCodeSize - This is the declared size of code values used for the
   /// current block, in bits.
-  unsigned CurCodeSize;
+  unsigned CurCodeSize = 2;
 
   /// BlockInfoCurBID - When emitting a BLOCKINFO_BLOCK, this is the currently
   /// selected BLOCK ID.
-  unsigned BlockInfoCurBID;
+  unsigned BlockInfoCurBID = 0;
 
   /// CurAbbrevs - Abbrevs installed at in this block.
   std::vector<std::shared_ptr<BitCodeAbbrev>> CurAbbrevs;
 
+  // Support for retrieving a section of the output, for purposes such as
+  // checksumming.
+  std::optional<size_t> BlockFlushingStartPos;
+
   struct Block {
     unsigned PrevCodeSize;
     size_t StartSizeWord;
@@ -77,47 +88,106 @@ class BitstreamWriter {
   void WriteWord(unsigned Value) {
     Value =
         support::endian::byte_swap<uint32_t, llvm::endianness::little>(Value);
-    Out.append(reinterpret_cast<const char *>(&Value),
-               reinterpret_cast<const char *>(&Value + 1));
+    Buffer.append(reinterpret_cast<const char *>(&Value),
+                  reinterpret_cast<const char *>(&Value + 1));
   }
 
-  uint64_t GetNumOfFlushedBytes() const { return FS ? FS->tell() : 0; }
+  uint64_t GetNumOfFlushedBytes() const {
+    return fdStream() ? fdStream()->tell() : 0;
+  }
 
-  size_t GetBufferOffset() const { return Out.size() + GetNumOfFlushedBytes(); }
+  size_t GetBufferOffset() const {
+    return Buffer.size() + GetNumOfFlushedBytes();
+  }
 
   size_t GetWordIndex() const {
     size_t Offset = GetBufferOffset();
     assert((Offset & 3) == 0 && "Not 32-bit aligned");
     return Offset / 4;
   }
 
-  /// If the related file stream supports reading, seeking and writing, flush
-  /// the buffer if its size is above a threshold.
-  void FlushToFile() {
-    if (!FS)
+  void flushAndClear() {
+    assert(FS);
+    assert(!Buffer.empty());
+    assert(!BlockFlushingStartPos &&
+           "a call to markAndBlockFlushing should have been paired with a "
+           "call to getMarkedBufferAndResumeFlushing");
+    FS->write(Buffer.data(), Buffer.size());
+    Buffer.clear();
+  }
+
+  /// If the related file stream is a raw_fd_stream, flush the buffer if its
+  /// size is above a threshold. If \p OnClosing is true, flushing happens
+  /// regardless of thresholds.
+  void FlushToFile(bool OnClosing = false) {
+    if (!FS || Buffer.empty())
       return;
-    if (Out.size() < FlushThreshold)
+    if (OnClosing)
+      return flushAndClear();
+    if (BlockFlushingStartPos)
       return;
-    FS->write((char *)&Out.front(), Out.size());
-    Out.clear();
+    if (fdStream() && Buffer.size() > FlushThreshold)
+      flushAndClear();
+  }
+
+  raw_fd_stream *fdStream() { return dyn_cast_or_null<raw_fd_stream>(FS); }
+
+  const raw_fd_stream *fdStream() const {
+    return dyn_cast_or_null<raw_fd_stream>(FS);
+  }
+
+  SmallVectorImpl<char> &getInternalBufferFromStream(raw_ostream &OutStream) {
+    if (auto *SV = dyn_cast<raw_svector_ostream>(&OutStream))
+      return SV->buffer();
+    return OwnBuffer;
   }
 
 public:
-  /// Create a BitstreamWriter that writes to Buffer \p O.
+  /// Create a BitstreamWriter over a raw_ostream \p OutStream.
+  /// If \p OutStream is a raw_svector_ostream, the BitstreamWriter will write
+  /// directly to the latter's buffer. In all other cases, the BitstreamWriter
+  /// will use an internal buffer and flush at the end of its lifetime.
   ///
-  /// \p FS is the file stream that \p O flushes to incrementally. If \p FS is
-  /// null, \p O does not flush incrementially, but writes to disk at the end.
+  /// In addition, if \p is a raw_fd_stream supporting seek, tell, and read
+  /// (besides write), the BitstreamWriter will also flush incrementally, when a
+  /// subblock is finished, and if the FlushThreshold is passed.
   ///
-  /// \p FlushThreshold is the threshold (unit M) to flush \p O if \p FS is
-  /// valid. Flushing only occurs at (sub)block boundaries.
-  BitstreamWriter(SmallVectorImpl<char> &O, raw_fd_stream *FS = nullptr,
-                  uint32_t FlushThreshold = 512)
-      : Out(O), FS(FS), FlushThreshold(uint64_t(FlushThreshold) << 20), CurBit(0),
-        CurValue(0), CurCodeSize(2) {}
+  /// NOTE: \p FlushThreshold's unit is MB.
+  BitstreamWriter(raw_ostream &OutStream, uint32_t FlushThreshold = 512)
+      : Buffer(getInternalBufferFromStream(OutStream)),
+        FS(!isa<raw_svector_ostream>(OutStream) ? &OutStream : nullptr),
+        FlushThreshold(uint64_t(FlushThreshold) << 20) {}
+
+  /// Convenience constructor for users that start with a vector - avoids
+  /// needing to wrap it in a raw_svector_ostream.
+  BitstreamWriter(SmallVectorImpl<char> &Buff)
+      : Buffer(Buff), FS(nullptr), FlushThreshold(0) {}
 
   ~BitstreamWriter() {
-    assert(CurBit == 0 && "Unflushed data remaining");
+    FlushToWord();
     assert(BlockScope.empty() && CurAbbrevs.empty() && "Block imbalance");
+    FlushToFile(/*OnClosing=*/true);
+  }
+
+  /// For scenarios where the user wants to access a section of the stream to
+  /// (for example) compute some checksum, disable flushing and remember the
+  /// position in the internal buffer where that happened. Must be paired with a
+  /// call to getMarkedBufferAndResumeFlushing.
+  void markAndBlockFlushing() {
+    assert(!BlockFlushingStartPos);
+    BlockFlushingStartPos = Buffer.size();
+  }
+
+  /// resumes flushing, but does not flush, and returns the section in the
+  /// internal buffer starting from the position marked with
+  /// markAndBlockFlushing. The return should be processed before any additional
+  /// calls to this object, because those may cause a flush and invalidate the
+  /// return.
+  StringRef getMarkedBufferAndResumeFlushing() {
+    assert(BlockFlushingStartPos);
+    size_t Start = *BlockFlushingStartPos;
+    BlockFlushingStartPos.reset();
+    return {&Buffer[Start], Buffer.size() - Start};
   }
 
   /// Retrieve the current position in the stream, in bits.
@@ -141,16 +211,19 @@ class BitstreamWriter {
     if (ByteNo >= NumOfFlushedBytes) {
       assert((!endian::readAtBitAlignment<uint8_t, llvm::endianness::little,
                                           unaligned>(
-                 &Out[ByteNo - NumOfFlushedBytes], StartBit)) &&
+                 &Buffer[ByteNo - NumOfFlushedBytes], StartBit)) &&
              "Expected to be patching over 0-value placeholders");
       endian::writeAtBitAlignment<uint8_t, llvm::endianness::little, unaligned>(
-          &Out[ByteNo - NumOfFlushedBytes], NewByte, StartBit);
+          &Buffer[ByteNo - NumOfFlushedBytes], NewByte, StartBit);
       return;
     }
 
+    // If we don't have a raw_fd_stream, GetNumOfFlushedBytes() should have
+    // returned 0, and we shouldn't be here.
+    assert(fdStream() != nullptr);
     // If the byte offset to backpatch is flushed, use seek to backfill data.
     // First, save the file position to restore later.
-    uint64_t CurPos = FS->tell();
+    uint64_t CurPos = fdStream()->tell();
 
     // Copy data to update into Bytes from the file FS and the buffer Out.
     char Bytes[3]; // Use one more byte to silence a warning from Visual C++.
@@ -159,19 +232,19 @@ class BitstreamWriter {
     size_t BytesFromBuffer = BytesNum - BytesFromDisk;
 
     // When unaligned, copy existing data into Bytes from the file FS and the
-    // buffer Out so that it can be updated before writing. For debug builds
+    // buffer Buffer so that it can be updated before writing. For debug builds
     // read bytes unconditionally in order to check that the existing value is 0
     // as expected.
 #ifdef NDEBUG
     if (StartBit)
 #endif
     {
-      FS->seek(ByteNo);
-      ssize_t BytesRead = FS->read(Bytes, BytesFromDisk);
+      fdStream()->seek(ByteNo);
+      ssize_t BytesRead = fdStream()->read(Bytes, BytesFromDisk);
       (void)BytesRead; // silence warning
       assert(BytesRead >= 0 && static_cast<size_t>(BytesRead) == BytesFromDisk);
       for (size_t i = 0; i < BytesFromBuffer; ++i)
-        Bytes[BytesFromDisk + i] = Out[i];
+        Bytes[BytesFromDisk + i] = Buffer[i];
       assert((!endian::readAtBitAlignment<uint8_t, llvm::endianness::little,
                                           unaligned>(Bytes, StartBit)) &&
              "Expected to be patching over 0-value placeholders");
@@ -182,13 +255,13 @@ class BitstreamWriter {
         Bytes, NewByte, StartBit);
 
     // Copy updated data back to the file FS and the buffer Out.
-    FS->seek(ByteNo);
-    FS->write(Bytes, BytesFromDisk);
+    fdStream()->seek(ByteNo);
+    fdStream()->write(Bytes, BytesFromDisk);
     for (size_t i = 0; i < BytesFromBuffer; ++i)
-      Out[i] = Bytes[BytesFromDisk + i];
+      Buffer[i] = Bytes[BytesFromDisk + i];
 
     // Restore the file position.
-    FS->seek(CurPos);
+    fdStream()->seek(CurPos);
   }
 
   void BackpatchHalfWord(uint64_t BitNo, uint16_t Val) {
@@ -481,11 +554,11 @@ class BitstreamWriter {
 
     // Emit literal bytes.
     assert(llvm::all_of(Bytes, [](UIntTy B) { return isUInt<8>(B); }));
-    Out.append(Bytes.begin(), Bytes.end());
+    Buffer.append(Bytes.begin(), Bytes.end());
 
     // Align end to 32-bits.
     while (GetBufferOffset() & 3)
-      Out.push_back(0);
+      Buffer.push_back(0);
   }
   void emitBlob(StringRef Bytes, bool ShouldEmitSize = true) {
     emitBlob(ArrayRef((const uint8_t *)Bytes.data(), Bytes.size()),

llvm/include/llvm/ProfileData/PGOCtxProfWriter.h

@@ -70,7 +70,7 @@ class PGOCtxProfileWriter final {
 public:
   PGOCtxProfileWriter(raw_fd_stream &Out,
                       std::optional<unsigned> VersionOverride = std::nullopt)
-      : Writer(Buff, &Out, 0) {
+      : Writer(Out, 0) {
     Writer.EnterSubblock(PGOCtxProfileBlockIDs::ProfileMetadataBlockID,
                          CodeLen);
     const auto Version = VersionOverride ? *VersionOverride : CurrentVersion;

llvm/include/llvm/Support/raw_ostream.h

@@ -55,6 +55,7 @@ class raw_ostream {
   enum class OStreamKind {
     OK_OStream,
     OK_FDStream,
+    OK_SVecStream,
   };
 
 private:
@@ -703,7 +704,11 @@ class raw_svector_ostream : public raw_pwrite_stream {
   ///
   /// \param O The vector to write to; this should generally have at least 128
   /// bytes free to avoid any extraneous memory overhead.
-  explicit raw_svector_ostream(SmallVectorImpl<char> &O) : OS(O) {
+  explicit raw_svector_ostream(SmallVectorImpl<char> &O)
+      : raw_pwrite_stream(false, raw_ostream::OStreamKind::OK_SVecStream),
+        OS(O) {
+    // FIXME: here and in a few other places, set directly to unbuffered in the
+    // ctor.
     SetUnbuffered();
   }
 
@@ -713,10 +718,13 @@ class raw_svector_ostream : public raw_pwrite_stream {
 
   /// Return a StringRef for the vector contents.
   StringRef str() const { return StringRef(OS.data(), OS.size()); }
+  SmallVectorImpl<char> &buffer() { return OS; }
 
   void reserveExtraSpace(uint64_t ExtraSize) override {
     OS.reserve(tell() + ExtraSize);
   }
+
+  static bool classof(const raw_ostream *OS);
 };
 
 /// A raw_ostream that discards all output.