src/stree.rs
//! S-tree static search index for sorted i64 keys.
//!
//! ORIGINAL AUTHOR: **Ragnar Groot Koerkamp** — `static-search-tree`
//! <https://github.com/RagnarGrootKoerkamp/static-search-tree>
//! Writeup: "Static search trees: 40× faster than binary search"
//! <https://curiouscoding.nl/posts/static-search-tree/>
//!
//! This module is a derivative work; the algorithm — the cache-line-sized
//! implicit S(+)-tree, the branchless SIMD node compare, and the software-
//! pipelined/prefetched batch traversal that hides memory latency — is entirely
//! Ragnar's. **Thank you, Ragnar Groot Koerkamp**, for the design and the
//! beautifully written explanation; every fast path here traces back to your
//! work, including the B-factor and batching trade-offs we re-measured for i64.
//! Licensed MIT — Copyright (c) 2025 Ragnar Groot Koerkamp. The upstream MIT
//! notice is retained in this crate's LICENSE file.
//!
//! Changes from the original (adapted for Rust 2024 edition):
//! - Key type: u32 → i64
//! - Node size: N=16 u32 (64 B) → N=8 i64 (64 B, one cache line). We also
//! benchmarked Ragnar's larger-node idea at i64 width (B=16, two lines) and
//! found it regresses the batched path — see the `const B` note below.
//! - Branching factor: B=16 → B=8
//! - SIMD: portable_simd (nightly) → std::arch AVX2 (stable)
//! `_mm256_cmpgt_epi64` compares 4 × i64 per 256-bit register; a `B/8`
//! compare *pair* covers a node, reduced by one packed movemask. The
//! check is hoisted once per query/batch into a
//! `#[target_feature(enable="avx2")]` wrapper so it inlines (not per node).
//! - **AVX-512 as a third kernel, selected at run time.** A B=8 node is 64 B
//! == one zmm register, so [`count_lt_avx512`] does in ONE
//! `_mm512_cmplt_epi64_mask` what AVX2 needs two compares plus two
//! movemasks for — on the same tree, same layout, same file format. Which
//! kernel runs is resolved once per process by [`crate::stree_simd`] (a
//! `OnceLock` ceiling from CPUID plus an atomic active selection), never
//! per node and no longer per query either. This is what lets ONE binary
//! be optimal on t14s (Zen 5, avx512f/bw/vl/dq) and correct on oden
//! (Zen 2, no AVX-512 at all): the 512-bit instructions are only reachable
//! through a `#[target_feature(enable="avx512f")]` wrapper that CPUID has
//! already cleared.
//! - Hugepage hint (`madvise(MADV_HUGEPAGE)`) on the in-RAM tree — cuts dTLB
//! misses in the RAM-bound regime, mirroring Ragnar's hugepage results.
//! - Return type: predecessor value → index in original array
//! (what NodeStore actually needs).
//!
//! Build: `STree64::new(sorted_ids)` — O(n).
//! Query: `find_exact(id)` → `Option<usize>` — O(log₉ n), ≈7 cache misses for 9B nodes.
//!
//! `STree64Mmap` — same tree but WITHOUT the leaf layer.
//! The sorted mmap (16-byte records: [i64 id][f32 lat][f32 lon]) IS the leaf layer.
//! Internal nodes only: ~94 MB for Europe (94M nodes), ~9 GB for planet (9B nodes).
//! Build reads every 8th ID from the mmap (stride scan, no full copy in RAM).
//! Lookup: navigate internal nodes → final linear scan of ≤9 mmap records (144 bytes).
// B=8 (one 64 B cache line) beats B=16 (two lines) for the *batched* lookup that
// production uses: benchmarked head-to-head (2026-06-11, see the bench_* tests
// below), B=16 regressed best-batch ns/q at every size (8 MB 40→47, 128 MB 68→82, 1 GB
// 104→120 default). Fewer tree levels don't pay off once P-way batching already
// hides miss latency — all that's left is B=16's doubled per-node load+compute.
// `count_lt_avx2` is written B-agnostically (loops B/8 AVX2 compare pairs) so flipping
// this constant is a one-line experiment if the workload ever shifts to serial.
const B: usize = 8; // elements per node = branching factor
const MAX64: i64 = i64::MAX; // sentinel for unused slots
/// The resolved-once kernel selector every public entry point below branches
/// on. See [`crate::stree_simd`] — CPUID runs on the first query and never
/// again; from then on the dispatch is a relaxed byte load. Imported under
/// `cfg(x86_64)` because that is the only arch with more than one kernel.
#[cfg(target_arch = "x86_64")]
use crate::stree_simd::SimdPath;
// ── Public type ────────────────────────────────────────────────────────────────
/// Immutable S-tree over a sorted `&[i64]` slice.
/// Stores a complete copy of the data internally.
pub struct STree64 {
tree: NodeBuf,
offsets: Vec<usize>,
n: usize,
}
impl STree64 {
/// Build from a **sorted** slice. Panics in debug if not sorted.
///
/// Nodes are 64-byte aligned (phase 0) — one cache line per node. The
/// old `Vec<[i64; B]>` landed at glibc's mmap base + 0x10, so every node
/// straddled two lines; measured 2026-08-10 (stree_lab, 1.34e8 keys,
/// `batch_stream::<128>`) that misalignment cost ~2× at prefetch
/// saturation, which is where the batch path lives.
pub fn new(vals: &[i64]) -> Self {
Self::new_phase(vals, 0)
}
/// [`new`](Self::new) with an explicit node-buffer cache-line phase — 0
/// aligned (default), 16 reproduces the old allocator reality for A/B.
pub fn new_phase(vals: &[i64], phase: usize) -> Self {
assert!(!vals.is_empty(), "STree64::new: empty input");
#[cfg(debug_assertions)]
for w in vals.windows(2) {
assert!(w[0] <= w[1], "STree64::new: input not sorted");
}
let n = vals.len();
let height = height(n);
let lsizes = layer_sizes(n, height);
let n_blocks: usize = lsizes.iter().sum();
let mut offsets = Vec::with_capacity(height);
let mut acc = 0;
for &ls in &lsizes {
offsets.push(acc);
acc += ls;
}
let mut tree = NodeBuf::alloc(n_blocks.max(1), phase);
// ── Leaf layer ────────────────────────────────────────────────────────
let ol = offsets[height - 1];
for (i, &val) in vals.iter().enumerate() {
tree[ol + i / B][i % B] = val;
}
if n % B != 0 {
tree[ol + n / B][n % B..].fill(MAX64);
}
// ── Internal layers (root … leaf−1, built bottom-up) ──────────────────
for h in (0..height - 1).rev() {
let oh = offsets[h];
tree[oh..oh + lsizes[h]]
.iter_mut()
.for_each(|nd| nd.fill(MAX64));
for i in 0..B * lsizes[h] {
let j = i % B;
let mut k = (i / B) * (B + 1) + j + 1;
for _ in h..height - 2 {
k *= B + 1;
}
tree[oh + i / B][j] = if k * B < n { tree[ol + k][0] } else { MAX64 };
}
}
advise_hugepage(&tree[..]);
Self { tree, offsets, n }
}
/// Find the index of `q` in the original sorted slice, or `None` if absent.
///
/// Internal routing uses `count_lt` (same as the original library). When all
/// elements in a leaf block are < q, `idx` equals B and the search overflows
/// into the next leaf block via `idx / B` — the same mechanism the original
/// S-tree uses with its `self.get(o + k + idx/N, idx%N)` expression.
#[inline]
pub fn find_exact(&self, q: i64) -> Option<usize> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.find_exact_avx512(q) },
SimdPath::Avx2 => return unsafe { self.find_exact_avx2(q) },
SimdPath::Scalar => {}
}
self.find_exact_impl(q, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_exact_avx2(&self, q: i64) -> Option<usize> {
self.find_exact_impl(q, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn find_exact_avx512(&self, q: i64) -> Option<usize> {
self.find_exact_impl(q, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn find_exact_impl<F: Fn(&[i64; B], i64) -> usize>(&self, q: i64, cnt: F) -> Option<usize> {
let height = self.offsets.len();
let mut k = 0usize;
for h in 0..height - 1 {
let o = self.offsets[h];
let jump = cnt(&self.tree[o + k], q);
k = k * (B + 1) + jump;
}
let o = self.offsets[height - 1];
let idx = cnt(&self.tree[o + k], q);
// idx can be 0..=B. When idx == B all leaf elements were < q,
// so the target is at the start of the NEXT leaf block (overflow by 1).
let block = k + idx / B;
let slot = idx % B;
let pos = block * B + slot;
// Guard both self.n (logical) and tree (physical) bounds.
if pos < self.n && self.tree[o + block][slot] == q {
Some(pos)
} else {
None
}
}
}
// ── STree64Mmap — internal-nodes-only tree, mmap is the leaf layer ────────────
//
// Mmap record layout: [i64 le id (8 B)][f32 le lat (4 B)][f32 le lon (4 B)] = 16 B.
// Build reads the first ID of each B-record leaf block directly from the mmap
// (stride scan at 128-byte intervals — cheap and cache-friendly).
// find_exact navigates internal nodes then does a ≤9-record linear scan in the mmap.
const MMAP_RECORD: usize = 16; // default bytes per record (AoS layout)
/// 64-byte-phase-controlled internal-node storage.
///
/// `Vec<[i64; B]>` is align-8 by type; a multi-MB allocation lands at glibc's
/// mmap chunk base + 0x10, so every 64-B node used to straddle two cache
/// lines — the untested H4 from oid_index.rs:110, now testable and fixed.
/// Over-allocate a plain `Vec<i64>` and start at the first 64-byte boundary
/// plus `phase` bytes: phase 0 is the aligned default, phase 16 reproduces
/// the old allocator reality for A/B measurement.
pub struct NodeBuf {
buf: Vec<i64>,
off: usize, // in i64 units
len: usize, // in nodes
}
impl NodeBuf {
fn alloc(n_nodes: usize, phase: usize) -> Self {
assert!(phase % 8 == 0 && phase < 64, "phase must be 8-multiple < 64");
let len = n_nodes * B + 16;
// MADV_HUGEPAGE must land BEFORE first touch: `vec![MAX64; len]`
// faults every 4 KB page at creation, and a hint issued after faulting
// waits on khugepaged, which never collapses anything on a benchmark's
// timescale — measured 2026-08-10: AnonHugePages stayed 0 kB against
// 4.9 GB of trees. `with_capacity` reserves without touching; advising
// the reservation makes the subsequent fill fault straight onto 2 MB
// pages (THP mode `madvise`, defrag `madvise`).
let mut buf: Vec<i64> = Vec::with_capacity(len);
advise_hugepage_raw(buf.as_ptr() as *const u8, len * 8);
buf.resize(len, MAX64);
let addr = buf.as_ptr() as usize;
let off = ((64 - (addr & 63)) & 63) / 8 + phase / 8;
NodeBuf { buf, off, len: n_nodes }
}
}
impl std::ops::Deref for NodeBuf {
type Target = [[i64; B]];
#[inline(always)]
fn deref(&self) -> &[[i64; B]] {
unsafe {
std::slice::from_raw_parts(self.buf.as_ptr().add(self.off) as *const [i64; B], self.len)
}
}
}
impl std::ops::DerefMut for NodeBuf {
#[inline(always)]
fn deref_mut(&mut self) -> &mut [[i64; B]] {
unsafe {
std::slice::from_raw_parts_mut(
self.buf.as_mut_ptr().add(self.off) as *mut [i64; B],
self.len,
)
}
}
}
pub struct STree64Mmap {
tree: NodeBuf,
offsets: Vec<usize>,
pub count: usize,
pub stride: usize, // bytes between consecutive i64 IDs in the data slice
}
impl STree64Mmap {
/// Build from a data slice with the default 16-byte AoS record layout
/// `[i64 id][f32 lat][f32 lon]`.
pub fn new(mmap: &[u8], count: usize) -> Self {
Self::new_with_stride(mmap, count, MMAP_RECORD)
}
/// Build from a data slice where each ID is `stride` bytes apart.
/// Use stride=8 for a pure ID column (Arrow IPC / SoA layout).
/// Use stride=16 for the legacy AoS `[i64 id][f32 lat][f32 lon]` layout.
///
/// Internal nodes are 64-byte aligned (phase 0) — one cache line per node.
pub fn new_with_stride(data: &[u8], count: usize, stride: usize) -> Self {
Self::new_with_stride_phase(data, count, stride, 0)
}
/// [`new_with_stride`](Self::new_with_stride) with an explicit node-buffer
/// cache-line phase — 0 for aligned (the default), 16 to reproduce the old
/// `Vec<[i64; B]>`-at-glibc's-mercy layout for A/B measurement (H4).
pub fn new_with_stride_phase(data: &[u8], count: usize, stride: usize, phase: usize) -> Self {
assert!(count > 0);
assert!(stride >= 8);
let h = height(count);
let lsizes = layer_sizes(count, h);
let ni = h - 1;
let n_blocks: usize = lsizes[..ni].iter().sum();
let mut offsets = Vec::with_capacity(ni);
let mut acc = 0usize;
for &ls in &lsizes[..ni] {
offsets.push(acc);
acc += ls;
}
let mut tree = NodeBuf::alloc(n_blocks.max(1), phase);
// Per-layer fill, bottom-up. Each layer's blocks are independent so
// they parallelise trivially; the layer barrier between iterations
// is algorithmic (parent nodes read child IDs from `data`, not from
// the upper-layer tree, so layers can in principle even run out of
// order — we keep bottom-up for cache friendliness).
//
// ROOT LAW #0 (2026-07-22): this was rayon, then a hand-rolled
// `std::thread::scope` pool that carved the layer into one static chunk
// per thread — a scoped work pool, which the law forbids just as hard as
// the rayon it replaced. It is now `gatling_scanlines`, which IS this
// shape: a `rows * stride` buffer where each row is a disjoint mutable
// slice claimed off the engine's atomic cursor.
//
// The ROW IS A BAND OF BLOCKS, NOT ONE BLOCK — measured, not assumed.
// One block per unit reads like the natural gatling unit, but a block is
// 64 bytes of work and the whole layer is `n/B` blocks (12.5 M on a 100 M
// -key tree), so the shared `fetch_add` and its cache line dominate:
// 0.038 s → 0.082 s, a 2x REGRESSION. Bands of
// `n_nodes / (4 * n_threads)` blocks put the cursor cost back in the
// noise while still giving 4 units per core to rebalance with — strictly
// more schedulable than the old one-static-chunk-per-thread carve-up.
// `min_rows_for_threads = 2` keeps the engine serial for a one-band layer.
const SMALL_LAYER_BLOCKS: usize = 4096;
let n_threads = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
.max(1);
for lvl in (0..ni).rev() {
let oh = offsets[lvl];
let n_nodes = lsizes[lvl];
let slice = &mut tree[oh..oh + n_nodes];
let fill_block = |blk: usize, node: &mut [i64]| {
node.fill(MAX64);
for j in 0..B {
let mut k = blk * (B + 1) + j + 1;
for _ in lvl..h - 2 {
k *= B + 1;
}
node[j] = if k * B < count {
id_at(data, k * B, stride)
} else {
MAX64
};
}
};
if n_nodes < SMALL_LAYER_BLOCKS {
for (blk, node) in slice.chunks_mut(1).enumerate() {
fill_block(blk, node.as_flattened_mut());
}
continue;
}
// Bands of `band` blocks; the ragged tail (< band blocks) is filled
// on this thread — it is at most one band, and the band count is
// 4 x n_threads, so the tail is < 1/(4*n_threads) of the layer.
let band = n_nodes.div_ceil(n_threads * 4).max(1);
let rows = n_nodes / band;
let flat = slice.as_flattened_mut();
let (head, tail) = flat.split_at_mut(rows * band * B);
gatling::gatling_forkjoin::gatling_scanlines(
head,
rows,
band * B,
0,
2,
|u, banded| {
for (i, node) in banded.chunks_mut(B).enumerate() {
fill_block(u * band + i, node);
}
},
);
for (i, node) in tail.chunks_mut(B).enumerate() {
fill_block(rows * band + i, node);
}
}
advise_hugepage(&tree[..]);
Self {
tree,
offsets,
count,
stride,
}
}
/// The internal-node buffer's cache-line phase (`ptr % 64`) — 0 means every
/// 64-B node occupies exactly one line.
pub fn tree_phase(&self) -> usize {
self.tree.as_ptr() as usize & 63
}
#[inline]
pub fn find_exact(&self, q: i64, mmap: &[u8]) -> Option<usize> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.find_exact_avx512(q, mmap) },
SimdPath::Avx2 => return unsafe { self.find_exact_avx2(q, mmap) },
SimdPath::Scalar => {}
}
self.find_exact_impl(q, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_exact_avx2(&self, q: i64, mmap: &[u8]) -> Option<usize> {
self.find_exact_impl(q, mmap, |n, qq| unsafe { count_lt_avx2(n, qq) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn find_exact_avx512(&self, q: i64, mmap: &[u8]) -> Option<usize> {
self.find_exact_impl(q, mmap, |n, qq| unsafe { count_lt_avx512(n, qq) })
}
#[inline(always)]
fn find_exact_impl<F: Fn(&[i64; B], i64) -> usize>(
&self,
q: i64,
mmap: &[u8],
cnt: F,
) -> Option<usize> {
let mut k = 0usize;
for &o in &self.offsets {
let jump = cnt(&self.tree[o + k], q);
k = k * (B + 1) + jump;
}
let stride = self.stride;
for i in 0..=B {
let pos = k * B + i;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => return Some(pos),
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
None
}
/// Tree-only routing: descend internal nodes to a leaf block index.
/// Skips the mmap leaf scan — useful when the caller wants to sort all
/// queries by leaf-block index before touching the mmap (to convert
/// random page faults into a sequential read pattern).
///
/// Returns the leaf block index `k` such that the candidate record range
/// is `[k*B, k*B + B + 1)` (the trailing +1 handles overflow when every
/// element in block `k` is `< q`).
#[inline]
pub fn route_to_block(&self, q: i64) -> usize {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.route_to_block_avx512(q) },
SimdPath::Avx2 => return unsafe { self.route_to_block_avx2(q) },
SimdPath::Scalar => {}
}
self.route_to_block_impl(q, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn route_to_block_avx2(&self, q: i64) -> usize {
self.route_to_block_impl(q, |n, qq| unsafe { count_lt_avx2(n, qq) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn route_to_block_avx512(&self, q: i64) -> usize {
self.route_to_block_impl(q, |n, qq| unsafe { count_lt_avx512(n, qq) })
}
#[inline(always)]
fn route_to_block_impl<F: Fn(&[i64; B], i64) -> usize>(&self, q: i64, cnt: F) -> usize {
let mut k = 0usize;
for &o in &self.offsets {
let jump = cnt(&self.tree[o + k], q);
k = k * (B + 1) + jump;
}
k
}
/// Batch lookup: resolve `ids` against the on-disk sorted record store.
/// Returns `out[i]` = position of `ids[i]` in the mmap, or `None` if absent.
///
/// Internally:
/// 1. Tree-only routing for every id (in-RAM, ~10 cache-line reads each)
/// 2. Sort `(leaf_block, original_index)` pairs — sequential mmap walk
/// 3. `madvise(WILLNEED)` on the distinct page range so the kernel can
/// issue parallel readahead for the upcoming leaf scans
/// 4. Linear leaf scan in sorted order; scatter results back
///
/// For a slot with millions of node-refs spread across a 144 GB coord
/// mmap, this converts millions of random page faults into a near-
/// sequential read pattern + parallel readahead. Expected speedup on
/// cold pages: ~4-8× vs N serial `find_exact` calls.
pub fn lookup_batch(&self, ids: &[i64], mmap: &[u8]) -> Vec<Option<usize>> {
let n = ids.len();
if n == 0 {
return Vec::new();
}
// 1. Route every id through the tree → (leaf_block, original_index).
let mut routed: Vec<(usize, u32)> = (0..n)
.map(|i| (self.route_to_block(ids[i]), i as u32))
.collect();
// 2. Sort by leaf block so the mmap walk is sequential.
routed.sort_unstable_by_key(|&(k, _)| k);
// 3. Best-effort readahead hint over the touched range.
// On Linux this is a non-blocking call to advise the page cache;
// failures are silently ignored (this is purely a performance hint).
if let (Some(&(lo_k, _)), Some(&(hi_k, _))) = (routed.first(), routed.last()) {
let stride = self.stride;
let lo_byte = lo_k * B * stride;
let hi_byte = ((hi_k + 1) * B + 1).min(self.count) * stride;
if hi_byte > lo_byte {
advise_willneed(&mmap[lo_byte..hi_byte]);
}
}
// 4. Scan leaves in sorted order, scatter into `out` by original index.
let stride = self.stride;
let mut out = vec![None; n];
for (k, orig) in routed {
let q = ids[orig as usize];
for i in 0..=B {
let pos = k * B + i;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[orig as usize] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
}
out
}
/// Pipelined batched lookup using software-pipelined internal traversal.
/// Const generic P sets the number of in-flight queries to overlap memory
/// latency. This is Ragnar's key trick to increase memory-level parallelism.
pub fn lookup_batch_pipeline<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.lookup_batch_pipeline_avx512::<P>(ids, mmap) },
SimdPath::Avx2 => return unsafe { self.lookup_batch_pipeline_avx2::<P>(ids, mmap) },
SimdPath::Scalar => {}
}
self.lookup_batch_pipeline_impl::<P, _>(ids, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_pipeline_avx2<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_pipeline_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_pipeline_avx512<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_pipeline_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn lookup_batch_pipeline_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
ids: &[i64],
mmap: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
if n == 0 {
return Vec::new();
}
// 1. Route queries in batches of P through internal nodes with prefetch.
let mut routed: Vec<(usize, u32)> = Vec::with_capacity(n);
let mut i = 0usize;
let offs = &self.offsets;
let n_levels = offs.len();
while i + P <= n {
let mut q = [0i64; P];
for j in 0..P {
q[j] = ids[i + j];
}
let mut kk = [0usize; P];
// Traverse all internal layers except the last with prefetch
for h in 0..n_levels.saturating_sub(1) {
let o = offs[h];
let o2 = offs[h + 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], q[j]);
kk[j] = kk[j] * (B + 1) + jump;
prefetch_index(&self.tree, o2 + kk[j]);
}
}
// Final internal level (no next tree level to prefetch)
if n_levels > 0 {
let o = offs[n_levels - 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], q[j]);
kk[j] = kk[j] * (B + 1) + jump;
}
}
for j in 0..P {
routed.push((kk[j], (i + j) as u32));
}
i += P;
}
// Remainder
while i < n {
routed.push((self.route_to_block(ids[i]), i as u32));
i += 1;
}
// 2. Sort & madvise as before
routed.sort_unstable_by_key(|&(k, _)| k);
let stride = self.stride;
if let (Some(&(lo_k, _)), Some(&(hi_k, _))) = (routed.first(), routed.last()) {
let lo_byte = lo_k * B * stride;
let hi_byte = ((hi_k + 1) * B + 1).min(self.count) * stride;
if hi_byte > lo_byte {
advise_willneed(&mmap[lo_byte..hi_byte]);
}
}
// 3. linear leaf scan, scatter results
let mut out = vec![None; n];
for (k, orig) in routed {
let q = ids[orig as usize];
for t in 0..=B {
let pos = k * B + t;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[orig as usize] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
_ => {}
}
}
}
out
}
/// Single-phase pipelined batch — the in-RAM-section shape (2026-08-10).
///
/// [`lookup_batch_pipeline`](Self::lookup_batch_pipeline) is two-phase:
/// route everything, `sort_unstable` all routed pairs, `madvise`, then a
/// separate serial leaf scan. That shape exists for the cold 144 GB OSM
/// mmap, where sorting converts random page faults into sequential
/// readahead. For an in-RAM keyspace (the `__gunnar_oid__` section) the
/// sort is pure overhead and the leaf misses sit OUTSIDE the pipeline.
///
/// This walk keeps everything inside it: route P queries level-
/// synchronously with next-level prefetch, prefetch the leaf lines at the
/// last internal level, scan leaves immediately while the next chunk's
/// prefetches would otherwise idle, write results positionally. No routed
/// Vec, no sort, no scatter.
pub fn lookup_batch_fused<const P: usize>(&self, ids: &[i64], mmap: &[u8]) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.lookup_batch_fused_avx512::<P>(ids, mmap) },
SimdPath::Avx2 => return unsafe { self.lookup_batch_fused_avx2::<P>(ids, mmap) },
SimdPath::Scalar => {}
}
self.lookup_batch_fused_impl::<P, _>(ids, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_fused_avx2<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_fused_avx512<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx512(n, q) })
}
/// [`lookup_batch_fused`](Self::lookup_batch_fused) with one more stage of
/// software pipelining: chunk k's leaf scan runs AFTER chunk k+1's
/// internal routing has been issued, so the leaf prefetches get a full
/// chunk's worth of routing work to complete under instead of being
/// consumed at the chunk boundary they were issued at.
pub fn lookup_batch_fused2<const P: usize>(&self, ids: &[i64], mmap: &[u8]) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.lookup_batch_fused2_avx512::<P>(ids, mmap) },
SimdPath::Avx2 => return unsafe { self.lookup_batch_fused2_avx2::<P>(ids, mmap) },
SimdPath::Scalar => {}
}
self.lookup_batch_fused2_impl::<P, _>(ids, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_fused2_avx2<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused2_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_fused2_avx512<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused2_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn lookup_batch_fused2_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
ids: &[i64],
mmap: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
let mut out = vec![None; n];
let n_levels = self.offsets.len();
let stride = self.stride;
// Route one chunk of P: internal levels with next-level prefetch, leaf
// lines prefetched at the last level. Identical to the fused walk's
// routing stage.
let route = |base: usize, kk: &mut [usize; P]| {
*kk = [0usize; P];
for h in 0..n_levels.saturating_sub(1) {
let o = self.offsets[h];
let o2 = self.offsets[h + 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[base + j]);
kk[j] = kk[j] * (B + 1) + jump;
prefetch_index(&self.tree, o2 + kk[j]);
}
}
if n_levels > 0 {
let o = self.offsets[n_levels - 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[base + j]);
kk[j] = kk[j] * (B + 1) + jump;
let byte = (kk[j] * B).min(self.count - 1) * stride;
let base_p = unsafe { mmap.as_ptr().add(byte.min(mmap.len() - 1)) };
prefetch_ptr(base_p);
prefetch_ptr(base_p.wrapping_add(64));
}
}
};
let scan = |base: usize, kk: &[usize; P], out: &mut Vec<Option<usize>>| {
for j in 0..P {
let q = ids[base + j];
let k = kk[j];
for t in 0..=B {
let pos = k * B + t;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[base + j] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
}
};
let mut i = 0usize;
let mut prev = [0usize; P];
let mut have_prev = false;
while i + P <= n {
let mut kk = [0usize; P];
route(i, &mut kk);
if have_prev {
scan(i - P, &prev, &mut out);
}
prev = kk;
have_prev = true;
i += P;
}
if have_prev {
scan(i - P, &prev, &mut out);
}
while i < n {
out[i] = self.find_exact(ids[i], mmap);
i += 1;
}
out
}
#[inline(always)]
fn lookup_batch_fused_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
ids: &[i64],
mmap: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
let mut out = vec![None; n];
let n_levels = self.offsets.len();
let stride = self.stride;
let mut i = 0usize;
while i + P <= n {
let mut kk = [0usize; P];
for h in 0..n_levels.saturating_sub(1) {
let o = self.offsets[h];
let o2 = self.offsets[h + 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[i + j]);
kk[j] = kk[j] * (B + 1) + jump;
prefetch_index(&self.tree, o2 + kk[j]);
}
}
if n_levels > 0 {
let o = self.offsets[n_levels - 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], ids[i + j]);
kk[j] = kk[j] * (B + 1) + jump;
// Candidate range is B+1 records from kk[j]*B — at stride 8
// that spans ≤ 72 B; two lines cover it at any phase.
let byte = (kk[j] * B).min(self.count - 1) * stride;
let base = unsafe { mmap.as_ptr().add(byte.min(mmap.len() - 1)) };
prefetch_ptr(base);
prefetch_ptr(base.wrapping_add(64));
}
}
for j in 0..P {
let q = ids[i + j];
let k = kk[j];
for t in 0..=B {
let pos = k * B + t;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[i + j] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
std::cmp::Ordering::Less => {}
}
}
}
i += P;
}
while i < n {
out[i] = self.find_exact(ids[i], mmap);
i += 1;
}
out
}
}
/// Best-effort `madvise(MADV_HUGEPAGE)` over the tree's backing store. The tree
/// is the random-access hot structure: every query touches one cache line per
/// level (≈9 dependent loads for a planet-scale i64 tree), and with 4 KB pages
/// each of those is also a TLB miss. Hinting transparent hugepages (2 MB) cuts
/// the page count — and thus dTLB pressure — ~512×, the single biggest lever for
/// the RAM-bound regime (cf. Ragnar's hugepage results). On hosts with
/// `transparent_hugepage=madvise` (the common default) this hint is *required*
/// for the tree to get hugepages at all. Best-effort: errors ignored, no-op off
/// Linux. Alignment isn't forced — the kernel promotes the 2 MB-aligned interior
/// of the allocation, which is the bulk of a large tree.
#[inline]
fn advise_hugepage<T>(slice: &[T]) {
advise_hugepage_raw(slice.as_ptr() as *const u8, std::mem::size_of_val(slice));
}
/// Clear this process's `PR_SET_THP_DISABLE` flag, returning whether THP is
/// enabled afterwards. Sandboxed/containerized parents (observed 2026-08-10:
/// an agent harness) set the flag and children inherit it, which makes every
/// `MADV_HUGEPAGE` in this module a silent no-op regardless of alignment or
/// ordering — `/proc/self/status` shows `THP_enabled: 0` while PID 1 shows 1.
/// A process may clear its own flag; call this once at startup in binaries
/// that want the trees on 2 MB pages. Library code never calls it — flipping
/// process-wide state is the binary's decision.
pub fn thp_enable_for_process() -> bool {
#[cfg(target_os = "linux")]
unsafe {
const PR_SET_THP_DISABLE: libc::c_int = 41;
let _ = libc::prctl(PR_SET_THP_DISABLE, 0usize, 0usize, 0usize, 0usize);
std::fs::read_to_string("/proc/self/status")
.map(|s| s.lines().any(|l| l.starts_with("THP_enabled") && l.contains('1')))
.unwrap_or(false)
}
#[cfg(not(target_os = "linux"))]
{
false
}
}
/// `MADV_HUGEPAGE` over a raw range. **Order matters**: to be effective this
/// must run before the range is first touched — pages faulted before the hint
/// only reach 2 MB via khugepaged, i.e. effectively never (see [`NodeBuf`]).
/// Public so callers that own the leaf storage (a keys column, a section byte
/// buffer) can put IT on hugepages too — the leaf is one touch per query and
/// its dTLB miss is as real as the tree's.
#[inline]
pub fn advise_hugepage_raw(ptr: *const u8, bytes: usize) {
#[cfg(target_os = "linux")]
unsafe {
// `madvise` demands a PAGE-ALIGNED address and glibc's large-alloc
// pointers are mmap base + 0x10 — so the naive call fails EINVAL and,
// errors being ignored by design, every hugepage hint this module ever
// issued was a silent no-op (found 2026-08-10: AnonHugePages 0 kB
// under 4.9 GB of advised trees). Round inward to page boundaries;
// the ≤4 KB lost at each edge cannot carry a 2 MB page anyway.
let start = (ptr as usize + 4095) & !4095;
let end = (ptr as usize).saturating_add(bytes) & !4095;
if end > start {
let _ = libc::madvise(start as *mut libc::c_void, end - start, libc::MADV_HUGEPAGE);
}
}
#[cfg(not(target_os = "linux"))]
{
let _ = (ptr, bytes);
}
}
/// Best-effort `madvise(MADV_WILLNEED)` over a mmap slice. Linux-only; on
/// other platforms this is a no-op. Errors are silently ignored — the
/// kernel may reject the hint (e.g. unaligned, beyond mapping) but the
/// subsequent reads still succeed via normal page faults.
#[inline]
fn advise_willneed(slice: &[u8]) {
#[cfg(target_os = "linux")]
unsafe {
let _ = libc::madvise(
slice.as_ptr() as *mut libc::c_void,
slice.len(),
libc::MADV_WILLNEED,
);
}
#[cfg(not(target_os = "linux"))]
{
let _ = slice;
}
}
/// Prefetch the given cacheline into L1 cache (pointer form).
#[inline]
fn prefetch_ptr<T>(ptr: *const T) {
#[cfg(target_arch = "x86_64")]
unsafe {
std::arch::x86_64::_mm_prefetch(ptr as *const i8, std::arch::x86_64::_MM_HINT_T0);
}
#[cfg(target_arch = "x86")]
unsafe {
std::arch::x86::_mm_prefetch(ptr as *const i8, std::arch::x86::_MM_HINT_T0);
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))]
{
let _ = ptr;
}
}
/// Prefetch the given cacheline by indexing into a slice.
#[inline]
fn prefetch_index<T>(s: &[T], index: usize) {
let ptr = unsafe { s.as_ptr().add(index) } as *const T;
prefetch_ptr(ptr);
}
#[inline]
fn id_at(data: &[u8], idx: usize, stride: usize) -> i64 {
let off = idx * stride;
i64::from_le_bytes(data[off..off + 8].try_into().unwrap())
}
// ── Tree shape helpers ─────────────────────────────────────────────────────────
fn blocks(n: usize) -> usize {
n.div_ceil(B)
}
fn prev_keys(n: usize) -> usize {
blocks(n).div_ceil(B + 1) * B
}
fn height(n: usize) -> usize {
if n <= B { 1 } else { height(prev_keys(n)) + 1 }
}
fn layer_size(mut n: usize, h: usize, height: usize) -> usize {
for _ in h..height - 1 {
n = prev_keys(n);
}
n
}
fn layer_sizes(n: usize, height: usize) -> Vec<usize> {
(0..height)
.map(|h| layer_size(n, h, height).div_ceil(B))
.collect()
}
// ── SIMD node comparison ───────────────────────────────────────────────────────
/// Count elements in `node` that are strictly less than `q`.
/// Equivalently: the index of the first element >= q (predecessor step).
///
/// Three implementations, selected **once per process** ([`crate::stree_simd`])
/// and dispatched once per query or batch by the public entry points — never
/// per node. The hot traversal loops are generic
/// over the counter (`Cnt: Fn`) and the AVX2 variant runs inside a
/// `#[target_feature(enable = "avx2")]` wrapper, so [`count_lt_avx2`] inlines
/// straight into the loop with zero call/dispatch overhead. The scalar variant
/// auto-vectorises too when the surrounding crate is built with AVX2 in its
/// baseline (`-C target-cpu=native`), and is the correctness fallback on
/// non-AVX2 CPUs. Per-node `is_x86_feature_detected!` (the old hot-path cost)
/// is gone, and as of the runtime-dispatch pass so is the per-call one:
/// `batch_stream` used to re-run the detect once per P-query chunk, and now
/// loads one already-resolved byte instead.
///
/// `Cnt` here is a function/closure `(&[i64; B], i64) -> usize`.
#[inline(always)]
fn count_lt_scalar(node: &[i64; B], q: i64) -> usize {
node.iter().filter(|&&x| x < q).count()
}
/// Scalar alias for the deliberately-unoptimised `batch_no_prefetch` baseline,
/// which exists only as the "no SIMD multiversioning, no prefetch" reference
/// point in benches. Everything on the production path routes through the
/// AVX2-multiversioned entry points instead.
#[inline(always)]
fn count_lt(node: &[i64; B], q: i64) -> usize {
count_lt_scalar(node, q)
}
/// AVX2 path: two 256-bit `_mm256_cmpgt_epi64` per 8 i64 — a whole
/// one-cache-line B=8 node — reduced by **one** `vpackssdw` + **one**
/// `vmovmskps` + **one** `popcnt`. `B` must be a multiple of 8 (it is: 8, or
/// 16 for the two-cache-line experiment). The loop is fully unrolled by the
/// optimiser.
///
/// ## Why the pair is reduced together (2026-08-11, measured on oden)
///
/// The obvious per-register reduction — `_mm256_movemask_epi8` on each compare,
/// `popcnt` each, add, `/8` — is what this function used to do, and on Zen 2 it
/// **lost to the portable fallback**: `count_lt_scalar` auto-vectorises under
/// `-C target-cpu=native`, and LLVM's own reduction was *shorter* than the
/// hand-written one. `vpmovmskb` yields 8 duplicate bits per i64 lane, so that
/// shape needs two vector→GPR moves, two popcounts, an add and a shift.
///
/// `vpackssdw` collapses the two 4×qword masks into 8 dword lanes first, so a
/// single `vmovmskps` yields exactly one bit per key and `popcnt` *is* the
/// whole reduction — no `/8`. Five µops where there were eight. The pack
/// scrambles lane order across the 128-bit halves; that is deliberate and
/// harmless, because a population count does not care which bit is which.
///
/// The cost of the old shape was not the `vpcmpgtq` (both shapes issue the
/// identical pair) and not a downclock (the fallback issues 256-bit compares
/// too, at the same measured frequency). It was µop count in a dependency
/// chain, and on the un-pipelined `find_exact` path it was amplified up to
/// 1.58x on a 1 GB tree by squeezing the out-of-order window that supplies the
/// serial path's only memory-level parallelism.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn count_lt_avx2(node: &[i64; B], q: i64) -> usize {
use std::arch::x86_64::*;
// Rust 2024: explicit unsafe block required even inside unsafe fn.
unsafe {
let q_v = _mm256_set1_epi64x(q);
let mut bits = 0u32;
let mut i = 0;
while i < B {
let lo = _mm256_cmpgt_epi64(
q_v,
_mm256_loadu_si256(node.as_ptr().add(i) as *const __m256i),
);
let hi = _mm256_cmpgt_epi64(
q_v,
_mm256_loadu_si256(node.as_ptr().add(i + 4) as *const __m256i),
);
// 8 qword lanes (0 or -1) → 8 dword lanes (0 or -1): one bit per key.
let packed = _mm256_packs_epi32(lo, hi);
bits += (_mm256_movemask_ps(_mm256_castsi256_ps(packed)) as u32).count_ones();
i += 8;
}
bits as usize
}
}
/// AVX-512 path: `B * 8 == 64` bytes == 512 bits, so a one-cache-line B=8 node
/// is **exactly one zmm register** — one `_mm512_cmplt_epi64_mask` where
/// [`count_lt_avx2`] needs two `cmpgt` + two `movemask` + a `/8`, and the mask
/// is already one bit per key so `popcnt` is the whole reduction.
///
/// Same tree, same file format, nothing to migrate: this is a pure reader-side
/// instruction swap on the geometry that is already built. Kept B-agnostic
/// (`B/8` iterations) exactly like the AVX2 sibling, so flipping `B` to 16
/// stays a one-line experiment. `B` must be a multiple of 8 — it is.
///
/// See [`crate::stree512`]'s machine note before expecting a speedup: on t14s
/// AVX-512 issues on a 256-bit datapath, so this decodes to ~two 256-bit µops.
/// Fewer instructions, not necessarily fewer cycles.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
#[inline]
unsafe fn count_lt_avx512(node: &[i64; B], q: i64) -> usize {
use std::arch::x86_64::*;
unsafe {
let q_v = _mm512_set1_epi64(q);
let mut bits = 0u32;
let mut i = 0;
while i < B {
let v = _mm512_loadu_epi64(node.as_ptr().add(i));
bits += _mm512_cmplt_epi64_mask(v, q_v).count_ones();
i += 8;
}
bits as usize
}
}
// ── STree64 batch methods (Ragnar-style pipelined) ─────────────────────────────
impl STree64 {
/// Pipelined batch search — process P queries lock-step through the tree,
/// issuing prefetch for the next level's node while other queries overlap
/// their cache-miss latency. This is the core of Ragnar's throughput trick.
///
/// Returns `[Option<usize>; P]` — index in the original sorted slice or None.
pub fn batch_prefetch<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.batch_prefetch_avx512(queries) },
SimdPath::Avx2 => return unsafe { self.batch_prefetch_avx2(queries) },
SimdPath::Scalar => {}
}
self.batch_prefetch_impl(queries, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn batch_prefetch_avx2<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
self.batch_prefetch_impl(queries, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn batch_prefetch_avx512<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
self.batch_prefetch_impl(queries, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn batch_prefetch_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
queries: &[i64; P],
cnt: F,
) -> [Option<usize>; P] {
let height = self.offsets.len();
let mut k = [0usize; P];
// Internal levels with prefetch
for h in 0..height - 1 {
let o = self.offsets[h];
let o2 = self.offsets[h + 1];
for j in 0..P {
let jump = cnt(&self.tree[o + k[j]], queries[j]);
k[j] = k[j] * (B + 1) + jump;
// Prefetch next-level node for query j
prefetch_index(&self.tree, o2 + k[j]);
}
}
// Leaf level — no prefetch needed, just resolve
let o = self.offsets[height - 1];
let mut results = [None; P];
for j in 0..P {
let idx = cnt(&self.tree[o + k[j]], queries[j]);
let block = k[j] + idx / B;
let slot = idx % B;
let pos = block * B + slot;
if pos < self.n && self.tree[o + block][slot] == queries[j] {
results[j] = Some(pos);
}
}
results
}
/// Batch search without prefetch (baseline for comparison).
pub fn batch_no_prefetch<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
let height = self.offsets.len();
let mut k = [0usize; P];
for h in 0..height - 1 {
let o = self.offsets[h];
for j in 0..P {
let jump = count_lt(&self.tree[o + k[j]], queries[j]);
k[j] = k[j] * (B + 1) + jump;
}
}
let o = self.offsets[height - 1];
let mut results = [None; P];
for j in 0..P {
let idx = count_lt(&self.tree[o + k[j]], queries[j]);
let block = k[j] + idx / B;
let slot = idx % B;
let pos = block * B + slot;
if pos < self.n && self.tree[o + block][slot] == queries[j] {
results[j] = Some(pos);
}
}
results
}
/// Pointer-based pipelined batch (Ragnar's `batch_ptr` style).
/// Uses raw pointers and byte offsets for the tightest inner loop.
/// Safety: tree must be non-empty and offsets correct (guaranteed by constructor).
pub fn batch_ptr<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.batch_ptr_avx512(queries) },
SimdPath::Avx2 => return unsafe { self.batch_ptr_avx2(queries) },
SimdPath::Scalar => {}
}
self.batch_ptr_impl(queries, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn batch_ptr_avx2<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
self.batch_ptr_impl(queries, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn batch_ptr_avx512<const P: usize>(&self, queries: &[i64; P]) -> [Option<usize>; P] {
self.batch_ptr_impl(queries, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn batch_ptr_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
queries: &[i64; P],
cnt: F,
) -> [Option<usize>; P] {
let height = self.offsets.len();
let mut k = [0usize; P];
let offsets: Vec<*const [i64; B]> = self
.offsets
.iter()
.map(|&o| unsafe { self.tree.as_ptr().add(o) })
.collect();
// Internal levels with pointer-based prefetch
for h in 0..height - 1 {
let o = offsets[h];
let o2 = offsets[h + 1];
for j in 0..P {
let node = unsafe { &*o.add(k[j]) };
let jump = cnt(node, queries[j]);
k[j] = k[j] * (B + 1) + jump;
prefetch_ptr(unsafe { o2.add(k[j]) });
}
}
// Leaf level
let o = offsets[height - 1];
let mut results = [None; P];
for j in 0..P {
let node = unsafe { &*o.add(k[j]) };
let idx = cnt(node, queries[j]);
let block = k[j] + idx / B;
let slot = idx % B;
let pos = block * B + slot;
if pos < self.n {
let leaf_node = unsafe { &*o.add(block) };
if leaf_node[slot] == queries[j] {
results[j] = Some(pos);
}
}
}
results
}
/// Streaming batch: process a slice of arbitrary length, chunked into
/// batches of P internally with prefetch. Returns Vec<Option<usize>>.
pub fn batch_stream<const P: usize>(&self, queries: &[i64]) -> Vec<Option<usize>> {
let n = queries.len();
let mut out = Vec::with_capacity(n);
let mut i = 0;
while i + P <= n {
let chunk: &[i64; P] = queries[i..i + P].try_into().unwrap();
let results = self.batch_prefetch(chunk);
out.extend_from_slice(&results);
i += P;
}
// Remainder: single queries
for j in i..n {
out.push(self.find_exact(queries[j]));
}
out
}
}
// ── STree64Mmap pointer-based pipelined batch ──────────────────────────────────
impl STree64Mmap {
/// Pointer-based pipelined batch traversal (Ragnar's `batch_ptr` for mmap tree).
/// Uses raw pointer arithmetic for maximum throughput on the in-RAM tree,
/// then sorts positions and does sequential mmap leaf scan.
pub fn lookup_batch_ptr<const P: usize>(&self, ids: &[i64], mmap: &[u8]) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.lookup_batch_ptr_avx512::<P>(ids, mmap) },
SimdPath::Avx2 => return unsafe { self.lookup_batch_ptr_avx2::<P>(ids, mmap) },
SimdPath::Scalar => {}
}
self.lookup_batch_ptr_impl::<P, _>(ids, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_ptr_avx2<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_ptr_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_ptr_avx512<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_ptr_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn lookup_batch_ptr_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
ids: &[i64],
mmap: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
if n == 0 {
return Vec::new();
}
let offsets: Vec<*const [i64; B]> = self
.offsets
.iter()
.map(|&o| unsafe { self.tree.as_ptr().add(o) })
.collect();
let n_levels = offsets.len();
// 1. Route queries in batches of P using pointer-based prefetch
let mut routed: Vec<(usize, u32)> = Vec::with_capacity(n);
let mut i = 0usize;
while i + P <= n {
let mut kk = [0usize; P];
for h in 0..n_levels - 1 {
let o = offsets[h];
let o2 = offsets[h + 1];
for j in 0..P {
let node = unsafe { &*o.add(kk[j]) };
let jump = cnt(node, ids[i + j]);
kk[j] = kk[j] * (B + 1) + jump;
prefetch_ptr(unsafe { o2.add(kk[j]) });
}
}
// Last internal level (no next-level prefetch, but prefetch mmap leaf)
let o_last = offsets[n_levels - 1];
for j in 0..P {
let node = unsafe { &*o_last.add(kk[j]) };
let jump = cnt(node, ids[i + j]);
kk[j] = kk[j] * (B + 1) + jump;
routed.push((kk[j], (i + j) as u32));
}
i += P;
}
// Remainder
while i < n {
routed.push((self.route_to_block(ids[i]), i as u32));
i += 1;
}
// 2. Sort + madvise
routed.sort_unstable_by_key(|&(k, _)| k);
let stride = self.stride;
if let (Some(&(lo_k, _)), Some(&(hi_k, _))) = (routed.first(), routed.last()) {
let lo_byte = lo_k * B * stride;
let hi_byte = ((hi_k + 1) * B + 1).min(self.count) * stride;
if hi_byte > lo_byte {
advise_willneed(&mmap[lo_byte..hi_byte]);
}
}
// 3. Leaf scan
let mut out = vec![None; n];
for (k, orig) in routed {
let q = ids[orig as usize];
for t in 0..=B {
let pos = k * B + t;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[orig as usize] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
_ => {}
}
}
}
out
}
/// Interleaved batch: process two groups of P queries with half-tree
/// interleaving so that one group's cache misses hide behind the other's
/// compute. Inspired by Ragnar's `batch_interleave_half`.
///
/// This processes `ids` in chunks of 2*P, interleaving the first half of
/// group B's traversal with the second half of group A's traversal.
pub fn lookup_batch_interleave<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => {
return unsafe { self.lookup_batch_interleave_avx512::<P>(ids, mmap) };
}
SimdPath::Avx2 => return unsafe { self.lookup_batch_interleave_avx2::<P>(ids, mmap) },
SimdPath::Scalar => {}
}
self.lookup_batch_interleave_impl::<P, _>(ids, mmap, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_interleave_avx2<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_interleave_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx2(n, q) })
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn lookup_batch_interleave_avx512<const P: usize>(
&self,
ids: &[i64],
mmap: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_interleave_impl::<P, _>(ids, mmap, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn lookup_batch_interleave_impl<const P: usize, F: Fn(&[i64; B], i64) -> usize>(
&self,
ids: &[i64],
mmap: &[u8],
cnt: F,
) -> Vec<Option<usize>> {
let n = ids.len();
if n == 0 {
return Vec::new();
}
let n_levels = self.offsets.len();
if n_levels < 2 {
// Tree too shallow for interleaving; fall back to pipeline
return self.lookup_batch_pipeline::<P>(ids, mmap);
}
let half = n_levels / 2;
let offsets: Vec<*const [i64; B]> = self
.offsets
.iter()
.map(|&o| unsafe { self.tree.as_ptr().add(o) })
.collect();
let mut routed: Vec<(usize, u32)> = Vec::with_capacity(n);
let mut i = 0usize;
while i + 2 * P <= n {
let mut k1 = [0usize; P];
let mut k2 = [0usize; P];
// First half of group 1
for h in 0..half {
let o = offsets[h];
let o2 = offsets[h + 1];
for j in 0..P {
let node = unsafe { &*o.add(k1[j]) };
let jump = cnt(node, ids[i + j]);
k1[j] = k1[j] * (B + 1) + jump;
prefetch_ptr(unsafe { o2.add(k1[j]) });
}
}
// Interleaved: second half of group 1 + first half of group 2
let g1_remaining = n_levels - half;
let interleave_iters = g1_remaining.max(half);
for step in 0..interleave_iters {
let g1_level = half + step;
let g2_level = step;
if g1_level < n_levels {
let o1 = offsets[g1_level];
let o12 = if g1_level + 1 < n_levels {
offsets[g1_level + 1]
} else {
o1
};
for j in 0..P {
let node = unsafe { &*o1.add(k1[j]) };
let jump = cnt(node, ids[i + j]);
k1[j] = k1[j] * (B + 1) + jump;
if g1_level + 1 < n_levels {
prefetch_ptr(unsafe { o12.add(k1[j]) });
}
}
}
if g2_level < half {
let o2 = offsets[g2_level];
let o22 = offsets[g2_level + 1];
for j in 0..P {
let node = unsafe { &*o2.add(k2[j]) };
let jump = cnt(node, ids[i + P + j]);
k2[j] = k2[j] * (B + 1) + jump;
prefetch_ptr(unsafe { o22.add(k2[j]) });
}
}
}
// Finish group 2's second half
for h in half..n_levels {
let o = offsets[h];
let o2 = if h + 1 < n_levels { offsets[h + 1] } else { o };
for j in 0..P {
let node = unsafe { &*o.add(k2[j]) };
let jump = cnt(node, ids[i + P + j]);
k2[j] = k2[j] * (B + 1) + jump;
if h + 1 < n_levels {
prefetch_ptr(unsafe { o2.add(k2[j]) });
}
}
}
for j in 0..P {
routed.push((k1[j], (i + j) as u32));
}
for j in 0..P {
routed.push((k2[j], (i + P + j) as u32));
}
i += 2 * P;
}
// Remainder: fall back to route_to_block
while i < n {
routed.push((self.route_to_block(ids[i]), i as u32));
i += 1;
}
// Sort + madvise + leaf scan (same as other batch methods)
routed.sort_unstable_by_key(|&(k, _)| k);
let stride = self.stride;
if let (Some(&(lo_k, _)), Some(&(hi_k, _))) = (routed.first(), routed.last()) {
let lo_byte = lo_k * B * stride;
let hi_byte = ((hi_k + 1) * B + 1).min(self.count) * stride;
if hi_byte > lo_byte {
advise_willneed(&mmap[lo_byte..hi_byte]);
}
}
let mut out = vec![None; n];
for (k, orig) in routed {
let q = ids[orig as usize];
for t in 0..=B {
let pos = k * B + t;
if pos >= self.count {
break;
}
match id_at(mmap, pos, stride).cmp(&q) {
std::cmp::Ordering::Equal => {
out[orig as usize] = Some(pos);
break;
}
std::cmp::Ordering::Greater => break,
_ => {}
}
}
}
out
}
}
// ═══════════════════════════════════════════════════════════════════════════════
// TESTS — extensive correctness + micro-bench comparisons
// ═══════════════════════════════════════════════════════════════════════════════
#[cfg(test)]
mod tests {
use super::*;
use rand::rngs::StdRng;
use rand::{RngExt, SeedableRng};
use std::time::Instant;
/// Every kernel this CPU can run must agree with `binary_search`, on both
/// the in-RAM tree and the mmap-leaf tree, serial and batched.
///
/// Scalar is included deliberately: it is the oracle's peer, so an AVX-512
/// mask/popcount slip cannot hide behind "AVX2 agreed". On oden the
/// AVX-512 arm is skipped by the same `<= detected()` clamp production
/// uses, so this test also documents what that box actually exercises.
#[test]
fn every_simd_path_agrees_across_kernels() {
crate::stree_simd::with_each_path(|p| {
let mut rng = StdRng::seed_from_u64(0xC0DE_1234);
for &n in &[1usize, 7, 9, 73, 5000, 40_000] {
let mut ids: Vec<i64> = (0..n).map(|_| rng.random::<i64>()).collect();
ids.extend_from_slice(&[i64::MIN, i64::MAX, -1, 0]);
ids.sort_unstable();
ids.dedup();
let data = build_ids_only(&ids);
let ram = STree64::new(&ids);
let mm = STree64Mmap::new_with_stride(&data, ids.len(), 8);
for (i, &k) in ids.iter().enumerate() {
assert_eq!(ram.find_exact(k), Some(i), "{p:?} ram n={n} k={k}");
assert_eq!(mm.find_exact(k, &data), Some(i), "{p:?} mmap n={n} k={k}");
}
let fused = mm.lookup_batch_fused::<16>(&ids, &data);
let fused2 = mm.lookup_batch_fused2::<16>(&ids, &data);
let pipe = mm.lookup_batch_pipeline::<16>(&ids, &data);
let ptr = mm.lookup_batch_ptr::<16>(&ids, &data);
let stream = ram.batch_stream::<16>(&ids);
for i in 0..ids.len() {
assert_eq!(fused[i], Some(i), "{p:?} fused n={n} i={i}");
assert_eq!(fused2[i], Some(i), "{p:?} fused2 n={n} i={i}");
assert_eq!(pipe[i], Some(i), "{p:?} pipeline n={n} i={i}");
assert_eq!(ptr[i], Some(i), "{p:?} batch_ptr n={n} i={i}");
assert_eq!(stream[i], Some(i), "{p:?} stream n={n} i={i}");
}
// Absent keys must miss on every kernel too.
let absent: Vec<i64> =
ids.iter().map(|k| k.wrapping_add(1)).filter(|k| ids.binary_search(k).is_err()).collect();
assert!(
mm.lookup_batch_fused::<16>(&absent, &data).iter().all(Option::is_none),
"{p:?} absent hit n={n}"
);
}
});
}
// ── Helpers ────────────────────────────────────────────────────────────────
fn build_mmap_from_ids(ids: &[i64]) -> Vec<u8> {
let mut buf = Vec::with_capacity(ids.len() * MMAP_RECORD);
for (i, &id) in ids.iter().enumerate() {
buf.extend_from_slice(&id.to_le_bytes());
// Fake lat/lon using index for verifiability
let lat = (i as f32) * 0.001;
let lon = (i as f32) * -0.002;
buf.extend_from_slice(&lat.to_le_bytes());
buf.extend_from_slice(&lon.to_le_bytes());
}
buf
}
/// Build a pure ID column buffer (stride=8) — simulates Arrow IPC SoA layout.
fn build_ids_only(ids: &[i64]) -> Vec<u8> {
let mut buf = Vec::with_capacity(ids.len() * 8);
for &id in ids {
buf.extend_from_slice(&id.to_le_bytes());
}
buf
}
fn gen_sorted_sparse_ids(n: usize, seed: u64) -> Vec<i64> {
let mut rng = StdRng::seed_from_u64(seed);
let mut ids = Vec::with_capacity(n);
let mut cur: i64 = 1;
for _ in 0..n {
cur += rng.random_range(1..20) as i64;
ids.push(cur);
}
ids
}
fn gen_queries(ids: &[i64], n_hit: usize, n_miss: usize, seed: u64) -> Vec<i64> {
let mut rng = StdRng::seed_from_u64(seed);
let mut qs = Vec::with_capacity(n_hit + n_miss);
for _ in 0..n_hit {
qs.push(ids[rng.random_range(0..ids.len())]);
}
let max_id = ids.last().copied().unwrap_or(1000);
for _ in 0..n_miss {
qs.push(rng.random_range(0..max_id * 2) * 3 + 2); // likely misses (odd stride)
}
qs
}
// ── STree64 basic tests ───────────────────────────────────────────────────
#[test]
fn round_trip_small() {
let vals: Vec<i64> = vec![1, 3, 5, 7, 9, 11, 13, 15, 100, 200];
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i), "failed for val {v}");
}
assert_eq!(tree.find_exact(2), None);
assert_eq!(tree.find_exact(0), None);
assert_eq!(tree.find_exact(201), None);
}
#[test]
fn round_trip_large() {
let vals: Vec<i64> = (0..100_000).map(|i| i * 3).collect();
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i));
}
assert_eq!(tree.find_exact(1), None);
assert_eq!(tree.find_exact(299_999), None);
}
#[test]
fn planet_scale_ids() {
let vals: Vec<i64> = (0..1000).map(|i| i as i64 * 12_000_000).collect();
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i));
}
}
#[test]
fn single_element() {
let tree = STree64::new(&[42]);
assert_eq!(tree.find_exact(42), Some(0));
assert_eq!(tree.find_exact(41), None);
assert_eq!(tree.find_exact(43), None);
}
#[test]
fn exactly_b_elements() {
let vals: Vec<i64> = (1..=B as i64).collect();
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i));
}
assert_eq!(tree.find_exact(0), None);
assert_eq!(tree.find_exact(B as i64 + 1), None);
}
#[test]
fn b_plus_one_elements() {
let vals: Vec<i64> = (1..=B as i64 + 1).collect();
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i));
}
}
#[test]
fn negative_ids() {
let vals: Vec<i64> = (-500..500).collect();
let tree = STree64::new(&vals);
for (i, &v) in vals.iter().enumerate() {
assert_eq!(tree.find_exact(v), Some(i));
}
assert_eq!(tree.find_exact(-501), None);
assert_eq!(tree.find_exact(500), None);
}
#[test]
fn duplicates() {
let vals: Vec<i64> = vec![1, 1, 2, 2, 3, 3, 4, 4, 5, 5];
let tree = STree64::new(&vals);
// find_exact should find at least one occurrence
assert!(tree.find_exact(1).is_some());
assert!(tree.find_exact(5).is_some());
assert_eq!(tree.find_exact(6), None);
}
// ── STree64 batch tests ───────────────────────────────────────────────────
#[test]
fn batch_prefetch_correctness() {
let vals = gen_sorted_sparse_ids(50_000, 100);
let tree = STree64::new(&vals);
let qs = gen_queries(&vals, 500, 500, 200);
// Process in chunks of 16
for chunk in qs.chunks(16) {
if chunk.len() == 16 {
let arr: [i64; 16] = chunk.try_into().unwrap();
let batch_res = tree.batch_prefetch(&arr);
let serial_res: Vec<_> = chunk.iter().map(|&q| tree.find_exact(q)).collect();
for (j, (&batch, &serial)) in batch_res.iter().zip(serial_res.iter()).enumerate() {
assert_eq!(
batch, serial,
"batch_prefetch mismatch at idx {j}, query={}",
chunk[j]
);
}
}
}
}
#[test]
fn batch_no_prefetch_correctness() {
let vals = gen_sorted_sparse_ids(50_000, 101);
let tree = STree64::new(&vals);
let qs = gen_queries(&vals, 500, 500, 201);
for chunk in qs.chunks(16) {
if chunk.len() == 16 {
let arr: [i64; 16] = chunk.try_into().unwrap();
let batch_res = tree.batch_no_prefetch(&arr);
let serial_res: Vec<_> = chunk.iter().map(|&q| tree.find_exact(q)).collect();
for (j, (&batch, &serial)) in batch_res.iter().zip(serial_res.iter()).enumerate() {
assert_eq!(batch, serial, "batch_no_prefetch mismatch at idx {j}");
}
}
}
}
#[test]
fn batch_ptr_correctness() {
let vals = gen_sorted_sparse_ids(50_000, 102);
let tree = STree64::new(&vals);
let qs = gen_queries(&vals, 500, 500, 202);
for chunk in qs.chunks(16) {
if chunk.len() == 16 {
let arr: [i64; 16] = chunk.try_into().unwrap();
let batch_res = tree.batch_ptr(&arr);
let serial_res: Vec<_> = chunk.iter().map(|&q| tree.find_exact(q)).collect();
for (j, (&batch, &serial)) in batch_res.iter().zip(serial_res.iter()).enumerate() {
assert_eq!(batch, serial, "batch_ptr mismatch at idx {j}");
}
}
}
}
#[test]
fn batch_stream_correctness() {
let vals = gen_sorted_sparse_ids(50_000, 103);
let tree = STree64::new(&vals);
let qs = gen_queries(&vals, 1000, 1000, 203);
let stream_res = tree.batch_stream::<16>(&qs);
let serial_res: Vec<_> = qs.iter().map(|&q| tree.find_exact(q)).collect();
assert_eq!(stream_res, serial_res);
// P=8
let stream8 = tree.batch_stream::<8>(&qs);
assert_eq!(stream8, serial_res);
// P=32
let stream32 = tree.batch_stream::<32>(&qs);
assert_eq!(stream32, serial_res);
}
#[test]
fn batch_stream_various_sizes() {
let vals = gen_sorted_sparse_ids(10_000, 104);
let tree = STree64::new(&vals);
// test sizes that are NOT multiples of P
for size in [
0, 1, 7, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 255, 256, 1000,
] {
let qs = gen_queries(&vals, size / 2 + 1, size / 2, 300 + size as u64);
let qs = &qs[..size.min(qs.len())];
let stream_res = tree.batch_stream::<16>(qs);
let serial_res: Vec<_> = qs.iter().map(|&q| tree.find_exact(q)).collect();
assert_eq!(
stream_res, serial_res,
"batch_stream mismatch for size={size}"
);
}
}
// ── STree64Mmap basic tests ───────────────────────────────────────────────
#[test]
fn mmap_find_exact_basic() {
let ids = gen_sorted_sparse_ids(10_000, 50);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
for (i, &id) in ids.iter().enumerate() {
assert_eq!(
tree.find_exact(id, &mmap),
Some(i),
"find_exact failed for id={id} at pos={i}"
);
}
// Missing ids
assert_eq!(tree.find_exact(0, &mmap), None);
assert_eq!(tree.find_exact(ids.last().unwrap() + 1, &mmap), None);
}
#[test]
fn mmap_find_exact_single() {
let ids = vec![999i64];
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, 1);
assert_eq!(tree.find_exact(999, &mmap), Some(0));
assert_eq!(tree.find_exact(998, &mmap), None);
}
// ── STree64Mmap batch correctness ─────────────────────────────────────────
#[test]
fn mmap_lookup_batch_vs_serial() {
let ids = gen_sorted_sparse_ids(20_000, 60);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 2_500, 2_500, 61);
let batch = tree.lookup_batch(&qs, &mmap);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
assert_eq!(batch, serial);
}
#[test]
fn mmap_pipeline_vs_serial_all_p() {
let ids = gen_sorted_sparse_ids(20_000, 70);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 2_500, 2_500, 71);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let p4 = tree.lookup_batch_pipeline::<4>(&qs, &mmap);
let p8 = tree.lookup_batch_pipeline::<8>(&qs, &mmap);
let p16 = tree.lookup_batch_pipeline::<16>(&qs, &mmap);
let p32 = tree.lookup_batch_pipeline::<32>(&qs, &mmap);
assert_eq!(p4, serial, "pipeline P=4 mismatch");
assert_eq!(p8, serial, "pipeline P=8 mismatch");
assert_eq!(p16, serial, "pipeline P=16 mismatch");
assert_eq!(p32, serial, "pipeline P=32 mismatch");
}
#[test]
fn mmap_batch_ptr_vs_serial() {
let ids = gen_sorted_sparse_ids(20_000, 80);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 2_500, 2_500, 81);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let ptr8 = tree.lookup_batch_ptr::<8>(&qs, &mmap);
let ptr16 = tree.lookup_batch_ptr::<16>(&qs, &mmap);
let ptr32 = tree.lookup_batch_ptr::<32>(&qs, &mmap);
assert_eq!(ptr8, serial, "batch_ptr P=8 mismatch");
assert_eq!(ptr16, serial, "batch_ptr P=16 mismatch");
assert_eq!(ptr32, serial, "batch_ptr P=32 mismatch");
}
#[test]
fn mmap_interleave_vs_serial() {
let ids = gen_sorted_sparse_ids(20_000, 90);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 2_500, 2_500, 91);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let interleave8 = tree.lookup_batch_interleave::<8>(&qs, &mmap);
let interleave16 = tree.lookup_batch_interleave::<16>(&qs, &mmap);
assert_eq!(interleave8, serial, "interleave P=8 mismatch");
assert_eq!(interleave16, serial, "interleave P=16 mismatch");
}
// ── Edge cases ────────────────────────────────────────────────────────────
#[test]
fn mmap_batch_empty_queries() {
let ids = gen_sorted_sparse_ids(1000, 110);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
assert_eq!(tree.lookup_batch(&[], &mmap), Vec::<Option<usize>>::new());
assert_eq!(
tree.lookup_batch_pipeline::<16>(&[], &mmap),
Vec::<Option<usize>>::new()
);
assert_eq!(
tree.lookup_batch_ptr::<16>(&[], &mmap),
Vec::<Option<usize>>::new()
);
assert_eq!(
tree.lookup_batch_interleave::<8>(&[], &mmap),
Vec::<Option<usize>>::new()
);
}
#[test]
fn mmap_batch_single_query() {
let ids = gen_sorted_sparse_ids(1000, 111);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let hit = vec![ids[500]];
// The generator steps by 1..20, so `ids[500] + 1` CAN be an id. The rand
// 0.10 seeded stream (W3, 2026-10-02) produced exactly that and flipped this
// assertion; walk forward to the first value that is genuinely absent.
let mut absent = ids[500] + 1;
while ids.binary_search(&absent).is_ok() {
absent += 1;
}
let miss = vec![absent];
assert_eq!(
tree.lookup_batch_pipeline::<4>(&hit, &mmap),
vec![Some(500)]
);
assert_eq!(tree.lookup_batch_pipeline::<4>(&miss, &mmap), vec![None]);
assert_eq!(tree.lookup_batch_ptr::<4>(&hit, &mmap), vec![Some(500)]);
}
#[test]
fn mmap_batch_all_hits() {
let ids = gen_sorted_sparse_ids(5_000, 120);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs: Vec<i64> = ids.iter().copied().collect();
let expected: Vec<_> = (0..ids.len()).map(Some).collect();
assert_eq!(tree.lookup_batch(&qs, &mmap), expected);
assert_eq!(tree.lookup_batch_pipeline::<16>(&qs, &mmap), expected);
assert_eq!(tree.lookup_batch_ptr::<16>(&qs, &mmap), expected);
}
#[test]
fn mmap_batch_all_misses() {
let ids = gen_sorted_sparse_ids(5_000, 130);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
// Queries guaranteed to miss (even ids when stored are odd-spaced)
let qs: Vec<i64> = (0..1000).map(|i| ids.last().unwrap() + i + 1).collect();
let expected: Vec<_> = vec![None; 1000];
assert_eq!(tree.lookup_batch(&qs, &mmap), expected);
assert_eq!(tree.lookup_batch_pipeline::<16>(&qs, &mmap), expected);
assert_eq!(tree.lookup_batch_ptr::<16>(&qs, &mmap), expected);
}
#[test]
fn mmap_batch_remainder_handling() {
// Test with sizes that leave various remainders for P
let ids = gen_sorted_sparse_ids(5_000, 140);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
for qlen in [1, 3, 7, 15, 16, 17, 31, 32, 33, 63, 100] {
let qs = gen_queries(&ids, qlen, 0, 140 + qlen as u64);
let qs = &qs[..qlen.min(qs.len())];
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let p16 = tree.lookup_batch_pipeline::<16>(qs, &mmap);
let ptr16 = tree.lookup_batch_ptr::<16>(qs, &mmap);
assert_eq!(p16, serial, "pipeline remainder mismatch for qlen={qlen}");
assert_eq!(
ptr16, serial,
"batch_ptr remainder mismatch for qlen={qlen}"
);
}
}
// ── Larger stress tests ───────────────────────────────────────────────────
#[test]
fn stress_100k_ids_10k_queries() {
let ids = gen_sorted_sparse_ids(100_000, 200);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 5_000, 5_000, 201);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let batch = tree.lookup_batch(&qs, &mmap);
let pipe16 = tree.lookup_batch_pipeline::<16>(&qs, &mmap);
let ptr16 = tree.lookup_batch_ptr::<16>(&qs, &mmap);
let interleave8 = tree.lookup_batch_interleave::<8>(&qs, &mmap);
assert_eq!(batch, serial);
assert_eq!(pipe16, serial);
assert_eq!(ptr16, serial);
assert_eq!(interleave8, serial);
}
#[test]
fn stress_osm_like_ids() {
// Simulate OSM-like node IDs: start at ~1B, gaps of 1-100
let mut rng = StdRng::seed_from_u64(300);
let n = 50_000;
let mut ids = Vec::with_capacity(n);
let mut cur: i64 = 1_000_000_000;
for _ in 0..n {
cur += rng.random_range(1..100) as i64;
ids.push(cur);
}
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 5_000, 5_000, 301);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
assert_eq!(tree.lookup_batch_pipeline::<16>(&qs, &mmap), serial);
assert_eq!(tree.lookup_batch_ptr::<16>(&qs, &mmap), serial);
assert_eq!(tree.lookup_batch_interleave::<8>(&qs, &mmap), serial);
}
// ── Stride=8 (Arrow IPC / SoA) tests ─────────────────────────────────────
#[test]
fn stride8_find_exact() {
let ids = gen_sorted_sparse_ids(10_000, 500);
let id_buf = build_ids_only(&ids);
let tree = STree64Mmap::new_with_stride(&id_buf, ids.len(), 8);
for (i, &id) in ids.iter().enumerate() {
assert_eq!(tree.find_exact(id, &id_buf), Some(i));
}
assert_eq!(tree.find_exact(ids[0] - 1, &id_buf), None);
assert_eq!(tree.find_exact(*ids.last().unwrap() + 1, &id_buf), None);
}
#[test]
fn stride8_batch_all_methods() {
let ids = gen_sorted_sparse_ids(50_000, 510);
let id_buf = build_ids_only(&ids);
let tree = STree64Mmap::new_with_stride(&id_buf, ids.len(), 8);
let qs = gen_queries(&ids, 5_000, 5_000, 511);
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &id_buf)).collect();
let batch = tree.lookup_batch(&qs, &id_buf);
let pipe = tree.lookup_batch_pipeline::<16>(&qs, &id_buf);
let ptr = tree.lookup_batch_ptr::<16>(&qs, &id_buf);
let interleave = tree.lookup_batch_interleave::<8>(&qs, &id_buf);
assert_eq!(batch, serial, "stride=8 lookup_batch");
assert_eq!(pipe, serial, "stride=8 pipeline");
assert_eq!(ptr, serial, "stride=8 ptr");
assert_eq!(interleave, serial, "stride=8 interleave");
}
#[test]
fn stride8_all_hits() {
let ids = gen_sorted_sparse_ids(5_000, 520);
let id_buf = build_ids_only(&ids);
let tree = STree64Mmap::new_with_stride(&id_buf, ids.len(), 8);
let expected: Vec<_> = (0..ids.len()).map(Some).collect();
assert_eq!(tree.lookup_batch(&ids, &id_buf), expected);
assert_eq!(tree.lookup_batch_pipeline::<8>(&ids, &id_buf), expected);
}
// ── Micro-benchmark (prints timing, not gated) ────────────────────────────
#[test]
fn bench_compare_all_methods() {
// 500K records, 100K queries — run with `cargo test --release -- --ignored`
let ids = gen_sorted_sparse_ids(500_000, 400);
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 50_000, 50_000, 401);
let n = qs.len();
// Warm up
let _ = tree.lookup_batch(&qs, &mmap);
let t0 = Instant::now();
let _r1 = tree.lookup_batch(&qs, &mmap);
let d1 = t0.elapsed();
let t0 = Instant::now();
let _r2 = tree.lookup_batch_pipeline::<16>(&qs, &mmap);
let d2 = t0.elapsed();
let t0 = Instant::now();
let _r3 = tree.lookup_batch_pipeline::<32>(&qs, &mmap);
let d3 = t0.elapsed();
let t0 = Instant::now();
let _r4 = tree.lookup_batch_ptr::<16>(&qs, &mmap);
let d4 = t0.elapsed();
let t0 = Instant::now();
let _r5 = tree.lookup_batch_ptr::<32>(&qs, &mmap);
let d5 = t0.elapsed();
let t0 = Instant::now();
let _r6 = tree.lookup_batch_interleave::<8>(&qs, &mmap);
let d6 = t0.elapsed();
let t0 = Instant::now();
let _r7 = tree.lookup_batch_interleave::<16>(&qs, &mmap);
let d7 = t0.elapsed();
// Serial baseline
let t0 = Instant::now();
let _r0: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
let d0 = t0.elapsed();
eprintln!(
"\n=== STree64Mmap batch benchmark ({n} queries, {} records) ===",
ids.len()
);
eprintln!(
" serial find_exact: {:>8.3} ms ({:.0} ns/query)",
d0.as_secs_f64() * 1000.0,
d0.as_nanos() as f64 / n as f64
);
eprintln!(
" lookup_batch (sort+adv): {:>8.3} ms ({:.0} ns/query)",
d1.as_secs_f64() * 1000.0,
d1.as_nanos() as f64 / n as f64
);
eprintln!(
" pipeline P=16: {:>8.3} ms ({:.0} ns/query)",
d2.as_secs_f64() * 1000.0,
d2.as_nanos() as f64 / n as f64
);
eprintln!(
" pipeline P=32: {:>8.3} ms ({:.0} ns/query)",
d3.as_secs_f64() * 1000.0,
d3.as_nanos() as f64 / n as f64
);
eprintln!(
" batch_ptr P=16: {:>8.3} ms ({:.0} ns/query)",
d4.as_secs_f64() * 1000.0,
d4.as_nanos() as f64 / n as f64
);
eprintln!(
" batch_ptr P=32: {:>8.3} ms ({:.0} ns/query)",
d5.as_secs_f64() * 1000.0,
d5.as_nanos() as f64 / n as f64
);
eprintln!(
" interleave P=8: {:>8.3} ms ({:.0} ns/query)",
d6.as_secs_f64() * 1000.0,
d6.as_nanos() as f64 / n as f64
);
eprintln!(
" interleave P=16: {:>8.3} ms ({:.0} ns/query)",
d7.as_secs_f64() * 1000.0,
d7.as_nanos() as f64 / n as f64
);
}
#[test]
fn bench_stree64_batch_variants() {
// In-RAM tree (STree64) batch comparison — 1M records, 500K queries
let ids = gen_sorted_sparse_ids(1_000_000, 500);
let tree = STree64::new(&ids);
let qs = gen_queries(&ids, 250_000, 250_000, 501);
let n = qs.len();
// Warm up
let _ = tree.batch_stream::<16>(&qs);
let t0 = Instant::now();
let _r0: Vec<_> = qs.iter().map(|&q| tree.find_exact(q)).collect();
let d0 = t0.elapsed();
let t0 = Instant::now();
let _r1 = tree.batch_stream::<8>(&qs);
let d1 = t0.elapsed();
let t0 = Instant::now();
let _r2 = tree.batch_stream::<16>(&qs);
let d2 = t0.elapsed();
let t0 = Instant::now();
let _r3 = tree.batch_stream::<32>(&qs);
let d3 = t0.elapsed();
eprintln!(
"\n=== STree64 batch benchmark ({n} queries, {} records) ===",
ids.len()
);
eprintln!(
" serial find_exact: {:>8.3} ms ({:.0} ns/query)",
d0.as_secs_f64() * 1000.0,
d0.as_nanos() as f64 / n as f64
);
eprintln!(
" batch_stream P=8: {:>8.3} ms ({:.0} ns/query)",
d1.as_secs_f64() * 1000.0,
d1.as_nanos() as f64 / n as f64
);
eprintln!(
" batch_stream P=16: {:>8.3} ms ({:.0} ns/query)",
d2.as_secs_f64() * 1000.0,
d2.as_nanos() as f64 / n as f64
);
eprintln!(
" batch_stream P=32: {:>8.3} ms ({:.0} ns/query)",
d3.as_secs_f64() * 1000.0,
d3.as_nanos() as f64 / n as f64
);
}
#[test]
fn bench_compare_1m_records_osm_ids() {
// Simulate planet-like conditions: 1M records with OSM-style IDs
let mut rng = StdRng::seed_from_u64(600);
let n = 1_000_000;
let mut ids = Vec::with_capacity(n);
let mut cur: i64 = 100_000_000;
for _ in 0..n {
cur += rng.random_range(1..1000) as i64;
ids.push(cur);
}
let mmap = build_mmap_from_ids(&ids);
let tree = STree64Mmap::new(&mmap, ids.len());
let qs = gen_queries(&ids, 100_000, 100_000, 601);
let n_q = qs.len();
// Verify correctness first
let serial: Vec<_> = qs.iter().map(|&q| tree.find_exact(q, &mmap)).collect();
assert_eq!(tree.lookup_batch_pipeline::<16>(&qs, &mmap), serial);
assert_eq!(tree.lookup_batch_ptr::<16>(&qs, &mmap), serial);
// Benchmark
let t0 = Instant::now();
let _ = qs
.iter()
.map(|&q| tree.find_exact(q, &mmap))
.collect::<Vec<_>>();
let d_serial = t0.elapsed();
let t0 = Instant::now();
let _ = tree.lookup_batch(&qs, &mmap);
let d_batch = t0.elapsed();
let t0 = Instant::now();
let _ = tree.lookup_batch_pipeline::<16>(&qs, &mmap);
let d_pipe = t0.elapsed();
let t0 = Instant::now();
let _ = tree.lookup_batch_ptr::<16>(&qs, &mmap);
let d_ptr = t0.elapsed();
eprintln!("\n=== 1M OSM-like records, {n_q} queries ===");
eprintln!(
" serial: {:>8.3} ms ({:.0} ns/q)",
d_serial.as_secs_f64() * 1000.0,
d_serial.as_nanos() as f64 / n_q as f64
);
eprintln!(
" batch: {:>8.3} ms ({:.0} ns/q)",
d_batch.as_secs_f64() * 1000.0,
d_batch.as_nanos() as f64 / n_q as f64
);
eprintln!(
" pipeline: {:>8.3} ms ({:.0} ns/q)",
d_pipe.as_secs_f64() * 1000.0,
d_pipe.as_nanos() as f64 / n_q as f64
);
eprintln!(
" batch_ptr: {:>8.3} ms ({:.0} ns/q)",
d_ptr.as_secs_f64() * 1000.0,
d_ptr.as_nanos() as f64 / n_q as f64
);
}
}