lz_map2.c

#include "rt/ustd.h"
#include "rt/fileio.h"

// https://en.wikipedia.org/wiki/Ukkonen%27s_algorithm
// https://en.wikipedia.org/wiki/Two-way_string-matching_algorithm
// https://en.wikipedia.org/wiki/LZ77_and_LZ78

#define LZ_VERBOSE
#undef  LZ_VERBOSE

// in DEBUG lz_linear_search() + verify is incredibly slow
#undef  LZ_ALL_TESTS
#define LZ_ALL_TESTS

#ifdef LZ_VERBOSE
#define trace(...) printf(__VA_ARGS__)
#else
#define trace(...) ((void)0)
#endif

enum { min_window = 4, max_window = 1u << 16, min_len = 2, max_len = 254 };

struct map {
    const void** p; // p[n]
    size_t       n; // number of entries in the map (max_window * 4)
    uint32_t     m; // mask 0xFFFFFF for 3 bytes and 0xFFFFFFFFu for 4 bytes
};

struct lz {
    size_t   prev[max_window]; // previous `i` of 4 bytes entry
    size_t   map2[1u << (sizeof(uint16_t) * 8)]; // `i` + 1 of 2 bytes
    uint64_t prev_count; // number of times `prev` was used
    size_t prev_max;     // maximum length of prev[] chain
};

static void map_init(struct map* m, void** p, size_t n, size_t b) {
    assert(16 < n && n <= (1u << 24) && 3 <= b && b <= 4);
    if ( !(16 < n && n <= (1u << 24) && 3 <= b && b <= 4) ) {
        exit(1); // this is not a generic map implementation
    }
    memset(m, 0, sizeof(*m));
    m->p = p;
    m->n = n;
    m->m = b == 3 ? 0x00FFFFFFu : 0xFFFFFFFFu;
    memset(m->p, 0, n * sizeof(m->p[0]));
}

static inline uint32_t map_hash_key(uint32_t k) {
    k ^= k >> 13; // Bob Jenkins' hash
    k *= 0x85EBCA6Bu;
    k ^= k >> 16;
    return k;
}

static inline size_t map_hash(struct map* m, uint32_t k) {
    return (size_t)(map_hash_key(k) % m->n);
}

static inline bool map_get(struct map* m, uint32_t k,
                                  const size_t s, size_t* ix) {
    size_t i = s; // start
    const uint32_t mask = m->m;
    while (m->p[i]) {
        if ((mask & *((uint32_t*)m->p[i])) == k) {
            *ix = i;
            return true;
        } else {
            i = (i + 1) % m->n;
            if (i == s) { return false; }
        }
    }
    return false;
}

static inline void map_remove(struct map* m, size_t i) {
    m->p[i] = 0;
    size_t x = i; // next
    for (;;) {
        x = (x + 1) % m->n;
        if (!m->p[x]) { break; }
        const uint32_t b4 = *(uint32_t*)m->p[x];
        size_t h = map_hash(m, b4 & m->m);
        // Check if `h` lies within [i, x), accounting for wrap-around:
        const bool can_move = i <= x ? x < h || h <= i :
                                       x < h && h <= i;
        if (can_move) {
            m->p[i] = m->p[x];
            m->p[x] = 0;
            i = x;
        }
    }
}

static inline bool map_put(struct map* m, const void* p,
                      const size_t h, uint32_t k) {
    const uint32_t mask = m->m;
    size_t i = h;
    while (m->p[i]) {
        if ((mask & *((uint32_t*)m->p[i])) == k) {
            m->p[i] = p;
            return true;
        } else {
            i = (i + 1) % m->n;
            assert(i != h); // in the way map is used should not happen
            if (i == h) { return false; }
        }
    }
    m->p[i] = p;
    return true;
}

static inline bool map_put3(struct map* m, const void* p, uint32_t b3) {
    assert((b3 & ~0xFFFFFFu) == 0 && m->m == 0x00FFFFFFu);
    return map_put(m, p, map_hash(m, b3), b3);
}

static inline bool map_put4(struct map* m, const void* p, uint32_t b4) {
    assert(m->m == 0xFFFFFFFFu);
    return map_put(m, p, map_hash(m, b4), b4);
}

static inline bool map_get3(struct map* m, uint32_t b3, size_t* ix) {
    assert((b3 & ~0xFFFFFFu) == 0 && m->m == 0x00FFFFFFu);
    return map_get(m, b3, map_hash(m, b3), ix);
}

static inline bool map_get4(struct map* m, uint32_t b4, size_t* ix) {
    assert(m->m == 0xFFFFFFFFu);
    return map_get(m, b4, map_hash(m, b4), ix);
}


static void lz_init(struct lz *lz) {
    memset(lz->prev, 0, sizeof(lz->prev));
    memset(lz->map2, 0, sizeof(lz->map2));
    lz->prev_count = 0;
    lz->prev_max   = 0;
}

static void lz_insert(struct lz* lz, const uint8_t* in, const size_t n,
                      const size_t w, const size_t i) {
    assert(min_window <= w && w <= max_window);
    assert(((w - 1) & w) == 0); // window is power of 2
    if (i < n - 2) {
        const uint16_t b2 = *((uint16_t*)(in + i));
        size_t pp = lz->map2[b2];
        bool seen = pp != 0;
        lz->map2[b2] = i + 1;
        const size_t index = i % w;
        if (!seen) {
            lz->prev[index] = 0;
        } else {
            pp--; // because map2[] keeps +1 values
            // `i`, `pp` and `w` are unsigned,
            // if `i` may is less than `w` pp <= i - w is incorrect
            if (i <= w || i - w <= pp && pp < i) {
                lz->prev[index] = i - pp;
            } else {
                lz->prev[index] = 0;
            }
        }
    }
}

// both lz_find and lz_linear function search longest match
// at a shortest distance from position `i` and return
// ml: [2..max_len] inclusive
// md: [1..window] inclusive

static inline void lz_find(struct lz* lz, const uint8_t* in, const size_t n,
                           const size_t w, const size_t i,
                           size_t *ml, size_t *md) {
    assert(min_window <= w && w <= max_window);
    assert(((w - 1) & w) == 0); // window is power of 2
    assert(*ml == 0); // caller's responsibility
    assert(*md == 0);
    size_t len = 0;
    size_t dst = 0;
    const uint16_t b2 = *((uint16_t*)(in + i));
    size_t p = lz->map2[b2];
    if (p > 0) {
        p--; // because map2[] keep position + 1
        size_t max_k = n - i > max_len ? max_len : n - i;
        size_t k = 2; // start with at least 4
        while (k < max_k && in[p + k] == in[i + k]) { k++; }
        if (i - p <= w) { len = k; dst = i - p; }
        size_t index = p & (w - 1); // same as p % w for w = 2^x
        size_t d = lz->prev[index];
        size_t prev_chain = 0;
        while (len < max_len && 0 < d && d <= p && p - d < i && i <= p - d + w) {
            p -= d;
            if (len == 2 || memcmp(in + p + 2, in + i + 2, len - 2) == 0) {
                k = len; // because with `len` bytes are the same
                while (k < max_k && in[p + k] == in[i + k]) { k++; }
                if (k > len) { len = k; dst = i - p; }
            }
            index = p & (w - 1); // same as p % w for w = 2^x
            d = lz->prev[index];
            prev_chain++;
        }
        lz->prev_count += prev_chain;
        lz->prev_max = max(lz->prev_max, prev_chain);
    }
    if (len > 0) {
        *ml = len;
        *md = dst;
    }
}

static void lz_linear(const uint8_t* in, const size_t n,
                      const size_t w, const size_t i,
                      size_t *ml, size_t *md) {
    assert(*ml == 0); // caller's responsibility
    assert(*md == 0);
    if (1 <= i && i < n - 2) {
        size_t len = 0;
        size_t dst = 0;
        size_t j = i - 1;
        size_t min_j = i >= w ? i - w : 0;
        for (;;) {
            size_t max_k = n - i > max_len ? max_len : n - i;
            size_t k = 0;
            while (k < max_k && in[j + k] == in[i + k]) { k++; }
            if (k >= 2 && k > len) {
                len = k;
                dst = i - j;
                if (len == max_len) { break; }
            }
            if (j == min_j) { break; }
            j--;
        }
        *ml = len;
        *md = dst;
    }
}

// tests:

static uint64_t seed = 1; // random generator seed/state

static int test(struct lz* lz, const uint8_t* in, const size_t n,
                const size_t w, bool verbose) {
    assert(min_window <= w && w <= max_window);
    lz_init(lz);
    #ifdef DEBUG
    if (verbose) {
        printf("\"%.*s\": %zu\n ", (int)n, in, n);
        for (size_t i = 0; i < n; i++) {
            printf("%d", i % 10);
        }
        printf("\n ");
        for (size_t i = 0; i < n; i++) {
            printf("%c", i % 10 == 0 ? '0' + (i / 10) % 10 : 0x20);
        }
        printf("\n");
    }
    #else
    (void)verbose;
    #endif
    size_t i = 0;
    lz_insert(lz, in, n, w, i);
    while (i < n) {
        if (1 <= i && i < n - 2) {
            size_t ml1 = 0, md1 = 0;
            // TODO: accumulate last 2 bytes in a register (shifting and pass to find and insert)
            lz_find(lz, in, n, w, i, &ml1, &md1);
            #ifdef DEBUG
                // TODO: this is incredibly slow
                //       need an option (compile or runtime)
                //       to turn it on/off
                size_t ml2 = 0, md2 = 0;
                lz_linear(in, n, w, i, &ml2, &md2);
                if (ml1 > 0 || ml2 > 0) {
                    if (ml1 != ml2 || md1 != md2) {
                        printf("[%zd] longest len:dst %3zd:%zd \"%.2s\" \n",
                               i, ml1, md1, in + i);
                        printf("[%zd] linear  len:dst %3zd:%zd \"%.2s\" \n",
                               i, ml2, md2, in + i);
                    }
                    assert(ml1 == ml2 && md1 == md2);
                    if (ml1 != ml2 || md1 != md2) { return 1; }
                    size_t next_i = i + ml1;
                    while (i < next_i) {
                        lz_insert(lz, in, n, w, i);
                        i++;
                    }
                } else {
                    lz_insert(lz, in, n, w, i);
                    i++;
                }
            #else
                if (ml1 > 0) {
                    size_t next_i = i + ml1;
                    while (i < next_i) {
                        lz_insert(lz, in, n, w, i);
                        i++;
                    }
                } else {
                    lz_insert(lz, in, n, w, i);
                    i++;
                }
            #endif
        } else {
            i++;
        }
    }
    return 0;
}

static int test0(struct lz* lz) {
    const char* in = "_abc;abd:abe_";
    const size_t n = strlen(in);
    const size_t w = 8;
    printf("n: %6zd window: %5zd\n", n, w);
    return test(lz, (const uint8_t*)in, n, w, true);
}

static int test1(struct lz* lz) {
    const char* in = "aaaaa";
    const size_t n = strlen(in);
    const size_t w = min_window;
    printf("n: %6zd window: %5zd\n", n, w);
    return test(lz, (const uint8_t*)in, n, w, true);
}

static int test2(struct lz* lz) {
    const char* in = "a_aa_aa_aaa_aaaa_aaaaa_aaaa_aaa_aa_a";
    const size_t n = strlen(in);
    size_t max_w = n < max_window ? n : max_window;
    for (size_t w = min_window; w < max_w; w += w) {
        printf("n: %6zd window: %5zd\n", n, w);
        if (test(lz, (const uint8_t*)in, n, w, true)) { return 1; }
    }
    return 0;
}

static int test3(struct lz* lz) {
//  enum { N = 16, M = 256 * 1024, T = 1 }; // experimental
    enum { N = 16, M = 256, T = 32 };
    int pl[N] = { 0 };
    for (int t = 0; t < T; t++) {
        for (int i = 0; i < 16; i++) {
            pl[i] = 2 + (int)(rand64(&seed) * 2);
        }
        uint8_t in[M] = {0};
        int k = 0;
        int j = 0;
        while (k < M - 2) {
            for (int i = 0; i < pl[j] && k < M - 2; i++) {
                in[k++] = 'a';
            }
            j = (j + 1) % N;
            if (k < M - 2) { in[k++] = '_'; }
        }
        assert(in[sizeof(in) - 1] == 0);
        const size_t n = strlen((char*)in);
        size_t max_w = n < max_window ? n : max_window;
        for (size_t w = min_window; w < max_w; w += w) {
            trace("n: %6zd window: %5zd\n", n, w);
            if (test(lz, (const uint8_t*)in, n, w, false)) { return 1; }
        }
    }
    return 0;
}

static int test4(struct lz* lz) {
    #if !defined(DEBUG) || defined(LZ_ALL_TESTS)
    #ifdef DEBUG
        enum { n = (max_window + 1) / 2, w = max_window / 8 };
    #else
        enum { n = 32 * 1024 * 1024, w = max_window };
    #endif
    static uint8_t in[n];
    memset(in, 0x00, n);
    printf("n: %6d window: %5d\n", n, w);
    uint64_t t = nanoseconds();
    int r = test(lz, in, n, w, false);
    t = nanoseconds() - t;
    printf("Time: %6.3fs Throughput: %7.3f MiB/s ",
            t / 1.0e9, n / (t / 1.0e9) / (1024 * 1024));
    double prev_per_byte = (double)lz->prev_count / (double)n;
    printf("prev[] per byte: %7.3f max: %zd\n", prev_per_byte, lz->prev_max);
    // 32MB on Mac Book Air 2024 M3
    // RELEASE Time:  0.102s Throughput: 312.227 MiB/s
    return r;
    #else
    return lz == 0;
    #endif
}

static int test5(struct lz* lz) {
    #if !defined(DEBUG) || defined(LZ_ALL_TESTS)
    #ifdef DEBUG
        enum { n = 128 * 1024};
    #else
        enum { n = 16 * 1024 * 1024};
    #endif
    static uint8_t in[n];
    for (int i = 0; i < n; i++) {
        in[i] = (uint8_t)(256 * rand64(&seed));
    }
    printf("n: %6d window: %5d\n", n, max_window);
    uint64_t t = nanoseconds();
    int r = test(lz, in, n, max_window, false);
    t = nanoseconds() - t;
    printf("Time: %6.3fs Throughput: %7.3f MiB/s ",
            t / 1.0e9, n / (t / 1.0e9) / (1024 * 1024));
    double prev_per_byte = (double)lz->prev_count / (double)n;
    printf("prev[] per byte: %7.3f max: %zd\n", prev_per_byte, lz->prev_max);
    // 32MB on Mac Book Air 2024 M3
    // RELEASE Time:  0.608s Throughput:  52.612 MiB/s
    return r;
    #else
    return lz == 0;
    #endif
}

static int test6(struct lz* lz) {
    #if !defined(DEBUG) || defined(LZ_ALL_TESTS)
    #ifdef DEBUG
        enum { n = 128 * 1024};
    #else
        enum { n = 32 * 1024 * 1024 };
    #endif
    static uint8_t in[n];
    for (int i = 0; i < n; i++) {
        in[i] = (uint8_t)(256 * rand64(&seed));
    }
    for (int i = 0; i < n / 256; i++) {
        int len = 2 + (int)(rand64(&seed) * 256);
        int pos = (int)(rand64(&seed) * (n - len));
        uint8_t b = (uint8_t)(rand64(&seed) * 256);
        for (int j = pos; j < pos + len; j++) {
            assert(j < n);
            in[j] = b;
        }
    }
    printf("n: %6d window: %5d\n", n, max_window);
    uint64_t t = nanoseconds();
    int r = test(lz, in, n, max_window, false);
    t = nanoseconds() - t;
    printf("Time: %6.3fs Throughput: %7.3f MiB/s ",
            t / 1.0e9, n / (t / 1.0e9) / (1024 * 1024));
    double prev_per_byte = (double)lz->prev_count / (double)n;
    printf("prev[] per byte: %7.3f max: %zd\n", prev_per_byte, lz->prev_max);
    // 64MB on Mac Book Air 2024 M3
    // RELEASE Time:  1.159s Throughput:  55.222 MiB/s
    return r;
    #else
    return lz == 0;
    #endif
}

static int test_file(struct lz* lz, const char* fn) {
    uint8_t* in = null;
    size_t n = 0;
    errno_t r = file_read_fully(fn, &in, &n);
    if (r != 0) { return r; }
    printf("\"%s\" %zd bytes\n", fn, n);
    uint64_t t = nanoseconds();
    r = test(lz, in, n, max_window, false);
    t = nanoseconds() - t;
    printf("Time: %6.3fs Throughput: %7.3f MiB/s ",
            t / 1.0e9, n / (t / 1.0e9) / (1024 * 1024));
    double prev_per_byte = (double)lz->prev_count / (double)n;
    printf("prev[] per byte: %7.3f max: %zd\n", prev_per_byte, lz->prev_max);
    free(in);
    return r;
}

int lz_map2_test(void) {
    static struct lz lz77;
    struct lz* lz = &lz77;
    return test0(lz) || test1(lz) || test2(lz) || test3(lz) ||
           test4(lz) || test5(lz) || test6(lz) ||
           test_file(lz, "test/bible.txt") ||
           test_file(lz, "test/mandrill.png");
}