<commit_before>#include "tdb_queue.h"
#include "tdb_profile.h"
#include "huffman.h"
#include "util.h"

#define MIN(a,b) ((a)>(b)?(b):(a))

struct hnode{
    uint32_t code;
    uint32_t num_bits;
    uint64_t symbol;
    uint64_t weight;
    struct hnode *left;
    struct hnode *right;
};

static void allocate_codewords(struct hnode *node,
                               uint32_t code,
                               uint32_t depth)
{
    if (node == NULL)
        return;
    if (depth < 16 && (node->right || node->left)){
        allocate_codewords(node->left, code, depth + 1);
        allocate_codewords(node->right, code | (1U << depth), depth + 1);
    }else{
        node->code = code;
        node->num_bits = depth;
    }
}

static struct hnode *pop_min_weight(struct hnode *symbols,
                                    uint32_t *num_symbols,
                                    struct tdb_queue *nodes)
{
    const struct hnode *n = (const struct hnode*)tdb_queue_peek(nodes);
    if (!*num_symbols || (n && n->weight < symbols[*num_symbols - 1].weight))
        return tdb_queue_pop(nodes);
    else if (*num_symbols)
        return &symbols[--*num_symbols];
    return NULL;
}

static int huffman_code(struct hnode *symbols, uint32_t num)
{
    struct tdb_queue *nodes = NULL;
    struct hnode *newnodes = NULL;
    uint32_t new_i = 0;

    if (!num)
        return 0;
    if (!(nodes = tdb_queue_new(num * 2)))
        return -1;
    if (!(newnodes = malloc(num * sizeof(struct hnode)))){
        tdb_queue_free(nodes);
        return -1;
    }

    /* construct the huffman tree bottom up */
    while (num || tdb_queue_length(nodes) > 1){
        struct hnode *new = &newnodes[new_i++];
        new->left = pop_min_weight(symbols, &num, nodes);
        new->right = pop_min_weight(symbols, &num, nodes);
        new->weight = (new->left ? new->left->weight: 0) +
                      (new->right ? new->right->weight: 0);
        tdb_queue_push(nodes, new);
    }
    /* allocate codewords top down (depth-first) */
    allocate_codewords(tdb_queue_pop(nodes), 0, 0);
    free(newnodes);
    tdb_queue_free(nodes);
    return 0;
}


static uint32_t sort_symbols(const Pvoid_t freqs,
                             uint64_t *totalfreq,
                             struct hnode *book)
{
    Word_t i;
    Word_t num_symbols;
    struct sortpair *pairs = sort_judyl(freqs, &num_symbols);

    *totalfreq = 0;
    for (i = 0; i < num_symbols; i++)
        *totalfreq += pairs[i].value;

    num_symbols = MIN(HUFF_CODEBOOK_SIZE, num_symbols);
    for (i = 0; i < num_symbols; i++){
        book[i].symbol = pairs[i].key;
        book[i].weight = pairs[i].value;
    }

    free(pairs);
    return (uint32_t)num_symbols;
}

static void print_codeword(const struct hnode *node)
{
    uint32_t j;
    for (j = 0; j < node->num_bits; j++)
        fprintf(stderr, "%u", (node->code & (1U << j) ? 1: 0));
}

static void output_stats(const struct hnode *book,
                         uint32_t num_symbols,
                         uint64_t tot)
{
    fprintf(stderr, "#codewords: %u\n", num_symbols);
    uint64_t cum = 0;
    uint32_t i;
    fprintf(stderr, "index) gramtype [field value] freq prob cum\n");

    for (i = 0; i < num_symbols; i++){

        long long unsigned int sym = book[i].symbol;
        long long unsigned int sym2 = sym >> 32;
        uint64_t f = book[i].weight;
        cum += f;

        fprintf(stderr, "%u) ", i);
        if (sym2 & 255){
            fprintf(stderr,
                    "bi [%llu %llu | %llu %llu] ",
                    sym & 255,
                    (sym >> 8) & ((1 << 24) - 1),
                    sym2 & 255,
                    sym2 >> 8);
        }else
            fprintf(stderr, "uni [%llu %llu] ", sym & 255, sym >> 8);

        fprintf(stderr, "%lu %2.3f %2.3f | ",
                f,
                100. * (double)f / (double)tot,
                100. * (double)cum / (double)tot);
        print_codeword(&book[i]);
        fprintf(stderr, "\n");
    }
}

static Pvoid_t make_codemap(struct hnode *nodes, uint32_t num_symbols)
{
    Pvoid_t codemap = NULL;

    uint32_t i = num_symbols;
    while (i--){
        Word_t *ptr;

        if (nodes[i].num_bits){
            JLI(ptr, codemap, nodes[i].symbol);
            *ptr = nodes[i].code | (nodes[i].num_bits << 16);
        }
    }
    return codemap;
}

struct field_stats *huff_field_stats(const uint64_t *field_cardinalities,
                                     uint32_t num_fields,
                                     uint32_t max_timestamp_delta)
{
    uint32_t i;
    struct field_stats *fstats;

    if (!(fstats = malloc(sizeof(struct field_stats) + num_fields * 4)))
        return NULL;

    fstats->field_id_bits = bits_needed(num_fields);
    fstats->field_bits[0] = bits_needed(max_timestamp_delta);
    for (i = 0; i < num_fields - 1; i++){
        fstats->field_bits[i + 1] = bits_needed(field_cardinalities[i]);
    }
    return fstats;
}

Pvoid_t huff_create_codemap(const Pvoid_t key_freqs)
{
    struct hnode *nodes;
    Pvoid_t codemap;
    uint64_t total_freq;
    uint32_t num_symbols;
    TDB_TIMER_DEF

    if (!(nodes = calloc(HUFF_CODEBOOK_SIZE, sizeof(struct hnode))))
        DIE("Could not allocate huffman codebook");

    TDB_TIMER_START
    num_symbols = sort_symbols(key_freqs, &total_freq, nodes);
    TDB_TIMER_END("huffman/sort_symbols")

    TDB_TIMER_START
    huffman_code(nodes, num_symbols);
    TDB_TIMER_END("huffman/huffman_code")

    if (getenv("DEBUG_HUFFMAN"))
        output_stats(nodes, num_symbols, total_freq);

    TDB_TIMER_START
    codemap = make_codemap(nodes, num_symbols);
    TDB_TIMER_END("huffman/make_codemap")

    free(nodes);
    return codemap;
}

static inline void encode_gram(const Pvoid_t codemap,
                               uint64_t gram,
                               char *buf,
                               uint64_t *offs,
                               const struct field_stats *fstats)
{
    const uint32_t field_id = gram & 255;
    const uint32_t value = (uint32_t)(gram >> 8);
    const uint32_t is_bigram = (gram >> 32) & 255;
    const uint32_t literal_bits = 1 + fstats->field_id_bits +
                                  fstats->field_bits[field_id];

    uint32_t huff_code, huff_bits;
    Word_t *ptr;

    JLG(ptr, codemap, gram);
    if (ptr){
        /* codeword: prefix code by an up bit */
        huff_code = 1 | (HUFF_CODE(*ptr) << 1);
        huff_bits = HUFF_BITS(*ptr) + 1;
    }
<commit_msg>enable debugging functions in huffman.c only with -DTDB_DEBUG_HUFFMAN<commit_after>#include "tdb_queue.h"
#include "tdb_profile.h"
#include "huffman.h"
#include "util.h"

#define MIN(a,b) ((a)>(b)?(b):(a))

struct hnode{
    uint32_t code;
    uint32_t num_bits;
    uint64_t symbol;
    uint64_t weight;
    struct hnode *left;
    struct hnode *right;
};

static void allocate_codewords(struct hnode *node,
                               uint32_t code,
                               uint32_t depth)
{
    if (node == NULL)
        return;
    if (depth < 16 && (node->right || node->left)){
        allocate_codewords(node->left, code, depth + 1);
        allocate_codewords(node->right, code | (1U << depth), depth + 1);
    }else{
        node->code = code;
        node->num_bits = depth;
    }
}

static struct hnode *pop_min_weight(struct hnode *symbols,
                                    uint32_t *num_symbols,
                                    struct tdb_queue *nodes)
{
    const struct hnode *n = (const struct hnode*)tdb_queue_peek(nodes);
    if (!*num_symbols || (n && n->weight < symbols[*num_symbols - 1].weight))
        return tdb_queue_pop(nodes);
    else if (*num_symbols)
        return &symbols[--*num_symbols];
    return NULL;
}

static int huffman_code(struct hnode *symbols, uint32_t num)
{
    struct tdb_queue *nodes = NULL;
    struct hnode *newnodes = NULL;
    uint32_t new_i = 0;

    if (!num)
        return 0;
    if (!(nodes = tdb_queue_new(num * 2)))
        return -1;
    if (!(newnodes = malloc(num * sizeof(struct hnode)))){
        tdb_queue_free(nodes);
        return -1;
    }

    /* construct the huffman tree bottom up */
    while (num || tdb_queue_length(nodes) > 1){
        struct hnode *new = &newnodes[new_i++];
        new->left = pop_min_weight(symbols, &num, nodes);
        new->right = pop_min_weight(symbols, &num, nodes);
        new->weight = (new->left ? new->left->weight: 0) +
                      (new->right ? new->right->weight: 0);
        tdb_queue_push(nodes, new);
    }
    /* allocate codewords top down (depth-first) */
    allocate_codewords(tdb_queue_pop(nodes), 0, 0);
    free(newnodes);
    tdb_queue_free(nodes);
    return 0;
}


static uint32_t sort_symbols(const Pvoid_t freqs,
                             uint64_t *totalfreq,
                             struct hnode *book)
{
    Word_t i;
    Word_t num_symbols;
    struct sortpair *pairs = sort_judyl(freqs, &num_symbols);

    *totalfreq = 0;
    for (i = 0; i < num_symbols; i++)
        *totalfreq += pairs[i].value;

    num_symbols = MIN(HUFF_CODEBOOK_SIZE, num_symbols);
    for (i = 0; i < num_symbols; i++){
        book[i].symbol = pairs[i].key;
        book[i].weight = pairs[i].value;
    }

    free(pairs);
    return (uint32_t)num_symbols;
}

#ifdef TDB_DEBUG_HUFFMAN
static void print_codeword(const struct hnode *node)
{
    uint32_t j;
    for (j = 0; j < node->num_bits; j++)
        fprintf(stderr, "%u", (node->code & (1U << j) ? 1: 0));
}

static void output_stats(const struct hnode *book,
                         uint32_t num_symbols,
                         uint64_t tot)
{
    fprintf(stderr, "#codewords: %u\n", num_symbols);
    uint64_t cum = 0;
    uint32_t i;
    fprintf(stderr, "index) gramtype [field value] freq prob cum\n");

    for (i = 0; i < num_symbols; i++){

        long long unsigned int sym = book[i].symbol;
        long long unsigned int sym2 = sym >> 32;
        uint64_t f = book[i].weight;
        cum += f;

        fprintf(stderr, "%u) ", i);
        if (sym2 & 255){
            fprintf(stderr,
                    "bi [%llu %llu | %llu %llu] ",
                    sym & 255,
                    (sym >> 8) & ((1 << 24) - 1),
                    sym2 & 255,
                    sym2 >> 8);
        }else
            fprintf(stderr, "uni [%llu %llu] ", sym & 255, sym >> 8);

        fprintf(stderr, "%lu %2.3f %2.3f | ",
                f,
                100. * (double)f / (double)tot,
                100. * (double)cum / (double)tot);
        print_codeword(&book[i]);
        fprintf(stderr, "\n");
    }
}
#endif

static Pvoid_t make_codemap(struct hnode *nodes, uint32_t num_symbols)
{
    Pvoid_t codemap = NULL;

    uint32_t i = num_symbols;
    while (i--){
        Word_t *ptr;

        if (nodes[i].num_bits){
            JLI(ptr, codemap, nodes[i].symbol);
            *ptr = nodes[i].code | (nodes[i].num_bits << 16);
        }
    }
    return codemap;
}

struct field_stats *huff_field_stats(const uint64_t *field_cardinalities,
                                     uint32_t num_fields,
                                     uint32_t max_timestamp_delta)
{
    uint32_t i;
    struct field_stats *fstats;

    if (!(fstats = malloc(sizeof(struct field_stats) + num_fields * 4)))
        return NULL;

    fstats->field_id_bits = bits_needed(num_fields);
    fstats->field_bits[0] = bits_needed(max_timestamp_delta);
    for (i = 0; i < num_fields - 1; i++){
        fstats->field_bits[i + 1] = bits_needed(field_cardinalities[i]);
    }
    return fstats;
}

Pvoid_t huff_create_codemap(const Pvoid_t key_freqs)
{
    struct hnode *nodes;
    Pvoid_t codemap;
    uint64_t total_freq;
    uint32_t num_symbols;
    TDB_TIMER_DEF

    if (!(nodes = calloc(HUFF_CODEBOOK_SIZE, sizeof(struct hnode))))
        DIE("Could not allocate huffman codebook");

    TDB_TIMER_START
    num_symbols = sort_symbols(key_freqs, &total_freq, nodes);
    TDB_TIMER_END("huffman/sort_symbols")

    TDB_TIMER_START
    huffman_code(nodes, num_symbols);
    TDB_TIMER_END("huffman/huffman_code")

#ifdef TDB_DEBUG_HUFFMAN
    if (getenv("TDB_DEBUG_HUFFMAN"))
        output_stats(nodes, num_symbols, total_freq);
#endif

    TDB_TIMER_START
    codemap = make_codemap(nodes, num_symbols);
    TDB_TIMER_END("huffman/make_codemap")

    free(nodes);
    return codemap;
}

static inline void encode_gram(const Pvoid_t codemap,
                               uint64_t gram,
                               char *buf,
                               uint64_t *offs,
                               const struct field_stats *fstats)
{
    const uint32_t field_id = gram & 255;
    const uint32_t value = (uint32_t)(gram >> 8);
    const uint32_t is_bigram = (gram >> 32) & 255;
    const uint32_t literal_bits = 1 + fstats->field_id_bits +
                                  fstats->field_bits[field_id];

    uint32_t huff_code, huff_bits;
    Word_t *ptr;

    JLG(ptr, codemap, gram);
    if (ptr){
        /* codeword: prefix code by an up bit */
        huff_code = 1 | (HUFF_CODE(*ptr) << 1);
        huff_bits = HUFF_BITS(*ptr) + 1;
    }

{'fraction_non_alphanumeric': 0.12391667933708378, 'fraction_numerical': 0.022654705792914703, 'mean_word_length': 2.1011315417256013, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 77, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5645324', 'n_tokens_mistral': 4851, 'n_tokens_neox': 4565, 'n_words': 1033}