[llvm] Unittests and usability for BitstreamWriter incremental flushing (PR #92983)
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Fri May 24 12:33:56 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/4] 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/4] 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/4] 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);
>From b4f6587bcbe73b8fa26d456ce491f4652b8ef1f2 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Fri, 24 May 2024 12:08:44 -0700
Subject: [PATCH 4/4] feedback
---
llvm/include/llvm/Bitstream/BitstreamWriter.h | 2 ++
llvm/unittests/Bitstream/BitstreamWriterTest.cpp | 2 ++
2 files changed, 4 insertions(+)
diff --git a/llvm/include/llvm/Bitstream/BitstreamWriter.h b/llvm/include/llvm/Bitstream/BitstreamWriter.h
index d0476784a1c48..627503884985b 100644
--- a/llvm/include/llvm/Bitstream/BitstreamWriter.h
+++ b/llvm/include/llvm/Bitstream/BitstreamWriter.h
@@ -93,6 +93,8 @@ class BitstreamWriter {
/// If the related file stream supports reading, seeking and writing, flush
/// the buffer if its size is above a threshold.
+ /// If \p OnClosing is true, and if a stream is provided, flushing happens
+ /// regardless of thresholds.
void FlushToFile(bool OnClosing = false) {
if (!FS || Out.empty())
return;
diff --git a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
index 53a7aa5ecd886..096363190af4a 100644
--- a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
+++ b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
@@ -63,6 +63,8 @@ TEST(BitstreamWriterTest, emitBlob4ByteAligned) {
class BitstreamWriterFlushTest : public ::testing::TestWithParam<int> {
protected:
+ // Any value after bitc::FIRST_APPLICATION_BLOCKID is good, but let's pick a
+ // distinctive one.
const unsigned BlkID = bitc::FIRST_APPLICATION_BLOCKID + 17;
void write(StringRef TestFilePath, int FlushThreshold,
More information about the llvm-commits
mailing list