quickbloom.c

// quickbloom.c -- the quickbloom SBBF implementation.
//
// One file, one ABI: the fastest single-key SBBF probe kernel we've
// measured, with two backends (AVX2 on x86_64, NEON on aarch64) sharing
// a single kernel body via a small SIMD shim. See the README's
// Performance section for the comparison against other Bloom designs.
//
// Split Block Bloom Filter (Apache Parquet spec):
//   - 256-bit blocks (8 x 32-bit words), K=8 (one bit per word)
//   - wyhash-style hash (128-bit multiply + xor-fold) on 16-byte
//     fast path, fasthash64 fallback for variable-length keys
//   - Power-of-2 nblocks with the Parquet-spec fastrange block
//     index, which on pow2 reduces to a single right shift
//   - SIMD mask compute: AVX2 (vpmullo + vpsrli + vpsllv) on x86_64,
//     NEON (vmulq + vshrq + vshlq) on aarch64. The bitset layout is
//     bit-identical across arches.
//   - 4-way unrolled bulk paths
//   - bulk_insert uses sequential load+OR+store per key so hardware
//     store-to-load forwarding handles aliased blocks correctly
//
// Supported: x86_64 with AVX2 + BMI2, or aarch64 with NEON (baseline
// ARMv8-A — Apple Silicon, AWS Graviton 2/3/4, etc.).

#if defined(__x86_64__) || defined(_M_X64)
#  define QB_ARCH_X86_64 1
#elif defined(__aarch64__) || defined(_M_ARM64)
#  define QB_ARCH_AARCH64 1
#else
#  error "quickbloom requires x86_64 (AVX2) or aarch64 (NEON)"
#endif

#if defined(QB_ARCH_X86_64) && !defined(__AVX2__)
#  error "quickbloom on x86_64 requires AVX2 — compile with at least -mavx2 -mbmi2"
#endif

// K_HASHES 8 -- documented for bench tooling that scans for this.
#define K_HASHES 8

#include <assert.h>
#include <math.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#if defined(QB_ARCH_X86_64)
#  include <immintrin.h>
#elif defined(QB_ARCH_AARCH64)
#  include <arm_neon.h>
#endif

static inline uint64_t hash16(const void* data) {
    uint64_t a, b;
    memcpy(&a, data, 8);
    memcpy(&b, (const uint8_t*)data + 8, 8);
    // XOR with mixing constants so structured inputs (notably, a key
    // with a zero half — e.g. a 64-bit ID padded into a 16-byte
    // buffer) don't degenerate to multiplication by zero. Constants
    // are SipHash's initial state words; any pair of high-entropy
    // primes works.
    a ^= 0x736f6d6570736575ULL;
    b ^= 0x646f72616e646f6dULL;
    __uint128_t r = (__uint128_t)a * b;
    return (uint64_t)r ^ (uint64_t)(r >> 64);
}

static inline uint64_t fasthash64_var(const void* data, size_t len) {
    const uint64_t M = 0x880355f21e6d1965ULL;
    const uint8_t* p = (const uint8_t*)data;
    uint64_t h = (uint64_t)len * M;
    size_t nblocks = len / 8;
    for (size_t i = 0; i < nblocks; i++) {
        uint64_t k;
        memcpy(&k, p + i * 8, 8);
        k *= M; k ^= k >> 23; k *= M;
        h ^= k; h *= M;
    }
    const uint8_t* tail = p + nblocks * 8;
    uint64_t t = 0;
    switch (len & 7) {
        case 7: t ^= (uint64_t)tail[6] << 48; __attribute__((fallthrough));
        case 6: t ^= (uint64_t)tail[5] << 40; __attribute__((fallthrough));
        case 5: t ^= (uint64_t)tail[4] << 32; __attribute__((fallthrough));
        case 4: t ^= (uint64_t)tail[3] << 24; __attribute__((fallthrough));
        case 3: t ^= (uint64_t)tail[2] << 16; __attribute__((fallthrough));
        case 2: t ^= (uint64_t)tail[1] << 8;  __attribute__((fallthrough));
        case 1: t ^= (uint64_t)tail[0];
                t *= M; t ^= t >> 23; t *= M; h ^= t; h *= M;
    }
    h ^= h >> 23; h *= 0x2127599bf4325c37ULL; h ^= h >> 47;
    return h;
}

static inline uint64_t bloom_hash(const void* data, size_t len) {
    if (__builtin_expect(len == 16, 1)) return hash16(data);
    return fasthash64_var(data, len);
}

#define SBBF_BLOCK_BYTES 32

// Spec-mandated salt vector. Lane 0 .. lane 7. The Parquet spec
// fixes this so the on-disk bitset is bit-identical across
// implementations (arrow-cpp, arrow-rs, Velox, DuckDB, Impala).
#define QB_SALT_L0 0x47b6137bu
#define QB_SALT_L1 0x44974d91u
#define QB_SALT_L2 0x8824ad5bu
#define QB_SALT_L3 0xa2b7289du
#define QB_SALT_L4 0x705495c7u
#define QB_SALT_L5 0x2df1424bu
#define QB_SALT_L6 0x9efc4947u
#define QB_SALT_L7 0x5c6bfb31u

typedef struct {
    uint32_t* bits;
    size_t   nblocks_mask;   // nblocks - 1 (nblocks is a power of two)
    uint32_t idx_shift;      // 32 - log2(nblocks); see block_for()
} bloom_t;

// ---------------------------------------------------------------
// SIMD shim. Each arch defines:
//   qb_block_t                opaque 256-bit block vector
//   qb_mask_for(h32)          per-lane bit mask from the 32-bit hash
//   qb_block_load(p)          load 32 aligned bytes from p
//   qb_block_store(p, v)      store 32 aligned bytes to p
//   qb_block_or(a, b)         bitwise OR
//   qb_testc(cur, m)          1 iff every set bit of m is set in cur
//                             (mirrors x86 _mm256_testc semantics)

#if defined(QB_ARCH_X86_64)

typedef __m256i qb_block_t;

static inline qb_block_t qb_mask_for(uint32_t h32) {
    const __m256i salt = _mm256_set_epi32(
        (int)QB_SALT_L7, (int)QB_SALT_L6, (int)QB_SALT_L5, (int)QB_SALT_L4,
        (int)QB_SALT_L3, (int)QB_SALT_L2, (int)QB_SALT_L1, (int)QB_SALT_L0);
    const __m256i hbcast = _mm256_set1_epi32((int)h32);
    const __m256i prod = _mm256_mullo_epi32(hbcast, salt);
    const __m256i shift = _mm256_srli_epi32(prod, 27);
    const __m256i ones = _mm256_set1_epi32(1);
    return _mm256_sllv_epi32(ones, shift);
}

static inline qb_block_t qb_block_load(const uint32_t* p) {
    return _mm256_load_si256((const __m256i*)p);
}

static inline void qb_block_store(uint32_t* p, qb_block_t v) {
    _mm256_store_si256((__m256i*)p, v);
}

static inline qb_block_t qb_block_or(qb_block_t a, qb_block_t b) {
    return _mm256_or_si256(a, b);
}

static inline int qb_testc(qb_block_t cur, qb_block_t m) {
    return _mm256_testc_si256(cur, m);
}

#elif defined(QB_ARCH_AARCH64)

// 256-bit block carried as two 128-bit NEON vectors. The struct is
// returned by value through the helper functions; clang/gcc keep
// both halves in q-registers so there's no spill.
typedef struct { uint32x4_t lo, hi; } qb_block_t;

static inline qb_block_t qb_mask_for(uint32_t h32) {
    static const uint32_t kSalt[8] = {
        QB_SALT_L0, QB_SALT_L1, QB_SALT_L2, QB_SALT_L3,
        QB_SALT_L4, QB_SALT_L5, QB_SALT_L6, QB_SALT_L7,
    };
    const uint32x4_t salt_lo = vld1q_u32(&kSalt[0]);
    const uint32x4_t salt_hi = vld1q_u32(&kSalt[4]);
    const uint32x4_t hb      = vdupq_n_u32(h32);
    const uint32x4_t prod_lo = vmulq_u32(hb, salt_lo);
    const uint32x4_t prod_hi = vmulq_u32(hb, salt_hi);
    // Upper 5 bits select the bit position within a 32-bit word.
    const uint32x4_t shift_lo = vshrq_n_u32(prod_lo, 27);
    const uint32x4_t shift_hi = vshrq_n_u32(prod_hi, 27);
    const uint32x4_t ones     = vdupq_n_u32(1);
    // vshlq_u32 takes a signed shift count per lane; reinterpret keeps
    // the bit pattern (values are in [0,31] so the sign bit is clear).
    qb_block_t out;
    out.lo = vshlq_u32(ones, vreinterpretq_s32_u32(shift_lo));
    out.hi = vshlq_u32(ones, vreinterpretq_s32_u32(shift_hi));
    return out;
}

static inline qb_block_t qb_block_load(const uint32_t* p) {
    qb_block_t out;
    out.lo = vld1q_u32(p);
    out.hi = vld1q_u32(p + 4);
    return out;
}

static inline void qb_block_store(uint32_t* p, qb_block_t v) {
    vst1q_u32(p,     v.lo);
    vst1q_u32(p + 4, v.hi);
}

static inline qb_block_t qb_block_or(qb_block_t a, qb_block_t b) {
    qb_block_t out;
    out.lo = vorrq_u32(a.lo, b.lo);
    out.hi = vorrq_u32(a.hi, b.hi);
    return out;
}

static inline int qb_testc(qb_block_t cur, qb_block_t m) {
    // x86 _mm256_testc(a, b) := ((~a) & b) == 0. NEON vbicq(b, a) is
    // b AND NOT a, which is the same thing. Combine the two halves
    // and reduce to a single zero check with vmaxvq.
    uint32x4_t lo = vbicq_u32(m.lo, cur.lo);
    uint32x4_t hi = vbicq_u32(m.hi, cur.hi);
    return vmaxvq_u32(vorrq_u32(lo, hi)) == 0;
}

#endif

static inline uint32_t* block_for(const bloom_t* b, uint64_t h64) {
    // Parquet spec: idx = ((h >> 32) * num_blocks) >> 32  (fastrange).
    // We require num_blocks to be a power of two, so the multiply-and-
    // shift collapses to a single right shift of the upper 32-bit half
    // of the hash. Cast through uint64_t so an idx_shift of 32 (which
    // occurs when nblocks == 1) is well-defined and yields 0 instead
    // of being UB.
    uint32_t h32 = (uint32_t)(h64 >> 32);
    size_t idx = (size_t)(((uint64_t)h32) >> b->idx_shift);
    return b->bits + idx * 8;
}

// Helper: log2(n) for n a power of two, used to compute idx_shift.
static inline uint32_t log2_pow2(size_t n) {
    uint32_t k = 0;
    while ((((size_t)1) << k) < n) k++;
    return k;
}

// Hard cap on nblocks. Block index is computed from the upper 32
// bits of the hash, so any nblocks > 2^32 leaves bits unreachable
// (silent quality degradation). We cap well below that — 2^31
// blocks = 64 GiB filter, more than any realistic deployment — and
// keep block_for's right-shift count in [1, 32], avoiding shift-by-
// out-of-range UB on absurd inputs.
#define QB_MAX_NBLOCKS ((size_t)1 << 31)

// qb_new returns NULL on allocation failure or if nbits is so large
// that the resulting filter would exceed QB_MAX_NBLOCKS. The
// returned pointer must be released with qb_free.
void* qb_new(size_t nbits) {
    size_t nblocks = (nbits + 255) / 256;
    if (nblocks == 0) nblocks = 1;
    if (nblocks > QB_MAX_NBLOCKS) return NULL;
    size_t pow2 = 1;
    while (pow2 < nblocks) pow2 <<= 1;
    nblocks = pow2;
    bloom_t* b = (bloom_t*)malloc(sizeof(bloom_t));
    if (!b) return NULL;
    size_t nbytes = nblocks * SBBF_BLOCK_BYTES;
    // 32-byte alignment satisfies AVX2 _mm256_load and is harmless on
    // NEON (which accepts any natural alignment). aligned_alloc is C11
    // and requires nbytes to be a multiple of the alignment — nbytes =
    // nblocks * 32 satisfies that by construction. Failing fast is
    // safer than a calloc fallback that can't satisfy the contract.
    void* mem = aligned_alloc(32, nbytes);
    if (!mem) {
        free(b);
        return NULL;
    }
    memset(mem, 0, nbytes);
    b->nblocks_mask = nblocks - 1;
    b->idx_shift = 32 - log2_pow2(nblocks);
    b->bits = (uint32_t*)mem;
    return b;
}

void qb_free(void* p) {
    if (!p) return;
    bloom_t* b = (bloom_t*)p;
    free(b->bits);
    free(b);
}

void qb_insert(void* p, const void* key, size_t len) {
    bloom_t* b = (bloom_t*)p;
    uint64_t h = bloom_hash(key, len);
    uint32_t* blk = block_for(b, h);
    qb_block_t m = qb_mask_for((uint32_t)h);
    qb_block_t cur = qb_block_load(blk);
    qb_block_store(blk, qb_block_or(cur, m));
}

int qb_contains(void* p, const void* key, size_t len) {
    bloom_t* b = (bloom_t*)p;
    uint64_t h = bloom_hash(key, len);
    uint32_t* blk = block_for(b, h);
    qb_block_t m = qb_mask_for((uint32_t)h);
    qb_block_t cur = qb_block_load(blk);
    return qb_testc(cur, m);
}

#define APPLY(P, M) do {                                          \
    qb_block_t c_ = qb_block_load(P);                             \
    qb_block_store((P), qb_block_or(c_, (M)));                    \
} while (0)

void qb_insert_bulk(void* p, const uint8_t* keys, size_t klen, size_t n) {
    bloom_t* b = (bloom_t*)p;
    size_t i = 0;
    if (klen == 16) {
        for (; i + 4 <= n; i += 4) {
            uint64_t h0 = hash16(keys + (i+0)*16);
            uint64_t h1 = hash16(keys + (i+1)*16);
            uint64_t h2 = hash16(keys + (i+2)*16);
            uint64_t h3 = hash16(keys + (i+3)*16);
            uint32_t* p0 = block_for(b, h0);
            uint32_t* p1 = block_for(b, h1);
            uint32_t* p2 = block_for(b, h2);
            uint32_t* p3 = block_for(b, h3);
            qb_block_t m0 = qb_mask_for((uint32_t)h0);
            qb_block_t m1 = qb_mask_for((uint32_t)h1);
            qb_block_t m2 = qb_mask_for((uint32_t)h2);
            qb_block_t m3 = qb_mask_for((uint32_t)h3);
            APPLY(p0, m0); APPLY(p1, m1); APPLY(p2, m2); APPLY(p3, m3);
        }
    }
    for (; i < n; i++) qb_insert(b, keys + i * klen, klen);
}

size_t qb_contains_bulk(void* p, const uint8_t* keys, size_t klen, size_t n) {
    bloom_t* b = (bloom_t*)p;
    size_t hits = 0;
    size_t i = 0;
    if (klen == 16) {
        for (; i + 4 <= n; i += 4) {
            uint64_t h0 = hash16(keys + (i+0)*16);
            uint64_t h1 = hash16(keys + (i+1)*16);
            uint64_t h2 = hash16(keys + (i+2)*16);
            uint64_t h3 = hash16(keys + (i+3)*16);
            uint32_t* p0 = block_for(b, h0);
            uint32_t* p1 = block_for(b, h1);
            uint32_t* p2 = block_for(b, h2);
            uint32_t* p3 = block_for(b, h3);
            qb_block_t m0 = qb_mask_for((uint32_t)h0);
            qb_block_t m1 = qb_mask_for((uint32_t)h1);
            qb_block_t m2 = qb_mask_for((uint32_t)h2);
            qb_block_t m3 = qb_mask_for((uint32_t)h3);
            qb_block_t c0 = qb_block_load(p0);
            qb_block_t c1 = qb_block_load(p1);
            qb_block_t c2 = qb_block_load(p2);
            qb_block_t c3 = qb_block_load(p3);
            hits += qb_testc(c0, m0);
            hits += qb_testc(c1, m1);
            hits += qb_testc(c2, m2);
            hits += qb_testc(c3, m3);
        }
    }
    for (; i < n; i++) {
        if (qb_contains(b, keys + i * klen, klen)) hits++;
    }
    return hits;
}

// Bench-fair single-key loop variants. These are pure for-loops over
// the single-key kernel — no 4-way unroll, no parallel block loads —
// so they're directly comparable to other Bloom libraries' bulk APIs
// (which are also for-loops over their single-key kernels). The
// bench harness aliases bloom_*_bulk to these for apples-to-apples
// cross-candidate comparison; user code that wants the throughput
// boost should call qb_*_bulk directly.
void qb_insert_loop(void* p, const uint8_t* keys, size_t klen, size_t n) {
    for (size_t i = 0; i < n; i++) qb_insert(p, keys + i * klen, klen);
}

size_t qb_contains_loop(void* p, const uint8_t* keys, size_t klen, size_t n) {
    size_t hits = 0;
    for (size_t i = 0; i < n; i++) {
        if (qb_contains(p, keys + i * klen, klen)) hits++;
    }
    return hits;
}

void qb_insert_prehash(void* p, uint64_t h) {
    bloom_t* b = (bloom_t*)p;
    uint32_t* blk = block_for(b, h);
    qb_block_t m = qb_mask_for((uint32_t)h);
    qb_block_t cur = qb_block_load(blk);
    qb_block_store(blk, qb_block_or(cur, m));
}

int qb_contains_prehash(void* p, uint64_t h) {
    bloom_t* b = (bloom_t*)p;
    uint32_t* blk = block_for(b, h);
    qb_block_t m = qb_mask_for((uint32_t)h);
    qb_block_t cur = qb_block_load(blk);
    return qb_testc(cur, m);
}

void qb_insert_prehash_bulk(void* p, const uint64_t* hashes, size_t n) {
    bloom_t* b = (bloom_t*)p;
    size_t i = 0;
    for (; i + 4 <= n; i += 4) {
        uint32_t* p0 = block_for(b, hashes[i+0]);
        uint32_t* p1 = block_for(b, hashes[i+1]);
        uint32_t* p2 = block_for(b, hashes[i+2]);
        uint32_t* p3 = block_for(b, hashes[i+3]);
        qb_block_t m0 = qb_mask_for((uint32_t)hashes[i+0]);
        qb_block_t m1 = qb_mask_for((uint32_t)hashes[i+1]);
        qb_block_t m2 = qb_mask_for((uint32_t)hashes[i+2]);
        qb_block_t m3 = qb_mask_for((uint32_t)hashes[i+3]);
        APPLY(p0, m0); APPLY(p1, m1); APPLY(p2, m2); APPLY(p3, m3);
    }
    for (; i < n; i++) qb_insert_prehash(p, hashes[i]);
}

size_t qb_contains_prehash_bulk(void* p, const uint64_t* hashes, size_t n) {
    bloom_t* b = (bloom_t*)p;
    size_t hits = 0;
    size_t i = 0;
    for (; i + 4 <= n; i += 4) {
        uint32_t* p0 = block_for(b, hashes[i+0]);
        uint32_t* p1 = block_for(b, hashes[i+1]);
        uint32_t* p2 = block_for(b, hashes[i+2]);
        uint32_t* p3 = block_for(b, hashes[i+3]);
        qb_block_t m0 = qb_mask_for((uint32_t)hashes[i+0]);
        qb_block_t m1 = qb_mask_for((uint32_t)hashes[i+1]);
        qb_block_t m2 = qb_mask_for((uint32_t)hashes[i+2]);
        qb_block_t m3 = qb_mask_for((uint32_t)hashes[i+3]);
        qb_block_t c0 = qb_block_load(p0);
        qb_block_t c1 = qb_block_load(p1);
        qb_block_t c2 = qb_block_load(p2);
        qb_block_t c3 = qb_block_load(p3);
        hits += qb_testc(c0, m0);
        hits += qb_testc(c1, m1);
        hits += qb_testc(c2, m2);
        hits += qb_testc(c3, m3);
    }
    for (; i < n; i++) {
        if (qb_contains_prehash(p, hashes[i])) hits++;
    }
    return hits;
}

// Bench-fair single-key loop variants (prehash). See qb_insert_loop
// / qb_contains_loop above for rationale — these are the for-loop
// equivalents over qb_*_prehash, used by the bench harness to
// compare against other Bloom libraries' for-loop bulk APIs.
void qb_insert_prehash_loop(void* p, const uint64_t* hashes, size_t n) {
    for (size_t i = 0; i < n; i++) qb_insert_prehash(p, hashes[i]);
}

size_t qb_contains_prehash_loop(void* p, const uint64_t* hashes, size_t n) {
    size_t hits = 0;
    for (size_t i = 0; i < n; i++) {
        if (qb_contains_prehash(p, hashes[i])) hits++;
    }
    return hits;
}

// ---------------------------------------------------------------
// Serialization. The on-disk layout is identical to the in-memory
// one (nblocks 32-byte blocks of 8 little-endian uint32 lanes), which
// is also the Parquet spec layout. Both x86_64 and aarch64 are
// little-endian in our supported configurations, so memcpy is
// sufficient in both directions; a future big-endian port would need
// per-lane byteswaps here.

// Maximum nbytes qb_deserialize will accept. Larger inputs are
// rejected with NULL so an attacker-controlled Parquet bloom header
// can't trigger a multi-gigabyte aligned_alloc. Override at compile
// time with -DQB_DESERIALIZE_MAX_BYTES=... if you need a larger
// ceiling. Default 1 GiB.
#ifndef QB_DESERIALIZE_MAX_BYTES
#define QB_DESERIALIZE_MAX_BYTES ((size_t)1 << 30)
#endif

size_t qb_serialized_size(void* p) {
    assert(p != NULL && "qb_serialized_size: filter pointer must be non-NULL");
    if (!p) return 0;
    bloom_t* b = (bloom_t*)p;
    return (b->nblocks_mask + 1) * SBBF_BLOCK_BYTES;
}

void qb_serialize(void* p, uint8_t* dst) {
    assert(p   != NULL && "qb_serialize: filter pointer must be non-NULL");
    assert(dst != NULL && "qb_serialize: dst buffer must be non-NULL");
    if (!p || !dst) return;
    bloom_t* b = (bloom_t*)p;
    size_t nbytes = (b->nblocks_mask + 1) * SBBF_BLOCK_BYTES;
    memcpy(dst, b->bits, nbytes);
}

void* qb_deserialize(const uint8_t* bytes, size_t nbytes) {
    if (!bytes) return NULL;
    if (nbytes == 0 || nbytes > QB_DESERIALIZE_MAX_BYTES) return NULL;
    if (nbytes % SBBF_BLOCK_BYTES != 0) return NULL;
    size_t nblocks = nbytes / SBBF_BLOCK_BYTES;
    // nblocks must be a power of two so the fastrange index reduces
    // to a shift (matches the qb_new layout).
    if ((nblocks & (nblocks - 1)) != 0) return NULL;

    bloom_t* b = (bloom_t*)malloc(sizeof(bloom_t));
    if (!b) return NULL;
    // See qb_new() for the alignment rationale. nbytes here is checked
    // above to be a multiple of SBBF_BLOCK_BYTES (32) before we get here.
    void* mem = aligned_alloc(32, nbytes);
    if (!mem) {
        free(b);
        return NULL;
    }
    memcpy(mem, bytes, nbytes);
    b->nblocks_mask = nblocks - 1;
    b->idx_shift = 32 - log2_pow2(nblocks);
    b->bits = (uint32_t*)mem;
    return b;
}

// ---- helpers ----

size_t qb_estimate_bits(size_t n, double fp) {
    // Smallest meaningful filter is one 256-bit block.
    if (n == 0) return 256;
    // Catch the programmer bug in debug builds; release stays
    // conservative (return n*32) so production callers don't crash
    // on a stray sentinel value coming out of a config parser.
    assert(fp > 0.0 && fp < 1.0
           && "qb_estimate_bits: fp must be strictly between 0 and 1");
    if (!(fp > 0.0) || !(fp < 1.0)) return n * 32;

    // Classical Bloom filter bit-budget:
    //   m = -n * ln(fp) / (ln 2)^2
    // For SBBF (Parquet K=8) this is an approximation: SBBF's actual
    // FP rate at this sizing is typically within ~2x of fp on small
    // filters and converges to fp as the filter grows. Callers who
    // need a specific FP target on a small filter should over-size
    // (e.g. multiply the result by 1.5).
    const double LN2_SQ = 0.4804530139182014;  // ln(2)^2
    double m = -((double)n) * log(fp) / LN2_SQ;
    if (m < 256.0) m = 256.0;
    return (size_t)(m + 0.5);
}