diff --git a/include/revng/Support/GzipStream.h b/include/revng/Support/GzipStream.h index 7e6aaaf72..7028e5aed 100644 --- a/include/revng/Support/GzipStream.h +++ b/include/revng/Support/GzipStream.h @@ -7,20 +7,61 @@ #include "llvm/ADT/ArrayRef.h" #include "llvm/Support/raw_ostream.h" +#include "revng/ADT/STLExtras.h" + +#include "zlib.h" + void gzipCompress(llvm::raw_ostream &OS, llvm::ArrayRef Buffer, int CompressionLevel = 3); -inline void gzipCompress(llvm::raw_ostream &OS, llvm::ArrayRef Buffer) { - return gzipCompress(OS, - { reinterpret_cast(Buffer.data()), - Buffer.size() }); +template +inline void +gzipCompress(llvm::raw_ostream &OS, T Buffer, int CompressionLevel = 3) { + gzipCompress(OS, + { reinterpret_cast(Buffer.data()), + Buffer.size() }, + CompressionLevel); +} + +template +inline llvm::SmallVector +gzipCompress(T Buffer, int CompressionLevel = 3) { + llvm::SmallVector Result; + llvm::raw_svector_ostream OS(Result); + gzipCompress(OS, Buffer); + return Result; } void gzipDecompress(llvm::raw_ostream &OS, llvm::ArrayRef Buffer); -inline void gzipDecompress(llvm::raw_ostream &OS, llvm::ArrayRef Buffer) { - return gzipDecompress(OS, - { reinterpret_cast(Buffer.data()), - Buffer.size() }); +template +inline void gzipDecompress(llvm::raw_ostream &OS, T Buffer) { + gzipDecompress(OS, + { reinterpret_cast(Buffer.data()), + Buffer.size() }); } + +template +inline llvm::SmallVector gzipDecompress(T Buffer) { + llvm::SmallVector Result; + llvm::raw_svector_ostream OS(Result); + gzipDecompress(OS, Buffer); + return Result; +} + +class GzipCompressedOstream : public llvm::raw_ostream { +private: + llvm::raw_ostream &OS; + llvm::SmallVector OutBuffer; + z_stream Stream; + +public: + GzipCompressedOstream(llvm::raw_ostream &OS, int CompressionLevel = 3); + ~GzipCompressedOstream() override; + void flush(); + +private: + void write_impl(const char *Ptr, size_t Size) override; + uint64_t current_pos() const override { return OS.tell(); } +}; diff --git a/lib/Support/GzipStream.cpp b/lib/Support/GzipStream.cpp index 13bc258ac..25d05b7f9 100644 --- a/lib/Support/GzipStream.cpp +++ b/lib/Support/GzipStream.cpp @@ -19,11 +19,10 @@ constexpr int WindowBits = 15 // 2**15 bytes (32k) of window + 16; // Magic offset for gzip template -static void zlibCopyStream(z_stream &Stream, - llvm::raw_ostream &OutputOS, - llvm::ArrayRef InputBuffer) { - llvm::SmallVector OutBuffer; - OutBuffer.resize_for_overwrite(OutputBufferSize); +static void zlibReadInput(z_stream &Stream, + llvm::raw_ostream &OutputOS, + llvm::ArrayRef InputBuffer, + llvm::SmallVector &OutBuffer) { const char *OutBufferPtr = reinterpret_cast(OutBuffer.data()); size_t RemainingInput = InputBuffer.size(); @@ -54,7 +53,13 @@ static void zlibCopyStream(z_stream &Stream, Stream.avail_out = OutBuffer.size(); } } +} +template +static void zlibFlushOutput(z_stream &Stream, + llvm::raw_ostream &OutputOS, + llvm::SmallVector &OutBuffer) { + const char *OutBufferPtr = reinterpret_cast(OutBuffer.data()); Stream.next_in = Z_NULL; Stream.avail_in = 0; @@ -94,7 +99,10 @@ void gzipCompress(llvm::raw_ostream &OutputOS, Strategy); revng_assert(RC == Z_OK); - zlibCopyStream(Stream, OutputOS, InputBuffer); + llvm::SmallVector OutBuffer; + OutBuffer.resize_for_overwrite(OutputBufferSize); + zlibReadInput(Stream, OutputOS, InputBuffer, OutBuffer); + zlibFlushOutput(Stream, OutputOS, OutBuffer); revng_assert(deflateEnd(&Stream) == Z_OK); } @@ -104,7 +112,42 @@ void gzipDecompress(llvm::raw_ostream &OutputOS, z_stream Stream = { .zalloc = Z_NULL, .zfree = Z_NULL, .opaque = Z_NULL }; revng_assert(inflateInit2(&Stream, WindowBits) == Z_OK); - zlibCopyStream(Stream, OutputOS, InputBuffer); + llvm::SmallVector OutBuffer; + OutBuffer.resize_for_overwrite(OutputBufferSize); + zlibReadInput(Stream, OutputOS, InputBuffer, OutBuffer); + zlibFlushOutput(Stream, OutputOS, OutBuffer); revng_assert(inflateEnd(&Stream) == Z_OK); } + +GzipCompressedOstream::GzipCompressedOstream(llvm::raw_ostream &OS, + int CompressionLevel) : + llvm::raw_ostream(), OS(OS), OutBuffer() { + revng_assert(CompressionLevel >= 1 and CompressionLevel <= 9); + Stream = { .zalloc = Z_NULL, .zfree = Z_NULL, .opaque = Z_NULL }; + + int RC = deflateInit2(&Stream, + CompressionLevel, + Z_DEFLATED, + WindowBits, + 8, + Z_DEFAULT_STRATEGY); + revng_assert(RC == Z_OK); + OutBuffer.resize_for_overwrite(OutputBufferSize); +} + +GzipCompressedOstream::~GzipCompressedOstream() { + flush(); + revng_assert(deflateEnd(&Stream) == Z_OK); +} + +void GzipCompressedOstream::flush() { + llvm::raw_ostream::flush(); + zlibFlushOutput(Stream, OS, OutBuffer); + OS.flush(); +} + +void GzipCompressedOstream::write_impl(const char *Ptr, size_t Size) { + llvm::ArrayRef Input(reinterpret_cast(Ptr), Size); + zlibReadInput(Stream, OS, Input, OutBuffer); +} diff --git a/lib/Support/GzipTarFile.cpp b/lib/Support/GzipTarFile.cpp index 1045a80a1..8e4f5608c 100644 --- a/lib/Support/GzipTarFile.cpp +++ b/lib/Support/GzipTarFile.cpp @@ -132,12 +132,12 @@ static llvm::SmallString writePaxHeader(llvm::StringRef Path, static void writeFileHeader(llvm::raw_ostream &OS, llvm::StringRef Path, size_t Size) { llvm::SmallString FileHeader = writePaxHeader(Path, Size); - gzipCompress(OS, { FileHeader.data(), FileHeader.size() }); + gzipCompress(OS, FileHeader); } static void compressedPadding(llvm::raw_ostream &OS, size_t Size) { llvm::SmallVector Buffer(Size, '\0'); - return gzipCompress(OS, { Buffer.data(), Buffer.size() }); + return gzipCompress(OS, Buffer); } namespace revng { @@ -152,7 +152,7 @@ OffsetDescriptor GzipTarWriter::append(llvm::StringRef Path, writeFileHeader(*OS, Path, Data.size()); Result.DataStart = OS->tell(); - gzipCompress(*OS, { Data.data(), Data.size() }); + gzipCompress(*OS, Data); Result.PaddingStart = OS->tell(); if (size_t Padding = computePadding(Data.size()); Padding % BlockSize != 0) diff --git a/tests/unit/GzipTarFile.cpp b/tests/unit/GzipTarFile.cpp index 6f1fb202f..5ae1ea389 100644 --- a/tests/unit/GzipTarFile.cpp +++ b/tests/unit/GzipTarFile.cpp @@ -13,19 +13,13 @@ bool init_unit_test(); #include "revng/UnitTestHelpers/UnitTestHelpers.h" -static std::string gzipDecompress(llvm::ArrayRef Buffer) { - llvm::SmallString<128> Output; - llvm::raw_svector_ostream OS(Output); - gzipDecompress(OS, Buffer); - return Output.str().str(); -} - static void checkOffset(llvm::SmallVector &Buffer, size_t Start, size_t Size, llvm::StringRef ExpectedValue) { - std::string Result = gzipDecompress({ Buffer.data() + Start, Size }); - BOOST_TEST(Result == ExpectedValue.str()); + llvm::ArrayRef ArrayRef{ Buffer.data() + Start, Size }; + llvm::SmallVector Result = gzipDecompress(ArrayRef); + BOOST_TEST(llvm::StringRef(Result.begin(), Result.size()) == ExpectedValue); } BOOST_AUTO_TEST_CASE(GzipTarFileTest) {