Skip to main content

hermes_core/structures/vector/quantization/
tq.rs

1//! TurboQuant (TQ): training-free dense-vector codec.
2//!
3//! Design: `docs/turboquant-quantization.md`. Per padded coordinate the codec
4//! stores a 3-bit scalar code (analytic Lloyd-Max levels for the unit-sphere
5//! coordinate density) plus a 1-bit QJL sign of the rotated stage-1 residual;
6//! a per-vector f32 `gamma = ‖residual‖₂` makes the inner-product estimator
7//! unbiased. Everything is derived from `(dim, codec constants)` — no trained
8//! artifacts and no cross-segment generations.
9//!
10//! References: Zandieh, Daliri et al., "TurboQuant: Online Vector
11//! Quantization with Near-optimal Distortion Rate" (arXiv 2504.19874);
12//! layout and LUT16 scoring follow the FastScan pattern (Faiss,
13//! mayflower/pg_turboquant, both MIT).
14
15/// Bumping this refuses to mix payloads across incompatible codec revisions.
16/// v2: padding-free 3-round rotation (sub-FWHT + signs + permutation per
17/// round) replaced the single-round power-of-two-padded FWHT — 768-dim codes
18/// shrank 33% and the codebook density now uses the true dimension.
19pub const TQ_CODEC_VERSION: u32 = 2;
20/// Bits per padded coordinate: 3-bit stage-1 code + 1-bit QJL sign.
21pub const TQ_BITS: u32 = 4;
22/// Vectors per scoring block; one lane per vector.
23pub const TQ_BLOCK_LANES: usize = 16;
24/// Smallest supported padded dimension. Below this the coordinate density
25/// exponent `(P-3)/2` degenerates and LUT rows would not fill a SIMD lane.
26pub const TQ_MIN_PADDED_DIM: usize = 8;
27
28const TQ_STAGE1_LEVELS: usize = 8;
29const TQ_STAGE1_SEED: u64 = 0x7154_5354_4147_4531; // "qTSTAGE1"
30const TQ_QJL_SEED: u64 = 0x7154_514a_4c53_4b31; // "qTQJLSK1"
31const TQ_LLOYD_GRID: usize = 8192;
32const TQ_LLOYD_MAX_ITERATIONS: usize = 64;
33const TQ_LLOYD_TOLERANCE: f64 = 1e-9;
34/// i16 lane accumulators are widened to i32 at least every this many
35/// dimensions: 128 * 127 = 16256 stays far from i16 saturation.
36#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
37const TQ_ACCUMULATE_CHUNK_DIMS: usize = 128;
38
39/// Code-layout dimension: the input dimension rounded up to an even count
40/// (two 4-bit coordinates per byte), floored at `TQ_MIN_PADDED_DIM`. Since
41/// codec v2 the rotation is padding-free, so this tracks the true dimension
42/// instead of the next power of two. Cheap; usable for header validation
43/// without building a codec.
44#[inline]
45pub fn tq_padded_dim(dim: usize) -> usize {
46    dim.next_multiple_of(2).max(TQ_MIN_PADDED_DIM)
47}
48
49/// Fingerprint every payload built for `dim` must carry (no codebook build).
50#[inline]
51pub fn tq_expected_fingerprint(dim: usize) -> u64 {
52    tq_fingerprint(dim, tq_padded_dim(dim))
53}
54
55/// Process-wide codec cache. A codec is a pure function of the dimension and
56/// costs a Lloyd solve to build; segment opens and merges share one instance
57/// per dimension instead of re-deriving it.
58pub fn tq_shared_codec(dim: usize) -> std::sync::Arc<TqCodec> {
59    static CODECS: std::sync::OnceLock<
60        std::sync::Mutex<rustc_hash::FxHashMap<usize, std::sync::Arc<TqCodec>>>,
61    > = std::sync::OnceLock::new();
62    let cache = CODECS.get_or_init(Default::default);
63    let mut guard = cache
64        .lock()
65        .unwrap_or_else(|poisoned| poisoned.into_inner());
66    std::sync::Arc::clone(
67        guard
68            .entry(dim)
69            .or_insert_with(|| std::sync::Arc::new(TqCodec::new(dim))),
70    )
71}
72
73/// Bytes of one scoring block: 16 f32 gammas + 16 packed nibble rows.
74#[inline]
75pub const fn tq_block_bytes(code_size: usize) -> usize {
76    TQ_BLOCK_LANES * (size_of::<f32>() + code_size)
77}
78
79/// Total codes-column bytes for `count` vectors (final block zero-padded).
80#[inline]
81pub const fn tq_codes_column_len(count: usize, code_size: usize) -> usize {
82    count.div_ceil(TQ_BLOCK_LANES) * tq_block_bytes(code_size)
83}
84
85/// Overflow-checked [`tq_codes_column_len`] for untrusted header values.
86#[inline]
87pub fn tq_codes_column_len_checked(count: usize, code_size: usize) -> Option<usize> {
88    count
89        .div_ceil(TQ_BLOCK_LANES)
90        .checked_mul(tq_block_bytes(code_size))
91}
92
93/// Bytes of one IVF-TQ scoring block: 16 f32 residual scales + 16 f32 gammas
94/// + 16 packed nibble rows.
95#[inline]
96pub const fn tq_ivf_block_bytes(code_size: usize) -> usize {
97    TQ_BLOCK_LANES * (2 * size_of::<f32>() + code_size)
98}
99
100/// Overflow-checked IVF-TQ codes-column length for untrusted header values.
101#[inline]
102pub fn tq_ivf_codes_column_len_checked(count: usize, code_size: usize) -> Option<usize> {
103    count
104        .div_ceil(TQ_BLOCK_LANES)
105        .checked_mul(tq_ivf_block_bytes(code_size))
106}
107
108#[inline]
109fn splitmix64(state: &mut u64) -> u64 {
110    *state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
111    let mut z = *state;
112    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
113    z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
114    z ^ (z >> 31)
115}
116
117/// Number of sign/sub-FWHT/permutation rounds in the padding-free rotation.
118/// One round mixes the largest power-of-two prefix; the permutations carry
119/// every coordinate through that prefix across rounds, so three rounds give
120/// near-uniform mixing for any dimension (pinned by the estimator tests).
121const TQ_ROTATION_ROUNDS: usize = 3;
122
123/// Seeded structured rotation, padding-free: per round, sign flips → a
124/// normalized FWHT over the largest power-of-two prefix (identity on the
125/// remainder) → a full random permutation. Every factor is orthonormal on
126/// `R^padded_dim`, so the composition preserves norms and inner products
127/// exactly; inputs shorter than `padded_dim` (odd dims round up by one) are
128/// zero-extended, which embeds them isometrically.
129#[derive(Debug, Clone)]
130pub struct TqRotation {
131    input_dim: usize,
132    padded_dim: usize,
133    /// Largest power of two ≤ `padded_dim`: the per-round FWHT span.
134    fwht_len: usize,
135    /// +1.0 / -1.0 per coordinate, one strip per round.
136    signs: Vec<f32>,
137    /// `output[i] = mixed[perm[i]]`, one strip per round.
138    perms: Vec<u32>,
139}
140
141impl TqRotation {
142    pub fn new(input_dim: usize, seed: u64) -> Self {
143        let padded_dim = tq_padded_dim(input_dim);
144        let fwht_len = if padded_dim.is_power_of_two() {
145            padded_dim
146        } else {
147            padded_dim.next_power_of_two() / 2
148        };
149        let mut state = seed;
150        let mut signs = Vec::with_capacity(TQ_ROTATION_ROUNDS * padded_dim);
151        let mut perms = Vec::with_capacity(TQ_ROTATION_ROUNDS * padded_dim);
152        for _ in 0..TQ_ROTATION_ROUNDS {
153            signs.extend((0..padded_dim).map(|_| {
154                if splitmix64(&mut state) & 1 == 1 {
155                    1.0f32
156                } else {
157                    -1.0f32
158                }
159            }));
160            let round_base = perms.len();
161            perms.extend(0..padded_dim as u32);
162            for i in (1..padded_dim).rev() {
163                let j = (splitmix64(&mut state) % (i as u64 + 1)) as usize;
164                perms.swap(round_base + i, round_base + j);
165            }
166        }
167        Self {
168            input_dim,
169            padded_dim,
170            fwht_len,
171            signs,
172            perms,
173        }
174    }
175
176    #[inline]
177    pub fn padded_dim(&self) -> usize {
178        self.padded_dim
179    }
180
181    /// Rotate `input` (length `input_dim`, or `padded_dim` for already-padded
182    /// residuals) into `output` (length `padded_dim`). `scratch` is reused
183    /// across calls to keep encoding allocation-free.
184    pub fn apply(&self, input: &[f32], scratch: &mut Vec<f32>, output: &mut [f32]) {
185        debug_assert!(input.len() == self.input_dim || input.len() == self.padded_dim);
186        debug_assert_eq!(output.len(), self.padded_dim);
187        let padded_dim = self.padded_dim;
188        scratch.clear();
189        scratch.resize(padded_dim, 0.0);
190        // Round 0 reads straight from the (implicitly zero-extended) input;
191        // later rounds ping-pong between `output` and `scratch`.
192        let signs = &self.signs[..padded_dim];
193        for (slot, (sign, index)) in scratch.iter_mut().zip(signs.iter().zip(0..)) {
194            *slot = input.get(index).copied().unwrap_or(0.0) * sign;
195        }
196        fwht_normalized(&mut scratch[..self.fwht_len]);
197        let perm = &self.perms[..padded_dim];
198        for (slot, &source) in output.iter_mut().zip(perm) {
199            *slot = scratch[source as usize];
200        }
201        for round in 1..TQ_ROTATION_ROUNDS {
202            let signs = &self.signs[round * padded_dim..(round + 1) * padded_dim];
203            for (slot, (&value, sign)) in scratch.iter_mut().zip(output.iter().zip(signs)) {
204                *slot = value * sign;
205            }
206            fwht_normalized(&mut scratch[..self.fwht_len]);
207            let perm = &self.perms[round * padded_dim..(round + 1) * padded_dim];
208            for (slot, &source) in output.iter_mut().zip(perm) {
209                *slot = scratch[source as usize];
210            }
211        }
212    }
213}
214
215/// In-place normalized fast Walsh-Hadamard transform (`len` a power of two).
216fn fwht_normalized(values: &mut [f32]) {
217    let len = values.len();
218    debug_assert!(len.is_power_of_two());
219    let mut step = 1;
220    while step < len {
221        let mut base = 0;
222        while base < len {
223            let (left_half, right_half) = values[base..base + step * 2].split_at_mut(step);
224            for (left, right) in left_half.iter_mut().zip(right_half.iter_mut()) {
225                let sum = *left + *right;
226                let difference = *left - *right;
227                *left = sum;
228                *right = difference;
229            }
230            base += step * 2;
231        }
232        step *= 2;
233    }
234    let scale = 1.0 / (len as f32).sqrt();
235    for value in values.iter_mut() {
236        *value *= scale;
237    }
238}
239
240/// Analytic 3-bit Lloyd-Max codebook for the marginal density of one
241/// coordinate of a uniform unit vector in `R^padded_dim`:
242/// `f(t) ∝ (1 - t²)^((padded_dim - 3) / 2)` on `[-1, 1]`.
243#[derive(Debug, Clone)]
244pub struct TqCodebook {
245    levels: [f32; TQ_STAGE1_LEVELS],
246    /// Decision boundaries between adjacent levels (midpoints).
247    boundaries: [f32; TQ_STAGE1_LEVELS - 1],
248}
249
250impl TqCodebook {
251    pub fn analytic(padded_dim: usize) -> Self {
252        assert!(
253            padded_dim >= TQ_MIN_PADDED_DIM,
254            "TQ codebook requires padded_dim >= {TQ_MIN_PADDED_DIM}, got {padded_dim}"
255        );
256        let exponent = (padded_dim as f64 - 3.0) / 2.0;
257        let cell = 2.0 / TQ_LLOYD_GRID as f64;
258        // Grid midpoints and their density weights over [-1, 1].
259        // Heap-allocated: two [f64; 8192] frames (~128 KiB) would risk stack
260        // overflow on constrained runtimes (WASM readers, worker threads).
261        let mut weights = vec![0.0f64; TQ_LLOYD_GRID];
262        let mut positions = vec![0.0f64; TQ_LLOYD_GRID];
263        for index in 0..TQ_LLOYD_GRID {
264            let t = -1.0 + (index as f64 + 0.5) * cell;
265            positions[index] = t;
266            let log_density = exponent * (1.0 - t * t).max(f64::MIN_POSITIVE).ln();
267            weights[index] = log_density.exp();
268        }
269
270        // Initialize boundaries at equal-mass quantiles.
271        let total_mass: f64 = weights.iter().sum();
272        let mut levels = [0.0f64; TQ_STAGE1_LEVELS];
273        let mut boundaries = [0.0f64; TQ_STAGE1_LEVELS - 1];
274        let mut accumulated = 0.0f64;
275        let mut next_boundary = 0usize;
276        for index in 0..TQ_LLOYD_GRID {
277            accumulated += weights[index];
278            while next_boundary < TQ_STAGE1_LEVELS - 1
279                && accumulated
280                    >= total_mass * (next_boundary as f64 + 1.0) / TQ_STAGE1_LEVELS as f64
281            {
282                boundaries[next_boundary] = positions[index];
283                next_boundary += 1;
284            }
285        }
286
287        // Lloyd-Max: centroids are density-weighted means of their cell,
288        // boundaries are midpoints of adjacent centroids.
289        for _ in 0..TQ_LLOYD_MAX_ITERATIONS {
290            let mut mass = [0.0f64; TQ_STAGE1_LEVELS];
291            let mut moment = [0.0f64; TQ_STAGE1_LEVELS];
292            let mut bucket = 0usize;
293            for index in 0..TQ_LLOYD_GRID {
294                let t = positions[index];
295                while bucket < TQ_STAGE1_LEVELS - 1 && t > boundaries[bucket] {
296                    bucket += 1;
297                }
298                mass[bucket] += weights[index];
299                moment[bucket] += weights[index] * t;
300            }
301            let mut shift = 0.0f64;
302            for level in 0..TQ_STAGE1_LEVELS {
303                if mass[level] > 0.0 {
304                    let updated = moment[level] / mass[level];
305                    shift = shift.max((updated - levels[level]).abs());
306                    levels[level] = updated;
307                }
308            }
309            for boundary in 0..TQ_STAGE1_LEVELS - 1 {
310                boundaries[boundary] = 0.5 * (levels[boundary] + levels[boundary + 1]);
311            }
312            if shift < TQ_LLOYD_TOLERANCE {
313                break;
314            }
315        }
316
317        // The density is even, so the optimal codebook is exactly symmetric;
318        // grid discretization leaves ~1e-4 asymmetry. Symmetrize so the
319        // central decision boundary is exactly zero.
320        for index in 0..TQ_STAGE1_LEVELS / 2 {
321            let magnitude = 0.5 * (levels[TQ_STAGE1_LEVELS - 1 - index] - levels[index]);
322            levels[index] = -magnitude;
323            levels[TQ_STAGE1_LEVELS - 1 - index] = magnitude;
324        }
325        for boundary in 0..TQ_STAGE1_LEVELS - 1 {
326            boundaries[boundary] = 0.5 * (levels[boundary] + levels[boundary + 1]);
327        }
328
329        Self {
330            levels: levels.map(|level| level as f32),
331            boundaries: boundaries.map(|boundary| boundary as f32),
332        }
333    }
334
335    /// 3-bit code of the nearest level.
336    #[inline]
337    pub fn encode_coordinate(&self, value: f32) -> u8 {
338        let mut code = 0u8;
339        for &boundary in &self.boundaries {
340            code += u8::from(value > boundary);
341        }
342        code
343    }
344}
345
346/// Complete TQ codec for one field dimension. Cheap to build (sub-millisecond)
347/// and immutable; share via `Arc` per open segment.
348#[derive(Debug, Clone)]
349pub struct TqCodec {
350    dim: usize,
351    padded_dim: usize,
352    stage1_rotation: TqRotation,
353    qjl_rotation: TqRotation,
354    codebook: TqCodebook,
355    /// `sqrt(π/2) / sqrt(padded_dim)`: QJL correction for an orthonormal sketch.
356    qjl_scale: f32,
357    fingerprint: u64,
358}
359
360impl TqCodec {
361    pub fn new(dim: usize) -> Self {
362        assert!(dim > 0, "TQ codec requires a non-zero dimension");
363        let stage1_rotation = TqRotation::new(dim, TQ_STAGE1_SEED);
364        let padded_dim = stage1_rotation.padded_dim();
365        let qjl_rotation = TqRotation::new(padded_dim, TQ_QJL_SEED);
366        debug_assert_eq!(qjl_rotation.padded_dim(), padded_dim);
367        let codebook = TqCodebook::analytic(padded_dim);
368        let qjl_scale = (std::f64::consts::PI / 2.0).sqrt() as f32 / (padded_dim as f32).sqrt();
369        let fingerprint = tq_fingerprint(dim, padded_dim);
370        Self {
371            dim,
372            padded_dim,
373            stage1_rotation,
374            qjl_rotation,
375            codebook,
376            qjl_scale,
377            fingerprint,
378        }
379    }
380
381    #[inline]
382    pub fn dim(&self) -> usize {
383        self.dim
384    }
385
386    #[inline]
387    pub fn padded_dim(&self) -> usize {
388        self.padded_dim
389    }
390
391    /// Logical bytes per vector (two 4-bit coordinates per byte).
392    #[inline]
393    pub fn code_size(&self) -> usize {
394        self.padded_dim / 2
395    }
396
397    /// Deterministic compatibility fingerprint carried as `quantizer_version`.
398    #[inline]
399    pub fn fingerprint(&self) -> u64 {
400        self.fingerprint
401    }
402
403    /// Heap footprint: two rotations (signs f32 + perm u32 per padded coord)
404    /// plus the fixed-size codebook.
405    pub fn estimated_memory_bytes(&self) -> usize {
406        2 * self.padded_dim * (size_of::<f32>() + size_of::<u32>()) + size_of::<TqCodebook>()
407    }
408
409    /// Encode one vector into `nibbles` (one 0..=15 value per padded
410    /// coordinate) and return `gamma`. The vector is normalized internally;
411    /// zero vectors encode as all-zero nibbles with `gamma = 0`.
412    pub fn encode_into(
413        &self,
414        vector: &[f32],
415        nibbles: &mut [u8],
416        scratch: &mut TqEncodeScratch,
417    ) -> f32 {
418        self.encode_residual_into(vector, nibbles, scratch).1
419    }
420
421    /// Encode one (possibly non-unit) vector as `scale · unit_direction` and
422    /// return `(scale = ‖vector‖₂, gamma)`. IVF leaves store centroid
423    /// residuals, whose norms carry ranking information; `scale` restores it
424    /// at score time. Zero vectors encode as all-zero nibbles with
425    /// `scale = gamma = 0`.
426    pub fn encode_residual_into(
427        &self,
428        vector: &[f32],
429        nibbles: &mut [u8],
430        scratch: &mut TqEncodeScratch,
431    ) -> (f32, f32) {
432        assert_eq!(vector.len(), self.dim, "TQ encode dimension mismatch");
433        assert_eq!(nibbles.len(), self.padded_dim, "TQ nibble buffer mismatch");
434        let norm = crate::structures::simd::dot_product_f32(vector, vector, vector.len()).sqrt();
435        if !norm.is_finite() || norm <= 0.0 {
436            nibbles.fill(0);
437            return (0.0, 0.0);
438        }
439        scratch.normalized.clear();
440        scratch
441            .normalized
442            .extend(vector.iter().map(|value| value / norm));
443
444        scratch.rotated.resize(self.padded_dim, 0.0);
445        let (normalized, rotated, fwht) =
446            (&scratch.normalized, &mut scratch.rotated, &mut scratch.fwht);
447        self.stage1_rotation.apply(normalized, fwht, rotated);
448
449        // Stage-1 codes and residual (in stage-1 rotated space).
450        scratch.residual.resize(self.padded_dim, 0.0);
451        let mut residual_norm_sq = 0.0f32;
452        for ((&value, nibble), residual_slot) in scratch
453            .rotated
454            .iter()
455            .zip(nibbles.iter_mut())
456            .zip(scratch.residual.iter_mut())
457        {
458            let code = self.codebook.encode_coordinate(value);
459            *nibble = code << 1;
460            let residual = value - self.codebook.levels[code as usize];
461            *residual_slot = residual;
462            residual_norm_sq += residual * residual;
463        }
464
465        // QJL sign bits of the rotated residual.
466        scratch.rotated_residual.resize(self.padded_dim, 0.0);
467        let (residual, rotated_residual, fwht) = (
468            &scratch.residual,
469            &mut scratch.rotated_residual,
470            &mut scratch.fwht,
471        );
472        self.qjl_rotation.apply(residual, fwht, rotated_residual);
473        for (nibble, &rotated) in nibbles.iter_mut().zip(scratch.rotated_residual.iter()) {
474            *nibble |= u8::from(rotated >= 0.0);
475        }
476        (norm, residual_norm_sq.sqrt())
477    }
478}
479
480/// Reusable per-thread encode buffers (hot-path allocation hygiene).
481#[derive(Debug, Default)]
482pub struct TqEncodeScratch {
483    normalized: Vec<f32>,
484    rotated: Vec<f32>,
485    residual: Vec<f32>,
486    rotated_residual: Vec<f32>,
487    fwht: Vec<f32>,
488}
489
490fn tq_fingerprint(dim: usize, padded_dim: usize) -> u64 {
491    let mut hash = 0xcbf2_9ce4_8422_2325u64; // FNV-1a offset basis
492    let mut mix = |bytes: &[u8]| {
493        for &byte in bytes {
494            hash ^= u64::from(byte);
495            hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
496        }
497    };
498    mix(b"hermes-tq");
499    mix(&TQ_CODEC_VERSION.to_le_bytes());
500    mix(&TQ_BITS.to_le_bytes());
501    mix(&(dim as u64).to_le_bytes());
502    mix(&(padded_dim as u64).to_le_bytes());
503    mix(&TQ_STAGE1_SEED.to_le_bytes());
504    mix(&TQ_QJL_SEED.to_le_bytes());
505    if hash == 0 { 1 } else { hash }
506}
507
508// ---------------------------------------------------------------------------
509// Query plan and block scoring
510// ---------------------------------------------------------------------------
511
512/// Per-query LUTs: `padded_dim × 16` i8 tables (globally-scaled
513/// quantizations) for the block kernels. The intermediate f32 tables are
514/// dropped after quantization — they are not read on the search path.
515pub struct TqQueryPlan {
516    padded_dim: usize,
517    fingerprint: u64,
518    /// Exact query identity used to validate query-global plan caches.
519    ///
520    /// Keep the IEEE-754 bits rather than a hash: a `DenseVectorQuery` is
521    /// cloneable and its public vector may be mutated while clones continue
522    /// sharing the same cache. Exact bits make stale LUT reuse impossible,
523    /// including for distinct NaN payloads and signed zero.
524    query_bits: Box<[u32]>,
525    base_lut_i8: Vec<i8>,
526    qjl_lut_i8: Vec<i8>,
527    base_dequant: f32,
528    qjl_dequant: f32,
529    /// LUT16 kernel resolved once per query instead of once per 16-vector
530    /// block (runtime feature detection is not free on x86_64).
531    kernel: lut16::Lut16Kernel,
532    /// Full-precision tables, retained for the reference estimator in tests.
533    #[cfg(test)]
534    reference_luts: (Vec<f32>, Vec<f32>),
535}
536
537impl std::fmt::Debug for TqQueryPlan {
538    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539        formatter
540            .debug_struct("TqQueryPlan")
541            .field("padded_dim", &self.padded_dim)
542            .field("fingerprint", &self.fingerprint)
543            .field("query_dim", &self.query_bits.len())
544            .finish()
545    }
546}
547
548impl TqQueryPlan {
549    pub fn build(codec: &TqCodec, query: &[f32]) -> Self {
550        assert_eq!(query.len(), codec.dim, "TQ query dimension mismatch");
551        let padded_dim = codec.padded_dim;
552        let norm = crate::structures::simd::dot_product_f32(query, query, query.len()).sqrt();
553        let inverse_norm = if norm.is_finite() && norm > 0.0 {
554            1.0 / norm
555        } else {
556            0.0
557        };
558        let normalized: Vec<f32> = query.iter().map(|value| value * inverse_norm).collect();
559        let mut fwht = Vec::with_capacity(padded_dim);
560        let mut rotated = vec![0.0f32; padded_dim];
561        codec
562            .stage1_rotation
563            .apply(&normalized, &mut fwht, &mut rotated);
564        let mut qjl_rotated = vec![0.0f32; padded_dim];
565        codec
566            .qjl_rotation
567            .apply(&rotated, &mut fwht, &mut qjl_rotated);
568
569        let mut base_lut = vec![0.0f32; padded_dim * 16];
570        let mut qjl_lut = vec![0.0f32; padded_dim * 16];
571        for dim in 0..padded_dim {
572            for nibble in 0..16 {
573                let level = codec.codebook.levels[nibble >> 1];
574                let sign = if nibble & 1 == 1 { 1.0 } else { -1.0 };
575                base_lut[dim * 16 + nibble] = rotated[dim] * level;
576                qjl_lut[dim * 16 + nibble] = sign * codec.qjl_scale * qjl_rotated[dim];
577            }
578        }
579        let (base_lut_i8, base_dequant) = quantize_lut(&base_lut);
580        let (qjl_lut_i8, qjl_dequant) = quantize_lut(&qjl_lut);
581        Self {
582            padded_dim,
583            fingerprint: codec.fingerprint,
584            query_bits: query
585                .iter()
586                .map(|value| value.to_bits())
587                .collect::<Vec<_>>()
588                .into_boxed_slice(),
589            base_lut_i8,
590            qjl_lut_i8,
591            base_dequant,
592            qjl_dequant,
593            kernel: lut16::Lut16Kernel::resolve(),
594            #[cfg(test)]
595            reference_luts: (base_lut, qjl_lut),
596        }
597    }
598
599    #[inline]
600    pub fn padded_dim(&self) -> usize {
601        self.padded_dim
602    }
603
604    #[inline]
605    pub fn fingerprint(&self) -> u64 {
606        self.fingerprint
607    }
608
609    /// Whether these LUTs were built for this exact query.
610    #[inline]
611    pub(crate) fn matches_query(&self, query: &[f32]) -> bool {
612        self.query_bits.len() == query.len()
613            && self
614                .query_bits
615                .iter()
616                .zip(query)
617                .all(|(&bits, value)| bits == value.to_bits())
618    }
619
620    /// Reference f32 estimator over one unpacked nibble row (test oracle for
621    /// the quantized block path).
622    #[cfg(test)]
623    pub(crate) fn estimate_row(&self, nibbles: &[u8], gamma: f32) -> f32 {
624        debug_assert_eq!(nibbles.len(), self.padded_dim);
625        let (base_lut, qjl_lut) = &self.reference_luts;
626        let mut base = 0.0f32;
627        let mut qjl = 0.0f32;
628        for (dim, &nibble) in nibbles.iter().enumerate() {
629            base += base_lut[dim * 16 + nibble as usize];
630            qjl += qjl_lut[dim * 16 + nibble as usize];
631        }
632        base + gamma * qjl
633    }
634}
635
636fn quantize_lut(values: &[f32]) -> (Vec<i8>, f32) {
637    let max_abs = values
638        .iter()
639        .fold(0.0f32, |acc, &value| acc.max(value.abs()));
640    if !max_abs.is_finite() || max_abs <= 0.0 {
641        return (vec![0i8; values.len()], 0.0);
642    }
643    let quantize_scale = 127.0 / max_abs;
644    let quantized = values
645        .iter()
646        .map(|&value| (value * quantize_scale).round().clamp(-127.0, 127.0) as i8)
647        .collect();
648    (quantized, max_abs / 127.0)
649}
650
651/// Score one block (16 lanes) into `scores`. `block` is
652/// `[16 × f32 gamma][padded_dim × 8 packed nibbles]`; lanes past the run's
653/// vector count hold zero padding and must be ignored by the caller.
654pub fn tq_score_block(plan: &TqQueryPlan, block: &[u8], scores: &mut [f32; TQ_BLOCK_LANES]) {
655    debug_assert_eq!(block.len(), tq_block_bytes(plan.padded_dim / 2));
656    let (gamma_bytes, nibble_bytes) = block.split_at(TQ_BLOCK_LANES * size_of::<f32>());
657    let mut base = [0i32; TQ_BLOCK_LANES];
658    let mut qjl = [0i32; TQ_BLOCK_LANES];
659    plan.kernel.accumulate(
660        &plan.base_lut_i8,
661        &plan.qjl_lut_i8,
662        nibble_bytes,
663        plan.padded_dim,
664        &mut base,
665        &mut qjl,
666    );
667    lut16::finish_block(
668        gamma_bytes,
669        &base,
670        &qjl,
671        plan.base_dequant,
672        plan.qjl_dequant,
673        scores,
674    );
675}
676
677/// Pack up to 16 nibble rows (+ gammas) into one block. Missing lanes are
678/// zero-filled. `rows` are `padded_dim`-length 0..=15 values.
679pub fn tq_pack_block(rows: &[&[u8]], gammas: &[f32], padded_dim: usize, output: &mut Vec<u8>) {
680    assert!(rows.len() <= TQ_BLOCK_LANES && rows.len() == gammas.len());
681    for lane in 0..TQ_BLOCK_LANES {
682        let gamma = gammas.get(lane).copied().unwrap_or(0.0);
683        output.extend_from_slice(&gamma.to_le_bytes());
684    }
685    pack_nibble_rows(rows, padded_dim, output);
686}
687
688/// Pack an IVF-TQ block: per-lane residual scales, gammas, then nibbles.
689pub fn tq_pack_ivf_block(
690    rows: &[&[u8]],
691    scales: &[f32],
692    gammas: &[f32],
693    padded_dim: usize,
694    output: &mut Vec<u8>,
695) {
696    assert!(rows.len() <= TQ_BLOCK_LANES && rows.len() == gammas.len());
697    assert_eq!(scales.len(), gammas.len());
698    for lane in 0..TQ_BLOCK_LANES {
699        let scale = scales.get(lane).copied().unwrap_or(0.0);
700        output.extend_from_slice(&scale.to_le_bytes());
701    }
702    for lane in 0..TQ_BLOCK_LANES {
703        let gamma = gammas.get(lane).copied().unwrap_or(0.0);
704        output.extend_from_slice(&gamma.to_le_bytes());
705    }
706    pack_nibble_rows(rows, padded_dim, output);
707}
708
709fn pack_nibble_rows(rows: &[&[u8]], padded_dim: usize, output: &mut Vec<u8>) {
710    for dim in 0..padded_dim {
711        for byte_index in 0..TQ_BLOCK_LANES / 2 {
712            let low = rows.get(byte_index).map_or(0, |row| row[dim] & 0x0F);
713            let high = rows
714                .get(byte_index + TQ_BLOCK_LANES / 2)
715                .map_or(0, |row| row[dim] & 0x0F);
716            output.push(low | (high << 4));
717        }
718    }
719}
720
721/// Score one IVF-TQ block: `score[lane] = cluster_dot + scale · (base +
722/// gamma · qjl)`, where `cluster_dot = ⟨normalized query, centroid⟩` is the
723/// probed cluster's shared contribution and `scale = ‖residual‖`.
724pub fn tq_score_ivf_block(
725    plan: &TqQueryPlan,
726    block: &[u8],
727    cluster_dot: f32,
728    scores: &mut [f32; TQ_BLOCK_LANES],
729) {
730    debug_assert_eq!(block.len(), tq_ivf_block_bytes(plan.padded_dim() / 2));
731    let lane_f32 = TQ_BLOCK_LANES * size_of::<f32>();
732    let (scale_bytes, rest) = block.split_at(lane_f32);
733    let (gamma_bytes, nibble_bytes) = rest.split_at(lane_f32);
734    let mut base = [0i32; TQ_BLOCK_LANES];
735    let mut qjl = [0i32; TQ_BLOCK_LANES];
736    plan.kernel.accumulate(
737        &plan.base_lut_i8,
738        &plan.qjl_lut_i8,
739        nibble_bytes,
740        plan.padded_dim,
741        &mut base,
742        &mut qjl,
743    );
744    lut16::finish_ivf_block(
745        scale_bytes,
746        gamma_bytes,
747        &base,
748        &qjl,
749        plan.base_dequant,
750        plan.qjl_dequant,
751        cluster_dot,
752        scores,
753    );
754}
755
756mod lut16 {
757    #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
758    use super::TQ_ACCUMULATE_CHUNK_DIMS;
759    use super::TQ_BLOCK_LANES;
760
761    /// LUT16 accumulation kernel, resolved once per query plan (mirrors
762    /// `simd::HammingKernel`) so the per-block path carries no runtime
763    /// feature detection.
764    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
765    pub(super) enum Lut16Kernel {
766        #[cfg(target_arch = "aarch64")]
767        Neon,
768        #[cfg(target_arch = "x86_64")]
769        Avx2,
770        #[cfg(target_arch = "x86_64")]
771        Ssse3,
772        /// Portable fallback (and the WASM path). NEON is baseline on
773        /// aarch64, so it is never resolved there.
774        #[cfg_attr(target_arch = "aarch64", allow(dead_code))]
775        Scalar,
776    }
777
778    impl Lut16Kernel {
779        pub(super) fn resolve() -> Self {
780            #[cfg(target_arch = "aarch64")]
781            {
782                Self::Neon
783            }
784            #[cfg(target_arch = "x86_64")]
785            {
786                if std::arch::is_x86_feature_detected!("avx2") {
787                    return Self::Avx2;
788                }
789                if std::arch::is_x86_feature_detected!("ssse3") {
790                    return Self::Ssse3;
791                }
792                Self::Scalar
793            }
794            #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
795            {
796                Self::Scalar
797            }
798        }
799
800        /// Accumulate both LUT sums for 16 lanes over all dimensions.
801        /// `nibble_bytes` is dimension-major: 8 bytes per dimension, byte `j`
802        /// holding lane `j` (low nibble) and lane `j + 8` (high nibble).
803        #[inline]
804        pub(super) fn accumulate(
805            self,
806            base_lut: &[i8],
807            qjl_lut: &[i8],
808            nibble_bytes: &[u8],
809            padded_dim: usize,
810            base: &mut [i32; TQ_BLOCK_LANES],
811            qjl: &mut [i32; TQ_BLOCK_LANES],
812        ) {
813            match self {
814                #[cfg(target_arch = "aarch64")]
815                Self::Neon => unsafe {
816                    accumulate_block_neon(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
817                },
818                #[cfg(target_arch = "x86_64")]
819                Self::Avx2 => unsafe {
820                    accumulate_block_avx2(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
821                },
822                #[cfg(target_arch = "x86_64")]
823                Self::Ssse3 => unsafe {
824                    accumulate_block_ssse3(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
825                },
826                Self::Scalar => {
827                    accumulate_block_scalar(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
828                }
829            }
830        }
831    }
832
833    /// Accumulate with the kernel resolved for this CPU (test/oracle entry
834    /// point; the search path resolves once per plan).
835    #[cfg(test)]
836    pub(super) fn accumulate_block(
837        base_lut: &[i8],
838        qjl_lut: &[i8],
839        nibble_bytes: &[u8],
840        padded_dim: usize,
841        base: &mut [i32; TQ_BLOCK_LANES],
842        qjl: &mut [i32; TQ_BLOCK_LANES],
843    ) {
844        Lut16Kernel::resolve().accumulate(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
845    }
846
847    /// Flat-TQ epilogue: `score[lane] = base·bd + gamma·qjl·qd`.
848    ///
849    /// The SIMD forms below use plain multiplies and adds in the same order
850    /// as the scalar reference — no FMA — so every path is bit-identical
851    /// (pinned by `epilogue_matches_scalar_reference_bit_exactly`).
852    #[inline]
853    pub(super) fn finish_block(
854        gamma_bytes: &[u8],
855        base: &[i32; TQ_BLOCK_LANES],
856        qjl: &[i32; TQ_BLOCK_LANES],
857        base_dequant: f32,
858        qjl_dequant: f32,
859        scores: &mut [f32; TQ_BLOCK_LANES],
860    ) {
861        assert!(gamma_bytes.len() >= TQ_BLOCK_LANES * 4);
862        #[cfg(target_arch = "aarch64")]
863        {
864            unsafe { finish_block_neon(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores) }
865            return;
866        }
867        #[cfg(target_arch = "x86_64")]
868        {
869            // SSE2 is baseline on x86_64.
870            finish_block_sse2(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores);
871            return;
872        }
873        #[allow(unreachable_code)]
874        finish_block_scalar(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores);
875    }
876
877    /// IVF-TQ epilogue: `score[lane] = cluster_dot + scale·(base·bd + gamma·qjl·qd)`.
878    #[inline]
879    #[allow(clippy::too_many_arguments)]
880    pub(super) fn finish_ivf_block(
881        scale_bytes: &[u8],
882        gamma_bytes: &[u8],
883        base: &[i32; TQ_BLOCK_LANES],
884        qjl: &[i32; TQ_BLOCK_LANES],
885        base_dequant: f32,
886        qjl_dequant: f32,
887        cluster_dot: f32,
888        scores: &mut [f32; TQ_BLOCK_LANES],
889    ) {
890        assert!(scale_bytes.len() >= TQ_BLOCK_LANES * 4);
891        assert!(gamma_bytes.len() >= TQ_BLOCK_LANES * 4);
892        #[cfg(target_arch = "aarch64")]
893        {
894            unsafe {
895                finish_ivf_block_neon(
896                    scale_bytes,
897                    gamma_bytes,
898                    base,
899                    qjl,
900                    base_dequant,
901                    qjl_dequant,
902                    cluster_dot,
903                    scores,
904                )
905            }
906            return;
907        }
908        #[cfg(target_arch = "x86_64")]
909        {
910            finish_ivf_block_sse2(
911                scale_bytes,
912                gamma_bytes,
913                base,
914                qjl,
915                base_dequant,
916                qjl_dequant,
917                cluster_dot,
918                scores,
919            );
920            return;
921        }
922        #[allow(unreachable_code)]
923        finish_ivf_block_scalar(
924            scale_bytes,
925            gamma_bytes,
926            base,
927            qjl,
928            base_dequant,
929            qjl_dequant,
930            cluster_dot,
931            scores,
932        );
933    }
934
935    // The scalar epilogues are the portable fallback and the exactness oracle
936    // for the SIMD forms; on SIMD targets only tests reach them.
937    #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
938    #[inline]
939    fn lane_f32(bytes: &[u8], lane: usize) -> f32 {
940        f32::from_le_bytes(
941            bytes[lane * 4..lane * 4 + 4]
942                .try_into()
943                .expect("lane slice is 4 bytes"),
944        )
945    }
946
947    /// Scalar reference for [`finish_block`].
948    #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
949    pub(super) fn finish_block_scalar(
950        gamma_bytes: &[u8],
951        base: &[i32; TQ_BLOCK_LANES],
952        qjl: &[i32; TQ_BLOCK_LANES],
953        base_dequant: f32,
954        qjl_dequant: f32,
955        scores: &mut [f32; TQ_BLOCK_LANES],
956    ) {
957        for lane in 0..TQ_BLOCK_LANES {
958            let gamma = lane_f32(gamma_bytes, lane);
959            scores[lane] =
960                base[lane] as f32 * base_dequant + gamma * qjl[lane] as f32 * qjl_dequant;
961        }
962    }
963
964    /// Scalar reference for [`finish_ivf_block`].
965    #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
966    #[allow(clippy::too_many_arguments)]
967    pub(super) fn finish_ivf_block_scalar(
968        scale_bytes: &[u8],
969        gamma_bytes: &[u8],
970        base: &[i32; TQ_BLOCK_LANES],
971        qjl: &[i32; TQ_BLOCK_LANES],
972        base_dequant: f32,
973        qjl_dequant: f32,
974        cluster_dot: f32,
975        scores: &mut [f32; TQ_BLOCK_LANES],
976    ) {
977        for lane in 0..TQ_BLOCK_LANES {
978            let scale = lane_f32(scale_bytes, lane);
979            let gamma = lane_f32(gamma_bytes, lane);
980            scores[lane] = cluster_dot
981                + scale
982                    * (base[lane] as f32 * base_dequant + gamma * qjl[lane] as f32 * qjl_dequant);
983        }
984    }
985
986    #[cfg(target_arch = "aarch64")]
987    #[target_feature(enable = "neon")]
988    unsafe fn finish_block_neon(
989        gamma_bytes: &[u8],
990        base: &[i32; TQ_BLOCK_LANES],
991        qjl: &[i32; TQ_BLOCK_LANES],
992        base_dequant: f32,
993        qjl_dequant: f32,
994        scores: &mut [f32; TQ_BLOCK_LANES],
995    ) {
996        use std::arch::aarch64::*;
997        // SAFETY: the caller asserted a gamma column of at least 16
998        // little-endian f32 values; `base`/`qjl`/`scores` are 16-lane arrays,
999        // so every 4-lane load/store below is in bounds. Byte loads avoid any
1000        // alignment assumption on the mmap-backed column.
1001        unsafe {
1002            let bd = vdupq_n_f32(base_dequant);
1003            let qd = vdupq_n_f32(qjl_dequant);
1004            for quarter in 0..4 {
1005                let gamma = vreinterpretq_f32_u8(vld1q_u8(gamma_bytes.as_ptr().add(quarter * 16)));
1006                let b = vcvtq_f32_s32(vld1q_s32(base.as_ptr().add(quarter * 4)));
1007                let j = vcvtq_f32_s32(vld1q_s32(qjl.as_ptr().add(quarter * 4)));
1008                let score = vaddq_f32(vmulq_f32(b, bd), vmulq_f32(vmulq_f32(gamma, j), qd));
1009                vst1q_f32(scores.as_mut_ptr().add(quarter * 4), score);
1010            }
1011        }
1012    }
1013
1014    #[cfg(target_arch = "aarch64")]
1015    #[target_feature(enable = "neon")]
1016    #[allow(clippy::too_many_arguments)]
1017    unsafe fn finish_ivf_block_neon(
1018        scale_bytes: &[u8],
1019        gamma_bytes: &[u8],
1020        base: &[i32; TQ_BLOCK_LANES],
1021        qjl: &[i32; TQ_BLOCK_LANES],
1022        base_dequant: f32,
1023        qjl_dequant: f32,
1024        cluster_dot: f32,
1025        scores: &mut [f32; TQ_BLOCK_LANES],
1026    ) {
1027        use std::arch::aarch64::*;
1028        // SAFETY: as `finish_block_neon`; the scale column is a second
1029        // 16 × f32 prefix of the same block, asserted by the caller.
1030        unsafe {
1031            let bd = vdupq_n_f32(base_dequant);
1032            let qd = vdupq_n_f32(qjl_dequant);
1033            let cd = vdupq_n_f32(cluster_dot);
1034            for quarter in 0..4 {
1035                let scale = vreinterpretq_f32_u8(vld1q_u8(scale_bytes.as_ptr().add(quarter * 16)));
1036                let gamma = vreinterpretq_f32_u8(vld1q_u8(gamma_bytes.as_ptr().add(quarter * 16)));
1037                let b = vcvtq_f32_s32(vld1q_s32(base.as_ptr().add(quarter * 4)));
1038                let j = vcvtq_f32_s32(vld1q_s32(qjl.as_ptr().add(quarter * 4)));
1039                let inner = vaddq_f32(vmulq_f32(b, bd), vmulq_f32(vmulq_f32(gamma, j), qd));
1040                vst1q_f32(
1041                    scores.as_mut_ptr().add(quarter * 4),
1042                    vaddq_f32(cd, vmulq_f32(scale, inner)),
1043                );
1044            }
1045        }
1046    }
1047
1048    #[cfg(target_arch = "x86_64")]
1049    fn finish_block_sse2(
1050        gamma_bytes: &[u8],
1051        base: &[i32; TQ_BLOCK_LANES],
1052        qjl: &[i32; TQ_BLOCK_LANES],
1053        base_dequant: f32,
1054        qjl_dequant: f32,
1055        scores: &mut [f32; TQ_BLOCK_LANES],
1056    ) {
1057        use std::arch::x86_64::*;
1058        // SAFETY: SSE2 is part of the x86_64 baseline; all loads/stores are
1059        // unaligned forms over the 16-lane arrays and the >=64-byte gamma
1060        // column asserted by the caller.
1061        unsafe {
1062            let bd = _mm_set1_ps(base_dequant);
1063            let qd = _mm_set1_ps(qjl_dequant);
1064            for quarter in 0..4 {
1065                let gamma = _mm_loadu_ps(gamma_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1066                let b = _mm_cvtepi32_ps(_mm_loadu_si128(base.as_ptr().add(quarter * 4).cast()));
1067                let j = _mm_cvtepi32_ps(_mm_loadu_si128(qjl.as_ptr().add(quarter * 4).cast()));
1068                let score = _mm_add_ps(_mm_mul_ps(b, bd), _mm_mul_ps(_mm_mul_ps(gamma, j), qd));
1069                _mm_storeu_ps(scores.as_mut_ptr().add(quarter * 4), score);
1070            }
1071        }
1072    }
1073
1074    #[cfg(target_arch = "x86_64")]
1075    #[allow(clippy::too_many_arguments)]
1076    fn finish_ivf_block_sse2(
1077        scale_bytes: &[u8],
1078        gamma_bytes: &[u8],
1079        base: &[i32; TQ_BLOCK_LANES],
1080        qjl: &[i32; TQ_BLOCK_LANES],
1081        base_dequant: f32,
1082        qjl_dequant: f32,
1083        cluster_dot: f32,
1084        scores: &mut [f32; TQ_BLOCK_LANES],
1085    ) {
1086        use std::arch::x86_64::*;
1087        // SAFETY: as `finish_block_sse2`, plus the >=64-byte scale column.
1088        unsafe {
1089            let bd = _mm_set1_ps(base_dequant);
1090            let qd = _mm_set1_ps(qjl_dequant);
1091            let cd = _mm_set1_ps(cluster_dot);
1092            for quarter in 0..4 {
1093                let scale = _mm_loadu_ps(scale_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1094                let gamma = _mm_loadu_ps(gamma_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1095                let b = _mm_cvtepi32_ps(_mm_loadu_si128(base.as_ptr().add(quarter * 4).cast()));
1096                let j = _mm_cvtepi32_ps(_mm_loadu_si128(qjl.as_ptr().add(quarter * 4).cast()));
1097                let inner = _mm_add_ps(_mm_mul_ps(b, bd), _mm_mul_ps(_mm_mul_ps(gamma, j), qd));
1098                _mm_storeu_ps(
1099                    scores.as_mut_ptr().add(quarter * 4),
1100                    _mm_add_ps(cd, _mm_mul_ps(scale, inner)),
1101                );
1102            }
1103        }
1104    }
1105
1106    /// Scalar fallback mirroring the SIMD integer arithmetic exactly
1107    /// (i8 lookups, i32 sums), so all paths agree bit-for-bit.
1108    pub(super) fn accumulate_block_scalar(
1109        base_lut: &[i8],
1110        qjl_lut: &[i8],
1111        nibble_bytes: &[u8],
1112        padded_dim: usize,
1113        base: &mut [i32; TQ_BLOCK_LANES],
1114        qjl: &mut [i32; TQ_BLOCK_LANES],
1115    ) {
1116        for dim in 0..padded_dim {
1117            let row = &nibble_bytes[dim * 8..dim * 8 + 8];
1118            let base_table = &base_lut[dim * 16..dim * 16 + 16];
1119            let qjl_table = &qjl_lut[dim * 16..dim * 16 + 16];
1120            for (lane, &byte) in row.iter().enumerate() {
1121                let low = (byte & 0x0F) as usize;
1122                let high = (byte >> 4) as usize;
1123                base[lane] += i32::from(base_table[low]);
1124                base[lane + 8] += i32::from(base_table[high]);
1125                qjl[lane] += i32::from(qjl_table[low]);
1126                qjl[lane + 8] += i32::from(qjl_table[high]);
1127            }
1128        }
1129    }
1130
1131    #[cfg(target_arch = "aarch64")]
1132    #[target_feature(enable = "neon")]
1133    unsafe fn accumulate_block_neon(
1134        base_lut: &[i8],
1135        qjl_lut: &[i8],
1136        nibble_bytes: &[u8],
1137        padded_dim: usize,
1138        base: &mut [i32; TQ_BLOCK_LANES],
1139        qjl: &mut [i32; TQ_BLOCK_LANES],
1140    ) {
1141        use std::arch::aarch64::*;
1142        unsafe {
1143            let mask = vdup_n_u8(0x0F);
1144            let mut base_lo_i32 = [vdupq_n_s32(0); 4];
1145            let mut qjl_lo_i32 = [vdupq_n_s32(0); 4];
1146            let mut dim = 0;
1147            while dim < padded_dim {
1148                let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1149                let mut base_acc = [vdupq_n_s16(0); 2];
1150                let mut qjl_acc = [vdupq_n_s16(0); 2];
1151                while dim < chunk_end {
1152                    let row = vld1_u8(nibble_bytes.as_ptr().add(dim * 8));
1153                    let low = vand_u8(row, mask);
1154                    let high = vshr_n_u8::<4>(row);
1155                    let lanes = vcombine_u8(low, high);
1156                    let base_table = vld1q_s8(base_lut.as_ptr().add(dim * 16));
1157                    let qjl_table = vld1q_s8(qjl_lut.as_ptr().add(dim * 16));
1158                    let base_values =
1159                        vreinterpretq_s8_u8(vqtbl1q_u8(vreinterpretq_u8_s8(base_table), lanes));
1160                    let qjl_values =
1161                        vreinterpretq_s8_u8(vqtbl1q_u8(vreinterpretq_u8_s8(qjl_table), lanes));
1162                    base_acc[0] = vaddw_s8(base_acc[0], vget_low_s8(base_values));
1163                    base_acc[1] = vaddw_s8(base_acc[1], vget_high_s8(base_values));
1164                    qjl_acc[0] = vaddw_s8(qjl_acc[0], vget_low_s8(qjl_values));
1165                    qjl_acc[1] = vaddw_s8(qjl_acc[1], vget_high_s8(qjl_values));
1166                    dim += 1;
1167                }
1168                for half in 0..2 {
1169                    base_lo_i32[half * 2] =
1170                        vaddw_s16(base_lo_i32[half * 2], vget_low_s16(base_acc[half]));
1171                    base_lo_i32[half * 2 + 1] =
1172                        vaddw_s16(base_lo_i32[half * 2 + 1], vget_high_s16(base_acc[half]));
1173                    qjl_lo_i32[half * 2] =
1174                        vaddw_s16(qjl_lo_i32[half * 2], vget_low_s16(qjl_acc[half]));
1175                    qjl_lo_i32[half * 2 + 1] =
1176                        vaddw_s16(qjl_lo_i32[half * 2 + 1], vget_high_s16(qjl_acc[half]));
1177                }
1178            }
1179            for quarter in 0..4 {
1180                vst1q_s32(base.as_mut_ptr().add(quarter * 4), base_lo_i32[quarter]);
1181                vst1q_s32(qjl.as_mut_ptr().add(quarter * 4), qjl_lo_i32[quarter]);
1182            }
1183        }
1184    }
1185
1186    /// AVX2: two dimensions per iteration. Adjacent dims' packed rows are
1187    /// contiguous (8 bytes each) and so are their 16-entry LUTs, so one 16-byte
1188    /// row load + one 32-byte LUT load + a 256-bit `vpshufb` covers both.
1189    #[cfg(target_arch = "x86_64")]
1190    #[target_feature(enable = "avx2")]
1191    unsafe fn accumulate_block_avx2(
1192        base_lut: &[i8],
1193        qjl_lut: &[i8],
1194        nibble_bytes: &[u8],
1195        padded_dim: usize,
1196        base: &mut [i32; TQ_BLOCK_LANES],
1197        qjl: &mut [i32; TQ_BLOCK_LANES],
1198    ) {
1199        use std::arch::x86_64::*;
1200        unsafe {
1201            let mask = _mm_set1_epi8(0x0F);
1202            let zero256 = _mm256_setzero_si256();
1203            let mut base_i32 = [zero256; 2];
1204            let mut qjl_i32 = [zero256; 2];
1205            let mut dim = 0;
1206            while dim < padded_dim {
1207                let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1208                let mut base_acc = zero256;
1209                let mut qjl_acc = zero256;
1210                while dim + 2 <= chunk_end {
1211                    // Bytes [dim*8, dim*8+16): rows for `dim` and `dim + 1`.
1212                    let rows = _mm_loadu_si128(nibble_bytes.as_ptr().add(dim * 8).cast());
1213                    let low = _mm_and_si128(rows, mask);
1214                    let high = _mm_and_si128(_mm_srli_epi16(rows, 4), mask);
1215                    // Lanes 0..16 of each dim: [low.q0 | high.q0], [low.q1 | high.q1].
1216                    let lanes_first = _mm_unpacklo_epi64(low, high);
1217                    let lanes_second = _mm_unpackhi_epi64(low, high);
1218                    let lanes = _mm256_inserti128_si256(
1219                        _mm256_castsi128_si256(lanes_first),
1220                        lanes_second,
1221                        1,
1222                    );
1223                    let base_tables = _mm256_loadu_si256(base_lut.as_ptr().add(dim * 16).cast());
1224                    let qjl_tables = _mm256_loadu_si256(qjl_lut.as_ptr().add(dim * 16).cast());
1225                    let base_values = _mm256_shuffle_epi8(base_tables, lanes);
1226                    let qjl_values = _mm256_shuffle_epi8(qjl_tables, lanes);
1227                    base_acc = _mm256_add_epi16(
1228                        base_acc,
1229                        _mm256_cvtepi8_epi16(_mm256_castsi256_si128(base_values)),
1230                    );
1231                    base_acc = _mm256_add_epi16(
1232                        base_acc,
1233                        _mm256_cvtepi8_epi16(_mm256_extracti128_si256(base_values, 1)),
1234                    );
1235                    qjl_acc = _mm256_add_epi16(
1236                        qjl_acc,
1237                        _mm256_cvtepi8_epi16(_mm256_castsi256_si128(qjl_values)),
1238                    );
1239                    qjl_acc = _mm256_add_epi16(
1240                        qjl_acc,
1241                        _mm256_cvtepi8_epi16(_mm256_extracti128_si256(qjl_values, 1)),
1242                    );
1243                    dim += 2;
1244                }
1245                // Odd remainder dim within the chunk.
1246                while dim < chunk_end {
1247                    let row = _mm_loadl_epi64(nibble_bytes.as_ptr().add(dim * 8).cast());
1248                    let low = _mm_and_si128(row, mask);
1249                    let high = _mm_and_si128(_mm_srli_epi16(row, 4), mask);
1250                    let lanes = _mm_unpacklo_epi64(low, high);
1251                    let base_table = _mm_loadu_si128(base_lut.as_ptr().add(dim * 16).cast());
1252                    let qjl_table = _mm_loadu_si128(qjl_lut.as_ptr().add(dim * 16).cast());
1253                    base_acc = _mm256_add_epi16(
1254                        base_acc,
1255                        _mm256_cvtepi8_epi16(_mm_shuffle_epi8(base_table, lanes)),
1256                    );
1257                    qjl_acc = _mm256_add_epi16(
1258                        qjl_acc,
1259                        _mm256_cvtepi8_epi16(_mm_shuffle_epi8(qjl_table, lanes)),
1260                    );
1261                    dim += 1;
1262                }
1263                for (accumulators, chunk) in [(&mut base_i32, base_acc), (&mut qjl_i32, qjl_acc)] {
1264                    accumulators[0] = _mm256_add_epi32(
1265                        accumulators[0],
1266                        _mm256_cvtepi16_epi32(_mm256_castsi256_si128(chunk)),
1267                    );
1268                    accumulators[1] = _mm256_add_epi32(
1269                        accumulators[1],
1270                        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(chunk, 1)),
1271                    );
1272                }
1273            }
1274            _mm256_storeu_si256(base.as_mut_ptr().cast(), base_i32[0]);
1275            _mm256_storeu_si256(base.as_mut_ptr().add(8).cast(), base_i32[1]);
1276            _mm256_storeu_si256(qjl.as_mut_ptr().cast(), qjl_i32[0]);
1277            _mm256_storeu_si256(qjl.as_mut_ptr().add(8).cast(), qjl_i32[1]);
1278        }
1279    }
1280
1281    #[cfg(target_arch = "x86_64")]
1282    #[target_feature(enable = "ssse3")]
1283    unsafe fn accumulate_block_ssse3(
1284        base_lut: &[i8],
1285        qjl_lut: &[i8],
1286        nibble_bytes: &[u8],
1287        padded_dim: usize,
1288        base: &mut [i32; TQ_BLOCK_LANES],
1289        qjl: &mut [i32; TQ_BLOCK_LANES],
1290    ) {
1291        use std::arch::x86_64::*;
1292        unsafe {
1293            let mask = _mm_set1_epi8(0x0F);
1294            let zero = _mm_setzero_si128();
1295            let mut base_i32 = [zero; 4];
1296            let mut qjl_i32 = [zero; 4];
1297            let mut dim = 0;
1298            while dim < padded_dim {
1299                let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1300                let mut base_acc = [zero; 2];
1301                let mut qjl_acc = [zero; 2];
1302                while dim < chunk_end {
1303                    let row = _mm_loadl_epi64(nibble_bytes.as_ptr().add(dim * 8).cast());
1304                    let low = _mm_and_si128(row, mask);
1305                    let high = _mm_and_si128(_mm_srli_epi16(row, 4), mask);
1306                    let lanes = _mm_unpacklo_epi64(low, high);
1307                    let base_table = _mm_loadu_si128(base_lut.as_ptr().add(dim * 16).cast());
1308                    let qjl_table = _mm_loadu_si128(qjl_lut.as_ptr().add(dim * 16).cast());
1309                    let base_values = _mm_shuffle_epi8(base_table, lanes);
1310                    let qjl_values = _mm_shuffle_epi8(qjl_table, lanes);
1311                    // Sign-extend i8 → i16 without SSE4.1: compare-based sign mask.
1312                    let base_sign = _mm_cmpgt_epi8(zero, base_values);
1313                    let qjl_sign = _mm_cmpgt_epi8(zero, qjl_values);
1314                    base_acc[0] =
1315                        _mm_add_epi16(base_acc[0], _mm_unpacklo_epi8(base_values, base_sign));
1316                    base_acc[1] =
1317                        _mm_add_epi16(base_acc[1], _mm_unpackhi_epi8(base_values, base_sign));
1318                    qjl_acc[0] = _mm_add_epi16(qjl_acc[0], _mm_unpacklo_epi8(qjl_values, qjl_sign));
1319                    qjl_acc[1] = _mm_add_epi16(qjl_acc[1], _mm_unpackhi_epi8(qjl_values, qjl_sign));
1320                    dim += 1;
1321                }
1322                for half in 0..2 {
1323                    let base_sign = _mm_cmpgt_epi16(zero, base_acc[half]);
1324                    let qjl_sign = _mm_cmpgt_epi16(zero, qjl_acc[half]);
1325                    base_i32[half * 2] = _mm_add_epi32(
1326                        base_i32[half * 2],
1327                        _mm_unpacklo_epi16(base_acc[half], base_sign),
1328                    );
1329                    base_i32[half * 2 + 1] = _mm_add_epi32(
1330                        base_i32[half * 2 + 1],
1331                        _mm_unpackhi_epi16(base_acc[half], base_sign),
1332                    );
1333                    qjl_i32[half * 2] = _mm_add_epi32(
1334                        qjl_i32[half * 2],
1335                        _mm_unpacklo_epi16(qjl_acc[half], qjl_sign),
1336                    );
1337                    qjl_i32[half * 2 + 1] = _mm_add_epi32(
1338                        qjl_i32[half * 2 + 1],
1339                        _mm_unpackhi_epi16(qjl_acc[half], qjl_sign),
1340                    );
1341                }
1342            }
1343            for quarter in 0..4 {
1344                _mm_storeu_si128(base.as_mut_ptr().add(quarter * 4).cast(), base_i32[quarter]);
1345                _mm_storeu_si128(qjl.as_mut_ptr().add(quarter * 4).cast(), qjl_i32[quarter]);
1346            }
1347        }
1348    }
1349}
1350
1351// ---------------------------------------------------------------------------
1352// Segment build support
1353// ---------------------------------------------------------------------------
1354
1355/// Streaming builder for one segment's TQ payload: doc/ordinal columns plus a
1356/// block-packed codes column ready for `ann_disk` serialization.
1357#[cfg(feature = "native")]
1358pub struct TqFlatBuilder {
1359    codec: std::sync::Arc<TqCodec>,
1360    pub doc_ids: Vec<u32>,
1361    pub ordinals: Vec<u16>,
1362    pub codes: Vec<u8>,
1363    pending_rows: Vec<Vec<u8>>,
1364    pending_gammas: Vec<f32>,
1365}
1366
1367#[cfg(feature = "native")]
1368impl TqFlatBuilder {
1369    pub fn new(codec: std::sync::Arc<TqCodec>) -> Self {
1370        Self {
1371            codec,
1372            doc_ids: Vec::new(),
1373            ordinals: Vec::new(),
1374            codes: Vec::new(),
1375            pending_rows: Vec::with_capacity(TQ_BLOCK_LANES),
1376            pending_gammas: Vec::with_capacity(TQ_BLOCK_LANES),
1377        }
1378    }
1379
1380    #[inline]
1381    pub fn codec(&self) -> &TqCodec {
1382        &self.codec
1383    }
1384
1385    #[inline]
1386    pub fn len(&self) -> usize {
1387        self.doc_ids.len()
1388    }
1389
1390    #[inline]
1391    pub fn is_empty(&self) -> bool {
1392        self.doc_ids.is_empty()
1393    }
1394
1395    /// Encode one contiguous `(labels, vectors)` batch in parallel while
1396    /// preserving input order (lane order must match the doc-ID column).
1397    pub fn add_batch(
1398        &mut self,
1399        labels: &[(u32, u16)],
1400        vectors: &[f32],
1401    ) -> Result<(), &'static str> {
1402        use rayon::prelude::*;
1403
1404        let dim = self.codec.dim();
1405        let vector_count = labels.len();
1406        let expected = vector_count
1407            .checked_mul(dim)
1408            .ok_or("TQ input size overflow")?;
1409        if vectors.len() != expected {
1410            return Err("TQ vector and label matrices are inconsistent");
1411        }
1412        let padded_dim = self.codec.padded_dim();
1413        let codec = std::sync::Arc::clone(&self.codec);
1414        // One contiguous nibble matrix instead of a Vec per vector: each
1415        // Rayon task writes its disjoint row range (allocation hygiene on
1416        // the ingest/merge path).
1417        let mut rows = vec![0u8; vector_count * padded_dim];
1418        let mut gammas = vec![0.0f32; vector_count];
1419        vectors
1420            .par_chunks_exact(dim)
1421            .zip(rows.par_chunks_exact_mut(padded_dim))
1422            .zip(gammas.par_iter_mut())
1423            .for_each_init(
1424                TqEncodeScratch::default,
1425                |scratch, ((vector, row), gamma)| {
1426                    *gamma = codec.encode_into(vector, row, scratch);
1427                },
1428            );
1429
1430        // Top up a carried-over partial block, pack full blocks straight from
1431        // the contiguous matrix (no per-row copies), and carry the tail.
1432        let mut index = 0;
1433        while index < vector_count && !self.pending_rows.is_empty() {
1434            let (doc_id, ordinal) = labels[index];
1435            self.doc_ids.push(doc_id);
1436            self.ordinals.push(ordinal);
1437            self.pending_rows
1438                .push(rows[index * padded_dim..(index + 1) * padded_dim].to_vec());
1439            self.pending_gammas.push(gammas[index]);
1440            if self.pending_rows.len() == TQ_BLOCK_LANES {
1441                self.flush_block();
1442            }
1443            index += 1;
1444        }
1445        while vector_count - index >= TQ_BLOCK_LANES {
1446            let row_refs: Vec<&[u8]> = (0..TQ_BLOCK_LANES)
1447                .map(|lane| {
1448                    let row = index + lane;
1449                    &rows[row * padded_dim..(row + 1) * padded_dim]
1450                })
1451                .collect();
1452            tq_pack_block(
1453                &row_refs,
1454                &gammas[index..index + TQ_BLOCK_LANES],
1455                padded_dim,
1456                &mut self.codes,
1457            );
1458            for &(doc_id, ordinal) in &labels[index..index + TQ_BLOCK_LANES] {
1459                self.doc_ids.push(doc_id);
1460                self.ordinals.push(ordinal);
1461            }
1462            index += TQ_BLOCK_LANES;
1463        }
1464        for row in index..vector_count {
1465            let (doc_id, ordinal) = labels[row];
1466            self.doc_ids.push(doc_id);
1467            self.ordinals.push(ordinal);
1468            self.pending_rows
1469                .push(rows[row * padded_dim..(row + 1) * padded_dim].to_vec());
1470            self.pending_gammas.push(gammas[row]);
1471        }
1472        Ok(())
1473    }
1474
1475    fn flush_block(&mut self) {
1476        let rows: Vec<&[u8]> = self.pending_rows.iter().map(Vec::as_slice).collect();
1477        tq_pack_block(
1478            &rows,
1479            &self.pending_gammas,
1480            self.codec.padded_dim(),
1481            &mut self.codes,
1482        );
1483        self.pending_rows.clear();
1484        self.pending_gammas.clear();
1485    }
1486
1487    /// Flush the trailing partial block (zero-padded lanes).
1488    pub fn finish(&mut self) {
1489        if !self.pending_rows.is_empty() {
1490            self.flush_block();
1491        }
1492        debug_assert_eq!(
1493            self.codes.len(),
1494            tq_codes_column_len(self.doc_ids.len(), self.codec.code_size())
1495        );
1496    }
1497}
1498
1499#[cfg(test)]
1500mod tests {
1501    use super::*;
1502
1503    fn seeded_unit_vector(dim: usize, seed: u64) -> Vec<f32> {
1504        // Box-Muller from splitmix64 for an isotropic direction.
1505        let mut state = seed;
1506        let mut values: Vec<f32> = (0..dim)
1507            .map(|_| {
1508                let a = (splitmix64(&mut state) >> 11) as f64 / (1u64 << 53) as f64;
1509                let b = (splitmix64(&mut state) >> 11) as f64 / (1u64 << 53) as f64;
1510                let gaussian = (-2.0 * (1.0 - a).max(f64::MIN_POSITIVE).ln()).sqrt()
1511                    * (2.0 * std::f64::consts::PI * b).cos();
1512                gaussian as f32
1513            })
1514            .collect();
1515        let norm = values.iter().map(|v| v * v).sum::<f32>().sqrt();
1516        values.iter_mut().for_each(|v| *v /= norm);
1517        values
1518    }
1519
1520    #[test]
1521    fn rotation_is_orthonormal_and_deterministic() {
1522        let rotation = TqRotation::new(100, 42);
1523        assert_eq!(rotation.padded_dim(), 100);
1524        let mut fwht_scratch = Vec::new();
1525        let mut probe = vec![0.0f32; 100];
1526        // Odd dims round up by one zero coordinate.
1527        assert_eq!(TqRotation::new(99, 42).padded_dim(), 100);
1528        TqRotation::new(99, 42).apply(&vec![1.0; 99], &mut fwht_scratch, &mut probe);
1529        let input = seeded_unit_vector(100, 7);
1530        let mut fwht = Vec::new();
1531        let mut output = vec![0.0f32; 100];
1532        rotation.apply(&input, &mut fwht, &mut output);
1533        let norm: f32 = output.iter().map(|v| v * v).sum();
1534        assert!(
1535            (norm - 1.0).abs() < 1e-5,
1536            "rotation must preserve norm, got {norm}"
1537        );
1538
1539        let mut second = vec![0.0f32; 100];
1540        TqRotation::new(100, 42).apply(&input, &mut fwht, &mut second);
1541        assert_eq!(output, second, "rotation must be deterministic");
1542
1543        // Distinct inputs keep their inner product (isometry).
1544        let other = seeded_unit_vector(100, 8);
1545        let mut other_rotated = vec![0.0f32; 100];
1546        rotation.apply(&other, &mut fwht, &mut other_rotated);
1547        let dot_before: f32 = input.iter().zip(&other).map(|(a, b)| a * b).sum();
1548        let dot_after: f32 = output.iter().zip(&other_rotated).map(|(a, b)| a * b).sum();
1549        assert!(
1550            (dot_before - dot_after).abs() < 1e-4,
1551            "rotation must preserve inner products: {dot_before} vs {dot_after}"
1552        );
1553    }
1554
1555    #[test]
1556    fn analytic_codebook_is_symmetric_and_monotonic() {
1557        for padded_dim in [8, 128, 1024] {
1558            let codebook = TqCodebook::analytic(padded_dim);
1559            let levels = &codebook.levels;
1560            for pair in levels.windows(2) {
1561                assert!(pair[0] < pair[1], "levels must be strictly increasing");
1562            }
1563            for index in 0..TQ_STAGE1_LEVELS / 2 {
1564                assert_eq!(
1565                    levels[index],
1566                    -levels[TQ_STAGE1_LEVELS - 1 - index],
1567                    "levels must be exactly symmetric for P={padded_dim}: {levels:?}"
1568                );
1569            }
1570            assert!(levels[TQ_STAGE1_LEVELS - 1] < 1.0);
1571            // Coordinates concentrate near ±1/sqrt(P); the top level must be
1572            // on that scale, not at the interval edge.
1573            let scale = 1.0 / (padded_dim as f32).sqrt();
1574            assert!(
1575                levels[TQ_STAGE1_LEVELS - 1] < 6.0 * scale,
1576                "top level {} is implausibly large for P={padded_dim}",
1577                levels[TQ_STAGE1_LEVELS - 1]
1578            );
1579        }
1580    }
1581
1582    #[test]
1583    fn encode_coordinate_matches_nearest_level() {
1584        let codebook = TqCodebook::analytic(256);
1585        for step in -1000i32..=1000 {
1586            let value = step as f32 / 1000.0;
1587            let code = codebook.encode_coordinate(value) as usize;
1588            let nearest = codebook
1589                .levels
1590                .iter()
1591                .enumerate()
1592                .min_by(|a, b| (a.1 - value).abs().total_cmp(&(b.1 - value).abs()))
1593                .unwrap()
1594                .0;
1595            assert_eq!(
1596                code, nearest,
1597                "value {value} coded {code}, nearest {nearest}"
1598            );
1599        }
1600    }
1601
1602    #[test]
1603    fn estimator_is_unbiased_and_tight() {
1604        let dim = 96;
1605        let codec = TqCodec::new(dim);
1606        let mut scratch = TqEncodeScratch::default();
1607        let mut nibbles = vec![0u8; codec.padded_dim()];
1608
1609        let pairs = 512;
1610        let mut signed_error_sum = 0.0f64;
1611        let mut squared_error_sum = 0.0f64;
1612        let mut stage1_signed_error_sum = 0.0f64;
1613        for pair in 0..pairs {
1614            let vector = seeded_unit_vector(dim, 1000 + pair);
1615            let query = seeded_unit_vector(dim, 900_000 + pair);
1616            let gamma = codec.encode_into(&vector, &mut nibbles, &mut scratch);
1617            let plan = TqQueryPlan::build(&codec, &query);
1618            let estimate = plan.estimate_row(&nibbles, gamma);
1619            let stage1_only = plan.estimate_row(&nibbles, 0.0);
1620            let truth: f32 = vector.iter().zip(&query).map(|(a, b)| a * b).sum();
1621            signed_error_sum += f64::from(estimate - truth);
1622            squared_error_sum += f64::from(estimate - truth).powi(2);
1623            stage1_signed_error_sum += f64::from(stage1_only - truth);
1624        }
1625        let mean_error = signed_error_sum / pairs as f64;
1626        let rmse = (squared_error_sum / pairs as f64).sqrt();
1627        let stage1_mean_error = stage1_signed_error_sum / pairs as f64;
1628        assert!(
1629            mean_error.abs() < 3e-3,
1630            "QJL-corrected estimator must be unbiased: mean error {mean_error}"
1631        );
1632        assert!(rmse < 0.05, "estimator RMSE too large: {rmse}");
1633        assert!(
1634            mean_error.abs() <= stage1_mean_error.abs() + 1e-4,
1635            "QJL correction must not increase bias: {mean_error} vs stage-1 {stage1_mean_error}"
1636        );
1637    }
1638
1639    #[test]
1640    fn block_scoring_matches_row_estimates() {
1641        let dim = 100;
1642        let codec = TqCodec::new(dim);
1643        let mut scratch = TqEncodeScratch::default();
1644        let query = seeded_unit_vector(dim, 3);
1645        let plan = TqQueryPlan::build(&codec, &query);
1646
1647        let lanes = 13; // deliberately partial block
1648        let mut rows = Vec::new();
1649        let mut gammas = Vec::new();
1650        for lane in 0..lanes {
1651            let vector = seeded_unit_vector(dim, 50 + lane as u64);
1652            let mut nibbles = vec![0u8; codec.padded_dim()];
1653            let gamma = codec.encode_into(&vector, &mut nibbles, &mut scratch);
1654            rows.push(nibbles);
1655            gammas.push(gamma);
1656        }
1657        let row_refs: Vec<&[u8]> = rows.iter().map(Vec::as_slice).collect();
1658        let mut block = Vec::new();
1659        tq_pack_block(&row_refs, &gammas, codec.padded_dim(), &mut block);
1660        assert_eq!(block.len(), tq_block_bytes(codec.code_size()));
1661
1662        let mut scores = [0.0f32; TQ_BLOCK_LANES];
1663        tq_score_block(&plan, &block, &mut scores);
1664        for lane in 0..lanes {
1665            let expected = plan.estimate_row(&rows[lane], gammas[lane]);
1666            // The block path uses i8 LUTs; agreement is approximate.
1667            assert!(
1668                (scores[lane] - expected).abs() < 0.02,
1669                "lane {lane}: block score {} vs row estimate {expected}",
1670                scores[lane]
1671            );
1672        }
1673    }
1674
1675    #[test]
1676    fn simd_accumulation_matches_scalar_reference_exactly() {
1677        let padded_dim = 192; // not a multiple of the widen chunk
1678        let mut state = 99u64;
1679        let mut base_lut = vec![0i8; padded_dim * 16];
1680        let mut qjl_lut = vec![0i8; padded_dim * 16];
1681        for value in base_lut.iter_mut().chain(qjl_lut.iter_mut()) {
1682            *value = (splitmix64(&mut state) as i32 % 255 - 127) as i8;
1683        }
1684        let mut nibble_bytes = vec![0u8; padded_dim * 8];
1685        for byte in nibble_bytes.iter_mut() {
1686            *byte = splitmix64(&mut state) as u8;
1687        }
1688
1689        let mut base_reference = [0i32; TQ_BLOCK_LANES];
1690        let mut qjl_reference = [0i32; TQ_BLOCK_LANES];
1691        lut16::accumulate_block_scalar(
1692            &base_lut,
1693            &qjl_lut,
1694            &nibble_bytes,
1695            padded_dim,
1696            &mut base_reference,
1697            &mut qjl_reference,
1698        );
1699        let mut base_dispatch = [0i32; TQ_BLOCK_LANES];
1700        let mut qjl_dispatch = [0i32; TQ_BLOCK_LANES];
1701        lut16::accumulate_block(
1702            &base_lut,
1703            &qjl_lut,
1704            &nibble_bytes,
1705            padded_dim,
1706            &mut base_dispatch,
1707            &mut qjl_dispatch,
1708        );
1709        assert_eq!(
1710            base_reference, base_dispatch,
1711            "base sums must match exactly"
1712        );
1713        assert_eq!(qjl_reference, qjl_dispatch, "qjl sums must match exactly");
1714    }
1715
1716    #[test]
1717    fn epilogue_matches_scalar_reference_bit_exactly() {
1718        let mut state = 4242u64;
1719        for trial in 0..256 {
1720            let mut base = [0i32; TQ_BLOCK_LANES];
1721            let mut qjl = [0i32; TQ_BLOCK_LANES];
1722            let mut scale_bytes = Vec::with_capacity(TQ_BLOCK_LANES * 4);
1723            let mut gamma_bytes = Vec::with_capacity(TQ_BLOCK_LANES * 4);
1724            for lane in 0..TQ_BLOCK_LANES {
1725                base[lane] = (splitmix64(&mut state) % 40_000) as i32 - 20_000;
1726                qjl[lane] = (splitmix64(&mut state) % 40_000) as i32 - 20_000;
1727                let scale = (splitmix64(&mut state) % 10_000) as f32 / 10_000.0;
1728                let gamma = (splitmix64(&mut state) % 10_000) as f32 / 10_000.0;
1729                // Include padded (zero) lanes and a negative scale.
1730                let scale = match lane {
1731                    15 => 0.0,
1732                    7 => -scale,
1733                    _ => scale,
1734                };
1735                scale_bytes.extend_from_slice(&scale.to_le_bytes());
1736                gamma_bytes.extend_from_slice(&gamma.to_le_bytes());
1737            }
1738            let base_dequant = 1e-4 + (trial as f32) * 1e-6;
1739            let qjl_dequant = 3e-5 + (trial as f32) * 1e-7;
1740            let cluster_dot = (trial as f32) / 256.0 - 0.5;
1741
1742            let mut expected = [0.0f32; TQ_BLOCK_LANES];
1743            let mut actual = [0.0f32; TQ_BLOCK_LANES];
1744            lut16::finish_ivf_block_scalar(
1745                &scale_bytes,
1746                &gamma_bytes,
1747                &base,
1748                &qjl,
1749                base_dequant,
1750                qjl_dequant,
1751                cluster_dot,
1752                &mut expected,
1753            );
1754            lut16::finish_ivf_block(
1755                &scale_bytes,
1756                &gamma_bytes,
1757                &base,
1758                &qjl,
1759                base_dequant,
1760                qjl_dequant,
1761                cluster_dot,
1762                &mut actual,
1763            );
1764            assert_eq!(
1765                expected.map(f32::to_bits),
1766                actual.map(f32::to_bits),
1767                "IVF epilogue must be bit-identical to the scalar reference (trial {trial})"
1768            );
1769
1770            lut16::finish_block_scalar(
1771                &gamma_bytes,
1772                &base,
1773                &qjl,
1774                base_dequant,
1775                qjl_dequant,
1776                &mut expected,
1777            );
1778            lut16::finish_block(
1779                &gamma_bytes,
1780                &base,
1781                &qjl,
1782                base_dequant,
1783                qjl_dequant,
1784                &mut actual,
1785            );
1786            assert_eq!(
1787                expected.map(f32::to_bits),
1788                actual.map(f32::to_bits),
1789                "flat epilogue must be bit-identical to the scalar reference (trial {trial})"
1790            );
1791        }
1792    }
1793
1794    #[test]
1795    fn fingerprint_pins_codec_constants() {
1796        let codec = TqCodec::new(768);
1797        assert_eq!(codec.padded_dim(), 768, "codec v2 is padding-free");
1798        assert_eq!(codec.code_size(), 384);
1799        assert_ne!(codec.fingerprint(), 0);
1800        assert_eq!(codec.fingerprint(), TqCodec::new(768).fingerprint());
1801        assert_ne!(codec.fingerprint(), TqCodec::new(769).fingerprint());
1802        // Golden value: the fingerprint is persisted as quantizer_version in
1803        // every TQ segment. Any change to the hash, seeds, or codec constants
1804        // MUST bump TQ_CODEC_VERSION — never silently re-derive.
1805        assert_eq!(
1806            codec.fingerprint(),
1807            GOLDEN_FINGERPRINT_768,
1808            "TQ fingerprint for dim 768 changed; existing segments would be \
1809             rejected. Bump TQ_CODEC_VERSION deliberately instead."
1810        );
1811    }
1812
1813    #[test]
1814    fn query_plan_identity_uses_exact_float_bits() {
1815        let codec = TqCodec::new(4);
1816        let query = [0.0, -0.0, 1.0, f32::from_bits(0x7fc0_0001)];
1817        let plan = TqQueryPlan::build(&codec, &query);
1818
1819        assert!(plan.matches_query(&query));
1820        assert!(!plan.matches_query(&query[..3]));
1821
1822        let mut different_zero = query;
1823        different_zero[1] = 0.0;
1824        assert!(!plan.matches_query(&different_zero));
1825
1826        let mut different_nan = query;
1827        different_nan[3] = f32::from_bits(0x7fc0_0002);
1828        assert!(!plan.matches_query(&different_nan));
1829    }
1830
1831    const GOLDEN_FINGERPRINT_768: u64 = 7026088428300072418;
1832
1833    #[cfg(feature = "native")]
1834    #[test]
1835    fn builder_packs_blocks_in_input_order() {
1836        let dim = 32;
1837        let codec = std::sync::Arc::new(TqCodec::new(dim));
1838        let mut builder = TqFlatBuilder::new(std::sync::Arc::clone(&codec));
1839        let count = 37; // two full blocks + partial
1840        let labels: Vec<(u32, u16)> = (0..count).map(|i| (i as u32, (i % 3) as u16)).collect();
1841        let mut vectors = Vec::new();
1842        for i in 0..count {
1843            vectors.extend(seeded_unit_vector(dim, 7_000 + i as u64));
1844        }
1845        builder.add_batch(&labels, &vectors).unwrap();
1846        builder.finish();
1847        assert_eq!(builder.doc_ids.len(), count);
1848        assert_eq!(
1849            builder.codes.len(),
1850            tq_codes_column_len(count, codec.code_size())
1851        );
1852
1853        // Every lane must score identically to encoding the row directly.
1854        let query = seeded_unit_vector(dim, 1);
1855        let plan = TqQueryPlan::build(&codec, &query);
1856        let block_bytes = tq_block_bytes(codec.code_size());
1857        let mut scratch = TqEncodeScratch::default();
1858        let mut scores = [0.0f32; TQ_BLOCK_LANES];
1859        for index in 0..count {
1860            let block = &builder.codes[(index / TQ_BLOCK_LANES) * block_bytes..][..block_bytes];
1861            tq_score_block(&plan, block, &mut scores);
1862            let mut nibbles = vec![0u8; codec.padded_dim()];
1863            let gamma = codec.encode_into(
1864                &vectors[index * dim..(index + 1) * dim],
1865                &mut nibbles,
1866                &mut scratch,
1867            );
1868            let expected = plan.estimate_row(&nibbles, gamma);
1869            assert!(
1870                (scores[index % TQ_BLOCK_LANES] - expected).abs() < 0.02,
1871                "vector {index} scored {} expected {expected}",
1872                scores[index % TQ_BLOCK_LANES]
1873            );
1874        }
1875    }
1876}