From 0572d0cd4b023ca3cdfe71026d196845c11eadc9 Mon Sep 17 00:00:00 2001 From: Connor McLaughlin Date: Thu, 18 Aug 2022 22:42:22 +1000 Subject: [PATCH] Common: Add ByteStream implementations for zstd --- src/common/CMakeLists.txt | 2 +- src/common/byte_stream.cpp | 325 ++++++++++++++++++++++++++++++++++++- src/common/byte_stream.h | 4 + src/common/common.props | 2 +- src/common/common.vcxproj | 1 + 5 files changed, 328 insertions(+), 6 deletions(-) diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt index 2acc72a94..ea086d198 100644 --- a/src/common/CMakeLists.txt +++ b/src/common/CMakeLists.txt @@ -58,7 +58,7 @@ add_library(common target_include_directories(common PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/..") target_include_directories(common PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/..") target_link_libraries(common PUBLIC fmt Threads::Threads vulkan-headers) -target_link_libraries(common PRIVATE stb libchdr zlib minizip "${CMAKE_DL_LIBS}") +target_link_libraries(common PRIVATE stb libchdr zlib minizip Zstd::Zstd "${CMAKE_DL_LIBS}") if(WIN32) target_sources(common PRIVATE diff --git a/src/common/byte_stream.cpp b/src/common/byte_stream.cpp index 98e8e3139..8c5c9b1f5 100644 --- a/src/common/byte_stream.cpp +++ b/src/common/byte_stream.cpp @@ -3,6 +3,8 @@ #include "file_system.h" #include "log.h" #include "string_util.h" +#include "zstd.h" +#include "zstd_errors.h" #include #include #include @@ -169,7 +171,10 @@ public: return true; } - u64 GetPosition() const override { return _ftelli64(m_pFile); } + u64 GetPosition() const override + { + return _ftelli64(m_pFile); + } u64 GetSize() const override { @@ -224,7 +229,10 @@ public: return true; } - u64 GetPosition() const override { return static_cast(ftello(m_pFile)); } + u64 GetPosition() const override + { + return static_cast(ftello(m_pFile)); + } u64 GetSize() const override { @@ -251,9 +259,15 @@ public: return true; } - virtual bool Commit() override { return true; } + virtual bool Commit() override + { + return true; + } - virtual bool Discard() override { return false; } + virtual bool Discard() override + { + return false; + } protected: FILE* m_pFile; @@ -1369,3 +1383,306 @@ bool ByteStream::WriteBinaryToStream(ByteStream* stream, const void* data, size_ return stream->Write2(data, static_cast(data_length)); } + +class ZstdCompressStream final : public ByteStream +{ +public: + ZstdCompressStream(ByteStream* dst_stream, int compression_level) : m_dst_stream(dst_stream) + { + m_cstream = ZSTD_createCStream(); + ZSTD_CCtx_setParameter(m_cstream, ZSTD_c_compressionLevel, compression_level); + } + + ~ZstdCompressStream() override + { + if (!m_done) + Compress(ZSTD_e_end); + + ZSTD_freeCStream(m_cstream); + } + + bool ReadByte(u8* pDestByte) override { return false; } + + u32 Read(void* pDestination, u32 ByteCount) override { return 0; } + + bool Read2(void* pDestination, u32 ByteCount, u32* pNumberOfBytesRead = nullptr) override { return false; } + + bool WriteByte(u8 SourceByte) override + { + if (m_input_buffer_wpos == INPUT_BUFFER_SIZE && !Compress(ZSTD_e_continue)) + return false; + + m_input_buffer[m_input_buffer_wpos++] = SourceByte; + return true; + } + + u32 Write(const void* pSource, u32 ByteCount) override + { + u32 remaining = ByteCount; + const u8* read_ptr = static_cast(pSource); + for (;;) + { + const u32 copy_size = std::min(INPUT_BUFFER_SIZE - m_input_buffer_wpos, remaining); + std::memcpy(&m_input_buffer[m_input_buffer_wpos], read_ptr, copy_size); + read_ptr += copy_size; + remaining -= copy_size; + m_input_buffer_wpos += copy_size; + if (remaining == 0 || !Compress(ZSTD_e_continue)) + break; + } + + return ByteCount - remaining; + } + + bool Write2(const void* pSource, u32 ByteCount, u32* pNumberOfBytesWritten = nullptr) override + { + const u32 bytes_written = Write(pSource, ByteCount); + if (pNumberOfBytesWritten) + *pNumberOfBytesWritten = bytes_written; + return (bytes_written == ByteCount); + } + + bool SeekAbsolute(u64 Offset) override { return false; } + + bool SeekRelative(s64 Offset) override { return (Offset == 0); } + + bool SeekToEnd() override { return false; } + + u64 GetPosition() const override { return m_position; } + + u64 GetSize() const override { return 0; } + + bool Flush() override { return Compress(ZSTD_e_flush); } + + bool Discard() override { return true; } + + bool Commit() override { return Compress(ZSTD_e_end); } + +private: + enum : u32 + { + INPUT_BUFFER_SIZE = 131072, + OUTPUT_BUFFER_SIZE = 65536, + }; + + bool Compress(ZSTD_EndDirective action) + { + if (m_errorState || m_done) + return false; + + ZSTD_inBuffer inbuf = {m_input_buffer, m_input_buffer_wpos, 0}; + + for (;;) + { + ZSTD_outBuffer outbuf = {m_output_buffer, OUTPUT_BUFFER_SIZE, 0}; + + const size_t ret = ZSTD_compressStream2(m_cstream, &outbuf, &inbuf, action); + if (ZSTD_isError(ret)) + { + Log_ErrorPrintf("ZSTD_compressStream2() failed: %u (%s)", static_cast(ZSTD_getErrorCode(ret)), + ZSTD_getErrorString(ZSTD_getErrorCode(ret))); + SetErrorState(); + return false; + } + + if (outbuf.pos > 0) + { + if (!m_dst_stream->Write2(m_output_buffer, static_cast(outbuf.pos))) + { + SetErrorState(); + return false; + } + + outbuf.pos = 0; + } + + if (action == ZSTD_e_end) + { + // break when compression output has finished + if (ret == 0) + { + m_done = true; + break; + } + } + else + { + // break when all input data is consumed + if (inbuf.pos == inbuf.size) + break; + } + } + + m_position += m_input_buffer_wpos; + m_input_buffer_wpos = 0; + return true; + } + + ByteStream* m_dst_stream; + ZSTD_CStream* m_cstream = nullptr; + u64 m_position = 0; + u32 m_input_buffer_wpos = 0; + bool m_done = false; + + u8 m_input_buffer[INPUT_BUFFER_SIZE]; + u8 m_output_buffer[OUTPUT_BUFFER_SIZE]; +}; + +std::unique_ptr ByteStream::CreateZstdCompressStream(ByteStream* src_stream, int compression_level) +{ + return std::make_unique(src_stream, compression_level); +} + +class ZstdDecompressStream final : public ByteStream +{ +public: + ZstdDecompressStream(ByteStream* src_stream, u32 compressed_size) + : m_src_stream(src_stream), m_bytes_remaining(compressed_size) + { + m_cstream = ZSTD_createDStream(); + m_in_buffer.src = m_input_buffer; + Decompress(); + } + + ~ZstdDecompressStream() override { ZSTD_freeDStream(m_cstream); } + + bool ReadByte(u8* pDestByte) override { return Read(pDestByte, sizeof(u8)) == sizeof(u8); } + + u32 Read(void* pDestination, u32 ByteCount) override + { + u8* write_ptr = static_cast(pDestination); + u32 remaining = ByteCount; + for (;;) + { + const u32 copy_size = std::min(m_output_buffer_wpos - m_output_buffer_rpos, remaining); + std::memcpy(write_ptr, &m_output_buffer[m_output_buffer_rpos], copy_size); + m_output_buffer_rpos += copy_size; + write_ptr += copy_size; + remaining -= copy_size; + if (remaining == 0 || !Decompress()) + break; + } + + return ByteCount - remaining; + } + + bool Read2(void* pDestination, u32 ByteCount, u32* pNumberOfBytesRead = nullptr) override + { + const u32 bytes_read = Read(pDestination, ByteCount); + if (pNumberOfBytesRead) + *pNumberOfBytesRead = bytes_read; + return (bytes_read == ByteCount); + } + + bool WriteByte(u8 SourceByte) override { return false; } + + u32 Write(const void* pSource, u32 ByteCount) override { return 0; } + + bool Write2(const void* pSource, u32 ByteCount, u32* pNumberOfBytesWritten = nullptr) override { return false; } + + bool SeekAbsolute(u64 Offset) override { return false; } + + bool SeekRelative(s64 Offset) override + { + if (Offset < 0) + return false; + else if (Offset == 0) + return true; + + s64 remaining = Offset; + for (;;) + { + const s64 skip = std::min(m_output_buffer_wpos - m_output_buffer_rpos, remaining); + remaining -= skip; + m_output_buffer_wpos += static_cast(skip); + if (remaining == 0) + return true; + else if (!Decompress()) + return false; + } + } + + bool SeekToEnd() override { return false; } + + u64 GetPosition() const override { return 0; } + + u64 GetSize() const override { return 0; } + + bool Flush() override { return true; } + + bool Discard() override { return true; } + + bool Commit() override { return true; } + +private: + enum : u32 + { + INPUT_BUFFER_SIZE = 65536, + OUTPUT_BUFFER_SIZE = 131072, + }; + + bool Decompress() + { + if (m_output_buffer_rpos != m_output_buffer_wpos) + { + const u32 move_size = m_output_buffer_wpos - m_output_buffer_rpos; + std::memmove(&m_output_buffer[0], &m_output_buffer[m_output_buffer_rpos], move_size); + m_output_buffer_rpos = move_size; + m_output_buffer_wpos = move_size; + } + else + { + m_output_buffer_rpos = 0; + m_output_buffer_wpos = 0; + } + + ZSTD_outBuffer outbuf = {m_output_buffer, OUTPUT_BUFFER_SIZE - m_output_buffer_wpos, 0}; + while (outbuf.pos == 0) + { + if (m_in_buffer.pos == m_in_buffer.size && !m_errorState) + { + const u32 requested_size = std::min(m_bytes_remaining, INPUT_BUFFER_SIZE); + const u32 bytes_read = m_src_stream->Read(m_input_buffer, requested_size); + m_in_buffer.size = bytes_read; + m_in_buffer.pos = 0; + m_bytes_remaining -= bytes_read; + if (bytes_read != requested_size || m_bytes_remaining == 0) + { + m_errorState = true; + break; + } + } + + size_t ret = ZSTD_decompressStream(m_cstream, &outbuf, &m_in_buffer); + if (ZSTD_isError(ret)) + { + Log_ErrorPrintf("ZSTD_decompressStream() failed: %u (%s)", static_cast(ZSTD_getErrorCode(ret)), + ZSTD_getErrorString(ZSTD_getErrorCode(ret))); + m_in_buffer.pos = m_in_buffer.size; + m_output_buffer_rpos = 0; + m_output_buffer_wpos = 0; + m_errorState = true; + return false; + } + } + + m_output_buffer_wpos = static_cast(outbuf.pos); + return true; + } + + ByteStream* m_src_stream; + ZSTD_DStream* m_cstream = nullptr; + ZSTD_inBuffer m_in_buffer = {}; + u32 m_output_buffer_rpos = 0; + u32 m_output_buffer_wpos = 0; + u32 m_bytes_remaining; + bool m_errorState = false; + + u8 m_input_buffer[INPUT_BUFFER_SIZE]; + u8 m_output_buffer[OUTPUT_BUFFER_SIZE]; +}; + +std::unique_ptr ByteStream::CreateZstdDecompressStream(ByteStream* src_stream, u32 compressed_size) +{ + return std::make_unique(src_stream, compressed_size); +} diff --git a/src/common/byte_stream.h b/src/common/byte_stream.h index 94ffa4268..6061ada43 100644 --- a/src/common/byte_stream.h +++ b/src/common/byte_stream.h @@ -117,6 +117,10 @@ public: // null memory stream static std::unique_ptr CreateNullStream(); + // zstd stream + static std::unique_ptr CreateZstdCompressStream(ByteStream* src_stream, int compression_level); + static std::unique_ptr CreateZstdDecompressStream(ByteStream* src_stream, u32 compressed_size); + // copies one stream's contents to another. rewinds source streams automatically, and returns it back to its old // position. static bool CopyStream(ByteStream* pDestinationStream, ByteStream* pSourceStream); diff --git a/src/common/common.props b/src/common/common.props index be2e123f0..ae0da01ca 100644 --- a/src/common/common.props +++ b/src/common/common.props @@ -11,7 +11,7 @@ $(RootBuildDir)glad\glad.lib;$(RootBuildDir)glslang\glslang.lib;%(AdditionalDependencies) - $(RootBuildDir)fmt\fmt.lib;$(RootBuildDir)zlib\zlib.lib;$(RootBuildDir)minizip\minizip.lib;$(RootBuildDir)lzma\lzma.lib;d3dcompiler.lib;d3d11.lib;%(AdditionalDependencies) + $(RootBuildDir)zstd\zstd.lib;$(RootBuildDir)fmt\fmt.lib;$(RootBuildDir)zlib\zlib.lib;$(RootBuildDir)minizip\minizip.lib;$(RootBuildDir)lzma\lzma.lib;d3dcompiler.lib;d3d11.lib;%(AdditionalDependencies) diff --git a/src/common/common.vcxproj b/src/common/common.vcxproj index 96952793b..260391222 100644 --- a/src/common/common.vcxproj +++ b/src/common/common.vcxproj @@ -235,6 +235,7 @@ $(IntDir)/%(RelativeDir)/ + %(AdditionalIncludeDirectories);$(SolutionDir)dep\zstd\lib