/* Granny "Oodle1" section decompressor (M0 T2).
 *
 * Ported (decode path only) from the leaked Granny 2.9.12 SDK:
 *   granny_oodle1_compression.cpp  +  radlz.c  +  radarith.c  +  arithbit.c
 * Used as algorithm spec; RAD's build machinery / compressor are NOT included.
 * See docs/reference/steps/M0-gr2-reader.md T2.
 */
#include "oodle1.h"
#include <stdlib.h>
#include <string.h>

typedef uint8_t  U8;
typedef uint16_t U16;
typedef uint32_t U32;
typedef int32_t  S32;
typedef uintptr_t UINTa;

#define radassert(x) ((void)0)

/* ================================================================= */
/* arithbit.c — byte-aligned range decoder                            */
/* ================================================================= */

typedef struct ARITHBITS {
    U32 low;
    U32 range;
    U8 *ptr;
    U32 underflow;      /* temp between Get()/Remove() to skip re-divide */
    U8 *start;          /* low bit used as scratch */
} ARITHBITS;

static void ArithBitsGetStart(ARITHBITS *ab, const void *ptr) {
    U32 buf;
    ab->ptr   = (U8 *)ptr + 1;
    buf       = ((const U8 *)ptr)[0];
    ab->start = (U8 *)((UINTa)ptr + (buf & 1));
    ab->low   = buf >> 1;
    ab->range = 1u << 7;
}

static void ArithBitsDecRenorm(ARITHBITS *ab) {
    U32 range = ab->range;
    if (range <= 0x800000u) {
        U32 low = ab->low;
        U32 buf = (U32)((UINTa)ab->start & 1u);
        U8 *ptr = ab->ptr;
        do {
            low = (low + low + buf) << 7;
            buf = *ptr++;
            low |= (buf >> 1);
            buf &= 1u;
            range <<= 8;
        } while (range <= 0x800000u);
        ab->low   = low;
        ab->start = (U8 *)(((UINTa)ab->start & ~(UINTa)1) + buf);
        ab->range = range;
        ab->ptr   = ptr;
    }
}

static U32 ArithBitsGet(ARITHBITS *ab, U32 scale) {
    U32 tmp;
    ArithBitsDecRenorm(ab);
    ab->underflow = ab->range / scale;
    tmp = ab->low / ab->underflow;
    return (tmp >= scale) ? (scale - 1) : tmp;
}

static void ArithBitsRemove(ARITHBITS *ab, U32 start, U32 range, U32 scale) {
    U32 tmp = ab->underflow * start;
    ab->low -= tmp;
    if ((start + range) < scale) ab->range = ab->underflow * range;
    else                         ab->range -= tmp;
}

static U32 ArithBitsGetValue(ARITHBITS *ab, U32 scale) {
    U32 tmp, div, start;
    ArithBitsDecRenorm(ab);
    div   = ab->range / scale;
    start = ab->low / div;
    if (start >= scale) start = scale - 1;
    tmp = div * start;
    ab->low -= tmp;
    if ((start + 1) < scale) ab->range = div;
    else                     ab->range -= tmp;
    return start;
}

static U32 ArithBitsGetBits(ARITHBITS *ab, U32 bits, U32 scale) {
    U32 tmp;
    ArithBitsDecRenorm(ab);
    ab->underflow = ab->range >> bits;
    tmp = ab->low / ab->underflow;
    return (tmp >= scale) ? (scale - 1) : tmp;
}

static U32 ArithBitsGetBitsValue(ARITHBITS *ab, U32 bits, U32 scale) {
    U32 tmp, div, start;
    ArithBitsDecRenorm(ab);
    div   = ab->range >> bits;
    start = ab->low / div;
    if (start >= scale) start = scale - 1;
    tmp = div * start;
    ab->low -= tmp;
    if ((start + 1) < scale) ab->range = div;
    else                     ab->range -= tmp;
    return start;
}

/* ================================================================= */
/* radarith.c — adaptive arithmetic modeller (decode path)            */
/* ================================================================= */

#define NORM_BITS       14
#define NORM_COUNT      (1 << NORM_BITS)
#define ADJ_SUMS        (NORM_COUNT << 1)
#define OVERFLOW_COUNT  ((NORM_BITS / 2) - 1)     /* 6 */
#define COUNTTYPE       U16

#define COUNTS_SIZE(v)  (sizeof(COUNTTYPE) * (((v) + 1 + 1 + 3) & ~3u))
#define VALUES_SIZE(v)  COUNTS_SIZE(v)

typedef struct ARITHDATA {
    COUNTTYPE  singles_tot;
    COUNTTYPE  update_tot;
    COUNTTYPE  update_max;
    COUNTTYPE  update_range;
    COUNTTYPE  rescale_tot;
    COUNTTYPE  singles_length;
    COUNTTYPE  summed_length;
    COUNTTYPE  unique_count;
    COUNTTYPE *values;
    COUNTTYPE *single_counts;
    U32        table_walks[NORM_BITS];
    COUNTTYPE  summed_counts[OVERFLOW_COUNT];  /* deliberately overflows into the gap before single_counts */
} ARITHDATA;
typedef ARITHDATA *ARITH;

#define Arith_was_escaped(val)                 (((UINTa)(val)) > 65536u)
#define Arith_set_decompressed_symbol(val, v)  (*((U16 *)(val)) = (U16)(v))

static U32 Arith_decompress_alloc_size(U32 unique_values) {
    return (U32)(sizeof(ARITHDATA) + COUNTS_SIZE(unique_values)
                + COUNTS_SIZE(unique_values) + VALUES_SIZE(unique_values));
}

static ARITH Arith_open(void *ptr, void *compress_temp_buf, U32 max_value, U32 unique_values) {
    ARITH a = (ARITH)ptr;
    U32 u;
    (void)compress_temp_buf;   /* always NULL on the decode path */
    if (!a) return a;
    memset(a, 0, Arith_decompress_alloc_size(unique_values));
    a->single_counts = (COUNTTYPE *)((char *)a + (sizeof(ARITHDATA) + COUNTS_SIZE(unique_values)));
    a->values        = (COUNTTYPE *)((char *)a->single_counts + COUNTS_SIZE(unique_values));
    a->unique_count  = (COUNTTYPE)unique_values;
    a->singles_tot   = 4;
    a->summed_counts[0] = ADJ_SUMS;
    a->summed_counts[1] = NORM_COUNT + ADJ_SUMS;
    a->summed_counts[2] = NORM_COUNT + ADJ_SUMS;
    a->summed_counts[3] = NORM_COUNT + ADJ_SUMS;
    a->summed_counts[4] = NORM_COUNT + ADJ_SUMS;
    a->summed_counts[5] = NORM_COUNT + ADJ_SUMS;
    a->single_counts[0] = 4;
    a->update_tot   = 8;
    a->update_range = 4;
    u = max_value * 32;
    if (u < 256)        u = 256;
    else if (u > 15160) u = 15160;
    a->rescale_tot = (COUNTTYPE)u;
    u = max_value * 2;
    if (u < 128) u = 128;
    else if (u >= (((U32)a->rescale_tot >> 1) - 32)) u = (U32)(((U32)a->rescale_tot >> 1) - 32);
    a->update_max = (COUNTTYPE)u;
    return a;
}

static void rescale_decompress(ARITH a) {
    unsigned long i;
    U32 max = 0, pos = (U32)-1;
    a->single_counts[0] >>= 1;
    a->singles_tot = a->single_counts[0];
    for (i = 1; i <= a->singles_length; i++) {
        while (a->single_counts[i] <= 1) {
            if (i < a->singles_length) {
                a->single_counts[i] = a->single_counts[a->singles_length];
                a->single_counts[a->singles_length] = 0;
                a->values[i] = a->values[a->singles_length];
                --a->singles_length;
            } else {
                a->single_counts[i] = 0;
                --a->singles_length;
                goto done;
            }
        }
        a->single_counts[i] >>= 1;
        a->singles_tot = (COUNTTYPE)(a->singles_tot + a->single_counts[i]);
        if (a->single_counts[i] > max) { max = a->single_counts[i]; pos = (U32)i; }
    }
done:
    if (max && a->singles_length) {
        U32 j = a->singles_length;
        if (pos != j) {
            U32 t;
            t = a->single_counts[j]; a->single_counts[j] = a->single_counts[pos]; a->single_counts[pos] = (COUNTTYPE)t;
            t = a->values[j];        a->values[j]        = a->values[pos];        a->values[pos]        = (COUNTTYPE)t;
        }
    }
    if ((a->singles_length != a->unique_count) && (a->single_counts[0] == 0)) {
        ++a->single_counts[0];
        ++a->singles_tot;
    }
}

static void update_counts(ARITH a) {
    U32 i, tot, adj;
    adj = (NORM_COUNT * 8) / a->singles_tot;
    tot = ((((U32)a->single_counts[0]) * adj) >> 3) + ADJ_SUMS;
    a->summed_counts[0] = ADJ_SUMS;
    i = 1;
    for (;;) {
        a->summed_counts[i] = (COUNTTYPE)tot;
        if (i > a->singles_length) break;
        tot += (((U32)a->single_counts[i]) * adj) >> 3;
        ++i;
    }
    tot = (U32)a->update_range << 1;
    if (tot > a->update_max) {
        a->update_tot = (COUNTTYPE)(a->singles_tot + a->update_max);
    } else {
        a->update_range = (COUNTTYPE)tot;
        a->update_tot   = (COUNTTYPE)(a->singles_tot + tot);
    }
    a->summed_length = a->singles_length;
    a->summed_counts[a->summed_length + 1] = NORM_COUNT + ADJ_SUMS;
}

static void build_walk_table(U32 *wtable, U32 v) {
    U32 num = 0, *m = wtable;
    ++v;
    do {
        v = (v + 1) >> 1;
        ++num;
        *m++ = v;
    } while (v > 4);
    num -= 1;
    if (num) {
        memmove(wtable + NORM_BITS - 1 - num + 1, wtable + 1, num * sizeof(wtable[0]));
    }
    wtable[1] = wtable[0];
    wtable[0] = num;
}

static UINTa Arith_decompress(ARITH a, ARITHBITS *ab) {
    U32 offset;
    S32 index;

    if (a->singles_tot >= a->update_tot) {
        if (a->update_tot >= a->rescale_tot) rescale_decompress(a);
        update_counts(a);
        a->summed_counts[a->summed_length + 2] = NORM_COUNT + ADJ_SUMS;
        a->summed_counts[a->summed_length + 3] = NORM_COUNT + ADJ_SUMS;
        a->summed_counts[a->summed_length + 4] = NORM_COUNT + ADJ_SUMS;
        a->summed_counts[a->summed_length + 5] = NORM_COUNT + ADJ_SUMS;
        a->summed_counts[a->summed_length + 6] = NORM_COUNT + ADJ_SUMS;
        build_walk_table(a->table_walks, a->summed_length);
    }

    offset = ArithBitsGetBits(ab, NORM_BITS, NORM_COUNT) + ADJ_SUMS;

    {
        index = (S32)a->table_walks[1];
        #define check_level(lev)                                            \
            if (a->summed_counts[index] > (COUNTTYPE)offset)                 \
                index -= (S32)a->table_walks[NORM_BITS - 1 - (lev)];         \
            else                                                            \
                index += (S32)a->table_walks[NORM_BITS - 1 - (lev)];
        switch (a->table_walks[0]) {
            case 11: check_level(10) /* fallthrough */
            case 10: check_level(9)  /* fallthrough */
            case 9:  check_level(8)  /* fallthrough */
            case 8:  check_level(7)  /* fallthrough */
            case 7:  check_level(6)  /* fallthrough */
            case 6:  check_level(5)  /* fallthrough */
            case 5:  check_level(4)  /* fallthrough */
            case 4:  check_level(3)  /* fallthrough */
            case 3:  check_level(2)  /* fallthrough */
            case 2:  check_level(1)  /* fallthrough */
            case 1:  check_level(0)  /* fallthrough */
            case 0:  break;
            default: break;
        }
        #undef check_level

        if (a->summed_counts[index] > (COUNTTYPE)offset) index -= 2;
        else                                            index += 2;

        if (a->summed_counts[index] > (COUNTTYPE)offset) {
            if (a->summed_counts[index - 1] > (COUNTTYPE)offset) index -= 2;
            else                                                --index;
        } else {
            if (a->summed_counts[index + 1] <= (COUNTTYPE)offset) ++index;
        }
    }

    ArithBitsRemove(ab, a->summed_counts[index] - ADJ_SUMS,
                    a->summed_counts[index + 1] - a->summed_counts[index], NORM_COUNT);

    ++a->single_counts[index];
    ++a->singles_tot;

    if (index <= 0) {
        if (a->singles_length == a->summed_length) {
        new_index:
            index = ++a->singles_length;
            a->single_counts[index] += 2;
            a->singles_tot += 2;
            if (a->singles_length == a->unique_count) {
                a->singles_tot = (COUNTTYPE)(a->singles_tot - a->single_counts[0]);
                a->single_counts[0] = 0;
            }
            return (UINTa)&a->values[index];
        } else {
            if (ArithBitsGetBitsValue(ab, 1, 2)) {
                index = (S32)(ArithBitsGetValue(ab, a->singles_length - a->summed_length)
                              + a->summed_length + 1);
                a->single_counts[index] += 2;
                a->singles_tot += 2;
            } else {
                goto new_index;
            }
        }
    }

    return a->values[index];
}

/* ================================================================= */
/* radlz.c — LZ decoder (decode path)                                 */
/* ================================================================= */

#define LOW_BITS 2
#define MED_BITS 8
#define TOP_BITS 8
#define LARGEST_POSSIBLE_OFFSET ((1 << (LOW_BITS + MED_BITS + TOP_BITS)) - 1)
#define GRANNY_OFFSET LARGEST_POSSIBLE_OFFSET

#define MAX_LENS 64
#define LMAX (1 << LOW_BITS)   /* 4   */
#define MMAX (1 << MED_BITS)   /* 256 */
#define GETT(v) ((v) >> (LOW_BITS + MED_BITS))
#define ADDRESS_MASK 3

static const U32 long_lengths[4] = { MAX_LENS * 2, MAX_LENS * 3, MAX_LENS * 4, MAX_LENS * 8 };

typedef struct LZ_HEADER {
    U32 max_offset_and_byte;
    U32 uniq_offset_and_byte;
    U32 uniq_lens;
} LZ_HEADER;

typedef struct LZDDATA {
    ARITH bytes[ADDRESS_MASK + 1];
    ARITH lens[MAX_LENS + 1];
    ARITH offsl, offst;
    ARITH offsm[256];
    U32   max_bytes, max_offs, max_offsL, max_offsM, max_offsT;
    U32   bytes_decompressed;
    U32   last_len;
} LZDDATA;
typedef LZDDATA *LZD;

static void get_max_offsets(U32 max_offset, U32 *low_max, U32 *med_max, U32 *high_max) {
    *low_max  = (max_offset >= LMAX) ? LMAX : (max_offset + 1);
    *med_max  = ((max_offset >> LOW_BITS) >= MMAX) ? MMAX : ((max_offset >> LOW_BITS) + 1);
    *high_max = GETT(max_offset) + 1;
}

static void copy_bytes(void *d, const void *s, U32 length) {
    U8 *dest = (U8 *)d;
    const U8 *src = (const U8 *)s;
    /* LZ back-refs always have src < dest, so the fast path never triggers here;
     * byte copy is required for correct RLE-style overlap. */
    if (length >= 4 && (src - dest) >= 4) {
        do {
            length -= 4;
            ((U32 *)dest)[0] = ((const U32 *)src)[0];
            dest += 4; src += 4;
        } while (length > 4);
        if (length == 0) return;
    }
    do { length--; *dest++ = *src++; } while (length);
}

static U32 LZ_decompress_alloc_size(U32 max_byte_value, U32 uniq_byte_values, U32 max_offset) {
    U32 size, h, low_max, med_max, high_max;
    (void)max_byte_value;
    size = (U32)sizeof(LZDDATA);
    get_max_offsets(max_offset, &low_max, &med_max, &high_max);
    h = Arith_decompress_alloc_size(uniq_byte_values);
    size += h * (ADDRESS_MASK + 1);
    size += Arith_decompress_alloc_size(MAX_LENS + 1) * (MAX_LENS + 1);
    size += Arith_decompress_alloc_size(low_max);
    size += high_max * Arith_decompress_alloc_size(med_max);
    size += Arith_decompress_alloc_size(high_max);
    return size;
}

static LZD LZ_decompress_open_from_header(void *ptr, const LZ_HEADER *h) {
    U32 i, j, size, uniq = 0;
    U8 *addr;
    LZD l = (LZD)ptr;

    l->max_bytes = h->max_offset_and_byte & 511;
    l->max_offs  = h->max_offset_and_byte >> 9;
    get_max_offsets(l->max_offs, &l->max_offsL, &l->max_offsM, &l->max_offsT);
    l->last_len = 0;
    l->bytes_decompressed = 0;

    j = h->uniq_offset_and_byte & 511;
    addr = (U8 *)(l + 1);
    size = Arith_decompress_alloc_size(j);
    for (i = 0; i <= ADDRESS_MASK; i++) {
        l->bytes[i] = Arith_open(addr, 0, l->max_bytes - 1, j);
        addr += size;
    }

    for (j = 0; j < 4; j++) {
        uniq = (h->uniq_lens >> ((3 - j) * 8)) & 255;
        size = Arith_decompress_alloc_size(uniq);
        for (i = 0; i < (MAX_LENS / 4); i++) {
            l->lens[(j * (MAX_LENS / 4)) + i] = Arith_open(addr, 0, MAX_LENS, uniq);
            addr += size;
        }
    }
    for (i = (MAX_LENS / 4) * 4; i <= MAX_LENS; i++) {
        l->lens[i] = Arith_open(addr, 0, MAX_LENS, uniq);  /* inherits last uniq/size */
        addr += size;
    }

    size = Arith_decompress_alloc_size(l->max_offsM);
    for (i = 0; i < l->max_offsT; i++) {
        l->offsm[i] = Arith_open(addr, 0, l->max_offsM - 1, l->max_offsM);
        addr += size;
    }

    l->offsl = Arith_open(addr, 0, l->max_offsL - 1, l->max_offsL);
    l->offst = Arith_open(addr + Arith_decompress_alloc_size(l->max_offsL),
                          0, l->max_offsT - 1, GETT(h->uniq_offset_and_byte >> 9) + 1);
    return l;
}

static U32 LZ_decompress(LZD l, ARITHBITS *ab, U8 *output) {
    UINTa v;
    U32 escaped;

    v = Arith_decompress(l->lens[l->last_len], ab);
    if (Arith_was_escaped(v)) {
        escaped = ArithBitsGetValue(ab, MAX_LENS + 1);
        Arith_set_decompressed_symbol(v, escaped);
        v = escaped;
    }
    l->last_len = (U32)v;

    if (v) {
        U32 len, off, max_ofs;
        UINTa m;

        max_ofs = l->max_offs;
        if (max_ofs > l->bytes_decompressed) max_ofs = l->bytes_decompressed;

        len = (v >= (MAX_LENS - 3)) ? long_lengths[v - (MAX_LENS - 3)] : (U32)(v + 1);

        v = Arith_decompress(l->offsl, ab);
        if (Arith_was_escaped(v)) {
            escaped = ArithBitsGetValue(ab, l->max_offsL);
            Arith_set_decompressed_symbol(v, escaped);
            v = escaped;
        }
        off = (U32)(v + 1);

        v = Arith_decompress(l->offst, ab);
        if (Arith_was_escaped(v)) {
            escaped = ArithBitsGetValue(ab, GETT(max_ofs) + 1);
            Arith_set_decompressed_symbol(v, escaped);
            v = escaped;
        }

        m = Arith_decompress(l->offsm[v], ab);
        if (Arith_was_escaped(m)) {
            U32 offsm_used;
            if (max_ofs >= (MMAX << LOW_BITS)) offsm_used = MMAX;
            else                               offsm_used = (max_ofs >> LOW_BITS) + 1;
            escaped = ArithBitsGetValue(ab, offsm_used);
            Arith_set_decompressed_symbol(m, escaped);
            m = escaped;
        }

        off = (U32)(off + ((U32)m << LOW_BITS) + ((U32)v << (LOW_BITS + MED_BITS)));
        l->bytes_decompressed += len;
        copy_bytes(output, output - off, len);
        return len;
    } else {
        v = Arith_decompress(l->bytes[((UINTa)output) & ADDRESS_MASK], ab);
        if (Arith_was_escaped(v)) {
            escaped = ArithBitsGetValue(ab, l->max_bytes);
            Arith_set_decompressed_symbol(v, escaped);
            v = escaped;
        }
        *output = (U8)v;
        ++l->bytes_decompressed;
        return 1;
    }
}

/* ================================================================= */
/* granny_oodle1_compression.cpp — the wrapper                        */
/* ================================================================= */

#define ALIGN32(x) (((x) + 3u) & ~3u)

static void reverse32(void *p, U32 nbytes) {
    U8 *b = (U8 *)p;
    U32 i;
    for (i = 0; i + 4 <= nbytes; i += 4) {
        U8 t0 = b[i], t1 = b[i + 1];
        b[i]     = b[i + 3];
        b[i + 1] = b[i + 2];
        b[i + 2] = t1;
        b[i + 3] = t0;
    }
}

int gr2_oodle1_decompress(int file_is_byte_reversed,
                          U32 comp_size, void *comp_bytes,
                          U32 stop0, U32 stop1, U32 stop2, void *out) {
    LZ_HEADER headers[3];
    ARITHBITS ab;
    U32 temp_size, size, b;
    U32 stops[3];
    void *temp;
    U8 *to;

    /* hanging-bit fix: zero-pad up to 32-bit alignment (caller also over-allocates). */
    {
        U32 rounded = ALIGN32(comp_size) - comp_size;
        while (rounded--) ((U8 *)comp_bytes)[comp_size + rounded] = 0;
    }

    if (comp_size < sizeof(headers)) return -1;
    memcpy(headers, comp_bytes, sizeof(headers));
    if (file_is_byte_reversed) reverse32(headers, sizeof(headers));

    ArithBitsGetStart(&ab, (const U8 *)comp_bytes + sizeof(headers));

    temp_size = LZ_decompress_alloc_size(255, 256, GRANNY_OFFSET);
    temp = malloc(temp_size);
    if (!temp) return -2;

    size = 0;
    stops[0] = stop0; stops[1] = stop1; stops[2] = stop2;
    to = (U8 *)out;
    for (b = 0; b < 3; ++b) {
        U32 stop = stops[b];
        LZD lz = LZ_decompress_open_from_header(temp, &headers[b]);
        U32 guard = 0;
        while (size < stop) {
            U32 len = LZ_decompress(lz, &ab, to);
            if (len == 0 || size + len > stop2 + 512) { free(temp); return -3; }  /* corruption guard */
            size += len;
            to   += len;
            if (++guard > stop2 + 16) { free(temp); return -4; }
        }
        if (size != stop) { free(temp); return -5; }
    }

    free(temp);
    return 0;
}

int gr2_oodle1_decompress_chunk(int file_is_byte_reversed,
                                U32 comp_size, void *comp_bytes,
                                U32 decompressed_size, void *out) {
    LZ_HEADER header;
    ARITHBITS ab;
    U32 temp_size, size;
    void *temp;
    U8 *to;

    if (comp_size < sizeof(header)) return -1;
    memcpy(&header, comp_bytes, sizeof(header));
    if (file_is_byte_reversed) reverse32(&header, sizeof(header));

    ArithBitsGetStart(&ab, (const U8 *)comp_bytes + sizeof(header));
    temp_size = LZ_decompress_alloc_size(255, 256, GRANNY_OFFSET);
    temp = malloc(temp_size);
    if (!temp) return -2;

    size = 0;
    to = (U8 *)out;
    {
        LZD lz = LZ_decompress_open_from_header(temp, &header);
        while (size < decompressed_size) {
            U32 len = LZ_decompress(lz, &ab, to);
            if (len == 0) { free(temp); return -3; }
            size += len;
            to   += len;
        }
    }
    free(temp);
    return (size == decompressed_size) ? 0 : -5;
}