/* The postings half of an inverted index: five representations, one AND.
 * Companion to retrieval.html.
 *
 *   cc -O2 -o postings postings.c && ./postings
 *
 * A Boolean conjunction over an inverted index is the oldest operation in
 * information retrieval and it is, in every implementation, a merge of two
 * sorted integer sequences. The algorithm has been settled since the 1960s.
 * What is not settled -- what every search engine still re-decides -- is how
 * those integers are laid out, and that decision is a memory-hierarchy
 * decision rather than an algorithmic one. All five variants here compute the
 * identical answer over the identical postings:
 *
 *   merge    two uint32 arrays, one linear pass. 4 bytes per posting.
 *   gallop   the same arrays, but the shorter list drives and each of its
 *            postings is located in the longer one by exponential search.
 *            Sublinear in the long list when the lists are skewed.
 *   vbyte    d-gaps under variable-byte coding, decoded end to end. Roughly
 *            one to two bytes per posting, and no way to skip: a compressed
 *            sequence has no random access.
 *   blocks   the same coding cut into 128-posting blocks, each with its last
 *            docid and its byte length in an uncompressed header. Skipping a
 *            block costs a header read, so compression and skipping coexist.
 *            This is, in outline, what Lucene does.
 *   bitmap   one bit per document, AND-ed a word at a time and counted with
 *            popcount. Size depends on the collection, not on the list, so it
 *            is enormous for a rare term and unbeatable for a common one.
 *
 * Every variant is checked against `merge` before it is timed. The counts must
 * agree exactly; a representation that is fast and wrong is not a data point.
 *
 * Two sweeps, because the two halves of the answer live on different axes.
 * The first holds the lists at equal length and grows both, which is the case
 * where bytes moved decide the outcome. The second pins one list at 1024
 * postings and grows the other, which is the case where the ability to skip
 * decides it.
 */

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

#define LGN       23                     /* documents in the collection */
#define NDOCS     (1u << LGN)
#define BLOCK     128                    /* postings per skippable block */
#define TRIALS    5

static double now(void)
{
    struct timespec t;
    clock_gettime(CLOCK_MONOTONIC, &t);
    return t.tv_sec + t.tv_nsec * 1e-9;
}

static unsigned long long rng_state;

static unsigned long long rnd(void)
{
    unsigned long long z = (rng_state += 0x9e3779b97f4a7c15ull);
    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull;
    z = (z ^ (z >> 27)) * 0x94d049bb133111ebull;
    return z ^ (z >> 31);
}

/* ------------------------------------------------------------ structures */

typedef struct {
    unsigned  *doc;      /* raw sorted docids                              */
    unsigned   len;
    unsigned char *vb;   /* d-gaps, variable-byte                          */
    size_t     vblen;
    unsigned char *bb;   /* the same, cut into blocks                      */
    size_t     bblen;
    unsigned  *blast;    /* last docid of each block                       */
    size_t    *boff;     /* byte offset of each block                      */
    unsigned   nblk;
    unsigned long long *bits;   /* one bit per document                    */
} List;

#define WORDS ((NDOCS + 63) / 64)

static unsigned char *vb_put(unsigned char *p, unsigned g)
{
    while (g >= 128) { *p++ = (unsigned char)((g & 127) | 128); g >>= 7; }
    *p++ = (unsigned char)g;
    return p;
}

/* A postings list of `len` docids drawn from NDOCS, strictly increasing.
   Gaps are uniform on [1, 2*mean-1], which gives the right mean spacing
   without the heavy tail of a real term; the tail would change the *bytes*
   each coding needs, and this benchmark is about what those bytes then cost,
   so a distribution with a stable byte count is the honest control. */
static void build(List *L, unsigned len, unsigned long long seed)
{
    rng_state = seed;
    L->doc = malloc((size_t)len * sizeof *L->doc);
    L->bits = calloc(WORDS, sizeof *L->bits);
    if (!L->doc || !L->bits) { fprintf(stderr, "oom\n"); exit(1); }

    unsigned mean = NDOCS / len;
    unsigned d = 0, i = 0;
    for (; i < len; i++) {
        unsigned gap = mean > 1 ? 1 + (unsigned)(rnd() % (2u * mean - 1)) : 1;
        if ((unsigned long long)d + gap >= NDOCS) break;
        d += gap;
        L->doc[i] = d;
        L->bits[d >> 6] |= 1ull << (d & 63);
    }
    L->len = i;

    /* Variable-byte, flat and blocked. Five bytes is the most a 32-bit gap
       can need, so the buffer cannot overrun. */
    L->vb = malloc((size_t)L->len * 5 + 8);
    L->bb = malloc((size_t)L->len * 5 + 8);
    L->nblk = (L->len + BLOCK - 1) / BLOCK;
    L->blast = malloc((size_t)(L->nblk + 1) * sizeof *L->blast);
    L->boff = malloc((size_t)(L->nblk + 1) * sizeof *L->boff);
    if (!L->vb || !L->bb || !L->blast || !L->boff) { fprintf(stderr, "oom\n"); exit(1); }

    unsigned char *p = L->vb;
    unsigned prev = 0;
    for (unsigned k = 0; k < L->len; k++) { p = vb_put(p, L->doc[k] - prev); prev = L->doc[k]; }
    L->vblen = (size_t)(p - L->vb);

    /* Blocked: gaps restart from zero at each block boundary, so a block can
       be decoded without having decoded the one before it. That restart is
       what skipping costs in bytes, and it is why the blocked encoding is
       always a little larger than the flat one. */
    p = L->bb;
    for (unsigned b = 0; b < L->nblk; b++) {
        L->boff[b] = (size_t)(p - L->bb);
        unsigned lo = b * BLOCK, hi = lo + BLOCK < L->len ? lo + BLOCK : L->len;
        unsigned base = 0;
        for (unsigned k = lo; k < hi; k++) { p = vb_put(p, L->doc[k] - base); base = L->doc[k]; }
        L->blast[b] = L->doc[hi - 1];
    }
    L->boff[L->nblk] = (size_t)(p - L->bb);
    L->bblen = (size_t)(p - L->bb);
}

static void freelist(List *L)
{
    free(L->doc); free(L->vb); free(L->bb);
    free(L->blast); free(L->boff); free(L->bits);
}

/* ------------------------------------------------------------ operations */

static unsigned isect_merge(const List *A, const List *B)
{
    unsigned i = 0, j = 0, c = 0;
    while (i < A->len && j < B->len) {
        unsigned a = A->doc[i], b = B->doc[j];
        c += a == b;
        i += a <= b;
        j += b <= a;
    }
    return c;
}

static unsigned isect_gallop(const List *A, const List *B)
{
    const List *S = A->len <= B->len ? A : B;
    const List *L = A->len <= B->len ? B : A;
    unsigned c = 0, lo = 0;
    for (unsigned i = 0; i < S->len && lo < L->len; i++) {
        unsigned x = S->doc[i];
        /* Double the reach until it overshoots, then bisect the last span.
           Starting from the previous answer keeps the whole scan linear in
           the short list, not in the long one. */
        unsigned step = 1, hi;
        while (lo + step < L->len && L->doc[lo + step] < x) { lo += step; step <<= 1; }
        hi = lo + step < L->len ? lo + step : L->len - 1;
        while (lo < hi) {
            unsigned mid = lo + (hi - lo) / 2;
            if (L->doc[mid] < x) lo = mid + 1; else hi = mid;
        }
        c += L->doc[lo] == x;
    }
    return c;
}

#define VB_GET(p, out)                                                        \
    do {                                                                      \
        unsigned shift_ = 0, byte_;                                           \
        (out) = 0;                                                            \
        do { byte_ = *(p)++; (out) |= (unsigned)(byte_ & 127) << shift_;      \
             shift_ += 7; } while (byte_ & 128);                              \
    } while (0)

static unsigned isect_vbyte(const List *A, const List *B)
{
    const unsigned char *pa = A->vb, *pb = B->vb;
    unsigned ia = 0, ib = 0, a = 0, b = 0, c = 0, g;
    if (A->len == 0 || B->len == 0) return 0;
    VB_GET(pa, g); a = g; ia = 1;
    VB_GET(pb, g); b = g; ib = 1;
    for (;;) {
        if (a == b) {
            c++;
            if (ia == A->len || ib == B->len) break;
            VB_GET(pa, g); a += g; ia++;
            VB_GET(pb, g); b += g; ib++;
        } else if (a < b) {
            if (ia == A->len) break;
            VB_GET(pa, g); a += g; ia++;
        } else {
            if (ib == B->len) break;
            VB_GET(pb, g); b += g; ib++;
        }
    }
    return c;
}

/* Decode one block into `out`, returning how many postings it held. */
static unsigned decode_block(const List *L, unsigned b, unsigned *out)
{
    const unsigned char *p = L->bb + L->boff[b];
    unsigned lo = b * BLOCK, hi = lo + BLOCK < L->len ? lo + BLOCK : L->len;
    unsigned base = 0, g;
    for (unsigned k = lo; k < hi; k++) {
        VB_GET(p, g);
        base += g;
        out[k - lo] = base;
    }
    return hi - lo;
}

static unsigned isect_blocks(const List *A, const List *B)
{
    const List *S = A->len <= B->len ? A : B;
    const List *L = A->len <= B->len ? B : A;
    unsigned sbuf[BLOCK], lbuf[BLOCK];
    unsigned c = 0, lb = 0, lcached = (unsigned)-1, lcount = 0, li = 0;

    for (unsigned sb = 0; sb < S->nblk; sb++) {
        unsigned scount = decode_block(S, sb, sbuf);
        for (unsigned k = 0; k < scount; k++) {
            unsigned x = sbuf[k];
            /* Headers only: a block whose last docid is below x cannot hold
               it, and is stepped over without touching its payload. */
            while (lb < L->nblk && L->blast[lb] < x) { lb++; lcached = (unsigned)-1; }
            if (lb == L->nblk) return c;
            if (lcached != lb) { lcount = decode_block(L, lb, lbuf); lcached = lb; li = 0; }
            while (li < lcount && lbuf[li] < x) li++;
            if (li < lcount) c += lbuf[li] == x;
        }
    }
    return c;
}

static unsigned isect_bitmap(const List *A, const List *B)
{
    unsigned c = 0;
    for (size_t w = 0; w < WORDS; w++)
        c += (unsigned)__builtin_popcountll(A->bits[w] & B->bits[w]);
    return c;
}

/* ---------------------------------------------------------------- timing */

static volatile unsigned long long sink;

typedef unsigned (*op)(const List *, const List *);

static double time_op(op f, const List *A, const List *B, long reps)
{
    double best = 1e30;
    for (int t = 0; t < TRIALS; t++) {
        double t0 = now();
        unsigned long long acc = 0;
        for (long r = 0; r < reps; r++) acc += f(A, B);
        double t1 = now();
        sink += acc;
        double ns = (t1 - t0) / (double)reps * 1e9;
        if (ns < best) best = ns;
    }
    return best;
}

static void row(const char *tag, unsigned la, unsigned lb)
{
    List A, B;
    build(&A, la, 0x51ed270b3f1e2d4full);
    build(&B, lb, 0xc0ffee1234567890ull);

    unsigned want = isect_merge(&A, &B);
    if (isect_gallop(&A, &B) != want || isect_vbyte(&A, &B) != want
        || isect_blocks(&A, &B) != want || isect_bitmap(&A, &B) != want) {
        fprintf(stderr, "disagreement at %u x %u: merge=%u gallop=%u vbyte=%u "
                        "blocks=%u bitmap=%u\n", la, lb, want,
                isect_gallop(&A, &B), isect_vbyte(&A, &B),
                isect_blocks(&A, &B), isect_bitmap(&A, &B));
        exit(1);
    }

    /* Work per intersection is linear in the lists for four of the five, so a
       fixed repeat count would spend the whole run on the largest row. */
    long reps = (long)(1u << 26) / (long)(A.len + B.len) + 2;
    double m = time_op(isect_merge, &A, &B, reps);
    double g = time_op(isect_gallop, &A, &B, reps);
    double v = time_op(isect_vbyte, &A, &B, reps);
    double k = time_op(isect_blocks, &A, &B, reps);
    double p = time_op(isect_bitmap, &A, &B, reps);

    /* Bytes per posting, summed over both lists. The raw column is 4 by
       construction rather than by measurement, and is printed so the three
       encodings can be read against each other in one place.

       The blocked column counts its headers. They are the whole point of the
       representation -- a block is skippable because its last docid and its
       extent are stored uncompressed -- so charging only for the payload
       would report the compression of the flat encoding under the name of a
       structure that does not have it. Here a header is `blast` plus `boff`,
       12 bytes on this target; a production encoding packs it smaller, and
       this program is not one. */
    double post = A.len + B.len;
    double head = (double)(sizeof *A.blast + sizeof *A.boff);
    double blocked = (double)(A.bblen + B.bblen) + head * (A.nblk + B.nblk);
    printf("%6s %9u %9u %8u %10.0f %10.0f %10.0f %10.0f %10.0f %8.2f %8.2f %8.2f\n",
           tag, A.len, B.len, want, m, g, v, k, p,
           4.0, (double)(A.vblen + B.vblen) / post, blocked / post);
    fflush(stdout);
    freelist(&A);
    freelist(&B);
}

int main(void)
{
    printf("# collection: %u documents, bitmap %zu KiB per list, block %d postings\n",
           NDOCS, (size_t)WORDS * 8 / 1024, BLOCK);
    printf("%6s %9s %9s %8s %10s %10s %10s %10s %10s %8s %8s %8s\n",
           "sweep", "|A|", "|B|", "|A&B|", "merge", "gallop", "vbyte",
           "blocks", "bitmap", "B/p:raw", "B/p:vb", "B/p:blk");

    /* Equal lengths: the bytes-moved regime. */
    for (unsigned lg = 10; lg <= 22; lg += 2) row("equal", 1u << lg, 1u << lg);
    /* One short list against a growing one: the skipping regime. */
    for (unsigned lg = 10; lg <= 22; lg += 2) row("skew", 1u << 10, 1u << lg);
    return 0;
}
