Jungfraujoch/tests/ZSTDCompressorTest.cpp

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include <bitshuffle/bitshuffle.h>
#include <bitshuffle_hperf/bitshuffle.h>
#include <zstd.h>
#include <iostream>
#include <random>

#include "../compression/JFJochCompressor.h"
#include "../compression/JFJochDecompress.h"
#include "../common/DiffractionExperiment.h"
#include "../writer/HDF5Objects.h"

TEST_CASE("JFjochZstdCompressor_Raw_Block","[ZSTD]") {
    uint8_t src[128*8];
    uint8_t dst[128*8+3];

    for (int i = 0; i < 128*8; i++)
        src[i] = i % 256;

    size_t dst_size = JFJochZstdCompressor::RawBlock(dst, src, 128 * 8, false);

    REQUIRE(dst_size == 128*8+3);
    int diff = 0;
    for (int i = 0; i < 128*8; i++) {
        if (dst[i+3] != i % 256) diff++;
    }
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_RLE_Block","[ZSTD]") {
    uint8_t dst[4];
    size_t dst_size = JFJochZstdCompressor::RLEBlock(dst, 0xAF, 128 * 8, false);

    REQUIRE(dst_size == 4);
    REQUIRE(dst[3] == 0xAF);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];

    for (int i = 0; i < 128 * 8; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128*8, 128*8);
    REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);

    REQUIRE(dst_size == 4 + 9);
    REQUIRE(dst[9 + 3] == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_1block","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX, ZSTD_BLOCKSIZE_MAX);
    REQUIRE(dst_size == 4 + 9);
    REQUIRE(dst[9 + 3] == 0);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_2blocks","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0;
        dst[i] = 0xFF;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
    REQUIRE(dst_size == 8 + 9);
    REQUIRE(dst[9 + 3] == 0);
    REQUIRE(dst[9 + 7] == 0);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_onlyRLE_0xFF_2blocks","[ZSTD]") {
    JFJochZstdCompressor compressor;
    uint8_t src[ZSTD_BLOCKSIZE_MAX * 2];
    uint8_t dst[ZSTD_BLOCKSIZE_MAX * 4];
    uint8_t decomp_buffer[ZSTD_BLOCKSIZE_MAX * 4];

    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) {
        src[i] = 0xFF;
        dst[i] = 0x0;
    }

    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, ZSTD_BLOCKSIZE_MAX * 2, ZSTD_BLOCKSIZE_MAX * 2);
    REQUIRE(dst_size == 8 + 9);
    REQUIRE(dst[9 + 3] == 0xFF);
    REQUIRE(dst[9 + 7] == 0xFF);

    REQUIRE(ZSTD_decompress(decomp_buffer, ZSTD_BLOCKSIZE_MAX * 4, dst, dst_size) == ZSTD_BLOCKSIZE_MAX * 2);
    size_t diff = 0;
    for (int i = 0; i < ZSTD_BLOCKSIZE_MAX * 2; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}


TEST_CASE("JFjochZstdCompressor_Frame_onlyRAW","[ZSTD]") {
    JFJochZstdCompressor compressor;

    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];
    for (int i = 0; i < 128 * 8; i++)
        src[i] = 1 + i % 128;
    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128*8);
    REQUIRE(ZSTD_decompress(decomp_buffer, 128 * 8 * 2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128 * 8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);

    REQUIRE(dst_size == 3 + 9 + 128 * 8);

    diff = 0;
    for (int i = 0; i < 128 * 8; i++) {
        if (dst[i + 3 + 9] != 1 + i % 128)
            diff++;
    }
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_mixed","[ZSTD]") {
    JFJochZstdCompressor compressor;

    uint8_t src[128 * 8];
    uint8_t dst[128 * 8 * 2];
    uint8_t decomp_buffer[128 * 8 * 2];

    for (int i = 0; i < 128*8; i++)
        src[i] = (i / 128 + 1) % 2 ;
    size_t dst_size = compressor.Compress(dst, (uint64_t *) src, 128 * 8, 128 * 8);
    REQUIRE(dst_size == 9 + 4 * 4 + 4 * (3+128));

    REQUIRE(ZSTD_decompress(decomp_buffer, 128*8*2, dst, dst_size) == 8 * 128);
    size_t diff = 0;
    for (int i = 0; i < 128*8; i++) if (src[i] != decomp_buffer[i]) diff++;
    REQUIRE(diff == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_simulated","[ZSTD]") {
    for(int seed = 0; seed < 200; seed++) {
        JFJochZstdCompressor compressor;

        // Predictable random number generator
        std::mt19937 g1(seed);
        std::uniform_int_distribution<int16_t> distribution;

        std::vector<int16_t> input(RAW_MODULE_SIZE);
        for (auto &i: input) i = distribution(g1);

        std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
        bshuf_bitshuffle(input.data(), input_shuffled.data(), RAW_MODULE_SIZE, 2, 4096);
        std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
        std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

        size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                              RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));

        REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                                input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

        REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
    }
}

TEST_CASE("JFjochZstdCompressor_Frame_zeroes","[ZSTD]") {
    JFJochZstdCompressor compressor;

    std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
    for (auto &i: input_shuffled) i = 0x0u;

    std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
    std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

    size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                          RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
    REQUIRE(dst_size == 9 + 4 * 8);
    REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                            input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

    REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}

TEST_CASE("JFjochZstdCompressor_Frame_ones","[ZSTD]") {
    JFJochZstdCompressor compressor;

    std::vector<char> input_shuffled(RAW_MODULE_SIZE * sizeof(uint16_t));
    for (auto &i: input_shuffled) i = 0xFFu;

    std::vector<uint8_t> input_compressed(ZSTD_compressBound(RAW_MODULE_SIZE * sizeof(uint16_t)));
    std::vector<uint8_t> output(RAW_MODULE_SIZE * sizeof(uint16_t) * 2);

    size_t dst_size = compressor.Compress(input_compressed.data(), (uint64_t *) input_shuffled.data(),
                                          RAW_MODULE_SIZE * sizeof(uint16_t), RAW_MODULE_SIZE * sizeof(uint16_t));
    REQUIRE(dst_size == 9 + 4 * 8);
    REQUIRE(ZSTD_decompress(output.data(), RAW_MODULE_SIZE * sizeof(uint16_t) * 2,
                            input_compressed.data(), dst_size) == RAW_MODULE_SIZE * sizeof(uint16_t));

    REQUIRE(memcmp(input_shuffled.data(), output.data(), RAW_MODULE_SIZE * sizeof(uint16_t)) == 0);
}

TEST_CASE("JFJochCompressor_JFJochDecompressor_ZSTD","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum());

    for (auto &i: image)
        i = 345;

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);

    auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
    tmp.resize(tmp_size);

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("JFJochCompressor_DestTooSmall_Throws","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_ZSTD).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum(), 345);

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    // A destination far too small for the compressed output must throw, not overflow it.
    std::vector<char> tiny(64);
    REQUIRE_THROWS_AS(compressor.Compress(tiny.data(), tiny.size(), image),
                      CompressionBufferTooSmallException);

    // A buffer sized to the worst case never throws.
    std::vector<char> big(MaxCompressedSize(x.GetCompressionAlgorithm(), x.GetPixelsNum(), sizeof(int32_t)));
    REQUIRE_NOTHROW(compressor.Compress(big.data(), big.size(), image));
}

TEST_CASE("JFJochCompressor_JFJochDecompressor_LZ4","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::BSHUF_LZ4).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum());

    for (auto &i: image)
        i = 5678;

    JFJochBitShuffleCompressor compressor(x.GetCompressionAlgorithm());

    std::vector<char> tmp(x.GetPixelsNum() * sizeof(int32_t) * 4 + 12);

    auto tmp_size = compressor.Compress(tmp.data(), tmp.size(), image);
    tmp.resize(tmp_size);

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), tmp, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("JFJochDecompressor_None","[ZSTD]") {
    DiffractionExperiment x(DetJF4M());
    x.Compression(CompressionAlgorithm::NO_COMPRESSION).BitDepthImage(32).PixelSigned(true);

    std::vector<int32_t> image(x.GetPixelsNum());
    for (auto &i: image)
        i = 578;

    std::vector<int32_t> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, x.GetCompressionAlgorithm(), image, x.GetPixelsNum()));
    REQUIRE(output.size() == x.GetPixelsNum());
    REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0);
}

TEST_CASE("Bitshuffle_ZSTD","[ZSTD]") {
    std::vector<int32_t> image(RAW_MODULE_SIZE * sizeof(int32_t));
    std::vector<char> compressed(bshuf_compress_zstd_bound(RAW_MODULE_SIZE, 4, 0));
    std::vector<int32_t> decompressed(RAW_MODULE_SIZE * sizeof(int32_t));
    for (int i = 0; i < RAW_MODULE_SIZE; i++)
        image[i] = i;

    auto out_size = bshuf_compress_zstd(image.data(), compressed.data(), RAW_MODULE_SIZE, 4, 0,0);
    REQUIRE(out_size > 0);
    REQUIRE(bshuf_decompress_zstd(compressed.data(), decompressed.data(), RAW_MODULE_SIZE, 4, 0) == out_size);

    REQUIRE(memcmp(image.data(), decompressed.data(), RAW_MODULE_SIZE*sizeof(uint32_t)) == 0);
}

// BitShuffleBlock.h picks between the two vendored bitshuffle implementations by architecture:
// bitshuffle_hperf on x86, the classic one on aarch64, where only the latter has SIMD. That is only
// safe as long as the two agree bit for bit - otherwise an ARM build would write files an x86 build
// cannot read. Both are always compiled in, so this holds them to it on every architecture.
TEST_CASE("Bitshuffle_hperf_matches_classic","[ZSTD]") {
    for (size_t elem_size : {1, 2, 4, 8}) {
        // Element counts the compressor actually produces (128 kB / 16 kB blocks), plus the
        // smallest legal block and a short trailing one.
        const size_t sizes[] = {8, 64, 16384 / elem_size, 131072 / elem_size};
        for (size_t size : sizes) {
            std::vector<char> image(size * elem_size), hperf(size * elem_size),
                    classic(size * elem_size), scratch(size * elem_size), decoded(size * elem_size);
            std::mt19937 gen(1234);
            for (auto &c : image)
                c = static_cast<char>(gen() & 0xff);

            REQUIRE(bitshuf_encode_block(hperf.data(), image.data(), scratch.data(), size, elem_size) >= 0);
            REQUIRE(bshuf_trans_bit_elem(image.data(), classic.data(), size, elem_size) >= 0);
            REQUIRE(hperf == classic);

            // Each implementation must also decode what the other wrote.
            REQUIRE(bitshuf_decode_block(decoded.data(), classic.data(), scratch.data(), size, elem_size) >= 0);
            REQUIRE(decoded == image);
            REQUIRE(bshuf_untrans_bit_elem(hperf.data(), decoded.data(), size, elem_size) >= 0);
            REQUIRE(decoded == image);
        }
    }
}
// ---- BSHUF_ZSTD_RLE_HUFF (RLE runs + adaptive Huffman literals, standard zstd frame) ----
// Round-trips through JFJochDecompress (which calls stock ZSTD_decompress per block) and checks
// the output never exceeds the advertised worst-case size. Run on good and adversarial data to
// confirm the algorithm always produces valid output, even when the ratio is poor.
template<class T>
static void RequireHuffRoundTrip(const std::vector<T> &image) {
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF);
    auto compressed = compressor.Compress(image);

    REQUIRE(compressed.size() <= (size_t) MaxCompressedSize(CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
                                                            image.size(), sizeof(T)));
    std::vector<T> output;
    REQUIRE_NOTHROW(JFJochDecompress(output, CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF,
                                     compressed, image.size()));
    REQUIRE(output.size() == image.size());
    REQUIRE(memcmp(image.data(), output.data(), image.size() * sizeof(T)) == 0);
}

TEST_CASE("ZstdHuff_PureZeros", "[ZSTD]") {
    std::vector<int16_t> image(200000, 0);            // all runs, no literals
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_Poisson10", "[ZSTD]") {
    std::vector<uint16_t> image(200000);
    std::mt19937 gen(12345);
    std::poisson_distribution<int> dist(10);
    for (auto &v : image) v = (uint16_t) dist(gen);
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_RandomIncompressible", "[ZSTD]") {
    std::vector<uint16_t> image(200000);              // Mersenne-Twister noise: ~zero compression
    std::mt19937 gen(98765);
    for (auto &v : image) v = (uint16_t) gen();
    RequireHuffRoundTrip(image);                      // must still be valid and within bound
}

TEST_CASE("ZstdHuff_MaskLike", "[ZSTD]") {
    std::vector<uint32_t> image(200000, 0);           // extreme sparsity like a pixel_mask
    std::mt19937 gen(555);
    for (auto &v : image) {
        uint32_t r = gen() % 100;
        if (r < 5) v = 1;
        else if (r == 5) v = 0x80000000u;             // module-gap style flag
    }
    RequireHuffRoundTrip(image);
}

TEST_CASE("ZstdHuff_BenchmarkImage", "[ZSTD]") {
    RegisterHDF5Filter();   // bitshuffle filter, needed to read the compressed benchmark dataset
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);
    auto dims = file_space.GetDimensions();
    std::vector<int16_t> image(dims[1] * dims[2]);
    std::vector<hsize_t> start = {0, 0, 0}, size = {1, dims[1], dims[2]};
    dataset.ReadVector(image, start, size);
    RequireHuffRoundTrip(image);
}

// On-the-record proof that a Compressed_Block with Huffman literals and Number_of_Sequences = 0 is
// a valid Zstandard block that outputs its literals. Per zstd_compression_format.md: "if
// (Number_of_Sequences == 0) ... Block's decompressed content is defined solely by the Literals
// Section content." We build exactly such a block and confirm stock ZSTD_decompress returns it.
TEST_CASE("ZstdHuff_LiteralsOnly_NbSeqZero", "[ZSTD]") {
    std::vector<uint8_t> lits(4096);
    std::mt19937 gen(7);
    for (auto &b : lits) b = (uint8_t)(gen() % 6 + 1);   // 6-symbol alphabet: Huffman-compressible, no 0x00/0xFF runs

    JFJochZstdHuffCompressor huff;
    std::vector<uint8_t> frame(lits.size() * 2 + 4096);
    size_t fsz = huff.Compress(frame.data(), (const uint64_t *) lits.data(), lits.size());

    // Frame = Magic(4) + Frame_Header_Descriptor(1) + Frame_Content_Size(4) = 9 bytes, then one Block.
    // Block_Header (3 bytes, little-endian): [ Last_Block:1 | Block_Type:2 | Block_Size:21 ].
    uint32_t bh = frame[9] | (frame[10] << 8) | (frame[11] << 16);
    bool     last_block = bh & 1;
    uint32_t block_type = (bh >> 1) & 3;
    uint32_t block_size = bh >> 3;
    REQUIRE(last_block);                            // a single block for this all-literal input
    REQUIRE(block_type == 2);                       // Compressed_Block (Huffman literals path, not Raw/RLE)
    REQUIRE(frame[12 + block_size - 1] == 0x00);    // its Sequences_Section is Number_of_Sequences = 0

    std::vector<uint8_t> out(lits.size());
    size_t r = ZSTD_decompress(out.data(), out.size(), frame.data(), fsz);
    REQUIRE_FALSE(ZSTD_isError(r));
    REQUIRE(r == lits.size());
    REQUIRE(memcmp(lits.data(), out.data(), lits.size()) == 0);
}