| // Copyright 2022 The Fuchsia Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include <lib/zxdump/zstd-writer.h> |
| #include <unistd.h> |
| #include <zircon/assert.h> |
| |
| #include <cerrno> |
| #include <thread> |
| |
| #include <zstd/zstd.h> |
| |
| namespace zxdump { |
| |
| using namespace std::literals; |
| |
| namespace { |
| |
| auto ErrnoError(std::string_view op, int error = errno) { |
| return fit::error{FdError{.op_ = op, .error_ = error}}; |
| } |
| |
| auto ZstdError(size_t result) { return fit::error{FdError{.op_ = ZSTD_getErrorName(result)}}; } |
| |
| } // namespace |
| |
| ZstdWriter::ZstdWriter(fbl::unique_fd fd) |
| : ctx_(ZSTD_createCStream()), |
| buffer_(new std::byte[ZSTD_CStreamOutSize()]), |
| fd_(std::move(fd)) { |
| auto set = [ctx = static_cast<ZSTD_CCtx*>(ctx_)](auto param, auto value) { |
| ZSTD_CCtx_setParameter(ctx, param, value); |
| }; |
| set(ZSTD_c_compressionLevel, 11); |
| set(ZSTD_c_enableLongDistanceMatching, 1); |
| set(ZSTD_c_nbWorkers, std::thread::hardware_concurrency()); |
| } |
| |
| ZstdWriter::~ZstdWriter() { |
| auto ctx = static_cast<ZSTD_CCtx*>(ctx_); |
| ZSTD_freeCStream(ctx); |
| } |
| |
| fit::result<FdError> ZstdWriter::Flush() { |
| ByteView out{buffer_.get(), buffer_pos_}; |
| buffer_pos_ = 0; |
| while (!out.empty()) { |
| ssize_t n = write(fd_.get(), out.data(), out.size()); |
| if (n < 0) { |
| return ErrnoError("write"sv); |
| } |
| if (n == 0) { |
| return ErrnoError("write returned zero"sv, EAGAIN); |
| } |
| out = out.subspan(n); |
| } |
| return fit::ok(); |
| } |
| |
| fit::result<FdError> ZstdWriter::Write(size_t offset, ByteView data) { |
| ZX_ASSERT(offset >= offset_); |
| ZX_ASSERT(!data.empty()); |
| ZX_DEBUG_ASSERT(data.data()); |
| |
| // If there are holes we have to feed zero bytes to the compressor. |
| while (offset > offset_) { |
| static constexpr std::byte kZero[32] = {}; |
| auto pad = ByteView{kZero, sizeof(kZero)}.subspan(0, offset - offset_); |
| auto result = Write(offset_, pad); |
| if (result.is_error()) { |
| return result.take_error(); |
| } |
| } |
| |
| ZSTD_inBuffer in = {data.data(), data.size(), 0}; |
| while (in.pos < in.size) { |
| ZSTD_outBuffer out = {buffer_.get(), ZSTD_CStreamOutSize(), buffer_pos_}; |
| auto ctx = static_cast<ZSTD_CCtx*>(ctx_); |
| size_t result = ZSTD_compressStream2(ctx, &out, &in, ZSTD_e_continue); |
| buffer_pos_ = out.pos; |
| if (ZSTD_isError(result)) { |
| return ZstdError(result); |
| } |
| if (in.pos < in.size) { |
| // Not all consumed yet, so flush the buffer. |
| auto result = Flush(); |
| if (result.is_error()) { |
| return result.take_error(); |
| } |
| } |
| } |
| |
| offset_ += in.pos; |
| return fit::ok(); |
| } |
| |
| fit::result<FdError> ZstdWriter::Finish() { |
| size_t compress_result; |
| do { |
| ZSTD_inBuffer in = {}; |
| ZSTD_outBuffer out = {buffer_.get(), ZSTD_CStreamOutSize(), buffer_pos_}; |
| auto ctx = static_cast<ZSTD_CCtx*>(ctx_); |
| compress_result = ZSTD_compressStream2(ctx, &out, &in, ZSTD_e_end); |
| buffer_pos_ = out.pos; |
| if (ZSTD_isError(compress_result)) { |
| return ZstdError(compress_result); |
| } |
| if (auto result = Flush(); result.is_error()) { |
| return result.take_error(); |
| } |
| } while (compress_result != 0); |
| return fit::ok(); |
| } |
| |
| } // namespace zxdump |