[llvm] Unittests and usability for BitstreamWriter incremental flushing (PR #92983)
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Wed May 22 10:12:49 PDT 2024
https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/92983
>From ae8bd286d5d851891ada8921bffa581551d0df61 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 21 May 2024 18:59:49 -0700
Subject: [PATCH 1/3] Unittests and usability for BitstreamWriter incremental
flushing
- added unittests for the raw_fd_stream output case.
- the ctor documentation seemed to suggest that if there was no file stream
provided, the buffer will be written at the end. That's incorrect. Fixed the
documentation, clarifying flushing behavior and highlighting a few other
details.
- flushing everything in the dtor. While flushing to the file stream would
happen on subblock boundaries, if the choice of format didn't feature subblocks
or didn't end with a subblock, the user would have the expectation of flushing
to the given stream, and have no API indication that's not the case.
---
llvm/include/llvm/Bitstream/BitstreamWriter.h | 28 ++++---
.../Bitstream/BitstreamWriterTest.cpp | 79 +++++++++++++++++++
2 files changed, 97 insertions(+), 10 deletions(-)
diff --git a/llvm/include/llvm/Bitstream/BitstreamWriter.h b/llvm/include/llvm/Bitstream/BitstreamWriter.h
index c726508cd5285..d0476784a1c48 100644
--- a/llvm/include/llvm/Bitstream/BitstreamWriter.h
+++ b/llvm/include/llvm/Bitstream/BitstreamWriter.h
@@ -93,10 +93,10 @@ class BitstreamWriter {
/// 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 FlushToFile(bool OnClosing = false) {
+ if (!FS || Out.empty())
return;
- if (Out.size() < FlushThreshold)
+ if (!OnClosing && Out.size() < FlushThreshold)
return;
FS->write((char *)&Out.front(), Out.size());
Out.clear();
@@ -104,20 +104,28 @@ class BitstreamWriter {
public:
/// Create a BitstreamWriter that writes to Buffer \p O.
+ /// The buffer is flushed incrementally to a raw_fd_stream \p FS, if provided.
///
- /// \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.
+ /// Flushing to FS happens:
+ /// - at \ref ExitBlock, and only if \O.size() >= \p FlushThreshold (MB)
+ /// - unconditionally at the BitstreamWriter destruction.
///
- /// \p FlushThreshold is the threshold (unit M) to flush \p O if \p FS is
- /// valid. Flushing only occurs at (sub)block boundaries.
+ /// NOTES:
+ /// - \p FlushThreshold's unit is MB.
+ /// - The provided buffer will always be resized as needed. Its initial size
+ /// has no bearing on flushing behavior (if \p FS is provided).
+ /// - \p FS is expected to support seek, tell, and read (besides write).
+ /// The BitstreamWriter makes no attempt to pre-validate FS, meaning errors
+ /// stemming from an unsupported \p FS will happen during writing.
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) {}
+ : Out(O), FS(FS), FlushThreshold(uint64_t(FlushThreshold) << 20),
+ CurBit(0), CurValue(0), CurCodeSize(2) {}
~BitstreamWriter() {
- assert(CurBit == 0 && "Unflushed data remaining");
+ FlushToWord();
assert(BlockScope.empty() && CurAbbrevs.empty() && "Block imbalance");
+ FlushToFile(/*OnClosing=*/true);
}
/// Retrieve the current position in the stream, in bits.
diff --git a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
index 054948e7e8b63..41b381af09394 100644
--- a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
+++ b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
@@ -9,6 +9,12 @@
#include "llvm/Bitstream/BitstreamWriter.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallString.h"
+#include "llvm/Bitstream/BitCodeEnums.h"
+#include "llvm/Bitstream/BitstreamReader.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Testing/Support/SupportHelpers.h"
+#include "gmock/gmock.h"
#include "gtest/gtest.h"
using namespace llvm;
@@ -55,4 +61,77 @@ TEST(BitstreamWriterTest, emitBlob4ByteAligned) {
EXPECT_EQ(StringRef("str0"), Buffer);
}
+class BitstreamWriterFlushTest :public ::testing::TestWithParam<int> {
+protected:
+ const unsigned BlkID = bitc::FIRST_APPLICATION_BLOCKID + 17;
+
+ void write(StringRef TestFilePath, int FlushThreshold,
+ llvm::function_ref<void(BitstreamWriter &)> Action) {
+ SmallString<64> Buffer;
+ std::error_code EC;
+ raw_fd_stream Out(TestFilePath, EC);
+ ASSERT_FALSE(EC);
+ BitstreamWriter W(Buffer, &Out, FlushThreshold);
+ Action(W);
+ }
+};
+
+TEST_P(BitstreamWriterFlushTest, simpleExample) {
+ llvm::unittest::TempFile TestFile("bitstream", "", "",
+ /*Unique*/ true);
+ write(TestFile.path(), GetParam(),
+ [&](BitstreamWriter &W) { W.EmitVBR(42, 2); });
+
+ ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
+ MemoryBuffer::getFile(TestFile.path());
+ ASSERT_TRUE(!!MB);
+ ASSERT_NE(*MB, nullptr);
+ BitstreamCursor Cursor((*MB)->getBuffer());
+ auto V = Cursor.ReadVBR(2);
+ EXPECT_TRUE(!!V);
+ EXPECT_EQ(*V, 42U);
+}
+
+TEST_P(BitstreamWriterFlushTest, subBlock) {
+ llvm::unittest::TempFile TestFile("bitstream", "", "",
+ /*Unique*/ true);
+ write(TestFile.path(), GetParam(), [&](BitstreamWriter &W) {
+ W.EnterSubblock(BlkID, 2);
+ W.EmitVBR(42, 2);
+ W.ExitBlock();
+ });
+ ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
+ MemoryBuffer::getFile(TestFile.path());
+ ASSERT_TRUE(!!MB);
+ ASSERT_NE(*MB, nullptr);
+ BitstreamCursor Cursor((*MB)->getBuffer());
+ auto Blk = Cursor.advance(BitstreamCursor::AF_DontAutoprocessAbbrevs);
+ ASSERT_TRUE(!!Blk);
+ EXPECT_EQ(Blk->Kind, BitstreamEntry::SubBlock);
+ EXPECT_EQ(Blk->ID, BlkID);
+ EXPECT_FALSE(Cursor.EnterSubBlock(BlkID));
+ auto V = Cursor.ReadVBR(2);
+ EXPECT_TRUE(!!V);
+ EXPECT_EQ(*V, 42U);
+ // ReadBlockEnd() returns false if it actually read the block end.
+ EXPECT_FALSE(Cursor.ReadBlockEnd());
+ EXPECT_TRUE(Cursor.AtEndOfStream());
+}
+
+TEST_P(BitstreamWriterFlushTest, blobRawRead) {
+ llvm::unittest::TempFile TestFile("bitstream", "", "",
+ /*Unique*/ true);
+ write(TestFile.path(), GetParam(), [&](BitstreamWriter &W) {
+ W.emitBlob("str", /* ShouldEmitSize */ false);
+ });
+
+ ErrorOr<std::unique_ptr<MemoryBuffer>> MB =
+ MemoryBuffer::getFile(TestFile.path());
+ ASSERT_TRUE(!!MB);
+ ASSERT_NE(*MB, nullptr);
+ EXPECT_EQ(StringRef("str\0", 4), (*MB)->getBuffer());
+}
+
+INSTANTIATE_TEST_SUITE_P(BitstreamWriterFlushCases, BitstreamWriterFlushTest,
+ ::testing::Values(0, 1 /*MB*/));
} // end namespace
>From 173486b18534afdd59da0b8dbe044cbbee919211 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 21 May 2024 20:41:17 -0700
Subject: [PATCH 2/3] clang-format
---
llvm/unittests/Bitstream/BitstreamWriterTest.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
index 41b381af09394..53a7aa5ecd886 100644
--- a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
+++ b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
@@ -61,7 +61,7 @@ TEST(BitstreamWriterTest, emitBlob4ByteAligned) {
EXPECT_EQ(StringRef("str0"), Buffer);
}
-class BitstreamWriterFlushTest :public ::testing::TestWithParam<int> {
+class BitstreamWriterFlushTest : public ::testing::TestWithParam<int> {
protected:
const unsigned BlkID = bitc::FIRST_APPLICATION_BLOCKID + 17;
>From 57ebf01058ebe28bff293837470b46fb35ecb6af Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Wed, 22 May 2024 10:12:23 -0700
Subject: [PATCH 3/3] fix
---
llvm/lib/Bitcode/Writer/BitcodeWriter.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index c4cea3d6eef2d..86c4293afb95d 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -5055,13 +5055,13 @@ void llvm::WriteBitcodeToFile(const Module &M, raw_ostream &Out,
Triple TT(M.getTargetTriple());
if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
Buffer.insert(Buffer.begin(), BWH_HeaderSize, 0);
-
- BitcodeWriter Writer(Buffer, dyn_cast<raw_fd_stream>(&Out));
- Writer.writeModule(M, ShouldPreserveUseListOrder, Index, GenerateHash,
- ModHash);
- Writer.writeSymtab();
- Writer.writeStrtab();
-
+ {
+ BitcodeWriter Writer(Buffer, dyn_cast<raw_fd_stream>(&Out));
+ Writer.writeModule(M, ShouldPreserveUseListOrder, Index, GenerateHash,
+ ModHash);
+ Writer.writeSymtab();
+ Writer.writeStrtab();
+ }
if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
emitDarwinBCHeaderAndTrailer(Buffer, TT);
More information about the llvm-commits
mailing list