lz4 Compress Test Utility

lz4_test.cpp

#include <vector>
#include <string>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <cstdint>
#include <fstream>

#include <lz4.h>
#include <lz4hc.h>
#include <lz4frame.h>

long file_size(const std::string& file_name) {
  //TODO: detect gzip and actually validate the uncompressed size?
  struct stat s;
  int rc = stat(file_name.c_str(), &s);
  return rc == 0 ? s.st_size : -1;
}

std::vector<char> read_file(const std::string& file_name, long size) {
  int fd = open(file_name.c_str(), O_RDONLY);
  if(fd == -1)
    throw std::runtime_error("Could not open: " + file_name);
  posix_fadvise(fd, 0, 0, 1);  // FDADVICE_SEQUENTIAL
  std::vector<char> in(size);
  if(read(fd, in.data(), size) != size)
    throw std::runtime_error("Could not open: " + file_name);
  close(fd);
  return in;
}

void write_file(const std::string& file_name, const std::vector<char>& out) {
  std::ofstream file(file_name, std::ios::binary | std::ios::trunc);
  file.write(out.data(), out.size());
}

enum class algorithm_t { DEFAULT, HIGH, FRAME };

std::vector<char> lzip(const std::string& file_name, int compression_level, algorithm_t algorithm) {
  auto in_size = file_size(file_name);
  auto in = read_file(file_name, in_size);

  LZ4F_preferences_t frame_prefs{LZ4F_frameInfo_t{}, compression_level};

  //prepair for destination in memory
  size_t max_compressed_size = algorithm == algorithm_t::FRAME ?
      LZ4F_compressFrameBound(in_size, &frame_prefs) : LZ4_compressBound(in_size);
  std::vector<char> out(max_compressed_size);

  //compress it in memory
  int compressed_size = 0;
  switch(algorithm) {
    case algorithm_t::DEFAULT:
      compressed_size = LZ4_compress_fast(in.data(), out.data(), in_size, max_compressed_size, compression_level);
      break;
    case algorithm_t::HIGH:
      compressed_size = LZ4_compress_HC(in.data(), out.data(), in_size, max_compressed_size, compression_level);
      break;
    case algorithm_t::FRAME:
      compressed_size = LZ4F_compressFrame(out.data(), max_compressed_size, in.data(), in_size, &frame_prefs);
      break;
    default:
      throw std::runtime_error("Wrong compression type");
  }
  if(compressed_size <= 0)
    throw std::runtime_error("Compression failed: " + std::to_string(compressed_size));
  out.resize(compressed_size);

  return out;
}

std::vector<char> lunzip(const std::string& file_name, size_t out_size, algorithm_t algorithm) {
  size_t in_size = file_size(file_name);
  auto in = read_file(file_name, in_size);
  auto* in_ptr = in.data();
  auto* in_end = in.data() + in_size;

  LZ4F_decompressionContext_t context = nullptr;
  LZ4F_frameInfo_t info;
  int hint = -1;

  auto decompressed_size = 0;
  std::vector<char> out(0);
  switch(algorithm) {
    case algorithm_t::DEFAULT:
    case algorithm_t::HIGH:
      out.resize(out_size);
      decompressed_size = LZ4_decompress_fast(in.data(), out.data(), out_size);
      break;
    case algorithm_t::FRAME:
      hint = LZ4F_createDecompressionContext(&context, LZ4F_VERSION);
      if(LZ4F_isError(hint))
        throw std::runtime_error("Decompression initialization failed " + std::string(LZ4F_getErrorName(hint)));

      /*
      out_size = LZ4F_getFrameInfo(context, &info, in_ptr, &in_size);
      if(LZ4F_isError(hint))
        throw std::runtime_error("Decompression header parsing failed " + std::string(LZ4F_getErrorName(hint)));
      in_ptr += in_size;
      in_size = in_end - in_ptr;
      */
      //cant get the info to tell us how much to reserve in total so we'll just have it known for now..
      out.resize(out_size);

      //decompress
      do {
        hint = LZ4F_decompress(context, out.data(), &out_size, in_ptr, &in_size, nullptr);
        if(LZ4F_isError(hint))
          throw std::runtime_error("Decompression of frame failed " + std::string(LZ4F_getErrorName(hint)));
        in_ptr += in_size;
        in_size = in_end - in_ptr;
      } while(hint && in_size);

      LZ4F_freeDecompressionContext(context);
      break;
    default:
      throw std::runtime_error("Wrong decompression algorithm");
  }

  if(decompressed_size < 0)
    throw std::runtime_error("Deompression failed: " + std::to_string(decompressed_size));

  return out;
}

int main(int argc, char** argv){
  for(int i = 1; i < argc; ++i) {
    //auto raw = read_file(argv[i], file_size(argv[1]));
    auto raw = lunzip(argv[i], 3601*3601*2, algorithm_t::FRAME);

    //auto raw = lzip(argv[i], 4, algorithm_t::FRAME);
    write_file(std::string(argv[i]) + ".lz4.raw", raw);
  }
  return 0;
}