/*
  infstats version 1.0, 21 June 2024

  Copyright (C) 2005-2024 Mark Adler

  This software is provided 'as-is', without any express or implied
  warranty.  In no event will the author be held liable for any damages
  arising from the use of this software.

  Permission is granted to anyone to use this software for any purpose,
  including commercial applications, and to alter it and redistribute it
  freely, subject to the following restrictions:

  1. The origin of this software must not be misrepresented; you must not
     claim that you wrote the original software. If you use this software
     in a product, an acknowledgment in the product documentation would be
     appreciated but is not required.
  2. Altered source versions must be plainly marked as such, and must not be
     misrepresented as being the original software.
  3. This notice may not be removed or altered from any source distribution.

  Mark Adler    madler@alumni.caltech.edu
 */

// Take as input the binary format generated by infgen -b or infgen -bb.
// Generate and show statistics on the compressed and uncompressed data.
//
// First the number of deflate streams are shown. The subsequent statistics are
// across all of the streams, as if they were a single stream. The statistics
// are broken up into static block, fixed block, and dynamic block groups.
//
// For each block type is the number of those blocks, the number of compressed
// bytes(bits), the number of uncompressed bytes, and the compression ratio of
// input bytes divided by output bytes. For fixed and dynamic blocks that is
// followed by the same thing, but broken out into literals ("lits") and
// matches ("mats"). bytes(bits) is the number of whole bytes followed by the
// remaining number of bits in 0..7.
//
// The subsequent lines for the fixed and dynamic block types are statistics
// for the bits per block ("bits"), the bits in the header ("head") for dynamic
// blocks only, the number of symbols per block ("syms"), the distribution of
// match lengths ("lens"), and the uncompressed bytes per block ("data"). Each
// line of statistics is the mean ("μ ="), standard deviation ("σ ="), and
// [minimum..maximum].
//
// Information about the compressed bits will not be included if infgen -b was
// used instead of infgen -bb.

#include <stdio.h>
#include <stdint.h>
#include <locale.h>
#include <math.h>

// Compile with -DTEST to write the uncompressed data to stdout, instead of
// the statistics.

#ifdef TEST
// Window for writing decompressed output.
static char window[32768];
static size_t pos = 0;
#endif

// Accumulated statistics of a given value.
typedef struct {
    uintmax_t num;          // count of values
    uintmax_t sum;          // sum of values
    double sqr;             // sum of the values squared
    uintmax_t min;          // minimum value
    uintmax_t max;          // maximum value
} tally_t;

// Initialize t.
static void tally_init(tally_t *t) {
    t->num = 0;
    t->sum = 0;
    t->sqr = 0.;
    t->min = -1;
    t->max = 0;
}

// Update the total, min, and max in stats with value.
static void tally(tally_t *t, uintmax_t value) {
    t->num++;
    t->sum += value;
    t->sqr += value * (double)value;
    if (value < t->min)
        t->min = value;
    if (value > t->max)
        t->max = value;
}

// Show the tally mean, standard deviation, min, and max on out, with the
// prefix pre.
static void show_tally(FILE *out, char const *pre, tally_t const *t) {
    if (t->num == 1)
        fprintf(out, "%s%'ju\n", pre, t->sum);
    else if (t->num > 1)
        fprintf(out, "%sμ = %'.1f, σ = %'.1f, in [%'ju..%'ju]\n",
                pre, t->sum / (double)t->num,
                sqrt((t->sqr - (t->sum * (double)t->sum) / t->num) /
                     (t->num - 1)),
                t->min, t->max);
}

typedef struct {
    // Temporary block-local statistics.
    uintmax_t lit;          // number of literals in current block
    uintmax_t mat;          // number of matches in current block
    tally_t writ;           // data bytes per symbol stats for current block

    // infgen -bb statistics.
    int detail;             // true if infgen -bb was used
    tally_t bits[3];        // bits in stats per block type
    tally_t head[3];        // header bits stats per block type
    uintmax_t litb[3];      // number of literal bits for fixed and dynamic
    uintmax_t matb[3];      // number of match bits for fixed and dynamic

    // Statistics valid for both infgen -b and infgen -bb.
    uintmax_t lits[3];      // number of literals per block type
    uintmax_t mats[3];      // number of matches per block type
    tally_t syms[3];        // number of symbols stats per block type
    tally_t lens[3];        // match length stats per block type
    tally_t data[3];        // bytes of uncompressed data stats per block type
    uintmax_t streams;      // number of deflate streams
} stat_t;

// Initialize the statistics in s.
static void init(stat_t *s) {
    s->detail = 0;
    for (int i = 0; i < 3; i++) {
        tally_init(&s->bits[i]);
        tally_init(&s->head[i]);
        s->litb[i] = 0;
        s->matb[i] = 0;
        s->lits[i] = 0;
        s->mats[i] = 0;
        tally_init(&s->syms[i]);
        tally_init(&s->lens[i]);
        tally_init(&s->data[i]);
    }
    s->streams = 0;
}

// Show the statistics in s on out.
static void show_stats(FILE *out, stat_t *s) {
    static const char *name[] = {"stored", "fixed", "dynamic"};
    fprintf(out, "%'ju stream%s\n", s->streams, s->streams == 1 ? "" : "s");
    for (int i = 0; i < 3; i++)
        if (s->data[i].num != 0) {
            fprintf(out, "%'ju %s block%s: ", s->data[i].num, name[i],
                    s->data[i].num == 1 ? "" : "s");
            if (s->detail)
                fprintf(out, "%'ju(%ju) -> ",
                        s->bits[i].sum >> 3, s->bits[i].sum & 7);
            fprintf(out, "%'ju", s->data[i].sum);
            if (s->detail)
                fprintf(out, " (%.3f)",
                        s->bits[i].sum / (8 * (double)s->data[i].sum));
            putc('\n', out);
            if (i > 0) {
                fprintf(out, "  lits: ");
                if (s->detail)
                    fprintf(out, "%'ju(%ju) -> ",
                            s->litb[i] >> 3, s->litb[i] & 7);
                fprintf(out, "%'ju", s->lits[i]);
                if (s->detail)
                    fprintf(out, " (%.3f)",
                            s->litb[i] / (8 * (double)s->lits[i]));
                putc('\n', out);
                fprintf(out, "  mats: ");
                if (s->detail)
                    fprintf(out, "%'ju(%ju) -> ",
                            s->matb[i] >> 3, s->matb[i] & 7);
                fprintf(out, "%'ju", s->lens[i].sum);
                if (s->detail)
                    fprintf(out, " (%.3f)",
                            s->matb[i] /
                            (8 * (double)(s->lens[i].sum)));
                putc('\n', out);
                show_tally(out, "  bits: ", &s->bits[i]);
            }
            if (i == 2)
                show_tally(out, "  head: ", &s->head[i]);
            if (i > 0) {
                show_tally(out, "  syms: ", &s->syms[i]);
                show_tally(out, "  lens: ", &s->lens[i]);
            }
            show_tally(out, "  data: ", &s->data[i]);
        }
}

// At the start of a new block. type == 0 is stored, 1, is fixed, 2 is dynamic.
// last is 1 if this block is the last block.
static void block(int type, int last, stat_t *s) {
    (void)type;
    (void)last;
    s->lit = 0;
    s->mat = 0;
    tally_init(&s->writ);
}

// head[] is the description of the dynamic block header. Verify the integrity
// of the description, returning 0 on success or 1 on failure. (This does not
// check for the completeness of the resulting Huffman codes.)
static int dynamic(uint8_t *head, stat_t *s) {
    (void)s;
    if (head[0] < 1 || head[0] > 30 ||
        head[1] < 1 || head[1] > 30 ||
        head[2] < 4 || head[2] > 19)
        return 1;
    int i = 3, k = head[2];
    do {
        if (head[i] < 1 || head[i] > 8)
            return 1;
        i++;
    } while (--k);
    k = 256 + head[0] + head[1];
    do {
        if (head[i] < 0)
            return 1;
        else if (head[i] <= 16)
            k--;
        else if (head[i] <= 20)
            k -= head[i] - 14;
        else if (head[i] <= 156)
            k -= head[i] - 18;
        else
            return 1;
        i++;
    } while (k > 0);
    if (k < 0 || head[i] != 0)
        return 1;
    return 0;
}

// Literal byte in 0..255.
static void literal(int lit, stat_t *s) {
#ifdef TEST
    (void)s;
    putc(lit, stdout);
    window[pos++] = lit;
    pos &= 0x7fff;
#else
    (void)lit;
    s->lit++;
    tally(&s->writ, 1);
#endif
}

// Match of previous bytes dist bytes back (1..32768) of length len (3..258).
// type is 1 for a fixed block or 2 for a dynamic block.
static void match(unsigned dist, int len, int type, stat_t *s) {
#ifdef TEST
    (void)s;
    do {
        literal(window[(pos - dist) & 0x7fff]);
    } while (--len);
#else
    (void)dist;
    s->mat++;
    tally(&s->writ, len);
    tally(&s->lens[type], len);
#endif
}

// Update bit counts at the end of a block (if -bb was used).
static void counts(uintmax_t bits, uintmax_t head, uintmax_t litb,
                   uintmax_t matb, int type, int last, stat_t *s) {
    (void)last;
    s->detail = 1;
    tally(&s->bits[type], bits);
    tally(&s->head[type], head);
    s->litb[type] += litb;
    s->matb[type] += matb;
}

// At the end of a deflate block of type type. last is 1 if this is the last
// block in the stream.
static void end(int type, int last, stat_t *s) {
    if (last)
        s->streams++;
    s->lits[type] += s->lit;
    s->mats[type] += s->mat;
    tally(&s->syms[type], s->writ.num);
    tally(&s->data[type], s->writ.sum);
}

// Read a variable-length integer from beg[0..end-beg-1]. Return the integer in
// *val. Return a pointer to what follows the integer, or NULL if the input
// ends before the variable-length integer does.
static inline uint8_t *getvar(uintmax_t *val, uint8_t *beg, uint8_t *end) {
    *val = 0;
    uintmax_t octet;
    int bits = 0;
    for (;;) {
        if (beg >= end)
            return NULL;
        octet = *beg++;
        if ((octet & 0x80) == 0)
            break;
        *val |= (octet & 0x7f) << bits;
        bits += 7;
    }
    *val |= octet << bits;
    return beg;
}

// Declare and get variable-length integer name from next[0..next-end-1],
// updating next.
#define GET(name, next, end) \
    uintmax_t name; \
    next = getvar(&name, next, end); \
    if (next == NULL) \
        return 6;                   // invalid bit counts

// Parse the output of infgen -b or -bb from in. The input is expected to
// represent a series of complete deflate streams. Call block() for the start
// of each new block or each end of a deflate stream, dynamic() with the
// description of a dynamic block header, literal() for each literal byte, and
// match() for each match. Return 0 on success, 1 if the input ended
// prematurely, or >1 if invalid input is encountered.
static int parse(FILE *in, stat_t *s) {
    // State.
    enum {
        TOP,        // next byte is distance high, low literal, or prefix
        LOW,        // next byte is distance low
        LEN,        // next byte is length
        BLK,        // next byte is block type or end, or high literal
        STO,        // next byte is stored leftover bits
        DYN,        // next byte continues the dynamic header description
        BIT,        // next byte continues the bit counts
        UNK,        // next byte continues an unknown zero-terminated block
        END         // next byte is stream-end leftover bits
    } st = TOP;

    // State transitions:
    //   TOP -> TOP, LOW, or BLK
    //     LOW -> LEN
    //       LEN -> TOP
    //     BLK -> TOP, STO, DYN, or END, and may set last
    //       STO -> TOP
    //       DYN -> DYN or TOP
    //       BIT -> BIT or TOP
    //       UNK -> UNK or TOP
    //       END -> TOP to start over, sets last to zero

    uint8_t info[512];              // bytes after 4, 5, or 9..63
    unsigned dist, have;
    int type = -1, last = 0, ch;
    while ((ch = getc(in)) != EOF) {
        // Parse one byte of input at a time, in the context of the state.
        switch (st) {
        case TOP:
            if (ch < 0x80) {
                // First byte of a match.
                dist = ch << 8;
                st = LOW;
            }
            else if (ch == 0xff)
                // Prefix for the next level of instructions.
                st = BLK;
            else {
                // Low literal.
                if (type == -1)
                    return 2;       // literal outside of a block
                literal(ch - 0x80, s);
            }
            break;
        case LOW:
            // Second byte of a match.
            dist |= ch;
            st = LEN;
            break;
        case LEN:
            // Third and final byte of a match.
            match(dist + 1, ch + 3, type, s);
            st = TOP;
            break;
        case BLK:
            have = 0;
            if (ch < 6) {
                // Start of a deflate block.
                if (last)
                    return 3;       // another block after the last block
                if (type != -1)
                    end(type, last, s);
                type = ch >> 1;
                last = ch & 1;
                block(type, last, s);
            }
            else if (type == -1)
                return 4;           // code outside of a block
            if (ch < 2)
                // Stored block.
                st = STO;
            else if (ch < 4)
                // Fixed block.
                st = TOP;
            else if (ch < 6)
                // Dynamic block.
                st = DYN;
            else if (ch == 8)
                // End of stream.
                st = END;
            else if (ch == 9)
                // Bit counts.
                st = BIT;
            else if (ch < 0x40)
                // Unknown instruction with following information. Skip over
                // the zero-terminated information.
                st = UNK;
            else if (ch < 0x7f)
                // Unknown instruction with nothing after it.
                st = TOP;
            else {
                // High literal.
                literal(ch, s);
                st = TOP;
            }
            break;
        case STO:
            // Ignore the leftover bits.
            st = TOP;
            break;
        case DYN:
        case BIT:
        case UNK:
            if (have < sizeof(info))
                info[have++] = ch;
            if (ch == 0) {
                // Zero-terminated block of information is complete.
                if (st == DYN) {
                    // Dynamic block header description.
                    if (info[have - 1] != 0 || dynamic(info, s))
                        return 5;   // invalid description
                }
                else if (st == BIT) {
                    // Bit counts at the end of the block.
                    if (info[have - 1] != 0)
                        return 6;   // invalid bit counts

                    // Get the bit counts. GET() will return 6 on error.
                    uint8_t *next = info;
                    uint8_t *end = next + have;
                    GET(bits, next, end);
                    GET(head, next, end);
                    GET(litb, next, end);
                    GET(matb, next, end);
                    if (next != end - 1)
                        return 6;   // invalid bit counts

                    // Update the statistics.
                    counts(bits + 9, head + 2, litb - 1, matb - 1,
                           type, last, s);
                }
                st = TOP;
            }
            break;
        case END:
            if (last != 1)
                return 7;           // ended in a non-last block
            end(type, last, s);
            type = -1;              // reset to start a new deflate stream
            last = 0;
            st = TOP;
            break;
        }
    }
    return type == -1 ? 0 : 1;  // 1 -> input ended prematurely
}

// Parse infgen -b or -bb output from stdin.
int main(void) {
    // Enable thousands separated by commas with single quote in format.
    setlocale(LC_NUMERIC, "");

    // Process the input and gather statistics.
    stat_t s;
    init(&s);
    int ret = parse(stdin, &s);
    if (ret)
        fprintf(stderr, "** %s input (%d)\n",
                ret == 1 ? "premature end of" : "invalid", ret);
#ifndef TEST
    else
        // Show the statistics.
        show_stats(stdout, &s);
#endif
    return ret != 0;
}