ClickHouse/src/IO/LzmaWriteBuffer.cpp

#include <IO/LzmaWriteBuffer.h>


namespace DB {
namespace ErrorCodes
{
    extern const int LZMA_STREAM_ENCODER_FAILED;
}

LzmaWriteBuffer::LzmaWriteBuffer(
        std::unique_ptr<WriteBuffer> out_,
        int compression_level,
        size_t buf_size,
        char * existing_memory,
        size_t alignment)
    : BufferWithOwnMemory<WriteBuffer>(buf_size, existing_memory, alignment)
    , out(std::move(out_))
{
    lstr = LZMA_STREAM_INIT;
    lstr.allocator = nullptr;
    lstr.next_in = nullptr;
    lstr.avail_in = 0;
    lstr.next_out = nullptr;
    lstr.avail_out = 0;
    
    // options for further compression
    lzma_options_lzma opt_lzma2;
    if (lzma_lzma_preset(&opt_lzma2, compression_level)) {
        throw Exception(std::string("lzma preset failed: ") + "; lzma version: " + LZMA_VERSION_STRING, ErrorCodes::LZMA_STREAM_ENCODER_FAILED);
    }

    lzma_filter filters[] = {
		{ .id = LZMA_FILTER_X86, .options = NULL },
		{ .id = LZMA_FILTER_LZMA2, .options = &opt_lzma2 },
        { .id = LZMA_VLI_UNKNOWN, .options = NULL },
    };
    lzma_ret ret = lzma_stream_encoder(&lstr, filters, LZMA_CHECK_CRC64);

    if (ret != LZMA_OK) 
        throw Exception(std::string("lzma stream encoder init failed: ") + std::to_string(ret) + "; lzma version: " + LZMA_VERSION_STRING, ErrorCodes::LZMA_STREAM_ENCODER_FAILED);
     
}

LzmaWriteBuffer::~LzmaWriteBuffer()
{
    try {
        finish();

        lzma_end(&lstr);
    } catch (...)
    {
        tryLogCurrentException(__PRETTY_FUNCTION__);
    }
}

void LzmaWriteBuffer::nextImpl()
{
    if (!offset())
        return;

    lstr.next_in = reinterpret_cast<unsigned char *>(working_buffer.begin());
    lstr.avail_in = offset();
    //std::cout << lstr.avail_in << std::endl;

    lzma_action action = LZMA_RUN;
    do {
        out->nextIfAtEnd();
        lstr.next_out = reinterpret_cast<unsigned char *>(out->position());
        lstr.avail_out = out->buffer().end() - out->position();
        //std::cout << lstr.avail_out << " BEFOR" << std::endl;


        lzma_ret ret = lzma_code(&lstr, action);
        out->position() = out->buffer().end() - lstr.avail_out;

        //std::cout << lstr.avail_out << " AFTER" << std::endl;

		//std::cout << ret << " RET IMPL" << std::endl;

        if (ret == LZMA_STREAM_END) {
                return;
        }
        
        if (ret != LZMA_OK)
            throw Exception(std::string("lzma stream encoding failed: ") + "; lzma version: " + LZMA_VERSION_STRING, ErrorCodes::LZMA_STREAM_ENCODER_FAILED);
    
        //std::cout << lstr.avail_in << " " << lstr.avail_out << std::endl;
    } while (lstr.avail_in > 0 || lstr.avail_out == 0);
}


void LzmaWriteBuffer::finish()
{
    if (finished)
        return;

    next();

    do {
        out->nextIfAtEnd();
        lstr.next_out = reinterpret_cast<unsigned char *>(out->position());
        lstr.avail_out = out->buffer().end() - out->position();

        lzma_ret ret = lzma_code(&lstr, LZMA_FINISH);
        out->position() = out->buffer().end() - lstr.avail_out;

		//std::cout << ret << " RET FIN" << std::endl;

        if (ret == LZMA_STREAM_END) {
            finished = true;
			return;
        }
        
        if (ret != LZMA_OK)
            throw Exception(std::string("lzma stream encoding failed: ") + "; lzma version: " + LZMA_VERSION_STRING, ErrorCodes::LZMA_STREAM_ENCODER_FAILED);
    
        //std::cout << lstr.avail_in << std::endl;
    } while (lstr.avail_out == 0);
}

}