src/stree_hi32.rs
//! Hi32-routing S-tree: **u32 high-half internal nodes over an i64 leaf
//! column** — the fanout-17 geometry without the two-cache-line node.
//!
//! Derivative of Ragnar Groot Koerkamp's static search tree
//! (<https://curiouscoding.nl/posts/static-search-tree/>), same as
//! [`crate::stree`]. The 2026-08-10 session measured that an i64 B=16 node
//! (128 B, two lines) loses its level-count win to doubled line traffic
//! ([`crate::stree512`], refuted). This variant gets fanout 17 AND one-line
//! nodes by routing on a **monotone 32-bit truncation** of the key:
//!
//! ```text
//! hi(q) = ((q as u64 ^ (1<<63)) >> 32) as u32 // order-preserving
//! ```
//!
//! Internal nodes hold `hi(first key of child leaf block)` — 16 × u32 = 64 B,
//! exactly Ragnar's original node, compared with the same two
//! `_mm256_cmpgt_epi32` as [`crate::stree32`]. The leaf stays the full-width
//! i64 column (8 keys / 64 B block).
//!
//! **64 B == 512 bits == one zmm register**, so this geometry gets the whole
//! node compare in a SINGLE `_mm512_cmplt_epu32_mask` where AVX2 needs two
//! biased compares plus two movemasks — see [`count_lt_avx512`]. Which kernel
//! runs is resolved once per process by [`crate::stree_simd`] (CPUID-clamped),
//! so the same binary is optimal on an AVX-512 box and correct on one without.
//!
//! **Truncation is predecessor-safe.** `count_lt` on `hi` values undercounts
//! (keys equal in the high half but `< q` in the low half are not counted),
//! so routing can only land AT or BEFORE the true block, never after. The
//! leaf scan therefore walks forward `while id < q`, unbounded by block —
//! for sha-derived (uniform) keys the expected overshoot is well under one
//! block; a pathological key set degrades speed, never correctness.
//!
//! Levels at 1.34e8 keys: log₁₇(16.8M blocks) ≈ 5.9 internal + leaf ≈ 7 line
//! touches vs the B=8 tree's ≈ 9. The internal tree also halves in bytes
//! (~66 MB vs ~153 MB), so more of its upper levels stay cache-resident.
const BI: usize = 16; // u32 slots per internal node = fanout - 1
const BL: usize = 8; // i64 keys per leaf block (one cache line)
const MAXU: u32 = u32::MAX;
/// Resolved-once kernel selector — see [`crate::stree_simd`].
#[cfg(target_arch = "x86_64")]
use crate::stree_simd::SimdPath;
#[inline(always)]
fn hi(q: i64) -> u32 {
(((q as u64) ^ (1u64 << 63)) >> 32) as u32
}
#[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())
}
#[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(not(target_arch = "x86_64"))]
{
let _ = ptr;
}
}
// ── shape (fanout 17 over leaf-block indices; same recurrences as stree32) ────
fn blocks(n: usize) -> usize {
n.div_ceil(BI)
}
fn prev_keys(n: usize) -> usize {
blocks(n).div_ceil(BI + 1) * BI
}
fn height(n: usize) -> usize {
if n <= BI { 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(BI)).collect()
}
// ── aligned node storage (same fix lineage as stree::NodeBuf) ─────────────────
struct Nodes {
buf: Vec<u32>,
off: usize,
len: usize,
}
impl Nodes {
fn alloc(n_nodes: usize) -> Self {
let len = n_nodes * BI + 32;
let mut buf: Vec<u32> = Vec::with_capacity(len);
crate::stree::advise_hugepage_raw(buf.as_ptr() as *const u8, len * 4);
buf.resize(len, MAXU);
let addr = buf.as_ptr() as usize;
let off = ((64 - (addr & 63)) & 63) / 4;
Nodes { buf, off, len: n_nodes }
}
}
impl std::ops::Deref for Nodes {
type Target = [[u32; BI]];
#[inline(always)]
fn deref(&self) -> &[[u32; BI]] {
unsafe {
std::slice::from_raw_parts(self.buf.as_ptr().add(self.off) as *const [u32; BI], self.len)
}
}
}
impl std::ops::DerefMut for Nodes {
#[inline(always)]
fn deref_mut(&mut self) -> &mut [[u32; BI]] {
unsafe {
std::slice::from_raw_parts_mut(
self.buf.as_mut_ptr().add(self.off) as *mut [u32; BI],
self.len,
)
}
}
}
// ── node compare (unsigned via sign-bias, as in stree32) ──────────────────────
#[inline(always)]
fn count_lt_scalar(node: &[u32; BI], q: u32) -> usize {
node.iter().filter(|&&x| x < q).count()
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn count_lt_avx2(node: &[u32; BI], q: u32) -> usize {
use std::arch::x86_64::*;
unsafe {
let bias = _mm256_set1_epi32(i32::MIN);
let q_v = _mm256_xor_si256(_mm256_set1_epi32(q as i32), bias);
let mut bits = 0u32;
let mut i = 0;
while i < BI {
// Pairwise reduction — see `crate::stree32::count_lt_avx2` for why the
// two compares are packed into one mask instead of two `vpmovmskb` +
// two `popcnt` + a `/4` (2026-08-11, measured on oden).
let lo = _mm256_cmpgt_epi32(
q_v,
_mm256_xor_si256(
_mm256_loadu_si256(node.as_ptr().add(i) as *const __m256i),
bias,
),
);
let hi = _mm256_cmpgt_epi32(
q_v,
_mm256_xor_si256(
_mm256_loadu_si256(node.as_ptr().add(i + 8) as *const __m256i),
bias,
),
);
let packed = _mm256_packs_epi32(lo, hi);
bits += (_mm256_movemask_epi8(packed) as u32).count_ones();
i += 16;
}
(bits / 2) as usize
}
}
/// AVX-512 path: **the whole node in one instruction.**
///
/// `BI * 4 == 64` bytes == 512 bits, so a hi32 internal node is exactly one
/// zmm register — the geometry that already wins on cache-line traffic also
/// happens to be the widest single compare x86 has. Where
/// [`count_lt_avx2`] needs two loads, two xors (sign-bias, because AVX2 has no
/// unsigned dword compare), two `cmpgt`, two `movemask` and a `/4`, this is one
/// load, one `cmplt_epu32_mask` and a `popcnt`: AVX-512 compares unsigned
/// natively, so the bias vector disappears too, and the mask is already one bit
/// per element rather than four bytes.
///
/// Honest caveat, from [`crate::stree512`]'s header: on t14s (Strix-class Zen 5
/// mobile) AVX-512 runs on a 256-bit datapath, so this decodes to ~two 256-bit
/// µops and the *compute* is not halved. The instruction-count win is real; the
/// throughput win may not be. Measure, do not assume.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
#[inline]
unsafe fn count_lt_avx512(node: &[u32; BI], q: u32) -> usize {
use std::arch::x86_64::*;
unsafe {
let v = _mm512_loadu_epi32(node.as_ptr() as *const i32);
let m: __mmask16 = _mm512_cmplt_epu32_mask(v, _mm512_set1_epi32(q as i32));
m.count_ones() as usize
}
}
// ── the tree ──────────────────────────────────────────────────────────────────
/// Internal-nodes-only Hi32 tree; the sorted i64 column in `data` (one key
/// every `stride` bytes) is the leaf layer.
pub struct STreeHi32Mmap {
tree: Nodes,
offsets: Vec<usize>,
pub count: usize,
pub stride: usize,
}
impl STreeHi32Mmap {
pub fn new_with_stride(data: &[u8], count: usize, stride: usize) -> Self {
assert!(count > 0);
assert!(stride >= 8);
let n_leaf_blocks = count.div_ceil(BL);
let h = height(n_leaf_blocks);
// `h` internal levels route over leaf-block indices; there is no leaf
// level inside the node buffer (mirrors STree64Mmap's ni = h-1 — but
// here the shape is computed over BLOCK indices, so every level is
// internal and all `h` levels are stored).
let lsizes = layer_sizes(n_leaf_blocks, h);
let n_nodes: usize = lsizes.iter().sum();
let mut offsets = Vec::with_capacity(h);
let mut acc = 0usize;
for &ls in &lsizes {
offsets.push(acc);
acc += ls;
}
let mut tree = Nodes::alloc(n_nodes.max(1));
// Bottom-up fill, same index algebra as the other trees but over
// BLOCK indices: node slot j at (level lvl, block blk) holds the hi32
// of the first key of leaf block k, k = child-subtree walk to depth.
for lvl in (0..h).rev() {
let oh = offsets[lvl];
for blk in 0..lsizes[lvl] {
let node = &mut tree[oh + blk];
for j in 0..BI {
let mut k = blk * (BI + 1) + j + 1;
for _ in lvl..h - 1 {
k *= BI + 1;
}
node[j] = if k * BL < count {
hi(id_at(data, k * BL, stride))
} else {
MAXU
};
}
}
}
crate::stree::advise_hugepage_raw(
tree[..].as_ptr() as *const u8,
std::mem::size_of_val(&tree[..]),
);
Self { tree, offsets, count, stride }
}
/// Forward scan from leaf block `k` while `id < q` — the truncation-safe
/// leaf resolve. Never bounded to one block: an equal-hi32 run can span
/// blocks, and routing lands at the run's first block.
#[inline(always)]
fn leaf_resolve(&self, k: usize, q: i64, data: &[u8]) -> Option<usize> {
let stride = self.stride;
let mut pos = k * BL;
while pos < self.count {
let id = id_at(data, pos, stride);
if id == q {
return Some(pos);
}
if id > q {
return None;
}
pos += 1;
}
None
}
/// Kernel selection happens **here**, once per query — never per node. The
/// resolved path is a relaxed byte load (see [`crate::stree_simd`]); the
/// chosen `#[target_feature]` wrapper then inlines the whole descent, so
/// the compare is a straight-line instruction inside the loop.
#[inline]
pub fn find_exact(&self, q: i64, data: &[u8]) -> Option<usize> {
#[cfg(target_arch = "x86_64")]
match crate::stree_simd::path() {
SimdPath::Avx512 => return unsafe { self.find_exact_avx512(q, data) },
SimdPath::Avx2 => return unsafe { self.find_exact_avx2(q, data) },
SimdPath::Scalar => {}
}
self.find_exact_impl(q, data, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_exact_avx2(&self, q: i64, data: &[u8]) -> Option<usize> {
self.find_exact_impl(q, data, |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, data: &[u8]) -> Option<usize> {
self.find_exact_impl(q, data, |n, qq| unsafe { count_lt_avx512(n, qq) })
}
#[inline(always)]
fn find_exact_impl<F: Fn(&[u32; BI], u32) -> usize>(
&self,
q: i64,
data: &[u8],
cnt: F,
) -> Option<usize> {
let hq = hi(q);
let mut k = 0usize;
for &o in &self.offsets {
let jump = cnt(&self.tree[o + k], hq);
k = k * (BI + 1) + jump;
}
self.leaf_resolve(k, q, data)
}
/// The fused single-phase batch walk, Hi32 geometry.
pub fn lookup_batch_fused<const P: usize>(&self, ids: &[i64], data: &[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, data) };
}
SimdPath::Avx2 => return unsafe { self.lookup_batch_fused_avx2::<P>(ids, data) },
SimdPath::Scalar => {}
}
self.lookup_batch_fused_impl::<P, _>(ids, data, count_lt_scalar)
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn lookup_batch_fused_avx2<const P: usize>(
&self,
ids: &[i64],
data: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, data, |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],
data: &[u8],
) -> Vec<Option<usize>> {
self.lookup_batch_fused_impl::<P, _>(ids, data, |n, q| unsafe { count_lt_avx512(n, q) })
}
#[inline(always)]
fn lookup_batch_fused_impl<const P: usize, F: Fn(&[u32; BI], u32) -> usize>(
&self,
ids: &[i64],
data: &[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];
let mut hq = [0u32; P];
for j in 0..P {
hq[j] = hi(ids[i + j]);
}
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]], hq[j]);
kk[j] = kk[j] * (BI + 1) + jump;
prefetch_ptr(&self.tree[o2 + kk[j]] as *const [u32; BI]);
}
}
if n_levels > 0 {
let o = self.offsets[n_levels - 1];
for j in 0..P {
let jump = cnt(&self.tree[o + kk[j]], hq[j]);
kk[j] = kk[j] * (BI + 1) + jump;
let byte = (kk[j] * BL).min(self.count - 1) * stride;
let base = unsafe { data.as_ptr().add(byte.min(data.len() - 1)) };
prefetch_ptr(base);
prefetch_ptr(base.wrapping_add(64));
}
}
for j in 0..P {
out[i + j] = self.leaf_resolve(kk[j], ids[i + j], data);
}
i += P;
}
while i < n {
out[i] = self.find_exact(ids[i], data);
i += 1;
}
out
}
}
#[cfg(test)]
mod tests {
use super::*;
fn keybuf(keys: &[i64]) -> Vec<u8> {
let mut v = Vec::with_capacity(keys.len() * 8);
for k in keys {
v.extend_from_slice(&k.to_le_bytes());
}
v
}
/// Same corpus, every kernel this CPU has. Scalar is the oracle's peer, so
/// an AVX-512 mask bug cannot hide behind "AVX2 agreed with AVX-512".
#[test]
fn every_simd_path_agrees_with_binary_search() {
crate::stree_simd::with_each_path(|p| {
let mut state = 0x0bad_f00d_1234_5678u64;
let mut next = move || {
state ^= state >> 12;
state ^= state << 25;
state ^= state >> 27;
state.wrapping_mul(0x2545_F491_4F6C_DD1D)
};
for &n in &[1usize, 8, 17, 4096, 60_000] {
let mut keys: Vec<i64> = (0..n).map(|_| next() as i64).collect();
// A hi32 value with the top bit of the biased key set exercises
// the unsigned compare: AVX2 needs the sign-bias xor, AVX-512
// compares epu32 natively. A signed-compare slip shows up here.
keys.push(i64::MIN);
keys.push(i64::MAX);
keys.push(-1);
keys.push(0);
keys.sort_unstable();
keys.dedup();
let data = keybuf(&keys);
let t = STreeHi32Mmap::new_with_stride(&data, keys.len(), 8);
for (i, &k) in keys.iter().enumerate() {
assert_eq!(t.find_exact(k, &data), Some(i), "{p:?} n={n} key={k}");
}
let b = t.lookup_batch_fused::<16>(&keys, &data);
for i in 0..keys.len() {
assert_eq!(b[i], Some(i), "{p:?} fused n={n} i={i}");
}
}
});
}
#[test]
fn matches_binary_search_including_hi32_collisions() {
let mut state = 0xfeed_beef_dead_c0deu64;
let mut next = move || {
state ^= state >> 12;
state ^= state << 25;
state ^= state >> 27;
state.wrapping_mul(0x2545_F491_4F6C_DD1D)
};
for &n in &[1usize, 8, 9, 17, 200, 4096, 200_000] {
let mut keys: Vec<i64> = (0..n).map(|_| next() as i64).collect();
// Force hi32 collision runs: clusters sharing the top 32 bits,
// differing only in the low half — the truncation's worst case.
if n >= 200 {
let base = keys[0] & !0xffff_ffffi64;
for c in 0..64.min(n / 3) {
keys[c] = base | (c as i64 * 7 + 1);
}
}
keys.sort_unstable();
keys.dedup();
let data = keybuf(&keys);
let t = STreeHi32Mmap::new_with_stride(&data, keys.len(), 8);
let qs: Vec<i64> =
(0..300.min(keys.len())).map(|_| keys[(next() as usize) % keys.len()]).collect();
for &q in &qs {
assert_eq!(t.find_exact(q, &data), keys.binary_search(&q).ok(), "serial n={n}");
}
let batch = t.lookup_batch_fused::<8>(&qs, &data);
for (i, &q) in qs.iter().enumerate() {
assert_eq!(batch[i], keys.binary_search(&q).ok(), "fused n={n} i={i}");
}
let absent: Vec<i64> = (0..200)
.map(|_| next() as i64)
.filter(|k| keys.binary_search(k).is_err())
.collect();
assert!(t.lookup_batch_fused::<8>(&absent, &data).iter().all(Option::is_none));
}
}
}