nordisk/znippy-zoomies

stree / stree32 — design

Class: design · Function: stree (the AVX2 static search tree). Strict: chosen shape only. The why (why two concrete types not a generic, how it relates to vann) is in stree-reasoning.md; the measured numbers are in benchmarks.md; usage in stree-guide.md.

Source: src/stree.rs (STree64, i64) and src/stree32.rs (STree32, u32 — Ragnar Groot Koerkamp's original width). Derived from static-search-tree (MIT; his copyright retained in LICENSE).

What it is

An exact-key search structure over a sorted slice of integer keys — "where is key q?" → Some(index) or None. It is an implicit S(+)-tree: the tree is a flat array, a node's children are found by arithmetic (child = k·(B+1)+i), never a stored pointer. Each node is exactly one 64-byte cache line holding B keys, so one line fetched compares one whole node. Three moves win, all about respecting the memory hierarchy: contiguous layout in traversal order, a branchless SIMD node compare, and a software-pipelined prefetched batch that hides the miss latency.

The two concrete types

STree64 serves OSM (node IDs exceed u32) and skade (i64 snapshot IDs). STree32 is Ragnar's original width for keys that fit in 32 bits — twice the key density per cache line, so log₁₇ levels instead of log₉, i.e. fewer dependent loads. Each width also ships an mmap-backed variant (STree64Mmap / STree32Mmap) that stores internal nodes only — the sorted on-disk column is the leaf layer (STree32Mmap::new_with_stride for AoS records whose leading bytes are the LE key).

Public surface (stree32)

STree64 mirrors this surface over i64.

The branchless node compare

count_lt(node, q) counts how many of a node's B keys are < q with one SIMD compare and a popcount(movemask)/4 — no branch, so no mispredict, so the pipeline never stalls on data. STree32 adds one wrinkle the u32 width needs: AVX2 has only a signed 32-bit greater-than, so count_lt biases both the query and each node lane by i32::MIN (XOR 0x8000_0000), mapping unsigned order onto signed. Each i32 lane yields 4 identical result bytes, hence the /4. The variant is selected once per query/batch (AVX2 vs count_lt_scalar fallback) by the public entry points, never per node; the hot loops are generic over Fn(&[u32; B], u32) -> usize and the AVX2 form inlines inside a #[target_feature(enable="avx2")] wrapper. count_lt_scalar auto-vectorises under -C target-cpu=native.

The batch that hides the miss (MLP)

A single query is latency-bound: every level is a dependent cache miss. batch_stream::<P> runs P queries in flight — each issues a prefetch for its next hop then yields, and another query's compute fills the gap, collapsing the wall clock from "sum of the misses" to "the compute, with misses overlapped." const P is the pipeline depth — the one parameter kept generic (P=16 pegs every core at 100 %). It is the MLP knob, not a key-width parametrization.

The static contract

Build is comparatively expensive (sort + pack the implicit layout) but fully deterministic; queries are cheap and read-only; the structure is frozen between rebuilds (the skade snapshot tree rebuilds only every ≥1024 commits). Correctness is pinned by the inject-and-assert tests in stree32.rs::tests — present/absent/boundary keys, the full u32::MAX range (the bias trick), leaf-overflow edges, a 200k round-trip, and batch_prefetch/batch_stream::<8|16|32> returning identical results to serial find_exact.