kko.cpp

// KKO - image data compression using dictionary-based compression algorithm
// Last change: 13.5.2025

#include <stdint.h>
#include <vector>
#include <stdexcept>
#include <iostream>
#include <getopt.h>
#include <fstream>

#include "utils.h"
#include "lzss.h"
#include "rle.h"
#include "blocks.h"


// parse arguments from command line, return them in a Params struct
struct Params parse_arguments(int argc, char **argv)
{
    Params params;
    std::string help_msg = "Usage: ./lz_codec [-c|-d] [-m] [-a] -i <input_file> -o <output_file> [-w <width>]";

    int c;
    while((c = getopt(argc, argv, "hcdmai:o:w:")) != -1)
    {
        switch (c)
        {
            case 'h':
                std::cout << help_msg << std::endl;
                exit(0);
            case 'c':
                params.compress = true;
                break;
            case 'd':
                params.decompress = true;
                break;
            case 'm':
                params.model = true;
                break;
            case 'a':
                params.adaptive = true;
                break;
            case 'i':
                params.input = std::string(optarg);
                break;
            case 'o':
                params.output = std::string(optarg);
                break;
            case 'w':
                params.width = atoi(optarg);
                break;
            default:
                std::cerr << help_msg << std::endl;
                exit(1);
        }
    }

    if(params.compress && params.decompress)
        throw std::runtime_error("wrong arguments: cannot compress and decompress at the same time");

    if(!params.compress && !params.decompress)
        throw std::runtime_error("wrong arguments: missing operation mode, use either -c or -d");

    if(params.input.length() == 0)
        throw std::runtime_error("wrong arguments: missing input file (-i)");

    if(params.output.length() == 0)
        throw std::runtime_error("wrong arguments: missing output file (-o)");

    if(params.compress && (params.width < 1 || params.width % 256 != 0))
        throw std::runtime_error("wrong arguments: width must be equal or larger than 1 and also be divisible by 256");

    return params;
}

// read binary file (given a filepath), return a vector of bytes
std::vector<uint8_t> read_file(std::string file_name)
{
    // open file
    std::ifstream file(file_name, std::ios::binary);
    if(!file)
    {
        std::cerr << "failed to open file: " << file_name << std::endl;
        exit(1);
    }

    // get size
    file.seekg(0, std::ios::end);
    std::streamsize size = file.tellg();
    file.seekg(0, std::ios::beg);

    // read into vector
    std::vector<uint8_t> raw_data(size);
    if(!file.read(reinterpret_cast<char*>(raw_data.data()), size))
    {
        std::cerr << "failed to read file: " + file_name << std::endl;
        exit(1);
    }

    return raw_data;
}


// load metadata from raw file, also remove the metadata bits (and padding bits at the end)
struct Params load(struct Params params, BitBuffer &bits)
{
    params.width = (ungamma_from_bitfield(bits) + 1) * 256;
    params.height = (ungamma_from_bitfield(bits) + 1) * 256;
    params.blocks = (ungamma_from_bitfield(bits));
    params.model = bits.read(0, 1) ? true : false;

    // next 3 bits are the number of padding bits
    params.padding_bits = bits.read(1, 3);

    //finish the byte
    uint64_t used_bits = gamma(params.width / 256 - 1) + gamma(params.height / 256 - 1) + gamma(params.blocks) + 1 + 3;
    uint64_t unused_bits = 8 - (used_bits % 8);
    bits.remove(4 + unused_bits);

    return params;
}


// save encoded data (+ metadata) to a file
bool save(struct Params params, BitBuffer bits)
{
    // create output file
    std::ofstream file(params.output, std::ios::binary);
    if(!file)
    {
        std::cerr << "failed to write output" << std::endl;
        exit(1);
    }

    // create header
    BitBuffer header;
    uint64_t width = (params.width / 256) - 1;
    uint64_t width_gamma = gamma(width);
    header.append(width+1, width_gamma);
    uint64_t height = (params.width / 256) - 1;
    uint64_t height_gamma = gamma(height);
    header.append(height+1, height_gamma);
    uint64_t blocks_gamma = gamma(params.blocks);
    header.append(params.blocks+1, blocks_gamma);
    header.append(params.model, 1);

    // calculate padding size
    uint64_t padding = (8 - ((bits.size()) % 8)) % 8;
    // add padding to the data
    header.append(padding, 3);

    // write header to file
    auto header_bytes = header.to_byte_vector();
    file.write(reinterpret_cast<const char*>(header_bytes.data()), header_bytes.size());

    // write data to file
    auto data_bytes = bits.to_byte_vector();
    file.write(reinterpret_cast<const char*>(data_bytes.data()), data_bytes.size());
    file.close();

    std::cout << "Saved as: " << params.output << "\nFilesize: " << (header_bytes.size() + data_bytes.size()) << " bytes" << std::endl;
    std::cout << "Final BPC: " << static_cast<float>((header_bytes.size() + data_bytes.size()) * 8) / (params.input_size_bytes) << std::endl;
    return true;
}


int main(int argc, char **argv) 
{
    // parse arguments
    Params params;
    try
    {
        params = parse_arguments(argc, argv);
    }
    catch(const std::exception& e)
    {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }

    // ###########
    // COMPRESSION
    // ###########
    if(params.compress)
    {
        // read data from file
        std::vector<uint8_t> raw_data = read_file(params.input);
        params.height = raw_data.size() / params.width;
        params.input_size_bytes = raw_data.size();

        // non-adaptive mode
        BitBuffer h_compressed_lzss;
        if(!params.adaptive)
        {
            h_compressed_lzss = compression_pipeline(params, raw_data);
            save(params, h_compressed_lzss);

        }

        // block compression
        BitBuffer block_compressed_lzss;
        if(params.adaptive)
        {
            block_compressed_lzss = block_compress(params, raw_data);
            params.blocks = BLOCK_SIZE;
            save(params, block_compressed_lzss);
        }
    }
    // #############
    // DECOMPRESSION
    // #############
    else if(params.decompress)
    {
        // read data from file
        std::vector<uint8_t> raw_data = read_file(params.input);

        BitBuffer bits;
        bits.load_from_byte_vector(raw_data);

        // load metadata
        params = load(params, bits);
        bits.adjust_padding(params.padding_bits);
        
        // decompress
        std::vector<uint8_t> decompressed;
        if(params.blocks > 0) // block compression was used
            decompressed = block_decompress(params, bits);
        else
            decompressed = decompression_pipeline(params, bits);

        // create output file and write the result
        std::ofstream file(params.output, std::ios::binary);
        if(!file)
        {
            std::cerr << "failed to write output" << std::endl;
            return 1;
        }
        file.write(reinterpret_cast<const char*>(decompressed.data()), decompressed.size());
        file.close();
        std::cout << "Saved as: " << params.output << "\nFilesize: " << decompressed.size() << " bytes" << std::endl;
    }
    
    return 0;
}