/* lzwh.c : LZW compression with Horspool's phased-in coding and accelerated loading */
/* Made by a machine. PUBLIC DOMAIN (CC0-1.0) */

#include "lzwh.h"

#include <stdlib.h>
#include <string.h>

/****************************************************************
 * Configuration
 ****************************************************************/

#define HASH_LOAD_FACTOR 4
#define CLEAR_CODE 256
#define FIRST_CODE 257

/****************************************************************
 * Bit I/O
 ****************************************************************/

struct bitwriter {
    uint8_t *buf;
    size_t cap;
    size_t pos;
    uint32_t bits;
    unsigned nbits;
};

static int
bw_init(struct bitwriter *w, size_t initial_cap)
{
    w->buf = malloc(initial_cap);
    if (!w->buf)
        return LZWH_ERR_NOMEM;
    w->cap = initial_cap;
    w->pos = 0;
    w->bits = 0;
    w->nbits = 0;
    return LZWH_OK;
}

static int
bw_ensure(struct bitwriter *w, size_t need)
{
    if (w->pos + need <= w->cap)
        return LZWH_OK;
    size_t newcap = w->cap * 2;
    if (newcap < w->pos + need)
        newcap = w->pos + need;
    uint8_t *p = realloc(w->buf, newcap);
    if (!p)
        return LZWH_ERR_NOMEM;
    w->buf = p;
    w->cap = newcap;
    return LZWH_OK;
}

static int
bw_put(struct bitwriter *w, uint32_t code, unsigned width)
{
    w->bits |= code << w->nbits;
    w->nbits += width;
    while (w->nbits >= 8) {
        if (bw_ensure(w, 1) != LZWH_OK)
            return LZWH_ERR_NOMEM;
        w->buf[w->pos++] = w->bits & 0xff;
        w->bits >>= 8;
        w->nbits -= 8;
    }
    return LZWH_OK;
}

static int
bw_flush(struct bitwriter *w)
{
    if (w->nbits > 0) {
        if (bw_ensure(w, 1) != LZWH_OK)
            return LZWH_ERR_NOMEM;
        w->buf[w->pos++] = w->bits & 0xff;
        w->bits = 0;
        w->nbits = 0;
    }
    return LZWH_OK;
}

struct bitreader {
    const uint8_t *buf;
    size_t len;
    size_t pos;
    uint32_t bits;
    unsigned nbits;
};

static void
br_init(struct bitreader *r, const uint8_t *buf, size_t len)
{
    r->buf = buf;
    r->len = len;
    r->pos = 0;
    r->bits = 0;
    r->nbits = 0;
}

static int
br_get(struct bitreader *r, unsigned width, uint32_t *out)
{
    while (r->nbits < width) {
        if (r->pos >= r->len)
            return LZWH_ERR_CORRUPT;
        r->bits |= (uint32_t)r->buf[r->pos++] << r->nbits;
        r->nbits += 8;
    }
    *out = r->bits & ((1u << width) - 1);
    r->bits >>= width;
    r->nbits -= width;
    return LZWH_OK;
}

/****************************************************************
 * Phased-in binary coding
 ****************************************************************/

/** Compute floor(log2(n)) for n >= 1. */
static unsigned
ilog2(uint32_t n)
{
    unsigned r = 0;
    while (n >>= 1)
        r++;
    return r;
}

/** Write a code using phased-in coding for a range of [0, n).
 *
 * LSB-first variant: the extra bit goes after the base value so that
 * the decoder can read b bits first and decide whether to read one more.
 */
static int
phased_write(struct bitwriter *w, uint32_t code, uint32_t n)
{
    if (n <= 1)
        return LZWH_OK;
    unsigned bits = ilog2(n);
    uint32_t threshold = (1u << (bits + 1)) - n;
    if (code < threshold)
        return bw_put(w, code, bits);
    uint32_t excess = code - threshold;
    int rc = bw_put(w, threshold + excess / 2, bits);
    if (rc != LZWH_OK)
        return rc;
    return bw_put(w, excess & 1, 1);
}

/** Read a code using phased-in coding for a range of [0, n). */
static int
phased_read(struct bitreader *r, uint32_t n, uint32_t *out)
{
    if (n <= 1) {
        *out = 0;
        return LZWH_OK;
    }
    unsigned bits = ilog2(n);
    uint32_t threshold = (1u << (bits + 1)) - n;
    uint32_t val;
    int rc = br_get(r, bits, &val);
    if (rc != LZWH_OK)
        return rc;
    if (val < threshold) {
        *out = val;
    } else {
        uint32_t extra;
        rc = br_get(r, 1, &extra);
        if (rc != LZWH_OK)
            return rc;
        *out = threshold + (val - threshold) * 2 + extra;
    }
    return LZWH_OK;
}

/****************************************************************
 * Compression hash table
 ****************************************************************/

struct ht_entry {
    uint32_t key;
    uint16_t code;
    uint8_t occupied;
};

struct comp_dict {
    struct ht_entry *table;
    uint32_t mask;
    uint32_t next_code;
    uint32_t max_code;
};

static uint32_t
ht_hash(uint32_t parent, uint8_t ch, uint32_t mask)
{
    uint32_t h = ((uint32_t)parent << 8) | ch;
    h = (h ^ (h >> 16)) * 0x45d9f3b;
    h = (h ^ (h >> 16));
    return h & mask;
}

static int
comp_dict_init(struct comp_dict *d, unsigned dict_bits)
{
    uint32_t dict_size = 1u << dict_bits;
    uint32_t ht_size = dict_size * HASH_LOAD_FACTOR;
    /* round up to power of 2 */
    uint32_t sz = 1;
    while (sz < ht_size)
        sz <<= 1;
    d->table = calloc(sz, sizeof(d->table[0]));
    if (!d->table)
        return LZWH_ERR_NOMEM;
    d->mask = sz - 1;
    d->next_code = FIRST_CODE;
    d->max_code = dict_size;
    return LZWH_OK;
}

static void
comp_dict_free(struct comp_dict *d)
{
    free(d->table);
}

static int
comp_dict_find(struct comp_dict *d, uint32_t parent, uint8_t ch,
               uint16_t *code)
{
    uint32_t key = (parent << 8) | ch;
    uint32_t idx = ht_hash(parent, ch, d->mask);
    for (;;) {
        struct ht_entry *e = &d->table[idx];
        if (!e->occupied)
            return 0;
        if (e->key == key) {
            *code = e->code;
            return 1;
        }
        idx = (idx + 1) & d->mask;
    }
}

static void
comp_dict_insert(struct comp_dict *d, uint32_t parent, uint8_t ch)
{
    if (d->next_code >= d->max_code)
        return;
    uint32_t key = (parent << 8) | ch;
    uint32_t idx = ht_hash(parent, ch, d->mask);
    while (d->table[idx].occupied)
        idx = (idx + 1) & d->mask;
    d->table[idx].key = key;
    d->table[idx].code = d->next_code++;
    d->table[idx].occupied = 1;
}

/****************************************************************
 * Decompression dictionary
 ****************************************************************/

struct decomp_dict {
    uint16_t *prefix;
    uint8_t *suffix;
    uint16_t *length;
    uint32_t next_code;
    uint32_t max_code;
};

static int
decomp_dict_init(struct decomp_dict *d, unsigned dict_bits)
{
    uint32_t dict_size = 1u << dict_bits;
    d->prefix = malloc(dict_size * sizeof(d->prefix[0]));
    d->suffix = malloc(dict_size * sizeof(d->suffix[0]));
    d->length = malloc(dict_size * sizeof(d->length[0]));
    if (!d->prefix || !d->suffix || !d->length) {
        free(d->prefix);
        free(d->suffix);
        free(d->length);
        return LZWH_ERR_NOMEM;
    }
    d->max_code = dict_size;
    d->next_code = FIRST_CODE;
    for (unsigned i = 0; i < 256; i++) {
        d->prefix[i] = UINT16_MAX;
        d->suffix[i] = i;
        d->length[i] = 1;
    }
    return LZWH_OK;
}

static void
decomp_dict_free(struct decomp_dict *d)
{
    free(d->prefix);
    free(d->suffix);
    free(d->length);
}

static void
decomp_dict_add(struct decomp_dict *d, uint16_t prefix, uint8_t ch)
{
    if (d->next_code >= d->max_code)
        return;
    uint32_t c = d->next_code++;
    d->prefix[c] = prefix;
    d->suffix[c] = ch;
    d->length[c] = d->length[prefix] + 1;
}

/** Decode a code into buf (written backwards from buf+maxlen). Returns
 * pointer to first byte of decoded string, or NULL on error. */
static uint8_t *
decomp_dict_decode(struct decomp_dict *d, uint32_t code,
                   uint8_t *buf, size_t maxlen)
{
    size_t pos = maxlen;
    uint32_t c = code;
    while (c >= 256) {
        if (pos == 0 || c >= d->next_code)
            return NULL;
        buf[--pos] = d->suffix[c];
        c = d->prefix[c];
    }
    if (pos == 0)
        return NULL;
    buf[--pos] = (uint8_t)c;
    return buf + pos;
}

/****************************************************************
 * Output buffer
 ****************************************************************/

struct outbuf {
    uint8_t *data;
    size_t len;
    size_t cap;
};

static int
outbuf_init(struct outbuf *o, size_t cap)
{
    o->data = malloc(cap);
    if (!o->data)
        return LZWH_ERR_NOMEM;
    o->len = 0;
    o->cap = cap;
    return LZWH_OK;
}

static int
outbuf_append(struct outbuf *o, const uint8_t *data, size_t len)
{
    if (o->len + len > o->cap) {
        size_t newcap = o->cap * 2;
        if (newcap < o->len + len)
            newcap = o->len + len;
        uint8_t *p = realloc(o->data, newcap);
        if (!p)
            return LZWH_ERR_NOMEM;
        o->data = p;
        o->cap = newcap;
    }
    memcpy(o->data + o->len, data, len);
    o->len += len;
    return LZWH_OK;
}

/****************************************************************
 * Compress
 ****************************************************************/

void
lzwh_params_init(struct lzwh_params *p)
{
    p->dict_bits = LZWH_DICT_BITS_DEFAULT;
    p->accel_max = LZWH_ACCEL_DEFAULT;
}

int
lzwh_compress(const uint8_t *src, size_t src_len,
              uint8_t **out, size_t *out_len,
              const struct lzwh_params *params)
{
    struct lzwh_params p;
    if (params)
        p = *params;
    else
        lzwh_params_init(&p);

    if (p.dict_bits < LZWH_DICT_BITS_MIN)
        p.dict_bits = LZWH_DICT_BITS_MIN;
    if (p.dict_bits > LZWH_DICT_BITS_MAX)
        p.dict_bits = LZWH_DICT_BITS_MAX;
    if (p.accel_max < 1)
        p.accel_max = 1;

    struct comp_dict dict;
    int rc = comp_dict_init(&dict, p.dict_bits);
    if (rc != LZWH_OK)
        return rc;

    struct bitwriter bw;
    rc = bw_init(&bw, src_len ? src_len : 64);
    if (rc != LZWH_OK) {
        comp_dict_free(&dict);
        return rc;
    }

    /* write header: original size as 4 bytes (little-endian) */
    bw_ensure(&bw, 4);
    bw.buf[0] = src_len & 0xff;
    bw.buf[1] = (src_len >> 8) & 0xff;
    bw.buf[2] = (src_len >> 16) & 0xff;
    bw.buf[3] = (src_len >> 24) & 0xff;
    bw.pos = 4;

    if (src_len == 0) {
        *out = bw.buf;
        *out_len = bw.pos;
        comp_dict_free(&dict);
        return LZWH_OK;
    }

    /*
     * Accelerated loading state. xscode tracks the previous code
     * whose extensions we may still be loading into the dictionary.
     * xsremain counts how many more extensions to add.
     *
     * emit_range tracks the phased-in coding range separately from
     * dict.next_code. Accel entries are added to the dictionary
     * between emits but the decompressor can only add them after
     * reading the next code, so the coding range must exclude them.
     */
    uint16_t xscode = 0;
    unsigned xsremain = 0;
    uint32_t emit_range = dict.next_code;

    uint16_t cur = src[0];
    size_t i = 1;
    while (i < src_len) {
        uint8_t ch = src[i];
        uint16_t next;
        if (comp_dict_find(&dict, cur, ch, &next) &&
            next < (uint16_t)emit_range) {
            if (xsremain > 0) {
                comp_dict_insert(&dict, xscode, ch);
                if (dict.next_code > FIRST_CODE)
                    xscode = dict.next_code - 1;
                xsremain--;
            }
            cur = next;
            i++;
        } else {
            rc = phased_write(&bw, cur, emit_range);
            if (rc != LZWH_OK)
                goto fail;

            comp_dict_insert(&dict, cur, ch);
            emit_range = dict.next_code;

            if (p.accel_max > 1 &&
                dict.next_code > FIRST_CODE) {
                xscode = dict.next_code - 1;
                xsremain = p.accel_max - 1;
            } else {
                xsremain = 0;
            }

            cur = ch;
            i++;
        }
    }

    /* emit final code */
    rc = phased_write(&bw, cur, emit_range);
    if (rc != LZWH_OK)
        goto fail;

    rc = bw_flush(&bw);
    if (rc != LZWH_OK)
        goto fail;

    *out = bw.buf;
    *out_len = bw.pos;
    comp_dict_free(&dict);
    return LZWH_OK;

fail:
    free(bw.buf);
    comp_dict_free(&dict);
    return rc;
}

/****************************************************************
 * Decompress
 ****************************************************************/

int
lzwh_decompress(const uint8_t *src, size_t src_len,
                uint8_t **out, size_t *out_len,
                const struct lzwh_params *params)
{
    struct lzwh_params p;
    if (params)
        p = *params;
    else
        lzwh_params_init(&p);

    if (p.dict_bits < LZWH_DICT_BITS_MIN)
        p.dict_bits = LZWH_DICT_BITS_MIN;
    if (p.dict_bits > LZWH_DICT_BITS_MAX)
        p.dict_bits = LZWH_DICT_BITS_MAX;
    if (p.accel_max < 1)
        p.accel_max = 1;

    if (src_len < 4)
        return LZWH_ERR_CORRUPT;

    uint32_t orig_len = src[0] | (src[1] << 8) |
                        (src[2] << 16) | (src[3] << 24);

    if (orig_len == 0) {
        *out = malloc(1);
        if (!*out)
            return LZWH_ERR_NOMEM;
        *out_len = 0;
        return LZWH_OK;
    }

    struct decomp_dict dict;
    int rc = decomp_dict_init(&dict, p.dict_bits);
    if (rc != LZWH_OK)
        return rc;

    struct outbuf ob;
    rc = outbuf_init(&ob, orig_len);
    if (rc != LZWH_OK) {
        decomp_dict_free(&dict);
        return rc;
    }

    struct bitreader br;
    br_init(&br, src + 4, src_len - 4);

    uint8_t tmpbuf[65536];

    /* read first code */
    uint32_t code;
    rc = phased_read(&br, dict.next_code, &code);
    if (rc != LZWH_OK)
        goto fail;

    if (code >= 256) {
        rc = LZWH_ERR_CORRUPT;
        goto fail;
    }

    uint8_t firstch = (uint8_t)code;
    outbuf_append(&ob, &firstch, 1);

    uint32_t prev = code;

    while (ob.len < orig_len) {
        uint32_t n = dict.next_code;
        uint32_t range = (n < dict.max_code) ? n + 1 : n;
        rc = phased_read(&br, range, &code);
        if (rc != LZWH_OK)
            goto fail;

        uint8_t *str;
        size_t slen;

        if (code < n) {
            str = decomp_dict_decode(&dict, code,
                                     tmpbuf, sizeof(tmpbuf));
            if (!str) {
                rc = LZWH_ERR_CORRUPT;
                goto fail;
            }
            slen = (tmpbuf + sizeof(tmpbuf)) - str;

            decomp_dict_add(&dict, prev, str[0]);
        } else if (code == n) {
            /* KwKwK case */
            str = decomp_dict_decode(&dict, prev,
                                     tmpbuf, sizeof(tmpbuf) - 1);
            if (!str) {
                rc = LZWH_ERR_CORRUPT;
                goto fail;
            }
            slen = (tmpbuf + sizeof(tmpbuf) - 1) - str;
            tmpbuf[sizeof(tmpbuf) - 1] = str[0];
            slen++;

            decomp_dict_add(&dict, prev, str[0]);
        } else {
            rc = LZWH_ERR_CORRUPT;
            goto fail;
        }

        if (p.accel_max > 1 && slen > 1) {
            uint16_t xs = dict.next_code - 1;
            unsigned n_extra = p.accel_max - 1;
            if (n_extra > slen - 1)
                n_extra = slen - 1;
            for (unsigned k = 0; k < n_extra; k++) {
                decomp_dict_add(&dict, xs, str[k + 1]);
                if (dict.next_code > FIRST_CODE)
                    xs = dict.next_code - 1;
            }
        }

        rc = outbuf_append(&ob, str, slen);
        if (rc != LZWH_OK)
            goto fail;

        prev = code;
    }

    *out = ob.data;
    *out_len = ob.len;
    decomp_dict_free(&dict);
    return LZWH_OK;

fail:
    free(ob.data);
    decomp_dict_free(&dict);
    return rc;
}