[llvm] 7cfffe7 - Unittests and usability for BitstreamWriter incremental flushing (#92983)
via llvm-commits
llvm-commits at lists.llvm.org
Thu May 30 12:26:03 PDT 2024
Author: Mircea Trofin
Date: 2024-05-30T12:25:59-07:00
New Revision: 7cfffe74eeb68fbb3fb9706ac7071f8caeeb6520
URL: https://github.com/llvm/llvm-project/commit/7cfffe74eeb68fbb3fb9706ac7071f8caeeb6520
DIFF: https://github.com/llvm/llvm-project/commit/7cfffe74eeb68fbb3fb9706ac7071f8caeeb6520.diff
LOG: Unittests and usability for BitstreamWriter incremental flushing (#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.
Added:
Modified:
llvm/include/llvm/Bitcode/BitcodeWriter.h
llvm/include/llvm/Bitstream/BitstreamWriter.h
llvm/include/llvm/ProfileData/PGOCtxProfWriter.h
llvm/include/llvm/Support/raw_ostream.h
llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
llvm/lib/Support/raw_ostream.cpp
llvm/unittests/Bitstream/BitstreamWriterTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/Bitcode/BitcodeWriter.h b/llvm/include/llvm/Bitcode/BitcodeWriter.h
index d1f9d57b6db6a..770e249290c3c 100644
--- a/llvm/include/llvm/Bitcode/BitcodeWriter.h
+++ b/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();
diff --git a/llvm/include/llvm/Bitstream/BitstreamWriter.h b/llvm/include/llvm/Bitstream/BitstreamWriter.h
index c726508cd5285..78f5eb4f364b6 100644
--- a/llvm/include/llvm/Bitstream/BitstreamWriter.h
+++ b/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,13 +88,17 @@ 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();
@@ -91,33 +106,88 @@ class BitstreamWriter {
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()),
diff --git a/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h b/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h
index edcf02c094697..e8b75d6f8cf4b 100644
--- a/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h
+++ b/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;
diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 696290a5d99cf..0951ffb19ffa1 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/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.
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index 046dad5721c4c..7d39c0db79fb1 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -260,9 +260,6 @@ class ModuleBitcodeWriterBase : public BitcodeWriterBase {
/// Class to manage the bitcode writing for a module.
class ModuleBitcodeWriter : public ModuleBitcodeWriterBase {
- /// Pointer to the buffer allocated by caller for bitcode writing.
- const SmallVectorImpl<char> &Buffer;
-
/// True if a module hash record should be written.
bool GenerateHash;
@@ -278,14 +275,13 @@ class ModuleBitcodeWriter : public ModuleBitcodeWriterBase {
public:
/// Constructs a ModuleBitcodeWriter object for the given Module,
/// writing to the provided \p Buffer.
- ModuleBitcodeWriter(const Module &M, SmallVectorImpl<char> &Buffer,
- StringTableBuilder &StrtabBuilder,
+ ModuleBitcodeWriter(const Module &M, StringTableBuilder &StrtabBuilder,
BitstreamWriter &Stream, bool ShouldPreserveUseListOrder,
const ModuleSummaryIndex *Index, bool GenerateHash,
ModuleHash *ModHash = nullptr)
: ModuleBitcodeWriterBase(M, StrtabBuilder, Stream,
ShouldPreserveUseListOrder, Index),
- Buffer(Buffer), GenerateHash(GenerateHash), ModHash(ModHash),
+ GenerateHash(GenerateHash), ModHash(ModHash),
BitcodeStartBit(Stream.GetCurrentBitNo()) {}
/// Emit the current module to the bitstream.
@@ -414,7 +410,7 @@ class ModuleBitcodeWriter : public ModuleBitcodeWriterBase {
writeFunction(const Function &F,
DenseMap<const Function *, uint64_t> &FunctionToBitcodeIndex);
void writeBlockInfo();
- void writeModuleHash(size_t BlockStartPos);
+ void writeModuleHash(StringRef View);
unsigned getEncodedSyncScopeID(SyncScope::ID SSID) {
return unsigned(SSID);
@@ -4819,13 +4815,13 @@ static void writeIdentificationBlock(BitstreamWriter &Stream) {
Stream.ExitBlock();
}
-void ModuleBitcodeWriter::writeModuleHash(size_t BlockStartPos) {
+void ModuleBitcodeWriter::writeModuleHash(StringRef View) {
// Emit the module's hash.
// MODULE_CODE_HASH: [5*i32]
if (GenerateHash) {
uint32_t Vals[5];
- Hasher.update(ArrayRef<uint8_t>((const uint8_t *)&(Buffer)[BlockStartPos],
- Buffer.size() - BlockStartPos));
+ Hasher.update(ArrayRef<uint8_t>(
+ reinterpret_cast<const uint8_t *>(View.data()), View.size()));
std::array<uint8_t, 20> Hash = Hasher.result();
for (int Pos = 0; Pos < 20; Pos += 4) {
Vals[Pos / 4] = support::endian::read32be(Hash.data() + Pos);
@@ -4844,7 +4840,9 @@ void ModuleBitcodeWriter::write() {
writeIdentificationBlock(Stream);
Stream.EnterSubblock(bitc::MODULE_BLOCK_ID, 3);
- size_t BlockStartPos = Buffer.size();
+ // We will want to write the module hash at this point. Block any flushing so
+ // we can have access to the whole underlying data later.
+ Stream.markAndBlockFlushing();
writeModuleVersion();
@@ -4895,7 +4893,7 @@ void ModuleBitcodeWriter::write() {
writeGlobalValueSymbolTable(FunctionToBitcodeIndex);
- writeModuleHash(BlockStartPos);
+ writeModuleHash(Stream.getMarkedBufferAndResumeFlushing());
Stream.ExitBlock();
}
@@ -4976,8 +4974,13 @@ static void writeBitcodeHeader(BitstreamWriter &Stream) {
Stream.Emit(0xD, 4);
}
-BitcodeWriter::BitcodeWriter(SmallVectorImpl<char> &Buffer, raw_fd_stream *FS)
- : Buffer(Buffer), Stream(new BitstreamWriter(Buffer, FS, FlushThreshold)) {
+BitcodeWriter::BitcodeWriter(SmallVectorImpl<char> &Buffer)
+ : Stream(new BitstreamWriter(Buffer)) {
+ writeBitcodeHeader(*Stream);
+}
+
+BitcodeWriter::BitcodeWriter(raw_ostream &FS)
+ : Stream(new BitstreamWriter(FS, FlushThreshold)) {
writeBitcodeHeader(*Stream);
}
@@ -5060,7 +5063,7 @@ void BitcodeWriter::writeModule(const Module &M,
assert(M.isMaterialized());
Mods.push_back(const_cast<Module *>(&M));
- ModuleBitcodeWriter ModuleWriter(M, Buffer, StrtabBuilder, *Stream,
+ ModuleBitcodeWriter ModuleWriter(M, StrtabBuilder, *Stream,
ShouldPreserveUseListOrder, Index,
GenerateHash, ModHash);
ModuleWriter.write();
@@ -5080,27 +5083,28 @@ void llvm::WriteBitcodeToFile(const Module &M, raw_ostream &Out,
bool ShouldPreserveUseListOrder,
const ModuleSummaryIndex *Index,
bool GenerateHash, ModuleHash *ModHash) {
- SmallVector<char, 0> Buffer;
- Buffer.reserve(256*1024);
-
- // If this is darwin or another generic macho target, reserve space for the
- // header.
+ auto Write = [&](BitcodeWriter &Writer) {
+ Writer.writeModule(M, ShouldPreserveUseListOrder, Index, GenerateHash,
+ ModHash);
+ Writer.writeSymtab();
+ Writer.writeStrtab();
+ };
Triple TT(M.getTargetTriple());
- if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
+ if (TT.isOSDarwin() || TT.isOSBinFormatMachO()) {
+ // If this is darwin or another generic macho target, reserve space for the
+ // header. Note that the header is computed *after* the output is known, so
+ // we currently explicitly use a buffer, write to it, and then subsequently
+ // flush to Out.
+ SmallVector<char, 256 * 1024> Buffer;
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();
-
- if (TT.isOSDarwin() || TT.isOSBinFormatMachO())
+ BitcodeWriter Writer(Buffer);
+ Write(Writer);
emitDarwinBCHeaderAndTrailer(Buffer, TT);
-
- // Write the generated bitstream to "Out".
- if (!Buffer.empty())
- Out.write((char *)&Buffer.front(), Buffer.size());
+ Out.write(Buffer.data(), Buffer.size());
+ } else {
+ BitcodeWriter Writer(Out);
+ Write(Writer);
+ }
}
void IndexBitcodeWriter::write() {
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 8cb7b5ac68ea7..e8d12357bfded 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -977,6 +977,10 @@ void raw_svector_ostream::pwrite_impl(const char *Ptr, size_t Size,
memcpy(OS.data() + Offset, Ptr, Size);
}
+bool raw_svector_ostream::classof(const raw_ostream *OS) {
+ return OS->get_kind() == OStreamKind::OK_SVecStream;
+}
+
//===----------------------------------------------------------------------===//
// raw_null_ostream
//===----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Bitstream/BitstreamWriterTest.cpp b/llvm/unittests/Bitstream/BitstreamWriterTest.cpp
index 054948e7e8b63..933965abdbadf 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,78 @@ TEST(BitstreamWriterTest, emitBlob4ByteAligned) {
EXPECT_EQ(StringRef("str0"), Buffer);
}
+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,
+ llvm::function_ref<void(BitstreamWriter &)> Action) {
+ std::error_code EC;
+ raw_fd_stream Out(TestFilePath, EC);
+ ASSERT_FALSE(EC);
+ BitstreamWriter W(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
More information about the llvm-commits
mailing list