Skip to main content

hermes_core/structures/
simd.rs

1//! Shared SIMD-accelerated functions for posting list compression
2//!
3//! This module provides platform-optimized implementations for common operations:
4//! - **Unpacking**: Convert packed 8/16/32-bit values to u32 arrays
5//! - **Delta decoding**: Prefix sum for converting deltas to absolute values
6//! - **Add one**: Increment all values in an array (for TF decoding)
7//!
8//! Supports:
9//! - **NEON** on aarch64 (Apple Silicon, ARM servers)
10//! - **SSE/SSE4.1** on x86_64 (Intel/AMD)
11//! - **Scalar fallback** for other architectures
12
13// ============================================================================
14// NEON intrinsics for aarch64 (Apple Silicon, ARM servers)
15// ============================================================================
16
17#[cfg(target_arch = "aarch64")]
18#[allow(unsafe_op_in_unsafe_fn)]
19mod neon {
20    use std::arch::aarch64::*;
21
22    /// SIMD unpack for 8-bit values using NEON
23    #[target_feature(enable = "neon")]
24    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
25        let chunks = count / 16;
26        let remainder = count % 16;
27
28        for chunk in 0..chunks {
29            let base = chunk * 16;
30            let in_ptr = input.as_ptr().add(base);
31
32            // Load 16 bytes
33            let bytes = vld1q_u8(in_ptr);
34
35            // Widen u8 -> u16 -> u32
36            let low8 = vget_low_u8(bytes);
37            let high8 = vget_high_u8(bytes);
38
39            let low16 = vmovl_u8(low8);
40            let high16 = vmovl_u8(high8);
41
42            let v0 = vmovl_u16(vget_low_u16(low16));
43            let v1 = vmovl_u16(vget_high_u16(low16));
44            let v2 = vmovl_u16(vget_low_u16(high16));
45            let v3 = vmovl_u16(vget_high_u16(high16));
46
47            let out_ptr = output.as_mut_ptr().add(base);
48            vst1q_u32(out_ptr, v0);
49            vst1q_u32(out_ptr.add(4), v1);
50            vst1q_u32(out_ptr.add(8), v2);
51            vst1q_u32(out_ptr.add(12), v3);
52        }
53
54        // Handle remainder
55        let base = chunks * 16;
56        for i in 0..remainder {
57            output[base + i] = input[base + i] as u32;
58        }
59    }
60
61    /// SIMD unpack for 16-bit values using NEON
62    #[target_feature(enable = "neon")]
63    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
64        let chunks = count / 8;
65        let remainder = count % 8;
66
67        for chunk in 0..chunks {
68            let base = chunk * 8;
69            let in_ptr = input.as_ptr().add(base * 2) as *const u16;
70
71            let vals = vld1q_u16(in_ptr);
72            let low = vmovl_u16(vget_low_u16(vals));
73            let high = vmovl_u16(vget_high_u16(vals));
74
75            let out_ptr = output.as_mut_ptr().add(base);
76            vst1q_u32(out_ptr, low);
77            vst1q_u32(out_ptr.add(4), high);
78        }
79
80        // Handle remainder
81        let base = chunks * 8;
82        for i in 0..remainder {
83            let idx = (base + i) * 2;
84            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
85        }
86    }
87
88    /// SIMD unpack for 32-bit values using NEON (fast copy)
89    #[target_feature(enable = "neon")]
90    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
91        let chunks = count / 4;
92        let remainder = count % 4;
93
94        let in_ptr = input.as_ptr() as *const u32;
95        let out_ptr = output.as_mut_ptr();
96
97        for chunk in 0..chunks {
98            let vals = vld1q_u32(in_ptr.add(chunk * 4));
99            vst1q_u32(out_ptr.add(chunk * 4), vals);
100        }
101
102        // Handle remainder
103        let base = chunks * 4;
104        for i in 0..remainder {
105            let idx = (base + i) * 4;
106            output[base + i] =
107                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
108        }
109    }
110
111    /// SIMD prefix sum for 4 u32 values using NEON
112    /// Input:  [a, b, c, d]
113    /// Output: [a, a+b, a+b+c, a+b+c+d]
114    #[inline]
115    #[target_feature(enable = "neon")]
116    unsafe fn prefix_sum_4(v: uint32x4_t) -> uint32x4_t {
117        // Step 1: shift by 1 and add
118        // [a, b, c, d] + [0, a, b, c] = [a, a+b, b+c, c+d]
119        let shifted1 = vextq_u32(vdupq_n_u32(0), v, 3);
120        let sum1 = vaddq_u32(v, shifted1);
121
122        // Step 2: shift by 2 and add
123        // [a, a+b, b+c, c+d] + [0, 0, a, a+b] = [a, a+b, a+b+c, a+b+c+d]
124        let shifted2 = vextq_u32(vdupq_n_u32(0), sum1, 2);
125        vaddq_u32(sum1, shifted2)
126    }
127
128    /// SIMD delta decode: convert deltas to absolute doc IDs
129    /// deltas[i] stores (gap - 1), output[i] = first + sum(gaps[0..i])
130    /// Uses NEON SIMD prefix sum for high throughput
131    #[target_feature(enable = "neon")]
132    pub unsafe fn delta_decode(
133        output: &mut [u32],
134        deltas: &[u32],
135        first_doc_id: u32,
136        count: usize,
137    ) {
138        if count == 0 {
139            return;
140        }
141
142        output[0] = first_doc_id;
143        if count == 1 {
144            return;
145        }
146
147        let ones = vdupq_n_u32(1);
148        let mut carry = vdupq_n_u32(first_doc_id);
149
150        let full_groups = (count - 1) / 4;
151        let remainder = (count - 1) % 4;
152
153        for group in 0..full_groups {
154            let base = group * 4;
155
156            // Load 4 deltas and add 1 (since we store gap-1)
157            let d = vld1q_u32(deltas[base..].as_ptr());
158            let gaps = vaddq_u32(d, ones);
159
160            // Compute prefix sum within the 4 elements
161            let prefix = prefix_sum_4(gaps);
162
163            // Add carry (broadcast last element of previous group)
164            let result = vaddq_u32(prefix, carry);
165
166            // Store result
167            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
168
169            // Update carry: broadcast the last element for next iteration
170            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
171        }
172
173        // Handle remainder
174        let base = full_groups * 4;
175        let mut scalar_carry = vgetq_lane_u32(carry, 0);
176        for j in 0..remainder {
177            scalar_carry = scalar_carry.wrapping_add(deltas[base + j]).wrapping_add(1);
178            output[base + j + 1] = scalar_carry;
179        }
180    }
181
182    /// SIMD add 1 to all values (for TF decoding: stored as tf-1)
183    #[target_feature(enable = "neon")]
184    pub unsafe fn add_one(values: &mut [u32], count: usize) {
185        let ones = vdupq_n_u32(1);
186        let chunks = count / 4;
187        let remainder = count % 4;
188
189        for chunk in 0..chunks {
190            let base = chunk * 4;
191            let ptr = values.as_mut_ptr().add(base);
192            let v = vld1q_u32(ptr);
193            let result = vaddq_u32(v, ones);
194            vst1q_u32(ptr, result);
195        }
196
197        let base = chunks * 4;
198        for i in 0..remainder {
199            values[base + i] += 1;
200        }
201    }
202
203    /// Fused unpack 8-bit + delta decode using NEON
204    /// Processes 4 values at a time, fusing unpack and prefix sum
205    #[target_feature(enable = "neon")]
206    pub unsafe fn unpack_8bit_delta_decode(
207        input: &[u8],
208        output: &mut [u32],
209        first_value: u32,
210        count: usize,
211    ) {
212        output[0] = first_value;
213        if count <= 1 {
214            return;
215        }
216
217        let ones = vdupq_n_u32(1);
218        let mut carry = vdupq_n_u32(first_value);
219
220        let full_groups = (count - 1) / 4;
221        let remainder = (count - 1) % 4;
222
223        for group in 0..full_groups {
224            let base = group * 4;
225
226            // Load 4 bytes as a u32, then widen u8→u16→u32 via NEON
227            let raw = std::ptr::read_unaligned(input.as_ptr().add(base) as *const u32);
228            let bytes = vreinterpret_u8_u32(vdup_n_u32(raw));
229            let u16s = vmovl_u8(bytes); // 8×u8 → 8×u16 (only low 4 matter)
230            let d = vmovl_u16(vget_low_u16(u16s)); // 4×u16 → 4×u32
231
232            // Add 1 (since we store gap-1)
233            let gaps = vaddq_u32(d, ones);
234
235            // Compute prefix sum within the 4 elements
236            let prefix = prefix_sum_4(gaps);
237
238            // Add carry
239            let result = vaddq_u32(prefix, carry);
240
241            // Store result
242            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
243
244            // Update carry
245            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
246        }
247
248        // Handle remainder
249        let base = full_groups * 4;
250        let mut scalar_carry = vgetq_lane_u32(carry, 0);
251        for j in 0..remainder {
252            scalar_carry = scalar_carry
253                .wrapping_add(input[base + j] as u32)
254                .wrapping_add(1);
255            output[base + j + 1] = scalar_carry;
256        }
257    }
258
259    /// Fused unpack 16-bit + delta decode using NEON
260    #[target_feature(enable = "neon")]
261    pub unsafe fn unpack_16bit_delta_decode(
262        input: &[u8],
263        output: &mut [u32],
264        first_value: u32,
265        count: usize,
266    ) {
267        output[0] = first_value;
268        if count <= 1 {
269            return;
270        }
271
272        let ones = vdupq_n_u32(1);
273        let mut carry = vdupq_n_u32(first_value);
274
275        let full_groups = (count - 1) / 4;
276        let remainder = (count - 1) % 4;
277
278        for group in 0..full_groups {
279            let base = group * 4;
280            let in_ptr = input.as_ptr().add(base * 2) as *const u16;
281
282            // Load 4 u16 values and widen to u32
283            let vals = vld1_u16(in_ptr);
284            let d = vmovl_u16(vals);
285
286            // Add 1 (since we store gap-1)
287            let gaps = vaddq_u32(d, ones);
288
289            // Compute prefix sum within the 4 elements
290            let prefix = prefix_sum_4(gaps);
291
292            // Add carry
293            let result = vaddq_u32(prefix, carry);
294
295            // Store result
296            vst1q_u32(output[base + 1..].as_mut_ptr(), result);
297
298            // Update carry
299            carry = vdupq_n_u32(vgetq_lane_u32(result, 3));
300        }
301
302        // Handle remainder
303        let base = full_groups * 4;
304        let mut scalar_carry = vgetq_lane_u32(carry, 0);
305        for j in 0..remainder {
306            let idx = (base + j) * 2;
307            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
308            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
309            output[base + j + 1] = scalar_carry;
310        }
311    }
312
313    /// NEON Hamming distance: XOR + byte popcount + horizontal sum.
314    /// Processes 16 bytes per iteration (vs 8 for scalar u64 path).
315    #[target_feature(enable = "neon")]
316    pub unsafe fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
317        let len = a.len();
318        let chunks16 = len / 16;
319        let mut total = 0u32;
320
321        // Process 16 bytes at a time, flush u8 accumulators every 31 iters
322        // (vcntq_u8 returns 0-8 per lane; 31 * 8 = 248 ≤ 255, avoiding u8 overflow)
323        let mut i = 0;
324        while i < chunks16 {
325            let batch_end = (i + 31).min(chunks16);
326            let mut acc = vdupq_n_u8(0);
327            for j in i..batch_end {
328                let off = j * 16;
329                let va = vld1q_u8(a.as_ptr().add(off));
330                let vb = vld1q_u8(b.as_ptr().add(off));
331                let popcnt = vcntq_u8(veorq_u8(va, vb));
332                acc = vaddq_u8(acc, popcnt);
333            }
334            // Widen u8 -> u16 -> u32 -> u64 and horizontal sum
335            let sum64 = vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(acc)));
336            total += vgetq_lane_u64(sum64, 0) as u32 + vgetq_lane_u64(sum64, 1) as u32;
337            i = batch_end;
338        }
339
340        // Remainder bytes (< 16)
341        let base = chunks16 * 16;
342        for k in base..len {
343            total += (a[k] ^ b[k]).count_ones();
344        }
345
346        total
347    }
348
349    /// Check if NEON is available (always true on aarch64)
350    #[inline]
351    pub fn is_available() -> bool {
352        true
353    }
354}
355
356// ============================================================================
357// SSE intrinsics for x86_64 (Intel/AMD)
358// ============================================================================
359
360#[cfg(target_arch = "x86_64")]
361#[allow(unsafe_op_in_unsafe_fn)]
362mod sse {
363    use std::arch::x86_64::*;
364
365    /// SIMD unpack for 8-bit values using SSE
366    #[target_feature(enable = "sse2", enable = "sse4.1")]
367    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
368        let chunks = count / 16;
369        let remainder = count % 16;
370
371        for chunk in 0..chunks {
372            let base = chunk * 16;
373            let in_ptr = input.as_ptr().add(base);
374
375            let bytes = _mm_loadu_si128(in_ptr as *const __m128i);
376
377            // Zero extend u8 -> u32 using SSE4.1 pmovzx
378            let v0 = _mm_cvtepu8_epi32(bytes);
379            let v1 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 4));
380            let v2 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 8));
381            let v3 = _mm_cvtepu8_epi32(_mm_srli_si128(bytes, 12));
382
383            let out_ptr = output.as_mut_ptr().add(base);
384            _mm_storeu_si128(out_ptr as *mut __m128i, v0);
385            _mm_storeu_si128(out_ptr.add(4) as *mut __m128i, v1);
386            _mm_storeu_si128(out_ptr.add(8) as *mut __m128i, v2);
387            _mm_storeu_si128(out_ptr.add(12) as *mut __m128i, v3);
388        }
389
390        let base = chunks * 16;
391        for i in 0..remainder {
392            output[base + i] = input[base + i] as u32;
393        }
394    }
395
396    /// SIMD unpack for 16-bit values using SSE
397    #[target_feature(enable = "sse2", enable = "sse4.1")]
398    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
399        let chunks = count / 8;
400        let remainder = count % 8;
401
402        for chunk in 0..chunks {
403            let base = chunk * 8;
404            let in_ptr = input.as_ptr().add(base * 2);
405
406            let vals = _mm_loadu_si128(in_ptr as *const __m128i);
407            let low = _mm_cvtepu16_epi32(vals);
408            let high = _mm_cvtepu16_epi32(_mm_srli_si128(vals, 8));
409
410            let out_ptr = output.as_mut_ptr().add(base);
411            _mm_storeu_si128(out_ptr as *mut __m128i, low);
412            _mm_storeu_si128(out_ptr.add(4) as *mut __m128i, high);
413        }
414
415        let base = chunks * 8;
416        for i in 0..remainder {
417            let idx = (base + i) * 2;
418            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
419        }
420    }
421
422    /// SIMD unpack for 32-bit values using SSE (fast copy)
423    #[target_feature(enable = "sse2")]
424    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
425        let chunks = count / 4;
426        let remainder = count % 4;
427
428        let in_ptr = input.as_ptr() as *const __m128i;
429        let out_ptr = output.as_mut_ptr() as *mut __m128i;
430
431        for chunk in 0..chunks {
432            let vals = _mm_loadu_si128(in_ptr.add(chunk));
433            _mm_storeu_si128(out_ptr.add(chunk), vals);
434        }
435
436        // Handle remainder
437        let base = chunks * 4;
438        for i in 0..remainder {
439            let idx = (base + i) * 4;
440            output[base + i] =
441                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
442        }
443    }
444
445    /// SIMD prefix sum for 4 u32 values using SSE
446    /// Input:  [a, b, c, d]
447    /// Output: [a, a+b, a+b+c, a+b+c+d]
448    #[inline]
449    #[target_feature(enable = "sse2")]
450    unsafe fn prefix_sum_4(v: __m128i) -> __m128i {
451        // Step 1: shift by 1 element (4 bytes) and add
452        // [a, b, c, d] + [0, a, b, c] = [a, a+b, b+c, c+d]
453        let shifted1 = _mm_slli_si128(v, 4);
454        let sum1 = _mm_add_epi32(v, shifted1);
455
456        // Step 2: shift by 2 elements (8 bytes) and add
457        // [a, a+b, b+c, c+d] + [0, 0, a, a+b] = [a, a+b, a+b+c, a+b+c+d]
458        let shifted2 = _mm_slli_si128(sum1, 8);
459        _mm_add_epi32(sum1, shifted2)
460    }
461
462    /// SIMD delta decode using SSE with true SIMD prefix sum
463    #[target_feature(enable = "sse2", enable = "sse4.1")]
464    pub unsafe fn delta_decode(
465        output: &mut [u32],
466        deltas: &[u32],
467        first_doc_id: u32,
468        count: usize,
469    ) {
470        if count == 0 {
471            return;
472        }
473
474        output[0] = first_doc_id;
475        if count == 1 {
476            return;
477        }
478
479        let ones = _mm_set1_epi32(1);
480        let mut carry = _mm_set1_epi32(first_doc_id as i32);
481
482        let full_groups = (count - 1) / 4;
483        let remainder = (count - 1) % 4;
484
485        for group in 0..full_groups {
486            let base = group * 4;
487
488            // Load 4 deltas and add 1 (since we store gap-1)
489            let d = _mm_loadu_si128(deltas[base..].as_ptr() as *const __m128i);
490            let gaps = _mm_add_epi32(d, ones);
491
492            // Compute prefix sum within the 4 elements
493            let prefix = prefix_sum_4(gaps);
494
495            // Add carry (broadcast last element of previous group)
496            let result = _mm_add_epi32(prefix, carry);
497
498            // Store result
499            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
500
501            // Update carry: broadcast the last element for next iteration
502            carry = _mm_shuffle_epi32(result, 0xFF); // broadcast lane 3
503        }
504
505        // Handle remainder
506        let base = full_groups * 4;
507        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
508        for j in 0..remainder {
509            scalar_carry = scalar_carry.wrapping_add(deltas[base + j]).wrapping_add(1);
510            output[base + j + 1] = scalar_carry;
511        }
512    }
513
514    /// SIMD add 1 to all values using SSE
515    #[target_feature(enable = "sse2")]
516    pub unsafe fn add_one(values: &mut [u32], count: usize) {
517        let ones = _mm_set1_epi32(1);
518        let chunks = count / 4;
519        let remainder = count % 4;
520
521        for chunk in 0..chunks {
522            let base = chunk * 4;
523            let ptr = values.as_mut_ptr().add(base) as *mut __m128i;
524            let v = _mm_loadu_si128(ptr);
525            let result = _mm_add_epi32(v, ones);
526            _mm_storeu_si128(ptr, result);
527        }
528
529        let base = chunks * 4;
530        for i in 0..remainder {
531            values[base + i] += 1;
532        }
533    }
534
535    /// Fused unpack 8-bit + delta decode using SSE
536    #[target_feature(enable = "sse2", enable = "sse4.1")]
537    pub unsafe fn unpack_8bit_delta_decode(
538        input: &[u8],
539        output: &mut [u32],
540        first_value: u32,
541        count: usize,
542    ) {
543        output[0] = first_value;
544        if count <= 1 {
545            return;
546        }
547
548        let ones = _mm_set1_epi32(1);
549        let mut carry = _mm_set1_epi32(first_value as i32);
550
551        let full_groups = (count - 1) / 4;
552        let remainder = (count - 1) % 4;
553
554        for group in 0..full_groups {
555            let base = group * 4;
556
557            // Load 4 bytes (unaligned) and zero-extend to u32
558            let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
559                input.as_ptr().add(base) as *const i32
560            ));
561            let d = _mm_cvtepu8_epi32(bytes);
562
563            // Add 1 (since we store gap-1)
564            let gaps = _mm_add_epi32(d, ones);
565
566            // Compute prefix sum within the 4 elements
567            let prefix = prefix_sum_4(gaps);
568
569            // Add carry
570            let result = _mm_add_epi32(prefix, carry);
571
572            // Store result
573            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
574
575            // Update carry: broadcast the last element
576            carry = _mm_shuffle_epi32(result, 0xFF);
577        }
578
579        // Handle remainder
580        let base = full_groups * 4;
581        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
582        for j in 0..remainder {
583            scalar_carry = scalar_carry
584                .wrapping_add(input[base + j] as u32)
585                .wrapping_add(1);
586            output[base + j + 1] = scalar_carry;
587        }
588    }
589
590    /// Fused unpack 16-bit + delta decode using SSE
591    #[target_feature(enable = "sse2", enable = "sse4.1")]
592    pub unsafe fn unpack_16bit_delta_decode(
593        input: &[u8],
594        output: &mut [u32],
595        first_value: u32,
596        count: usize,
597    ) {
598        output[0] = first_value;
599        if count <= 1 {
600            return;
601        }
602
603        let ones = _mm_set1_epi32(1);
604        let mut carry = _mm_set1_epi32(first_value as i32);
605
606        let full_groups = (count - 1) / 4;
607        let remainder = (count - 1) % 4;
608
609        for group in 0..full_groups {
610            let base = group * 4;
611            let in_ptr = input.as_ptr().add(base * 2);
612
613            // Load 8 bytes (4 u16 values, unaligned) and zero-extend to u32
614            let vals = _mm_loadl_epi64(in_ptr as *const __m128i); // loadl_epi64 supports unaligned
615            let d = _mm_cvtepu16_epi32(vals);
616
617            // Add 1 (since we store gap-1)
618            let gaps = _mm_add_epi32(d, ones);
619
620            // Compute prefix sum within the 4 elements
621            let prefix = prefix_sum_4(gaps);
622
623            // Add carry
624            let result = _mm_add_epi32(prefix, carry);
625
626            // Store result
627            _mm_storeu_si128(output[base + 1..].as_mut_ptr() as *mut __m128i, result);
628
629            // Update carry: broadcast the last element
630            carry = _mm_shuffle_epi32(result, 0xFF);
631        }
632
633        // Handle remainder
634        let base = full_groups * 4;
635        let mut scalar_carry = _mm_extract_epi32(carry, 0) as u32;
636        for j in 0..remainder {
637            let idx = (base + j) * 2;
638            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
639            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
640            output[base + j + 1] = scalar_carry;
641        }
642    }
643
644    /// Check if SSE4.1 is available at runtime
645    #[inline]
646    pub fn is_available() -> bool {
647        is_x86_feature_detected!("sse4.1")
648    }
649}
650
651// ============================================================================
652// AVX2 intrinsics for x86_64 (Intel/AMD with 256-bit registers)
653// ============================================================================
654
655#[cfg(target_arch = "x86_64")]
656#[allow(unsafe_op_in_unsafe_fn)]
657mod avx2 {
658    use std::arch::x86_64::*;
659
660    /// AVX2 unpack for 8-bit values (processes 32 bytes at a time)
661    #[target_feature(enable = "avx2")]
662    pub unsafe fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
663        let chunks = count / 32;
664        let remainder = count % 32;
665
666        for chunk in 0..chunks {
667            let base = chunk * 32;
668            let in_ptr = input.as_ptr().add(base);
669
670            // Load 32 bytes (two 128-bit loads, then combine)
671            let bytes_lo = _mm_loadu_si128(in_ptr as *const __m128i);
672            let bytes_hi = _mm_loadu_si128(in_ptr.add(16) as *const __m128i);
673
674            // Zero extend first 16 bytes: u8 -> u32
675            let v0 = _mm256_cvtepu8_epi32(bytes_lo);
676            let v1 = _mm256_cvtepu8_epi32(_mm_srli_si128(bytes_lo, 8));
677            let v2 = _mm256_cvtepu8_epi32(bytes_hi);
678            let v3 = _mm256_cvtepu8_epi32(_mm_srli_si128(bytes_hi, 8));
679
680            let out_ptr = output.as_mut_ptr().add(base);
681            _mm256_storeu_si256(out_ptr as *mut __m256i, v0);
682            _mm256_storeu_si256(out_ptr.add(8) as *mut __m256i, v1);
683            _mm256_storeu_si256(out_ptr.add(16) as *mut __m256i, v2);
684            _mm256_storeu_si256(out_ptr.add(24) as *mut __m256i, v3);
685        }
686
687        // Handle remainder with SSE
688        let base = chunks * 32;
689        for i in 0..remainder {
690            output[base + i] = input[base + i] as u32;
691        }
692    }
693
694    /// AVX2 unpack for 16-bit values (processes 16 values at a time)
695    #[target_feature(enable = "avx2")]
696    pub unsafe fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
697        let chunks = count / 16;
698        let remainder = count % 16;
699
700        for chunk in 0..chunks {
701            let base = chunk * 16;
702            let in_ptr = input.as_ptr().add(base * 2);
703
704            // Load 32 bytes (16 u16 values)
705            let vals_lo = _mm_loadu_si128(in_ptr as *const __m128i);
706            let vals_hi = _mm_loadu_si128(in_ptr.add(16) as *const __m128i);
707
708            // Zero extend u16 -> u32
709            let v0 = _mm256_cvtepu16_epi32(vals_lo);
710            let v1 = _mm256_cvtepu16_epi32(vals_hi);
711
712            let out_ptr = output.as_mut_ptr().add(base);
713            _mm256_storeu_si256(out_ptr as *mut __m256i, v0);
714            _mm256_storeu_si256(out_ptr.add(8) as *mut __m256i, v1);
715        }
716
717        // Handle remainder
718        let base = chunks * 16;
719        for i in 0..remainder {
720            let idx = (base + i) * 2;
721            output[base + i] = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
722        }
723    }
724
725    /// AVX2 unpack for 32-bit values (fast copy, 8 values at a time)
726    #[target_feature(enable = "avx2")]
727    pub unsafe fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
728        let chunks = count / 8;
729        let remainder = count % 8;
730
731        let in_ptr = input.as_ptr() as *const __m256i;
732        let out_ptr = output.as_mut_ptr() as *mut __m256i;
733
734        for chunk in 0..chunks {
735            let vals = _mm256_loadu_si256(in_ptr.add(chunk));
736            _mm256_storeu_si256(out_ptr.add(chunk), vals);
737        }
738
739        // Handle remainder
740        let base = chunks * 8;
741        for i in 0..remainder {
742            let idx = (base + i) * 4;
743            output[base + i] =
744                u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
745        }
746    }
747
748    /// AVX2 add 1 to all values (8 values at a time)
749    #[target_feature(enable = "avx2")]
750    pub unsafe fn add_one(values: &mut [u32], count: usize) {
751        let ones = _mm256_set1_epi32(1);
752        let chunks = count / 8;
753        let remainder = count % 8;
754
755        for chunk in 0..chunks {
756            let base = chunk * 8;
757            let ptr = values.as_mut_ptr().add(base) as *mut __m256i;
758            let v = _mm256_loadu_si256(ptr);
759            let result = _mm256_add_epi32(v, ones);
760            _mm256_storeu_si256(ptr, result);
761        }
762
763        let base = chunks * 8;
764        for i in 0..remainder {
765            values[base + i] += 1;
766        }
767    }
768
769    /// AVX2 prefix sum for 8 u32 values (Hillis-Steele)
770    /// Input:  [a, b, c, d, e, f, g, h]
771    /// Output: [a, a+b, a+b+c, ..., a+b+c+d+e+f+g+h]
772    #[inline]
773    #[target_feature(enable = "avx2")]
774    unsafe fn prefix_sum_8(v: __m256i) -> __m256i {
775        // Step 1: intra-lane shift by 1 element (4 bytes) and add
776        let s1 = _mm256_slli_si256(v, 4);
777        let r1 = _mm256_add_epi32(v, s1);
778
779        // Step 2: intra-lane shift by 2 elements (8 bytes) and add
780        let s2 = _mm256_slli_si256(r1, 8);
781        let r2 = _mm256_add_epi32(r1, s2);
782
783        // Step 3: propagate lower lane sum to upper lane
784        // Broadcast element 3 (lower lane sum) within each lane
785        let lo_sum = _mm256_shuffle_epi32(r2, 0xFF);
786        // Duplicate lane 0 to both lanes
787        let carry = _mm256_permute2x128_si256(lo_sum, lo_sum, 0x00);
788        // Zero carry for lower lane, keep for upper
789        let carry_hi = _mm256_blend_epi32::<0xF0>(_mm256_setzero_si256(), carry);
790        _mm256_add_epi32(r2, carry_hi)
791    }
792
793    /// AVX2 fused unpack 8-bit + delta decode (processes 8 values at a time)
794    #[target_feature(enable = "avx2")]
795    pub unsafe fn unpack_8bit_delta_decode(
796        input: &[u8],
797        output: &mut [u32],
798        first_value: u32,
799        count: usize,
800    ) {
801        output[0] = first_value;
802        if count <= 1 {
803            return;
804        }
805
806        let ones = _mm256_set1_epi32(1);
807        let mut carry = _mm256_set1_epi32(first_value as i32);
808        let broadcast_idx = _mm256_set1_epi32(7);
809
810        let full_groups = (count - 1) / 8;
811        let remainder = (count - 1) % 8;
812
813        for group in 0..full_groups {
814            let base = group * 8;
815
816            // Load 8 bytes and zero-extend to 8×u32
817            let bytes = _mm_loadl_epi64(input.as_ptr().add(base) as *const __m128i);
818            let d = _mm256_cvtepu8_epi32(bytes);
819
820            // Add 1 (since we store gap-1)
821            let gaps = _mm256_add_epi32(d, ones);
822
823            // Compute prefix sum within 8 elements
824            let prefix = prefix_sum_8(gaps);
825
826            // Add carry from previous group
827            let result = _mm256_add_epi32(prefix, carry);
828
829            // Store 8 results
830            _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
831
832            // Update carry: broadcast element 7 to all positions
833            carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
834        }
835
836        // Handle remainder with scalar
837        let base = full_groups * 8;
838        let mut scalar_carry = _mm256_extract_epi32::<0>(carry) as u32;
839        for j in 0..remainder {
840            scalar_carry = scalar_carry
841                .wrapping_add(input[base + j] as u32)
842                .wrapping_add(1);
843            output[base + j + 1] = scalar_carry;
844        }
845    }
846
847    /// AVX2 fused unpack 16-bit + delta decode (processes 8 values at a time)
848    #[target_feature(enable = "avx2")]
849    pub unsafe fn unpack_16bit_delta_decode(
850        input: &[u8],
851        output: &mut [u32],
852        first_value: u32,
853        count: usize,
854    ) {
855        output[0] = first_value;
856        if count <= 1 {
857            return;
858        }
859
860        let ones = _mm256_set1_epi32(1);
861        let mut carry = _mm256_set1_epi32(first_value as i32);
862        let broadcast_idx = _mm256_set1_epi32(7);
863
864        let full_groups = (count - 1) / 8;
865        let remainder = (count - 1) % 8;
866
867        for group in 0..full_groups {
868            let base = group * 8;
869            let in_ptr = input.as_ptr().add(base * 2);
870
871            // Load 16 bytes (8 u16 values) and zero-extend to 8×u32
872            let vals = _mm_loadu_si128(in_ptr as *const __m128i);
873            let d = _mm256_cvtepu16_epi32(vals);
874
875            // Add 1 (since we store gap-1)
876            let gaps = _mm256_add_epi32(d, ones);
877
878            // Compute prefix sum within 8 elements
879            let prefix = prefix_sum_8(gaps);
880
881            // Add carry from previous group
882            let result = _mm256_add_epi32(prefix, carry);
883
884            // Store 8 results
885            _mm256_storeu_si256(output[base + 1..].as_mut_ptr() as *mut __m256i, result);
886
887            // Update carry: broadcast element 7 to all positions
888            carry = _mm256_permutevar8x32_epi32(result, broadcast_idx);
889        }
890
891        // Handle remainder with scalar
892        let base = full_groups * 8;
893        let mut scalar_carry = _mm256_extract_epi32::<0>(carry) as u32;
894        for j in 0..remainder {
895            let idx = (base + j) * 2;
896            let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
897            scalar_carry = scalar_carry.wrapping_add(delta).wrapping_add(1);
898            output[base + j + 1] = scalar_carry;
899        }
900    }
901
902    /// AVX2 Hamming distance using VPSHUFB-based popcount (Muła algorithm).
903    /// Processes 32 bytes per iteration with a nibble lookup table.
904    #[target_feature(enable = "avx2")]
905    pub unsafe fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
906        let len = a.len();
907        let chunks32 = len / 32;
908        let low_mask = _mm256_set1_epi8(0x0f);
909        // Nibble popcount lookup table: popcount(0..15)
910        let lookup = _mm256_setr_epi8(
911            0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2,
912            3, 3, 4,
913        );
914        let mut total = 0u64;
915
916        let mut i = 0;
917        while i < chunks32 {
918            // Accumulate in u8 lanes, flush every 31 iters to avoid overflow
919            // (nibble popcount gives 0-8 per lane; 31 * 8 = 248 ≤ 255)
920            let batch_end = (i + 31).min(chunks32);
921            let mut acc = _mm256_setzero_si256();
922            for j in i..batch_end {
923                let off = j * 32;
924                let va = _mm256_loadu_si256(a.as_ptr().add(off) as *const __m256i);
925                let vb = _mm256_loadu_si256(b.as_ptr().add(off) as *const __m256i);
926                let xored = _mm256_xor_si256(va, vb);
927                // VPSHUFB popcount: count bits per byte via nibble lookup
928                let lo = _mm256_and_si256(xored, low_mask);
929                let hi = _mm256_and_si256(_mm256_srli_epi16(xored, 4), low_mask);
930                let popcnt = _mm256_add_epi8(
931                    _mm256_shuffle_epi8(lookup, lo),
932                    _mm256_shuffle_epi8(lookup, hi),
933                );
934                acc = _mm256_add_epi8(acc, popcnt);
935            }
936            // Horizontal sum: u8 -> u64 via SAD against zero
937            let sad = _mm256_sad_epu8(acc, _mm256_setzero_si256());
938            total += _mm256_extract_epi64(sad, 0) as u64
939                + _mm256_extract_epi64(sad, 1) as u64
940                + _mm256_extract_epi64(sad, 2) as u64
941                + _mm256_extract_epi64(sad, 3) as u64;
942            i = batch_end;
943        }
944
945        // Remainder bytes (< 32)
946        let base = chunks32 * 32;
947        for k in base..len {
948            total += (a[k] ^ b[k]).count_ones() as u64;
949        }
950
951        total as u32
952    }
953
954    /// Check if AVX2 is available at runtime
955    #[inline]
956    pub fn is_available() -> bool {
957        is_x86_feature_detected!("avx2")
958    }
959}
960
961// ============================================================================
962// Scalar fallback implementations
963// ============================================================================
964
965#[allow(dead_code)]
966mod scalar {
967    /// Scalar unpack for 8-bit values
968    #[inline]
969    pub fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
970        for i in 0..count {
971            output[i] = input[i] as u32;
972        }
973    }
974
975    /// Scalar unpack for 16-bit values
976    #[inline]
977    pub fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
978        for (i, out) in output.iter_mut().enumerate().take(count) {
979            let idx = i * 2;
980            *out = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
981        }
982    }
983
984    /// Scalar unpack for 32-bit values
985    #[inline]
986    pub fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
987        for (i, out) in output.iter_mut().enumerate().take(count) {
988            let idx = i * 4;
989            *out = u32::from_le_bytes([input[idx], input[idx + 1], input[idx + 2], input[idx + 3]]);
990        }
991    }
992
993    /// Scalar delta decode
994    #[inline]
995    pub fn delta_decode(output: &mut [u32], deltas: &[u32], first_doc_id: u32, count: usize) {
996        if count == 0 {
997            return;
998        }
999
1000        output[0] = first_doc_id;
1001        let mut carry = first_doc_id;
1002
1003        for i in 0..count - 1 {
1004            carry = carry.wrapping_add(deltas[i]).wrapping_add(1);
1005            output[i + 1] = carry;
1006        }
1007    }
1008
1009    /// Scalar add 1 to all values
1010    #[inline]
1011    pub fn add_one(values: &mut [u32], count: usize) {
1012        for val in values.iter_mut().take(count) {
1013            *val += 1;
1014        }
1015    }
1016}
1017
1018// ============================================================================
1019// Public dispatch functions that select SIMD or scalar at runtime
1020// ============================================================================
1021
1022/// Unpack 8-bit packed values to u32 with SIMD acceleration
1023#[inline]
1024pub fn unpack_8bit(input: &[u8], output: &mut [u32], count: usize) {
1025    #[cfg(target_arch = "aarch64")]
1026    {
1027        if neon::is_available() {
1028            unsafe {
1029                neon::unpack_8bit(input, output, count);
1030            }
1031            return;
1032        }
1033    }
1034
1035    #[cfg(target_arch = "x86_64")]
1036    {
1037        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1038        if avx2::is_available() {
1039            unsafe {
1040                avx2::unpack_8bit(input, output, count);
1041            }
1042            return;
1043        }
1044        if sse::is_available() {
1045            unsafe {
1046                sse::unpack_8bit(input, output, count);
1047            }
1048            return;
1049        }
1050    }
1051
1052    scalar::unpack_8bit(input, output, count);
1053}
1054
1055/// Unpack 16-bit packed values to u32 with SIMD acceleration
1056#[inline]
1057pub fn unpack_16bit(input: &[u8], output: &mut [u32], count: usize) {
1058    #[cfg(target_arch = "aarch64")]
1059    {
1060        if neon::is_available() {
1061            unsafe {
1062                neon::unpack_16bit(input, output, count);
1063            }
1064            return;
1065        }
1066    }
1067
1068    #[cfg(target_arch = "x86_64")]
1069    {
1070        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1071        if avx2::is_available() {
1072            unsafe {
1073                avx2::unpack_16bit(input, output, count);
1074            }
1075            return;
1076        }
1077        if sse::is_available() {
1078            unsafe {
1079                sse::unpack_16bit(input, output, count);
1080            }
1081            return;
1082        }
1083    }
1084
1085    scalar::unpack_16bit(input, output, count);
1086}
1087
1088/// Unpack 32-bit packed values to u32 with SIMD acceleration
1089#[inline]
1090pub fn unpack_32bit(input: &[u8], output: &mut [u32], count: usize) {
1091    #[cfg(target_arch = "aarch64")]
1092    {
1093        if neon::is_available() {
1094            unsafe {
1095                neon::unpack_32bit(input, output, count);
1096            }
1097            return;
1098        }
1099    }
1100
1101    #[cfg(target_arch = "x86_64")]
1102    {
1103        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1104        if avx2::is_available() {
1105            unsafe {
1106                avx2::unpack_32bit(input, output, count);
1107            }
1108            return;
1109        }
1110        if sse::is_available() {
1111            unsafe {
1112                sse::unpack_32bit(input, output, count);
1113            }
1114            return;
1115        }
1116    }
1117
1118    scalar::unpack_32bit(input, output, count);
1119}
1120
1121/// Delta decode with SIMD acceleration
1122///
1123/// Converts delta-encoded values to absolute values.
1124/// Input: deltas[i] = value[i+1] - value[i] - 1 (gap minus one)
1125/// Output: absolute values starting from first_value
1126#[inline]
1127pub fn delta_decode(output: &mut [u32], deltas: &[u32], first_value: u32, count: usize) {
1128    #[cfg(target_arch = "aarch64")]
1129    {
1130        if neon::is_available() {
1131            unsafe {
1132                neon::delta_decode(output, deltas, first_value, count);
1133            }
1134            return;
1135        }
1136    }
1137
1138    #[cfg(target_arch = "x86_64")]
1139    {
1140        if sse::is_available() {
1141            unsafe {
1142                sse::delta_decode(output, deltas, first_value, count);
1143            }
1144            return;
1145        }
1146    }
1147
1148    scalar::delta_decode(output, deltas, first_value, count);
1149}
1150
1151/// Add 1 to all values with SIMD acceleration
1152///
1153/// Used for TF decoding where values are stored as (tf - 1)
1154#[inline]
1155pub fn add_one(values: &mut [u32], count: usize) {
1156    #[cfg(target_arch = "aarch64")]
1157    {
1158        if neon::is_available() {
1159            unsafe {
1160                neon::add_one(values, count);
1161            }
1162            return;
1163        }
1164    }
1165
1166    #[cfg(target_arch = "x86_64")]
1167    {
1168        // Prefer AVX2 (256-bit) over SSE (128-bit) when available
1169        if avx2::is_available() {
1170            unsafe {
1171                avx2::add_one(values, count);
1172            }
1173            return;
1174        }
1175        if sse::is_available() {
1176            unsafe {
1177                sse::add_one(values, count);
1178            }
1179            return;
1180        }
1181    }
1182
1183    scalar::add_one(values, count);
1184}
1185
1186/// Compute the number of bits needed to represent a value
1187#[inline]
1188pub fn bits_needed(val: u32) -> u8 {
1189    if val == 0 {
1190        0
1191    } else {
1192        32 - val.leading_zeros() as u8
1193    }
1194}
1195
1196// ============================================================================
1197// Rounded bitpacking for truly vectorized encoding/decoding
1198// ============================================================================
1199//
1200// Instead of using arbitrary bit widths (1-32), we round up to SIMD-friendly
1201// widths: 0, 8, 16, or 32 bits. This trades ~10-20% more space for much faster
1202// decoding since we can use direct SIMD widening instructions (pmovzx) without
1203// any bit-shifting or masking.
1204//
1205// Bit width mapping:
1206//   0      -> 0  (all zeros)
1207//   1-8    -> 8  (u8)
1208//   9-16   -> 16 (u16)
1209//   17-32  -> 32 (u32)
1210
1211/// Rounded bit width type for SIMD-friendly encoding
1212#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1213#[repr(u8)]
1214pub enum RoundedBitWidth {
1215    Zero = 0,
1216    Bits8 = 8,
1217    Bits16 = 16,
1218    Bits32 = 32,
1219}
1220
1221impl RoundedBitWidth {
1222    /// Round an exact bit width to the nearest SIMD-friendly width
1223    #[inline]
1224    pub fn from_exact(bits: u8) -> Self {
1225        match bits {
1226            0 => RoundedBitWidth::Zero,
1227            1..=8 => RoundedBitWidth::Bits8,
1228            9..=16 => RoundedBitWidth::Bits16,
1229            _ => RoundedBitWidth::Bits32,
1230        }
1231    }
1232
1233    /// Convert from stored u8 value (must be 0, 8, 16, or 32)
1234    #[inline]
1235    pub fn from_u8(bits: u8) -> Self {
1236        match bits {
1237            0 => RoundedBitWidth::Zero,
1238            8 => RoundedBitWidth::Bits8,
1239            16 => RoundedBitWidth::Bits16,
1240            32 => RoundedBitWidth::Bits32,
1241            _ => RoundedBitWidth::Bits32, // Fallback for invalid values
1242        }
1243    }
1244
1245    /// Get the byte size per value
1246    #[inline]
1247    pub fn bytes_per_value(self) -> usize {
1248        match self {
1249            RoundedBitWidth::Zero => 0,
1250            RoundedBitWidth::Bits8 => 1,
1251            RoundedBitWidth::Bits16 => 2,
1252            RoundedBitWidth::Bits32 => 4,
1253        }
1254    }
1255
1256    /// Get the raw bit width value
1257    #[inline]
1258    pub fn as_u8(self) -> u8 {
1259        self as u8
1260    }
1261}
1262
1263/// Round a bit width to the nearest SIMD-friendly width (0, 8, 16, or 32)
1264#[inline]
1265pub fn round_bit_width(bits: u8) -> u8 {
1266    RoundedBitWidth::from_exact(bits).as_u8()
1267}
1268
1269/// Pack values using rounded bit width (SIMD-friendly)
1270///
1271/// This is much simpler than arbitrary bitpacking since values are byte-aligned.
1272/// Returns the number of bytes written.
1273#[inline]
1274pub fn pack_rounded(values: &[u32], bit_width: RoundedBitWidth, output: &mut [u8]) -> usize {
1275    let count = values.len();
1276    match bit_width {
1277        RoundedBitWidth::Zero => 0,
1278        RoundedBitWidth::Bits8 => {
1279            for (i, &v) in values.iter().enumerate() {
1280                output[i] = v as u8;
1281            }
1282            count
1283        }
1284        RoundedBitWidth::Bits16 => {
1285            for (i, &v) in values.iter().enumerate() {
1286                let bytes = (v as u16).to_le_bytes();
1287                output[i * 2] = bytes[0];
1288                output[i * 2 + 1] = bytes[1];
1289            }
1290            count * 2
1291        }
1292        RoundedBitWidth::Bits32 => {
1293            for (i, &v) in values.iter().enumerate() {
1294                let bytes = v.to_le_bytes();
1295                output[i * 4] = bytes[0];
1296                output[i * 4 + 1] = bytes[1];
1297                output[i * 4 + 2] = bytes[2];
1298                output[i * 4 + 3] = bytes[3];
1299            }
1300            count * 4
1301        }
1302    }
1303}
1304
1305/// Unpack values using rounded bit width with SIMD acceleration
1306///
1307/// This is the fast path - no bit manipulation needed, just widening.
1308#[inline]
1309pub fn unpack_rounded(input: &[u8], bit_width: RoundedBitWidth, output: &mut [u32], count: usize) {
1310    match bit_width {
1311        RoundedBitWidth::Zero => {
1312            for out in output.iter_mut().take(count) {
1313                *out = 0;
1314            }
1315        }
1316        RoundedBitWidth::Bits8 => unpack_8bit(input, output, count),
1317        RoundedBitWidth::Bits16 => unpack_16bit(input, output, count),
1318        RoundedBitWidth::Bits32 => unpack_32bit(input, output, count),
1319    }
1320}
1321
1322/// Fused unpack + delta decode using rounded bit width
1323///
1324/// Combines unpacking and prefix sum in a single pass for better cache utilization.
1325#[inline]
1326pub fn unpack_rounded_delta_decode(
1327    input: &[u8],
1328    bit_width: RoundedBitWidth,
1329    output: &mut [u32],
1330    first_value: u32,
1331    count: usize,
1332) {
1333    match bit_width {
1334        RoundedBitWidth::Zero => {
1335            // All deltas are 0, meaning gaps of 1
1336            let mut val = first_value;
1337            for out in output.iter_mut().take(count) {
1338                *out = val;
1339                val = val.wrapping_add(1);
1340            }
1341        }
1342        RoundedBitWidth::Bits8 => unpack_8bit_delta_decode(input, output, first_value, count),
1343        RoundedBitWidth::Bits16 => unpack_16bit_delta_decode(input, output, first_value, count),
1344        RoundedBitWidth::Bits32 => {
1345            // Unpack count-1 deltas from input, then prefix sum to absolute values
1346            if count > 0 {
1347                output[0] = first_value;
1348                let mut carry = first_value;
1349                for i in 0..count - 1 {
1350                    let idx = i * 4;
1351                    let delta = u32::from_le_bytes([
1352                        input[idx],
1353                        input[idx + 1],
1354                        input[idx + 2],
1355                        input[idx + 3],
1356                    ]);
1357                    carry = carry.wrapping_add(delta).wrapping_add(1);
1358                    output[i + 1] = carry;
1359                }
1360            }
1361        }
1362    }
1363}
1364
1365// ============================================================================
1366// Fused operations for better cache utilization
1367// ============================================================================
1368
1369/// Fused unpack 8-bit + delta decode in a single pass
1370///
1371/// This avoids writing the intermediate unpacked values to memory,
1372/// improving cache utilization for large blocks.
1373#[inline]
1374pub fn unpack_8bit_delta_decode(input: &[u8], output: &mut [u32], first_value: u32, count: usize) {
1375    if count == 0 {
1376        return;
1377    }
1378
1379    output[0] = first_value;
1380    if count == 1 {
1381        return;
1382    }
1383
1384    #[cfg(target_arch = "aarch64")]
1385    {
1386        if neon::is_available() {
1387            unsafe {
1388                neon::unpack_8bit_delta_decode(input, output, first_value, count);
1389            }
1390            return;
1391        }
1392    }
1393
1394    #[cfg(target_arch = "x86_64")]
1395    {
1396        if avx2::is_available() {
1397            unsafe {
1398                avx2::unpack_8bit_delta_decode(input, output, first_value, count);
1399            }
1400            return;
1401        }
1402        if sse::is_available() {
1403            unsafe {
1404                sse::unpack_8bit_delta_decode(input, output, first_value, count);
1405            }
1406            return;
1407        }
1408    }
1409
1410    // Scalar fallback
1411    let mut carry = first_value;
1412    for i in 0..count - 1 {
1413        carry = carry.wrapping_add(input[i] as u32).wrapping_add(1);
1414        output[i + 1] = carry;
1415    }
1416}
1417
1418/// Fused unpack 16-bit + delta decode in a single pass
1419#[inline]
1420pub fn unpack_16bit_delta_decode(input: &[u8], output: &mut [u32], first_value: u32, count: usize) {
1421    if count == 0 {
1422        return;
1423    }
1424
1425    output[0] = first_value;
1426    if count == 1 {
1427        return;
1428    }
1429
1430    #[cfg(target_arch = "aarch64")]
1431    {
1432        if neon::is_available() {
1433            unsafe {
1434                neon::unpack_16bit_delta_decode(input, output, first_value, count);
1435            }
1436            return;
1437        }
1438    }
1439
1440    #[cfg(target_arch = "x86_64")]
1441    {
1442        if avx2::is_available() {
1443            unsafe {
1444                avx2::unpack_16bit_delta_decode(input, output, first_value, count);
1445            }
1446            return;
1447        }
1448        if sse::is_available() {
1449            unsafe {
1450                sse::unpack_16bit_delta_decode(input, output, first_value, count);
1451            }
1452            return;
1453        }
1454    }
1455
1456    // Scalar fallback
1457    let mut carry = first_value;
1458    for i in 0..count - 1 {
1459        let idx = i * 2;
1460        let delta = u16::from_le_bytes([input[idx], input[idx + 1]]) as u32;
1461        carry = carry.wrapping_add(delta).wrapping_add(1);
1462        output[i + 1] = carry;
1463    }
1464}
1465
1466/// Fused unpack + delta decode for arbitrary bit widths
1467///
1468/// Combines unpacking and prefix sum in a single pass, avoiding intermediate buffer.
1469/// Uses SIMD-accelerated paths for 8/16-bit widths, scalar for others.
1470#[inline]
1471pub fn unpack_delta_decode(
1472    input: &[u8],
1473    bit_width: u8,
1474    output: &mut [u32],
1475    first_value: u32,
1476    count: usize,
1477) {
1478    if count == 0 {
1479        return;
1480    }
1481
1482    output[0] = first_value;
1483    if count == 1 {
1484        return;
1485    }
1486
1487    // Fast paths for SIMD-friendly bit widths
1488    match bit_width {
1489        0 => {
1490            // All zeros = consecutive doc IDs (gap of 1)
1491            let mut val = first_value;
1492            for item in output.iter_mut().take(count).skip(1) {
1493                val = val.wrapping_add(1);
1494                *item = val;
1495            }
1496        }
1497        8 => unpack_8bit_delta_decode(input, output, first_value, count),
1498        16 => unpack_16bit_delta_decode(input, output, first_value, count),
1499        32 => {
1500            // 32-bit: unpack inline and delta decode
1501            let mut carry = first_value;
1502            for i in 0..count - 1 {
1503                let idx = i * 4;
1504                let delta = u32::from_le_bytes([
1505                    input[idx],
1506                    input[idx + 1],
1507                    input[idx + 2],
1508                    input[idx + 3],
1509                ]);
1510                carry = carry.wrapping_add(delta).wrapping_add(1);
1511                output[i + 1] = carry;
1512            }
1513        }
1514        _ => {
1515            // Generic bit width: fused unpack + delta decode
1516            let mask = (1u64 << bit_width) - 1;
1517            let bit_width_usize = bit_width as usize;
1518            let mut bit_pos = 0usize;
1519            let input_ptr = input.as_ptr();
1520            let mut carry = first_value;
1521
1522            for i in 0..count - 1 {
1523                let byte_idx = bit_pos >> 3;
1524                let bit_offset = bit_pos & 7;
1525
1526                // SAFETY: Caller guarantees input has enough data
1527                let word = unsafe { (input_ptr.add(byte_idx) as *const u64).read_unaligned() };
1528                let delta = ((word >> bit_offset) & mask) as u32;
1529
1530                carry = carry.wrapping_add(delta).wrapping_add(1);
1531                output[i + 1] = carry;
1532                bit_pos += bit_width_usize;
1533            }
1534        }
1535    }
1536}
1537
1538// ============================================================================
1539// Sparse Vector SIMD Functions
1540// ============================================================================
1541
1542/// Dequantize UInt8 weights to f32 with SIMD acceleration
1543///
1544/// Computes: output[i] = input[i] as f32 * scale + min_val
1545#[inline]
1546pub fn dequantize_uint8(input: &[u8], output: &mut [f32], scale: f32, min_val: f32, count: usize) {
1547    #[cfg(target_arch = "aarch64")]
1548    {
1549        if neon::is_available() {
1550            unsafe {
1551                dequantize_uint8_neon(input, output, scale, min_val, count);
1552            }
1553            return;
1554        }
1555    }
1556
1557    #[cfg(target_arch = "x86_64")]
1558    {
1559        if sse::is_available() {
1560            unsafe {
1561                dequantize_uint8_sse(input, output, scale, min_val, count);
1562            }
1563            return;
1564        }
1565    }
1566
1567    // Scalar fallback
1568    for i in 0..count {
1569        output[i] = input[i] as f32 * scale + min_val;
1570    }
1571}
1572
1573#[cfg(target_arch = "aarch64")]
1574#[target_feature(enable = "neon")]
1575#[allow(unsafe_op_in_unsafe_fn)]
1576unsafe fn dequantize_uint8_neon(
1577    input: &[u8],
1578    output: &mut [f32],
1579    scale: f32,
1580    min_val: f32,
1581    count: usize,
1582) {
1583    use std::arch::aarch64::*;
1584
1585    let scale_v = vdupq_n_f32(scale);
1586    let min_v = vdupq_n_f32(min_val);
1587
1588    let chunks = count / 16;
1589    let remainder = count % 16;
1590
1591    for chunk in 0..chunks {
1592        let base = chunk * 16;
1593        let in_ptr = input.as_ptr().add(base);
1594
1595        // Load 16 bytes
1596        let bytes = vld1q_u8(in_ptr);
1597
1598        // Widen u8 -> u16 -> u32 -> f32
1599        let low8 = vget_low_u8(bytes);
1600        let high8 = vget_high_u8(bytes);
1601
1602        let low16 = vmovl_u8(low8);
1603        let high16 = vmovl_u8(high8);
1604
1605        // Process 4 values at a time
1606        let u32_0 = vmovl_u16(vget_low_u16(low16));
1607        let u32_1 = vmovl_u16(vget_high_u16(low16));
1608        let u32_2 = vmovl_u16(vget_low_u16(high16));
1609        let u32_3 = vmovl_u16(vget_high_u16(high16));
1610
1611        // Convert to f32 and apply scale + min_val
1612        let f32_0 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_0), scale_v);
1613        let f32_1 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_1), scale_v);
1614        let f32_2 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_2), scale_v);
1615        let f32_3 = vfmaq_f32(min_v, vcvtq_f32_u32(u32_3), scale_v);
1616
1617        let out_ptr = output.as_mut_ptr().add(base);
1618        vst1q_f32(out_ptr, f32_0);
1619        vst1q_f32(out_ptr.add(4), f32_1);
1620        vst1q_f32(out_ptr.add(8), f32_2);
1621        vst1q_f32(out_ptr.add(12), f32_3);
1622    }
1623
1624    // Handle remainder
1625    let base = chunks * 16;
1626    for i in 0..remainder {
1627        output[base + i] = input[base + i] as f32 * scale + min_val;
1628    }
1629}
1630
1631#[cfg(target_arch = "x86_64")]
1632#[target_feature(enable = "sse2", enable = "sse4.1")]
1633#[allow(unsafe_op_in_unsafe_fn)]
1634unsafe fn dequantize_uint8_sse(
1635    input: &[u8],
1636    output: &mut [f32],
1637    scale: f32,
1638    min_val: f32,
1639    count: usize,
1640) {
1641    use std::arch::x86_64::*;
1642
1643    let scale_v = _mm_set1_ps(scale);
1644    let min_v = _mm_set1_ps(min_val);
1645
1646    let chunks = count / 4;
1647    let remainder = count % 4;
1648
1649    for chunk in 0..chunks {
1650        let base = chunk * 4;
1651
1652        // Load 4 bytes as a single i32 and zero-extend u8→u32 via SSE4.1
1653        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
1654            input.as_ptr().add(base) as *const i32
1655        ));
1656        let ints = _mm_cvtepu8_epi32(bytes);
1657        let floats = _mm_cvtepi32_ps(ints);
1658
1659        // Apply scale and min_val: result = floats * scale + min_val
1660        let scaled = _mm_add_ps(_mm_mul_ps(floats, scale_v), min_v);
1661
1662        _mm_storeu_ps(output.as_mut_ptr().add(base), scaled);
1663    }
1664
1665    // Handle remainder
1666    let base = chunks * 4;
1667    for i in 0..remainder {
1668        output[base + i] = input[base + i] as f32 * scale + min_val;
1669    }
1670}
1671
1672/// Compute dot product of two f32 arrays with SIMD acceleration
1673#[inline]
1674pub fn dot_product_f32(a: &[f32], b: &[f32], count: usize) -> f32 {
1675    assert!(
1676        count <= a.len() && count <= b.len(),
1677        "dot_product_f32 count {count} exceeds input lengths ({}, {})",
1678        a.len(),
1679        b.len()
1680    );
1681    #[cfg(target_arch = "aarch64")]
1682    {
1683        if neon::is_available() {
1684            return unsafe { dot_product_f32_neon(a, b, count) };
1685        }
1686    }
1687
1688    #[cfg(target_arch = "x86_64")]
1689    {
1690        if is_x86_feature_detected!("avx512f") {
1691            return unsafe { dot_product_f32_avx512(a, b, count) };
1692        }
1693        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1694            return unsafe { dot_product_f32_avx2(a, b, count) };
1695        }
1696        if sse::is_available() {
1697            return unsafe { dot_product_f32_sse(a, b, count) };
1698        }
1699    }
1700
1701    // Scalar fallback
1702    let mut sum = 0.0f32;
1703    for i in 0..count {
1704        sum += a[i] * b[i];
1705    }
1706    sum
1707}
1708
1709#[cfg(target_arch = "aarch64")]
1710#[target_feature(enable = "neon")]
1711#[allow(unsafe_op_in_unsafe_fn)]
1712unsafe fn dot_product_f32_neon(a: &[f32], b: &[f32], count: usize) -> f32 {
1713    use std::arch::aarch64::*;
1714
1715    let chunks16 = count / 16;
1716    let remainder = count % 16;
1717
1718    let mut acc0 = vdupq_n_f32(0.0);
1719    let mut acc1 = vdupq_n_f32(0.0);
1720    let mut acc2 = vdupq_n_f32(0.0);
1721    let mut acc3 = vdupq_n_f32(0.0);
1722
1723    for c in 0..chunks16 {
1724        let base = c * 16;
1725        acc0 = vfmaq_f32(
1726            acc0,
1727            vld1q_f32(a.as_ptr().add(base)),
1728            vld1q_f32(b.as_ptr().add(base)),
1729        );
1730        acc1 = vfmaq_f32(
1731            acc1,
1732            vld1q_f32(a.as_ptr().add(base + 4)),
1733            vld1q_f32(b.as_ptr().add(base + 4)),
1734        );
1735        acc2 = vfmaq_f32(
1736            acc2,
1737            vld1q_f32(a.as_ptr().add(base + 8)),
1738            vld1q_f32(b.as_ptr().add(base + 8)),
1739        );
1740        acc3 = vfmaq_f32(
1741            acc3,
1742            vld1q_f32(a.as_ptr().add(base + 12)),
1743            vld1q_f32(b.as_ptr().add(base + 12)),
1744        );
1745    }
1746
1747    let acc = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
1748    let mut sum = vaddvq_f32(acc);
1749
1750    let base = chunks16 * 16;
1751    for i in 0..remainder {
1752        sum += a[base + i] * b[base + i];
1753    }
1754
1755    sum
1756}
1757
1758#[cfg(target_arch = "x86_64")]
1759#[target_feature(enable = "avx2", enable = "fma")]
1760#[allow(unsafe_op_in_unsafe_fn)]
1761unsafe fn dot_product_f32_avx2(a: &[f32], b: &[f32], count: usize) -> f32 {
1762    use std::arch::x86_64::*;
1763
1764    let chunks32 = count / 32;
1765    let remainder = count % 32;
1766
1767    let mut acc0 = _mm256_setzero_ps();
1768    let mut acc1 = _mm256_setzero_ps();
1769    let mut acc2 = _mm256_setzero_ps();
1770    let mut acc3 = _mm256_setzero_ps();
1771
1772    for c in 0..chunks32 {
1773        let base = c * 32;
1774        acc0 = _mm256_fmadd_ps(
1775            _mm256_loadu_ps(a.as_ptr().add(base)),
1776            _mm256_loadu_ps(b.as_ptr().add(base)),
1777            acc0,
1778        );
1779        acc1 = _mm256_fmadd_ps(
1780            _mm256_loadu_ps(a.as_ptr().add(base + 8)),
1781            _mm256_loadu_ps(b.as_ptr().add(base + 8)),
1782            acc1,
1783        );
1784        acc2 = _mm256_fmadd_ps(
1785            _mm256_loadu_ps(a.as_ptr().add(base + 16)),
1786            _mm256_loadu_ps(b.as_ptr().add(base + 16)),
1787            acc2,
1788        );
1789        acc3 = _mm256_fmadd_ps(
1790            _mm256_loadu_ps(a.as_ptr().add(base + 24)),
1791            _mm256_loadu_ps(b.as_ptr().add(base + 24)),
1792            acc3,
1793        );
1794    }
1795
1796    let acc = _mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3));
1797
1798    // Horizontal sum: 256-bit → 128-bit → scalar
1799    let hi = _mm256_extractf128_ps(acc, 1);
1800    let lo = _mm256_castps256_ps128(acc);
1801    let sum128 = _mm_add_ps(lo, hi);
1802    let shuf = _mm_shuffle_ps(sum128, sum128, 0b10_11_00_01);
1803    let sums = _mm_add_ps(sum128, shuf);
1804    let shuf2 = _mm_movehl_ps(sums, sums);
1805    let final_sum = _mm_add_ss(sums, shuf2);
1806
1807    let mut sum = _mm_cvtss_f32(final_sum);
1808
1809    let base = chunks32 * 32;
1810    for i in 0..remainder {
1811        sum += a[base + i] * b[base + i];
1812    }
1813
1814    sum
1815}
1816
1817#[cfg(target_arch = "x86_64")]
1818#[target_feature(enable = "sse")]
1819#[allow(unsafe_op_in_unsafe_fn)]
1820unsafe fn dot_product_f32_sse(a: &[f32], b: &[f32], count: usize) -> f32 {
1821    use std::arch::x86_64::*;
1822
1823    let chunks = count / 4;
1824    let remainder = count % 4;
1825
1826    let mut acc = _mm_setzero_ps();
1827
1828    for chunk in 0..chunks {
1829        let base = chunk * 4;
1830        let va = _mm_loadu_ps(a.as_ptr().add(base));
1831        let vb = _mm_loadu_ps(b.as_ptr().add(base));
1832        acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
1833    }
1834
1835    // Horizontal sum: [a, b, c, d] -> a + b + c + d
1836    let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01); // [b, a, d, c]
1837    let sums = _mm_add_ps(acc, shuf); // [a+b, a+b, c+d, c+d]
1838    let shuf2 = _mm_movehl_ps(sums, sums); // [c+d, c+d, ?, ?]
1839    let final_sum = _mm_add_ss(sums, shuf2); // [a+b+c+d, ?, ?, ?]
1840
1841    let mut sum = _mm_cvtss_f32(final_sum);
1842
1843    // Handle remainder
1844    let base = chunks * 4;
1845    for i in 0..remainder {
1846        sum += a[base + i] * b[base + i];
1847    }
1848
1849    sum
1850}
1851
1852#[cfg(target_arch = "x86_64")]
1853#[target_feature(enable = "avx512f")]
1854#[allow(unsafe_op_in_unsafe_fn)]
1855unsafe fn dot_product_f32_avx512(a: &[f32], b: &[f32], count: usize) -> f32 {
1856    use std::arch::x86_64::*;
1857
1858    let chunks64 = count / 64;
1859    let remainder = count % 64;
1860
1861    let mut acc0 = _mm512_setzero_ps();
1862    let mut acc1 = _mm512_setzero_ps();
1863    let mut acc2 = _mm512_setzero_ps();
1864    let mut acc3 = _mm512_setzero_ps();
1865
1866    for c in 0..chunks64 {
1867        let base = c * 64;
1868        acc0 = _mm512_fmadd_ps(
1869            _mm512_loadu_ps(a.as_ptr().add(base)),
1870            _mm512_loadu_ps(b.as_ptr().add(base)),
1871            acc0,
1872        );
1873        acc1 = _mm512_fmadd_ps(
1874            _mm512_loadu_ps(a.as_ptr().add(base + 16)),
1875            _mm512_loadu_ps(b.as_ptr().add(base + 16)),
1876            acc1,
1877        );
1878        acc2 = _mm512_fmadd_ps(
1879            _mm512_loadu_ps(a.as_ptr().add(base + 32)),
1880            _mm512_loadu_ps(b.as_ptr().add(base + 32)),
1881            acc2,
1882        );
1883        acc3 = _mm512_fmadd_ps(
1884            _mm512_loadu_ps(a.as_ptr().add(base + 48)),
1885            _mm512_loadu_ps(b.as_ptr().add(base + 48)),
1886            acc3,
1887        );
1888    }
1889
1890    let acc = _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3));
1891    let mut sum = _mm512_reduce_add_ps(acc);
1892
1893    let base = chunks64 * 64;
1894    for i in 0..remainder {
1895        sum += a[base + i] * b[base + i];
1896    }
1897
1898    sum
1899}
1900
1901#[cfg(target_arch = "x86_64")]
1902#[target_feature(enable = "avx512f")]
1903#[allow(unsafe_op_in_unsafe_fn)]
1904unsafe fn fused_dot_norm_avx512(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1905    use std::arch::x86_64::*;
1906
1907    let chunks64 = count / 64;
1908    let remainder = count % 64;
1909
1910    let mut d0 = _mm512_setzero_ps();
1911    let mut d1 = _mm512_setzero_ps();
1912    let mut d2 = _mm512_setzero_ps();
1913    let mut d3 = _mm512_setzero_ps();
1914    let mut n0 = _mm512_setzero_ps();
1915    let mut n1 = _mm512_setzero_ps();
1916    let mut n2 = _mm512_setzero_ps();
1917    let mut n3 = _mm512_setzero_ps();
1918
1919    for c in 0..chunks64 {
1920        let base = c * 64;
1921        let vb0 = _mm512_loadu_ps(b.as_ptr().add(base));
1922        d0 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base)), vb0, d0);
1923        n0 = _mm512_fmadd_ps(vb0, vb0, n0);
1924        let vb1 = _mm512_loadu_ps(b.as_ptr().add(base + 16));
1925        d1 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 16)), vb1, d1);
1926        n1 = _mm512_fmadd_ps(vb1, vb1, n1);
1927        let vb2 = _mm512_loadu_ps(b.as_ptr().add(base + 32));
1928        d2 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 32)), vb2, d2);
1929        n2 = _mm512_fmadd_ps(vb2, vb2, n2);
1930        let vb3 = _mm512_loadu_ps(b.as_ptr().add(base + 48));
1931        d3 = _mm512_fmadd_ps(_mm512_loadu_ps(a.as_ptr().add(base + 48)), vb3, d3);
1932        n3 = _mm512_fmadd_ps(vb3, vb3, n3);
1933    }
1934
1935    let acc_dot = _mm512_add_ps(_mm512_add_ps(d0, d1), _mm512_add_ps(d2, d3));
1936    let acc_norm = _mm512_add_ps(_mm512_add_ps(n0, n1), _mm512_add_ps(n2, n3));
1937    let mut dot = _mm512_reduce_add_ps(acc_dot);
1938    let mut norm = _mm512_reduce_add_ps(acc_norm);
1939
1940    let base = chunks64 * 64;
1941    for i in 0..remainder {
1942        dot += a[base + i] * b[base + i];
1943        norm += b[base + i] * b[base + i];
1944    }
1945
1946    (dot, norm)
1947}
1948
1949// ============================================================================
1950// Batched Cosine Similarity for Dense Vector Search
1951// ============================================================================
1952
1953/// Fused dot-product + self-norm in a single pass (SIMD accelerated).
1954///
1955/// Returns (dot(a, b), dot(b, b)) — i.e. the dot product of a·b and ||b||².
1956/// Loads `b` only once (halves memory bandwidth vs two separate dot products).
1957#[inline]
1958fn fused_dot_norm(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1959    #[cfg(target_arch = "aarch64")]
1960    {
1961        if neon::is_available() {
1962            return unsafe { fused_dot_norm_neon(a, b, count) };
1963        }
1964    }
1965
1966    #[cfg(target_arch = "x86_64")]
1967    {
1968        if is_x86_feature_detected!("avx512f") {
1969            return unsafe { fused_dot_norm_avx512(a, b, count) };
1970        }
1971        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
1972            return unsafe { fused_dot_norm_avx2(a, b, count) };
1973        }
1974        if sse::is_available() {
1975            return unsafe { fused_dot_norm_sse(a, b, count) };
1976        }
1977    }
1978
1979    // Scalar fallback
1980    let mut dot = 0.0f32;
1981    let mut norm_b = 0.0f32;
1982    for i in 0..count {
1983        dot += a[i] * b[i];
1984        norm_b += b[i] * b[i];
1985    }
1986    (dot, norm_b)
1987}
1988
1989#[cfg(target_arch = "aarch64")]
1990#[target_feature(enable = "neon")]
1991#[allow(unsafe_op_in_unsafe_fn)]
1992unsafe fn fused_dot_norm_neon(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
1993    use std::arch::aarch64::*;
1994
1995    let chunks16 = count / 16;
1996    let remainder = count % 16;
1997
1998    let mut d0 = vdupq_n_f32(0.0);
1999    let mut d1 = vdupq_n_f32(0.0);
2000    let mut d2 = vdupq_n_f32(0.0);
2001    let mut d3 = vdupq_n_f32(0.0);
2002    let mut n0 = vdupq_n_f32(0.0);
2003    let mut n1 = vdupq_n_f32(0.0);
2004    let mut n2 = vdupq_n_f32(0.0);
2005    let mut n3 = vdupq_n_f32(0.0);
2006
2007    for c in 0..chunks16 {
2008        let base = c * 16;
2009        let va0 = vld1q_f32(a.as_ptr().add(base));
2010        let vb0 = vld1q_f32(b.as_ptr().add(base));
2011        d0 = vfmaq_f32(d0, va0, vb0);
2012        n0 = vfmaq_f32(n0, vb0, vb0);
2013        let va1 = vld1q_f32(a.as_ptr().add(base + 4));
2014        let vb1 = vld1q_f32(b.as_ptr().add(base + 4));
2015        d1 = vfmaq_f32(d1, va1, vb1);
2016        n1 = vfmaq_f32(n1, vb1, vb1);
2017        let va2 = vld1q_f32(a.as_ptr().add(base + 8));
2018        let vb2 = vld1q_f32(b.as_ptr().add(base + 8));
2019        d2 = vfmaq_f32(d2, va2, vb2);
2020        n2 = vfmaq_f32(n2, vb2, vb2);
2021        let va3 = vld1q_f32(a.as_ptr().add(base + 12));
2022        let vb3 = vld1q_f32(b.as_ptr().add(base + 12));
2023        d3 = vfmaq_f32(d3, va3, vb3);
2024        n3 = vfmaq_f32(n3, vb3, vb3);
2025    }
2026
2027    let acc_dot = vaddq_f32(vaddq_f32(d0, d1), vaddq_f32(d2, d3));
2028    let acc_norm = vaddq_f32(vaddq_f32(n0, n1), vaddq_f32(n2, n3));
2029    let mut dot = vaddvq_f32(acc_dot);
2030    let mut norm = vaddvq_f32(acc_norm);
2031
2032    let base = chunks16 * 16;
2033    for i in 0..remainder {
2034        dot += a[base + i] * b[base + i];
2035        norm += b[base + i] * b[base + i];
2036    }
2037
2038    (dot, norm)
2039}
2040
2041#[cfg(target_arch = "x86_64")]
2042#[target_feature(enable = "avx2", enable = "fma")]
2043#[allow(unsafe_op_in_unsafe_fn)]
2044unsafe fn fused_dot_norm_avx2(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2045    use std::arch::x86_64::*;
2046
2047    let chunks32 = count / 32;
2048    let remainder = count % 32;
2049
2050    let mut d0 = _mm256_setzero_ps();
2051    let mut d1 = _mm256_setzero_ps();
2052    let mut d2 = _mm256_setzero_ps();
2053    let mut d3 = _mm256_setzero_ps();
2054    let mut n0 = _mm256_setzero_ps();
2055    let mut n1 = _mm256_setzero_ps();
2056    let mut n2 = _mm256_setzero_ps();
2057    let mut n3 = _mm256_setzero_ps();
2058
2059    for c in 0..chunks32 {
2060        let base = c * 32;
2061        let vb0 = _mm256_loadu_ps(b.as_ptr().add(base));
2062        d0 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base)), vb0, d0);
2063        n0 = _mm256_fmadd_ps(vb0, vb0, n0);
2064        let vb1 = _mm256_loadu_ps(b.as_ptr().add(base + 8));
2065        d1 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 8)), vb1, d1);
2066        n1 = _mm256_fmadd_ps(vb1, vb1, n1);
2067        let vb2 = _mm256_loadu_ps(b.as_ptr().add(base + 16));
2068        d2 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 16)), vb2, d2);
2069        n2 = _mm256_fmadd_ps(vb2, vb2, n2);
2070        let vb3 = _mm256_loadu_ps(b.as_ptr().add(base + 24));
2071        d3 = _mm256_fmadd_ps(_mm256_loadu_ps(a.as_ptr().add(base + 24)), vb3, d3);
2072        n3 = _mm256_fmadd_ps(vb3, vb3, n3);
2073    }
2074
2075    let acc_dot = _mm256_add_ps(_mm256_add_ps(d0, d1), _mm256_add_ps(d2, d3));
2076    let acc_norm = _mm256_add_ps(_mm256_add_ps(n0, n1), _mm256_add_ps(n2, n3));
2077
2078    // Horizontal sums: 256→128→scalar
2079    let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2080    let lo_d = _mm256_castps256_ps128(acc_dot);
2081    let sum_d = _mm_add_ps(lo_d, hi_d);
2082    let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2083    let sums_d = _mm_add_ps(sum_d, shuf_d);
2084    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2085    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2086
2087    let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2088    let lo_n = _mm256_castps256_ps128(acc_norm);
2089    let sum_n = _mm_add_ps(lo_n, hi_n);
2090    let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2091    let sums_n = _mm_add_ps(sum_n, shuf_n);
2092    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2093    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2094
2095    let base = chunks32 * 32;
2096    for i in 0..remainder {
2097        dot += a[base + i] * b[base + i];
2098        norm += b[base + i] * b[base + i];
2099    }
2100
2101    (dot, norm)
2102}
2103
2104#[cfg(target_arch = "x86_64")]
2105#[target_feature(enable = "sse")]
2106#[allow(unsafe_op_in_unsafe_fn)]
2107unsafe fn fused_dot_norm_sse(a: &[f32], b: &[f32], count: usize) -> (f32, f32) {
2108    use std::arch::x86_64::*;
2109
2110    let chunks = count / 4;
2111    let remainder = count % 4;
2112
2113    let mut acc_dot = _mm_setzero_ps();
2114    let mut acc_norm = _mm_setzero_ps();
2115
2116    for chunk in 0..chunks {
2117        let base = chunk * 4;
2118        let va = _mm_loadu_ps(a.as_ptr().add(base));
2119        let vb = _mm_loadu_ps(b.as_ptr().add(base));
2120        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2121        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2122    }
2123
2124    // Horizontal sums
2125    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2126    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2127    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2128    let final_d = _mm_add_ss(sums_d, shuf2_d);
2129    let mut dot = _mm_cvtss_f32(final_d);
2130
2131    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2132    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2133    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2134    let final_n = _mm_add_ss(sums_n, shuf2_n);
2135    let mut norm = _mm_cvtss_f32(final_n);
2136
2137    let base = chunks * 4;
2138    for i in 0..remainder {
2139        dot += a[base + i] * b[base + i];
2140        norm += b[base + i] * b[base + i];
2141    }
2142
2143    (dot, norm)
2144}
2145
2146/// Fast approximate reciprocal square root: 1/sqrt(x).
2147///
2148/// Uses the IEEE 754 bit trick (Quake III) + one Newton-Raphson iteration
2149/// for ~23-bit precision — sufficient for cosine similarity scoring.
2150/// ~3-5x faster than `1.0 / x.sqrt()` on most architectures.
2151#[inline]
2152pub fn fast_inv_sqrt(x: f32) -> f32 {
2153    let half = 0.5 * x;
2154    let i = 0x5F37_5A86_u32.wrapping_sub(x.to_bits() >> 1);
2155    let y = f32::from_bits(i);
2156    let y = y * (1.5 - half * y * y); // first Newton-Raphson step
2157    y * (1.5 - half * y * y) // second step: ~23-bit precision
2158}
2159
2160/// Batch cosine similarity: query vs N contiguous vectors.
2161///
2162/// `vectors` is a contiguous buffer of `n * dim` floats (row-major).
2163/// `scores` must have length >= n.
2164///
2165/// Optimizations over calling `cosine_similarity` N times:
2166/// 1. Query norm computed once (not N times)
2167/// 2. Fused dot+norm kernel — each vector loaded once (halves bandwidth)
2168/// 3. No per-call overhead (branch prediction, function calls)
2169/// 4. Fast reciprocal square root (~3-5x faster than 1/sqrt)
2170#[inline]
2171pub fn batch_cosine_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
2172    let n = scores.len();
2173    let required = n
2174        .checked_mul(dim)
2175        .expect("batch cosine vector length overflow");
2176    assert_eq!(query.len(), dim, "batch cosine query dimension mismatch");
2177    assert!(
2178        vectors.len() >= required,
2179        "batch cosine vectors are truncated: need {required}, got {}",
2180        vectors.len()
2181    );
2182
2183    if dim == 0 || n == 0 {
2184        return;
2185    }
2186
2187    // Pre-compute query inverse norm once
2188    let norm_q_sq = dot_product_f32(query, query, dim);
2189    if norm_q_sq < f32::EPSILON {
2190        for s in scores.iter_mut() {
2191            *s = 0.0;
2192        }
2193        return;
2194    }
2195    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2196
2197    for i in 0..n {
2198        let vec = &vectors[i * dim..(i + 1) * dim];
2199        let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
2200        if norm_v_sq < f32::EPSILON {
2201            scores[i] = 0.0;
2202        } else {
2203            scores[i] = dot * inv_norm_q * fast_inv_sqrt(norm_v_sq);
2204        }
2205    }
2206}
2207
2208// ============================================================================
2209// f16 (IEEE 754 half-precision) conversion
2210// ============================================================================
2211
2212/// Convert f32 to f16 (IEEE 754 half-precision), stored as u16
2213#[inline]
2214pub fn f32_to_f16(value: f32) -> u16 {
2215    let bits = value.to_bits();
2216    let sign = (bits >> 16) & 0x8000;
2217    let exp = ((bits >> 23) & 0xFF) as i32;
2218    let mantissa = bits & 0x7F_FFFF;
2219
2220    if exp == 255 {
2221        // Inf/NaN
2222        return (sign | 0x7C00 | ((mantissa >> 13) & 0x3FF)) as u16;
2223    }
2224
2225    let exp16 = exp - 127 + 15;
2226
2227    if exp16 >= 31 {
2228        return (sign | 0x7C00) as u16; // overflow → infinity
2229    }
2230
2231    if exp16 <= 0 {
2232        if exp16 < -10 {
2233            return sign as u16; // too small → zero
2234        }
2235        let shift = (1 - exp16) as u32;
2236        let m = (mantissa | 0x80_0000) >> shift;
2237        // Round-to-nearest-even
2238        let round_bit = (m >> 12) & 1;
2239        let sticky = m & 0xFFF;
2240        let m13 = m >> 13;
2241        let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2242        return (sign | rounded) as u16;
2243    }
2244
2245    // Round-to-nearest-even for normal numbers
2246    let round_bit = (mantissa >> 12) & 1;
2247    let sticky = mantissa & 0xFFF;
2248    let m13 = mantissa >> 13;
2249    let rounded = m13 + (round_bit & (m13 | if sticky != 0 { 1 } else { 0 }));
2250    // Check if rounding caused mantissa overflow (carry into exponent)
2251    if rounded > 0x3FF {
2252        let exp16_inc = exp16 as u32 + 1;
2253        if exp16_inc >= 31 {
2254            return (sign | 0x7C00) as u16; // overflow → infinity
2255        }
2256        (sign | (exp16_inc << 10)) as u16
2257    } else {
2258        (sign | ((exp16 as u32) << 10) | rounded) as u16
2259    }
2260}
2261
2262/// Convert f16 (stored as u16) to f32
2263#[inline]
2264pub fn f16_to_f32(half: u16) -> f32 {
2265    let sign = ((half & 0x8000) as u32) << 16;
2266    let exp = ((half >> 10) & 0x1F) as u32;
2267    let mantissa = (half & 0x3FF) as u32;
2268
2269    if exp == 0 {
2270        if mantissa == 0 {
2271            return f32::from_bits(sign);
2272        }
2273        // Subnormal: normalize
2274        let mut e = 0u32;
2275        let mut m = mantissa;
2276        while (m & 0x400) == 0 {
2277            m <<= 1;
2278            e += 1;
2279        }
2280        return f32::from_bits(sign | ((127 - 15 + 1 - e) << 23) | ((m & 0x3FF) << 13));
2281    }
2282
2283    if exp == 31 {
2284        return f32::from_bits(sign | 0x7F80_0000 | (mantissa << 13));
2285    }
2286
2287    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (mantissa << 13))
2288}
2289
2290// ============================================================================
2291// uint8 scalar quantization for [-1, 1] range
2292// ============================================================================
2293
2294const U8_SCALE: f32 = 127.5;
2295const U8_INV_SCALE: f32 = 1.0 / 127.5;
2296
2297/// Quantize f32 in [-1, 1] to u8 [0, 255]
2298#[inline]
2299pub fn f32_to_u8_saturating(value: f32) -> u8 {
2300    ((value.clamp(-1.0, 1.0) + 1.0) * U8_SCALE) as u8
2301}
2302
2303/// Dequantize u8 [0, 255] to f32 in [-1, 1]
2304#[inline]
2305pub fn u8_to_f32(byte: u8) -> f32 {
2306    byte as f32 * U8_INV_SCALE - 1.0
2307}
2308
2309// ============================================================================
2310// Batch conversion (used during builder write)
2311// ============================================================================
2312
2313/// Batch convert f32 slice to f16 (stored as u16)
2314pub fn batch_f32_to_f16(src: &[f32], dst: &mut [u16]) {
2315    debug_assert_eq!(src.len(), dst.len());
2316    for (s, d) in src.iter().zip(dst.iter_mut()) {
2317        *d = f32_to_f16(*s);
2318    }
2319}
2320
2321/// Batch convert f32 slice to u8 with [-1,1] → [0,255] mapping
2322pub fn batch_f32_to_u8(src: &[f32], dst: &mut [u8]) {
2323    debug_assert_eq!(src.len(), dst.len());
2324    for (s, d) in src.iter().zip(dst.iter_mut()) {
2325        *d = f32_to_u8_saturating(*s);
2326    }
2327}
2328
2329// ============================================================================
2330// NEON-accelerated fused dot+norm for quantized vectors
2331// ============================================================================
2332
2333#[cfg(target_arch = "aarch64")]
2334#[allow(unsafe_op_in_unsafe_fn)]
2335mod neon_quant {
2336    use std::arch::aarch64::*;
2337
2338    /// Fused dot(query_f16, vec_f16) + norm(vec_f16) for f16 vectors on NEON.
2339    ///
2340    /// Both query and vectors are f16 (stored as u16). Uses hardware `vcvt_f32_f16`
2341    /// for SIMD f16→f32 conversion (replaces scalar bit manipulation), processes
2342    /// 8 elements per iteration with f32 accumulation for precision.
2343    #[allow(clippy::incompatible_msrv)]
2344    #[target_feature(enable = "neon")]
2345    pub unsafe fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2346        let chunks16 = dim / 16;
2347        let remainder = dim % 16;
2348
2349        // 2 accumulator pairs to hide FMA latency (processes 16 f16 per iteration)
2350        let mut acc_dot0 = vdupq_n_f32(0.0);
2351        let mut acc_dot1 = vdupq_n_f32(0.0);
2352        let mut acc_norm0 = vdupq_n_f32(0.0);
2353        let mut acc_norm1 = vdupq_n_f32(0.0);
2354
2355        for c in 0..chunks16 {
2356            let base = c * 16;
2357
2358            // First 8 f16 elements
2359            let v_raw0 = vld1q_u16(vec_f16.as_ptr().add(base));
2360            let v_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw0)));
2361            let v_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw0)));
2362            let q_raw0 = vld1q_u16(query_f16.as_ptr().add(base));
2363            let q_lo0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw0)));
2364            let q_hi0 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw0)));
2365
2366            acc_dot0 = vfmaq_f32(acc_dot0, q_lo0, v_lo0);
2367            acc_dot0 = vfmaq_f32(acc_dot0, q_hi0, v_hi0);
2368            acc_norm0 = vfmaq_f32(acc_norm0, v_lo0, v_lo0);
2369            acc_norm0 = vfmaq_f32(acc_norm0, v_hi0, v_hi0);
2370
2371            // Second 8 f16 elements (independent accumulator chain)
2372            let v_raw1 = vld1q_u16(vec_f16.as_ptr().add(base + 8));
2373            let v_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw1)));
2374            let v_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw1)));
2375            let q_raw1 = vld1q_u16(query_f16.as_ptr().add(base + 8));
2376            let q_lo1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw1)));
2377            let q_hi1 = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw1)));
2378
2379            acc_dot1 = vfmaq_f32(acc_dot1, q_lo1, v_lo1);
2380            acc_dot1 = vfmaq_f32(acc_dot1, q_hi1, v_hi1);
2381            acc_norm1 = vfmaq_f32(acc_norm1, v_lo1, v_lo1);
2382            acc_norm1 = vfmaq_f32(acc_norm1, v_hi1, v_hi1);
2383        }
2384
2385        // Combine accumulator pairs
2386        let mut dot = vaddvq_f32(vaddq_f32(acc_dot0, acc_dot1));
2387        let mut norm = vaddvq_f32(vaddq_f32(acc_norm0, acc_norm1));
2388
2389        // Handle remainder
2390        let base = chunks16 * 16;
2391        for i in 0..remainder {
2392            let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2393            let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2394            dot += q * v;
2395            norm += v * v;
2396        }
2397
2398        (dot, norm)
2399    }
2400
2401    /// Fused dot(query, vec) + norm(vec) for u8 vectors on NEON.
2402    /// Processes 16 u8 values per iteration using NEON widening chain.
2403    #[target_feature(enable = "neon")]
2404    pub unsafe fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2405        let scale = vdupq_n_f32(super::U8_INV_SCALE);
2406        let offset = vdupq_n_f32(-1.0);
2407
2408        let chunks16 = dim / 16;
2409        let remainder = dim % 16;
2410
2411        let mut acc_dot = vdupq_n_f32(0.0);
2412        let mut acc_norm = vdupq_n_f32(0.0);
2413
2414        for c in 0..chunks16 {
2415            let base = c * 16;
2416
2417            // Load 16 u8 values
2418            let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2419
2420            // Widen: 16×u8 → 2×8×u16 → 4×4×u32 → 4×4×f32
2421            let lo8 = vget_low_u8(bytes);
2422            let hi8 = vget_high_u8(bytes);
2423            let lo16 = vmovl_u8(lo8);
2424            let hi16 = vmovl_u8(hi8);
2425
2426            let f0 = vaddq_f32(
2427                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2428                offset,
2429            );
2430            let f1 = vaddq_f32(
2431                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2432                offset,
2433            );
2434            let f2 = vaddq_f32(
2435                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2436                offset,
2437            );
2438            let f3 = vaddq_f32(
2439                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2440                offset,
2441            );
2442
2443            let q0 = vld1q_f32(query.as_ptr().add(base));
2444            let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2445            let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2446            let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2447
2448            acc_dot = vfmaq_f32(acc_dot, q0, f0);
2449            acc_dot = vfmaq_f32(acc_dot, q1, f1);
2450            acc_dot = vfmaq_f32(acc_dot, q2, f2);
2451            acc_dot = vfmaq_f32(acc_dot, q3, f3);
2452
2453            acc_norm = vfmaq_f32(acc_norm, f0, f0);
2454            acc_norm = vfmaq_f32(acc_norm, f1, f1);
2455            acc_norm = vfmaq_f32(acc_norm, f2, f2);
2456            acc_norm = vfmaq_f32(acc_norm, f3, f3);
2457        }
2458
2459        let mut dot = vaddvq_f32(acc_dot);
2460        let mut norm = vaddvq_f32(acc_norm);
2461
2462        let base = chunks16 * 16;
2463        for i in 0..remainder {
2464            let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2465            dot += *query.get_unchecked(base + i) * v;
2466            norm += v * v;
2467        }
2468
2469        (dot, norm)
2470    }
2471
2472    /// Dot product only for f16 vectors on NEON (no norm — for unit_norm vectors).
2473    #[allow(clippy::incompatible_msrv)]
2474    #[target_feature(enable = "neon")]
2475    pub unsafe fn dot_product_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2476        let chunks8 = dim / 8;
2477        let remainder = dim % 8;
2478
2479        let mut acc = vdupq_n_f32(0.0);
2480
2481        for c in 0..chunks8 {
2482            let base = c * 8;
2483            let v_raw = vld1q_u16(vec_f16.as_ptr().add(base));
2484            let v_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(v_raw)));
2485            let v_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(v_raw)));
2486            let q_raw = vld1q_u16(query_f16.as_ptr().add(base));
2487            let q_lo = vcvt_f32_f16(vreinterpret_f16_u16(vget_low_u16(q_raw)));
2488            let q_hi = vcvt_f32_f16(vreinterpret_f16_u16(vget_high_u16(q_raw)));
2489            acc = vfmaq_f32(acc, q_lo, v_lo);
2490            acc = vfmaq_f32(acc, q_hi, v_hi);
2491        }
2492
2493        let mut dot = vaddvq_f32(acc);
2494        let base = chunks8 * 8;
2495        for i in 0..remainder {
2496            let v = super::f16_to_f32(*vec_f16.get_unchecked(base + i));
2497            let q = super::f16_to_f32(*query_f16.get_unchecked(base + i));
2498            dot += q * v;
2499        }
2500        dot
2501    }
2502
2503    /// Dot product only for u8 vectors on NEON (no norm — for unit_norm vectors).
2504    #[target_feature(enable = "neon")]
2505    pub unsafe fn dot_product_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2506        let scale = vdupq_n_f32(super::U8_INV_SCALE);
2507        let offset = vdupq_n_f32(-1.0);
2508        let chunks16 = dim / 16;
2509        let remainder = dim % 16;
2510
2511        let mut acc = vdupq_n_f32(0.0);
2512
2513        for c in 0..chunks16 {
2514            let base = c * 16;
2515            let bytes = vld1q_u8(vec_u8.as_ptr().add(base));
2516            let lo8 = vget_low_u8(bytes);
2517            let hi8 = vget_high_u8(bytes);
2518            let lo16 = vmovl_u8(lo8);
2519            let hi16 = vmovl_u8(hi8);
2520            let f0 = vaddq_f32(
2521                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(lo16))), scale),
2522                offset,
2523            );
2524            let f1 = vaddq_f32(
2525                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(lo16))), scale),
2526                offset,
2527            );
2528            let f2 = vaddq_f32(
2529                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_low_u16(hi16))), scale),
2530                offset,
2531            );
2532            let f3 = vaddq_f32(
2533                vmulq_f32(vcvtq_f32_u32(vmovl_u16(vget_high_u16(hi16))), scale),
2534                offset,
2535            );
2536            let q0 = vld1q_f32(query.as_ptr().add(base));
2537            let q1 = vld1q_f32(query.as_ptr().add(base + 4));
2538            let q2 = vld1q_f32(query.as_ptr().add(base + 8));
2539            let q3 = vld1q_f32(query.as_ptr().add(base + 12));
2540            acc = vfmaq_f32(acc, q0, f0);
2541            acc = vfmaq_f32(acc, q1, f1);
2542            acc = vfmaq_f32(acc, q2, f2);
2543            acc = vfmaq_f32(acc, q3, f3);
2544        }
2545
2546        let mut dot = vaddvq_f32(acc);
2547        let base = chunks16 * 16;
2548        for i in 0..remainder {
2549            let v = super::u8_to_f32(*vec_u8.get_unchecked(base + i));
2550            dot += *query.get_unchecked(base + i) * v;
2551        }
2552        dot
2553    }
2554}
2555
2556// ============================================================================
2557// Scalar fallback for fused dot+norm on quantized vectors
2558// ============================================================================
2559
2560#[allow(dead_code)]
2561fn fused_dot_norm_f16_scalar(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2562    let mut dot = 0.0f32;
2563    let mut norm = 0.0f32;
2564    for i in 0..dim {
2565        let v = f16_to_f32(vec_f16[i]);
2566        let q = f16_to_f32(query_f16[i]);
2567        dot += q * v;
2568        norm += v * v;
2569    }
2570    (dot, norm)
2571}
2572
2573#[allow(dead_code)]
2574fn fused_dot_norm_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2575    let mut dot = 0.0f32;
2576    let mut norm = 0.0f32;
2577    for i in 0..dim {
2578        let v = u8_to_f32(vec_u8[i]);
2579        dot += query[i] * v;
2580        norm += v * v;
2581    }
2582    (dot, norm)
2583}
2584
2585#[allow(dead_code)]
2586fn dot_product_f16_scalar(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2587    let mut dot = 0.0f32;
2588    for i in 0..dim {
2589        dot += f16_to_f32(query_f16[i]) * f16_to_f32(vec_f16[i]);
2590    }
2591    dot
2592}
2593
2594#[allow(dead_code)]
2595fn dot_product_u8_scalar(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2596    let mut dot = 0.0f32;
2597    for i in 0..dim {
2598        dot += query[i] * u8_to_f32(vec_u8[i]);
2599    }
2600    dot
2601}
2602
2603// ============================================================================
2604// x86_64 SSE4.1 quantized fused dot+norm
2605// ============================================================================
2606
2607#[cfg(target_arch = "x86_64")]
2608#[target_feature(enable = "sse2", enable = "sse4.1")]
2609#[allow(unsafe_op_in_unsafe_fn)]
2610unsafe fn fused_dot_norm_f16_sse(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2611    use std::arch::x86_64::*;
2612
2613    let chunks = dim / 4;
2614    let remainder = dim % 4;
2615
2616    let mut acc_dot = _mm_setzero_ps();
2617    let mut acc_norm = _mm_setzero_ps();
2618
2619    for chunk in 0..chunks {
2620        let base = chunk * 4;
2621        // Load 4 f16 values and convert to f32 using scalar conversion
2622        let v0 = f16_to_f32(*vec_f16.get_unchecked(base));
2623        let v1 = f16_to_f32(*vec_f16.get_unchecked(base + 1));
2624        let v2 = f16_to_f32(*vec_f16.get_unchecked(base + 2));
2625        let v3 = f16_to_f32(*vec_f16.get_unchecked(base + 3));
2626        let vb = _mm_set_ps(v3, v2, v1, v0);
2627
2628        let q0 = f16_to_f32(*query_f16.get_unchecked(base));
2629        let q1 = f16_to_f32(*query_f16.get_unchecked(base + 1));
2630        let q2 = f16_to_f32(*query_f16.get_unchecked(base + 2));
2631        let q3 = f16_to_f32(*query_f16.get_unchecked(base + 3));
2632        let va = _mm_set_ps(q3, q2, q1, q0);
2633
2634        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2635        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2636    }
2637
2638    // Horizontal sums
2639    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2640    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2641    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2642    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2643
2644    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2645    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2646    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2647    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2648
2649    let base = chunks * 4;
2650    for i in 0..remainder {
2651        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2652        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2653        dot += q * v;
2654        norm += v * v;
2655    }
2656
2657    (dot, norm)
2658}
2659
2660#[cfg(target_arch = "x86_64")]
2661#[target_feature(enable = "sse2", enable = "sse4.1")]
2662#[allow(unsafe_op_in_unsafe_fn)]
2663unsafe fn fused_dot_norm_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2664    use std::arch::x86_64::*;
2665
2666    let scale = _mm_set1_ps(U8_INV_SCALE);
2667    let offset = _mm_set1_ps(-1.0);
2668
2669    let chunks = dim / 4;
2670    let remainder = dim % 4;
2671
2672    let mut acc_dot = _mm_setzero_ps();
2673    let mut acc_norm = _mm_setzero_ps();
2674
2675    for chunk in 0..chunks {
2676        let base = chunk * 4;
2677
2678        // Load 4 bytes, zero-extend to i32, convert to f32, dequantize
2679        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2680            vec_u8.as_ptr().add(base) as *const i32
2681        ));
2682        let ints = _mm_cvtepu8_epi32(bytes);
2683        let floats = _mm_cvtepi32_ps(ints);
2684        let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2685
2686        let va = _mm_loadu_ps(query.as_ptr().add(base));
2687
2688        acc_dot = _mm_add_ps(acc_dot, _mm_mul_ps(va, vb));
2689        acc_norm = _mm_add_ps(acc_norm, _mm_mul_ps(vb, vb));
2690    }
2691
2692    // Horizontal sums
2693    let shuf_d = _mm_shuffle_ps(acc_dot, acc_dot, 0b10_11_00_01);
2694    let sums_d = _mm_add_ps(acc_dot, shuf_d);
2695    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2696    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2697
2698    let shuf_n = _mm_shuffle_ps(acc_norm, acc_norm, 0b10_11_00_01);
2699    let sums_n = _mm_add_ps(acc_norm, shuf_n);
2700    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2701    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2702
2703    let base = chunks * 4;
2704    for i in 0..remainder {
2705        let v = u8_to_f32(*vec_u8.get_unchecked(base + i));
2706        dot += *query.get_unchecked(base + i) * v;
2707        norm += v * v;
2708    }
2709
2710    (dot, norm)
2711}
2712
2713// ============================================================================
2714// x86_64 F16C + AVX + FMA accelerated f16 scoring
2715// ============================================================================
2716
2717#[cfg(target_arch = "x86_64")]
2718#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2719#[allow(unsafe_op_in_unsafe_fn)]
2720unsafe fn fused_dot_norm_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2721    use std::arch::x86_64::*;
2722
2723    let chunks16 = dim / 16;
2724    let remainder = dim % 16;
2725
2726    // 2 accumulator pairs to hide FMA latency (processes 16 f16 per iteration)
2727    let mut acc_dot0 = _mm256_setzero_ps();
2728    let mut acc_dot1 = _mm256_setzero_ps();
2729    let mut acc_norm0 = _mm256_setzero_ps();
2730    let mut acc_norm1 = _mm256_setzero_ps();
2731
2732    for c in 0..chunks16 {
2733        let base = c * 16;
2734
2735        // First 8 f16 elements
2736        let v_raw0 = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2737        let vb0 = _mm256_cvtph_ps(v_raw0);
2738        let q_raw0 = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2739        let qa0 = _mm256_cvtph_ps(q_raw0);
2740        acc_dot0 = _mm256_fmadd_ps(qa0, vb0, acc_dot0);
2741        acc_norm0 = _mm256_fmadd_ps(vb0, vb0, acc_norm0);
2742
2743        // Second 8 f16 elements (independent accumulator chain)
2744        let v_raw1 = _mm_loadu_si128(vec_f16.as_ptr().add(base + 8) as *const __m128i);
2745        let vb1 = _mm256_cvtph_ps(v_raw1);
2746        let q_raw1 = _mm_loadu_si128(query_f16.as_ptr().add(base + 8) as *const __m128i);
2747        let qa1 = _mm256_cvtph_ps(q_raw1);
2748        acc_dot1 = _mm256_fmadd_ps(qa1, vb1, acc_dot1);
2749        acc_norm1 = _mm256_fmadd_ps(vb1, vb1, acc_norm1);
2750    }
2751
2752    // Combine accumulator pairs
2753    let acc_dot = _mm256_add_ps(acc_dot0, acc_dot1);
2754    let acc_norm = _mm256_add_ps(acc_norm0, acc_norm1);
2755
2756    // Horizontal sum 256→128→scalar
2757    let hi_d = _mm256_extractf128_ps(acc_dot, 1);
2758    let lo_d = _mm256_castps256_ps128(acc_dot);
2759    let sum_d = _mm_add_ps(lo_d, hi_d);
2760    let shuf_d = _mm_shuffle_ps(sum_d, sum_d, 0b10_11_00_01);
2761    let sums_d = _mm_add_ps(sum_d, shuf_d);
2762    let shuf2_d = _mm_movehl_ps(sums_d, sums_d);
2763    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums_d, shuf2_d));
2764
2765    let hi_n = _mm256_extractf128_ps(acc_norm, 1);
2766    let lo_n = _mm256_castps256_ps128(acc_norm);
2767    let sum_n = _mm_add_ps(lo_n, hi_n);
2768    let shuf_n = _mm_shuffle_ps(sum_n, sum_n, 0b10_11_00_01);
2769    let sums_n = _mm_add_ps(sum_n, shuf_n);
2770    let shuf2_n = _mm_movehl_ps(sums_n, sums_n);
2771    let mut norm = _mm_cvtss_f32(_mm_add_ss(sums_n, shuf2_n));
2772
2773    let base = chunks16 * 16;
2774    for i in 0..remainder {
2775        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2776        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2777        dot += q * v;
2778        norm += v * v;
2779    }
2780
2781    (dot, norm)
2782}
2783
2784#[cfg(target_arch = "x86_64")]
2785#[target_feature(enable = "avx", enable = "f16c", enable = "fma")]
2786#[allow(unsafe_op_in_unsafe_fn)]
2787unsafe fn dot_product_f16_f16c(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2788    use std::arch::x86_64::*;
2789
2790    let chunks = dim / 8;
2791    let remainder = dim % 8;
2792    let mut acc = _mm256_setzero_ps();
2793
2794    for chunk in 0..chunks {
2795        let base = chunk * 8;
2796        let v_raw = _mm_loadu_si128(vec_f16.as_ptr().add(base) as *const __m128i);
2797        let vb = _mm256_cvtph_ps(v_raw);
2798        let q_raw = _mm_loadu_si128(query_f16.as_ptr().add(base) as *const __m128i);
2799        let qa = _mm256_cvtph_ps(q_raw);
2800        acc = _mm256_fmadd_ps(qa, vb, acc);
2801    }
2802
2803    let hi = _mm256_extractf128_ps(acc, 1);
2804    let lo = _mm256_castps256_ps128(acc);
2805    let sum = _mm_add_ps(lo, hi);
2806    let shuf = _mm_shuffle_ps(sum, sum, 0b10_11_00_01);
2807    let sums = _mm_add_ps(sum, shuf);
2808    let shuf2 = _mm_movehl_ps(sums, sums);
2809    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2810
2811    let base = chunks * 8;
2812    for i in 0..remainder {
2813        let v = f16_to_f32(*vec_f16.get_unchecked(base + i));
2814        let q = f16_to_f32(*query_f16.get_unchecked(base + i));
2815        dot += q * v;
2816    }
2817    dot
2818}
2819
2820#[cfg(target_arch = "x86_64")]
2821#[target_feature(enable = "sse2", enable = "sse4.1")]
2822#[allow(unsafe_op_in_unsafe_fn)]
2823unsafe fn dot_product_u8_sse(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2824    use std::arch::x86_64::*;
2825
2826    let scale = _mm_set1_ps(U8_INV_SCALE);
2827    let offset = _mm_set1_ps(-1.0);
2828    let chunks = dim / 4;
2829    let remainder = dim % 4;
2830    let mut acc = _mm_setzero_ps();
2831
2832    for chunk in 0..chunks {
2833        let base = chunk * 4;
2834        let bytes = _mm_cvtsi32_si128(std::ptr::read_unaligned(
2835            vec_u8.as_ptr().add(base) as *const i32
2836        ));
2837        let ints = _mm_cvtepu8_epi32(bytes);
2838        let floats = _mm_cvtepi32_ps(ints);
2839        let vb = _mm_add_ps(_mm_mul_ps(floats, scale), offset);
2840        let va = _mm_loadu_ps(query.as_ptr().add(base));
2841        acc = _mm_add_ps(acc, _mm_mul_ps(va, vb));
2842    }
2843
2844    let shuf = _mm_shuffle_ps(acc, acc, 0b10_11_00_01);
2845    let sums = _mm_add_ps(acc, shuf);
2846    let shuf2 = _mm_movehl_ps(sums, sums);
2847    let mut dot = _mm_cvtss_f32(_mm_add_ss(sums, shuf2));
2848
2849    let base = chunks * 4;
2850    for i in 0..remainder {
2851        dot += *query.get_unchecked(base + i) * u8_to_f32(*vec_u8.get_unchecked(base + i));
2852    }
2853    dot
2854}
2855
2856// ============================================================================
2857// Platform dispatch
2858// ============================================================================
2859
2860#[inline]
2861fn fused_dot_norm_f16(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> (f32, f32) {
2862    #[cfg(target_arch = "aarch64")]
2863    {
2864        return unsafe { neon_quant::fused_dot_norm_f16(query_f16, vec_f16, dim) };
2865    }
2866
2867    #[cfg(target_arch = "x86_64")]
2868    {
2869        if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2870            return unsafe { fused_dot_norm_f16_f16c(query_f16, vec_f16, dim) };
2871        }
2872        if sse::is_available() {
2873            return unsafe { fused_dot_norm_f16_sse(query_f16, vec_f16, dim) };
2874        }
2875    }
2876
2877    #[allow(unreachable_code)]
2878    fused_dot_norm_f16_scalar(query_f16, vec_f16, dim)
2879}
2880
2881#[inline]
2882fn fused_dot_norm_u8(query: &[f32], vec_u8: &[u8], dim: usize) -> (f32, f32) {
2883    #[cfg(target_arch = "aarch64")]
2884    {
2885        return unsafe { neon_quant::fused_dot_norm_u8(query, vec_u8, dim) };
2886    }
2887
2888    #[cfg(target_arch = "x86_64")]
2889    {
2890        if sse::is_available() {
2891            return unsafe { fused_dot_norm_u8_sse(query, vec_u8, dim) };
2892        }
2893    }
2894
2895    #[allow(unreachable_code)]
2896    fused_dot_norm_u8_scalar(query, vec_u8, dim)
2897}
2898
2899// ── Dot-product-only dispatch (for unit_norm vectors) ─────────────────────
2900
2901#[inline]
2902fn dot_product_f16_quant(query_f16: &[u16], vec_f16: &[u16], dim: usize) -> f32 {
2903    #[cfg(target_arch = "aarch64")]
2904    {
2905        return unsafe { neon_quant::dot_product_f16(query_f16, vec_f16, dim) };
2906    }
2907
2908    #[cfg(target_arch = "x86_64")]
2909    {
2910        if is_x86_feature_detected!("f16c") && is_x86_feature_detected!("fma") {
2911            return unsafe { dot_product_f16_f16c(query_f16, vec_f16, dim) };
2912        }
2913    }
2914
2915    #[allow(unreachable_code)]
2916    dot_product_f16_scalar(query_f16, vec_f16, dim)
2917}
2918
2919#[inline]
2920fn dot_product_u8_quant(query: &[f32], vec_u8: &[u8], dim: usize) -> f32 {
2921    #[cfg(target_arch = "aarch64")]
2922    {
2923        return unsafe { neon_quant::dot_product_u8(query, vec_u8, dim) };
2924    }
2925
2926    #[cfg(target_arch = "x86_64")]
2927    {
2928        if sse::is_available() {
2929            return unsafe { dot_product_u8_sse(query, vec_u8, dim) };
2930        }
2931    }
2932
2933    #[allow(unreachable_code)]
2934    dot_product_u8_scalar(query, vec_u8, dim)
2935}
2936
2937// ============================================================================
2938// Public batch cosine scoring for quantized vectors
2939// ============================================================================
2940
2941/// Batch cosine similarity: f32 query vs N contiguous f16 vectors.
2942///
2943/// `vectors_raw` is raw bytes: N vectors × dim × 2 bytes (f16 stored as u16).
2944/// Query is quantized to f16 once, then both query and vectors are scored in
2945/// f16 space using hardware SIMD conversion (8 elements/iteration on NEON).
2946/// Memory bandwidth is halved for both query and vector loads.
2947#[inline]
2948pub fn batch_cosine_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
2949    let n = scores.len();
2950    let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
2951    let required = n
2952        .checked_mul(vec_bytes)
2953        .expect("f16 batch byte length overflow");
2954    assert_eq!(
2955        query.len(),
2956        dim,
2957        "f16 batch cosine query dimension mismatch"
2958    );
2959    assert!(
2960        vectors_raw.len() >= required,
2961        "f16 batch cosine vectors are truncated: need {required} bytes, got {}",
2962        vectors_raw.len()
2963    );
2964    if required > 0 {
2965        assert!(
2966            (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
2967            "f16 batch cosine vectors are not 2-byte aligned"
2968        );
2969    }
2970    if dim == 0 || n == 0 {
2971        return;
2972    }
2973
2974    // Compute query inverse norm in f32 (full precision, before quantization)
2975    let norm_q_sq = dot_product_f32(query, query, dim);
2976    if norm_q_sq < f32::EPSILON {
2977        for s in scores.iter_mut() {
2978            *s = 0.0;
2979        }
2980        return;
2981    }
2982    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
2983
2984    // Quantize query to f16 once (O(dim)), reused for all N vector scorings
2985    let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
2986
2987    for i in 0..n {
2988        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
2989        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
2990
2991        let (dot, norm_v_sq) = fused_dot_norm_f16(&query_f16, f16_slice, dim);
2992        scores[i] = if norm_v_sq < f32::EPSILON {
2993            0.0
2994        } else {
2995            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
2996        };
2997    }
2998}
2999
3000/// Batch cosine similarity: f32 query vs N contiguous u8 vectors.
3001///
3002/// `vectors_raw` is raw bytes: N vectors × dim bytes (u8, mapping [-1,1]→[0,255]).
3003/// Converts u8→f32 using NEON widening chain (16 values/iteration), scores with FMA.
3004/// Memory bandwidth is quartered compared to f32 scoring.
3005#[inline]
3006pub fn batch_cosine_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3007    let n = scores.len();
3008    let required = n.checked_mul(dim).expect("u8 batch byte length overflow");
3009    assert_eq!(query.len(), dim, "u8 batch cosine query dimension mismatch");
3010    assert!(
3011        vectors_raw.len() >= required,
3012        "u8 batch cosine vectors are truncated: need {required} bytes, got {}",
3013        vectors_raw.len()
3014    );
3015    if dim == 0 || n == 0 {
3016        return;
3017    }
3018
3019    let norm_q_sq = dot_product_f32(query, query, dim);
3020    if norm_q_sq < f32::EPSILON {
3021        for s in scores.iter_mut() {
3022            *s = 0.0;
3023        }
3024        return;
3025    }
3026    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3027
3028    for i in 0..n {
3029        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3030
3031        let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3032        scores[i] = if norm_v_sq < f32::EPSILON {
3033            0.0
3034        } else {
3035            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3036        };
3037    }
3038}
3039
3040// ============================================================================
3041// Batch dot-product scoring for unit-norm vectors
3042// ============================================================================
3043
3044/// Batch dot-product scoring: f32 query vs N contiguous f32 unit-norm vectors.
3045///
3046/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3047/// Skips per-vector norm computation — ~40% less work than `batch_cosine_scores`.
3048#[inline]
3049pub fn batch_dot_scores(query: &[f32], vectors: &[f32], dim: usize, scores: &mut [f32]) {
3050    let n = scores.len();
3051    let required = n
3052        .checked_mul(dim)
3053        .expect("batch dot vector length overflow");
3054    assert_eq!(query.len(), dim, "batch dot query dimension mismatch");
3055    assert!(
3056        vectors.len() >= required,
3057        "batch dot vectors are truncated: need {required}, got {}",
3058        vectors.len()
3059    );
3060
3061    if dim == 0 || n == 0 {
3062        return;
3063    }
3064
3065    let norm_q_sq = dot_product_f32(query, query, dim);
3066    if norm_q_sq < f32::EPSILON {
3067        for s in scores.iter_mut() {
3068            *s = 0.0;
3069        }
3070        return;
3071    }
3072    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3073
3074    for i in 0..n {
3075        let vec = &vectors[i * dim..(i + 1) * dim];
3076        let dot = dot_product_f32(query, vec, dim);
3077        scores[i] = dot * inv_norm_q;
3078    }
3079}
3080
3081/// Batch dot-product scoring: f32 query vs N contiguous f16 unit-norm vectors.
3082///
3083/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3084/// Uses F16C/NEON hardware conversion + dot-only kernel.
3085#[inline]
3086pub fn batch_dot_scores_f16(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3087    let n = scores.len();
3088    let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3089    let required = n
3090        .checked_mul(vec_bytes)
3091        .expect("f16 batch byte length overflow");
3092    assert_eq!(query.len(), dim, "f16 batch dot query dimension mismatch");
3093    assert!(
3094        vectors_raw.len() >= required,
3095        "f16 batch dot vectors are truncated: need {required} bytes, got {}",
3096        vectors_raw.len()
3097    );
3098    if required > 0 {
3099        assert!(
3100            (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3101            "f16 batch dot vectors are not 2-byte aligned"
3102        );
3103    }
3104    if dim == 0 || n == 0 {
3105        return;
3106    }
3107
3108    let norm_q_sq = dot_product_f32(query, query, dim);
3109    if norm_q_sq < f32::EPSILON {
3110        for s in scores.iter_mut() {
3111            *s = 0.0;
3112        }
3113        return;
3114    }
3115    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3116
3117    let query_f16: Vec<u16> = query.iter().map(|&v| f32_to_f16(v)).collect();
3118    for i in 0..n {
3119        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3120        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3121        let dot = dot_product_f16_quant(&query_f16, f16_slice, dim);
3122        scores[i] = dot * inv_norm_q;
3123    }
3124}
3125
3126/// Batch dot-product scoring: f32 query vs N contiguous u8 unit-norm vectors.
3127///
3128/// For pre-normalized vectors (||v|| = 1), cosine = dot(q, v) / ||q||.
3129/// Uses NEON/SSE widening chain for u8→f32 conversion + dot-only kernel.
3130#[inline]
3131pub fn batch_dot_scores_u8(query: &[f32], vectors_raw: &[u8], dim: usize, scores: &mut [f32]) {
3132    let n = scores.len();
3133    let required = n.checked_mul(dim).expect("u8 batch byte length overflow");
3134    assert_eq!(query.len(), dim, "u8 batch dot query dimension mismatch");
3135    assert!(
3136        vectors_raw.len() >= required,
3137        "u8 batch dot vectors are truncated: need {required} bytes, got {}",
3138        vectors_raw.len()
3139    );
3140    if dim == 0 || n == 0 {
3141        return;
3142    }
3143
3144    let norm_q_sq = dot_product_f32(query, query, dim);
3145    if norm_q_sq < f32::EPSILON {
3146        for s in scores.iter_mut() {
3147            *s = 0.0;
3148        }
3149        return;
3150    }
3151    let inv_norm_q = fast_inv_sqrt(norm_q_sq);
3152
3153    for i in 0..n {
3154        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3155        let dot = dot_product_u8_quant(query, u8_slice, dim);
3156        scores[i] = dot * inv_norm_q;
3157    }
3158}
3159
3160// ============================================================================
3161// Precomputed-norm batch scoring (avoids redundant query norm + f16 conversion)
3162// ============================================================================
3163
3164/// Batch cosine: f32 query vs N f32 vectors, with precomputed `inv_norm_q`.
3165#[inline]
3166pub fn batch_cosine_scores_precomp(
3167    query: &[f32],
3168    vectors: &[f32],
3169    dim: usize,
3170    scores: &mut [f32],
3171    inv_norm_q: f32,
3172) {
3173    let n = scores.len();
3174    let required = n
3175        .checked_mul(dim)
3176        .expect("precomputed cosine vector length overflow");
3177    assert_eq!(
3178        query.len(),
3179        dim,
3180        "precomputed cosine query dimension mismatch"
3181    );
3182    assert!(
3183        vectors.len() >= required,
3184        "precomputed cosine vectors are truncated: need {required}, got {}",
3185        vectors.len()
3186    );
3187    for i in 0..n {
3188        let vec = &vectors[i * dim..(i + 1) * dim];
3189        let (dot, norm_v_sq) = fused_dot_norm(query, vec, dim);
3190        scores[i] = if norm_v_sq < f32::EPSILON {
3191            0.0
3192        } else {
3193            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3194        };
3195    }
3196}
3197
3198/// Batch cosine: precomputed `inv_norm_q` + `query_f16` vs N f16 vectors.
3199#[inline]
3200pub fn batch_cosine_scores_f16_precomp(
3201    query_f16: &[u16],
3202    vectors_raw: &[u8],
3203    dim: usize,
3204    scores: &mut [f32],
3205    inv_norm_q: f32,
3206) {
3207    let n = scores.len();
3208    let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3209    let required = n
3210        .checked_mul(vec_bytes)
3211        .expect("precomputed f16 cosine batch byte length overflow");
3212    assert_eq!(
3213        query_f16.len(),
3214        dim,
3215        "precomputed f16 cosine query dimension mismatch"
3216    );
3217    assert!(
3218        vectors_raw.len() >= required,
3219        "precomputed f16 cosine vectors are truncated: need {required} bytes, got {}",
3220        vectors_raw.len()
3221    );
3222    if required > 0 {
3223        assert!(
3224            (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3225            "precomputed f16 cosine vectors are not 2-byte aligned"
3226        );
3227    }
3228    for i in 0..n {
3229        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3230        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3231        let (dot, norm_v_sq) = fused_dot_norm_f16(query_f16, f16_slice, dim);
3232        scores[i] = if norm_v_sq < f32::EPSILON {
3233            0.0
3234        } else {
3235            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3236        };
3237    }
3238}
3239
3240/// Batch cosine: precomputed `inv_norm_q` vs N u8 vectors.
3241#[inline]
3242pub fn batch_cosine_scores_u8_precomp(
3243    query: &[f32],
3244    vectors_raw: &[u8],
3245    dim: usize,
3246    scores: &mut [f32],
3247    inv_norm_q: f32,
3248) {
3249    let n = scores.len();
3250    let required = n
3251        .checked_mul(dim)
3252        .expect("precomputed u8 cosine batch byte length overflow");
3253    assert_eq!(
3254        query.len(),
3255        dim,
3256        "precomputed u8 cosine query dimension mismatch"
3257    );
3258    assert!(
3259        vectors_raw.len() >= required,
3260        "precomputed u8 cosine vectors are truncated: need {required} bytes, got {}",
3261        vectors_raw.len()
3262    );
3263    for i in 0..n {
3264        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3265        let (dot, norm_v_sq) = fused_dot_norm_u8(query, u8_slice, dim);
3266        scores[i] = if norm_v_sq < f32::EPSILON {
3267            0.0
3268        } else {
3269            dot * inv_norm_q * fast_inv_sqrt(norm_v_sq)
3270        };
3271    }
3272}
3273
3274/// Batch dot-product: precomputed `inv_norm_q` vs N f32 unit-norm vectors.
3275#[inline]
3276pub fn batch_dot_scores_precomp(
3277    query: &[f32],
3278    vectors: &[f32],
3279    dim: usize,
3280    scores: &mut [f32],
3281    inv_norm_q: f32,
3282) {
3283    let n = scores.len();
3284    let required = n
3285        .checked_mul(dim)
3286        .expect("precomputed dot vector length overflow");
3287    assert_eq!(query.len(), dim, "precomputed dot query dimension mismatch");
3288    assert!(
3289        vectors.len() >= required,
3290        "precomputed dot vectors are truncated: need {required}, got {}",
3291        vectors.len()
3292    );
3293    for i in 0..n {
3294        let vec = &vectors[i * dim..(i + 1) * dim];
3295        scores[i] = dot_product_f32(query, vec, dim) * inv_norm_q;
3296    }
3297}
3298
3299/// Batch dot-product: precomputed `inv_norm_q` + `query_f16` vs N f16 unit-norm vectors.
3300#[inline]
3301pub fn batch_dot_scores_f16_precomp(
3302    query_f16: &[u16],
3303    vectors_raw: &[u8],
3304    dim: usize,
3305    scores: &mut [f32],
3306    inv_norm_q: f32,
3307) {
3308    let n = scores.len();
3309    let vec_bytes = dim.checked_mul(2).expect("f16 vector size overflow");
3310    let required = n
3311        .checked_mul(vec_bytes)
3312        .expect("precomputed f16 dot batch byte length overflow");
3313    assert_eq!(
3314        query_f16.len(),
3315        dim,
3316        "precomputed f16 dot query dimension mismatch"
3317    );
3318    assert!(
3319        vectors_raw.len() >= required,
3320        "precomputed f16 dot vectors are truncated: need {required} bytes, got {}",
3321        vectors_raw.len()
3322    );
3323    if required > 0 {
3324        assert!(
3325            (vectors_raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>()),
3326            "precomputed f16 dot vectors are not 2-byte aligned"
3327        );
3328    }
3329    for i in 0..n {
3330        let raw = &vectors_raw[i * vec_bytes..(i + 1) * vec_bytes];
3331        let f16_slice = unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const u16, dim) };
3332        scores[i] = dot_product_f16_quant(query_f16, f16_slice, dim) * inv_norm_q;
3333    }
3334}
3335
3336/// Batch dot-product: precomputed `inv_norm_q` vs N u8 unit-norm vectors.
3337#[inline]
3338pub fn batch_dot_scores_u8_precomp(
3339    query: &[f32],
3340    vectors_raw: &[u8],
3341    dim: usize,
3342    scores: &mut [f32],
3343    inv_norm_q: f32,
3344) {
3345    let n = scores.len();
3346    let required = n
3347        .checked_mul(dim)
3348        .expect("precomputed u8 dot batch byte length overflow");
3349    assert_eq!(
3350        query.len(),
3351        dim,
3352        "precomputed u8 dot query dimension mismatch"
3353    );
3354    assert!(
3355        vectors_raw.len() >= required,
3356        "precomputed u8 dot vectors are truncated: need {required} bytes, got {}",
3357        vectors_raw.len()
3358    );
3359    for i in 0..n {
3360        let u8_slice = &vectors_raw[i * dim..(i + 1) * dim];
3361        scores[i] = dot_product_u8_quant(query, u8_slice, dim) * inv_norm_q;
3362    }
3363}
3364
3365/// Compute cosine similarity between two f32 vectors with SIMD acceleration
3366///
3367/// Returns dot(a,b) / (||a|| * ||b||), range [-1, 1]
3368/// Returns 0.0 if either vector has zero norm.
3369#[inline]
3370pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
3371    assert_eq!(a.len(), b.len(), "cosine vector dimension mismatch");
3372    let count = a.len();
3373
3374    if count == 0 {
3375        return 0.0;
3376    }
3377
3378    let dot = dot_product_f32(a, b, count);
3379    let norm_a = dot_product_f32(a, a, count);
3380    let norm_b = dot_product_f32(b, b, count);
3381
3382    let denom = (norm_a * norm_b).sqrt();
3383    if denom < f32::EPSILON {
3384        return 0.0;
3385    }
3386
3387    dot / denom
3388}
3389
3390// ============================================================================
3391// Hamming distance for binary dense vectors
3392// ============================================================================
3393
3394/// Compute Hamming distance between two packed-bit vectors.
3395/// Returns the number of differing bits.
3396///
3397/// Uses NEON intrinsics on aarch64, POPCNT on x86_64, scalar fallback otherwise.
3398#[inline]
3399pub fn hamming_distance(a: &[u8], b: &[u8]) -> u32 {
3400    assert_eq!(a.len(), b.len(), "Hamming vector byte length mismatch");
3401
3402    #[cfg(target_arch = "aarch64")]
3403    unsafe {
3404        neon::hamming_distance(a, b)
3405    }
3406
3407    #[cfg(target_arch = "x86_64")]
3408    {
3409        if avx2::is_available() {
3410            return unsafe { avx2::hamming_distance(a, b) };
3411        }
3412        hamming_distance_scalar(a, b)
3413    }
3414
3415    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3416    hamming_distance_scalar(a, b)
3417}
3418
3419/// Scalar Hamming distance using u64 chunks + count_ones().
3420/// On x86_64, count_ones() compiles to POPCNT when target-cpu supports it.
3421#[inline]
3422#[allow(dead_code)]
3423fn hamming_distance_scalar(a: &[u8], b: &[u8]) -> u32 {
3424    let len = a.len();
3425    let chunks = len / 8;
3426    let remainder = len % 8;
3427    let mut total = 0u32;
3428
3429    for i in 0..chunks {
3430        let off = i * 8;
3431        let va = unsafe { std::ptr::read_unaligned(a.as_ptr().add(off) as *const u64) };
3432        let vb = unsafe { std::ptr::read_unaligned(b.as_ptr().add(off) as *const u64) };
3433        total += (va ^ vb).count_ones();
3434    }
3435
3436    let base = chunks * 8;
3437    for i in 0..remainder {
3438        total += (a[base + i] ^ b[base + i]).count_ones();
3439    }
3440
3441    total
3442}
3443
3444/// Batch Hamming scoring: compute similarity scores for multiple binary vectors.
3445///
3446/// `query` and each vector in `db` are packed-bit vectors of `byte_len` bytes each.
3447/// `dim_bits` is the number of bits (dimensions) for normalization.
3448/// Score = 1.0 - hamming_distance / dim_bits (range [0.0, 1.0]).
3449pub fn batch_hamming_scores(
3450    query: &[u8],
3451    db: &[u8],
3452    byte_len: usize,
3453    dim_bits: usize,
3454    scores: &mut [f32],
3455) {
3456    let n = scores.len();
3457    let required = n
3458        .checked_mul(byte_len)
3459        .expect("Hamming batch byte length overflow");
3460    assert_eq!(query.len(), byte_len, "Hamming query byte length mismatch");
3461    assert!(
3462        db.len() >= required,
3463        "Hamming batch is truncated: need {required} bytes, got {}",
3464        db.len()
3465    );
3466
3467    if byte_len == 0 || n == 0 || dim_bits == 0 {
3468        return;
3469    }
3470
3471    let inv_dim = 1.0 / dim_bits as f32;
3472
3473    // Select the architecture kernel once for the whole batch. Calling
3474    // `hamming_distance` here used to repeat runtime feature detection for
3475    // every database vector in the hottest binary-scan loop.
3476    #[cfg(target_arch = "x86_64")]
3477    {
3478        if avx2::is_available() {
3479            for (i, score) in scores.iter_mut().enumerate() {
3480                let vector = &db[i * byte_len..(i + 1) * byte_len];
3481                let distance = unsafe { avx2::hamming_distance(query, vector) };
3482                *score = 1.0 - distance as f32 * inv_dim;
3483            }
3484        } else {
3485            for (i, score) in scores.iter_mut().enumerate() {
3486                let vector = &db[i * byte_len..(i + 1) * byte_len];
3487                let distance = hamming_distance_scalar(query, vector);
3488                *score = 1.0 - distance as f32 * inv_dim;
3489            }
3490        }
3491    }
3492
3493    #[cfg(target_arch = "aarch64")]
3494    {
3495        for (i, score) in scores.iter_mut().enumerate() {
3496            let vector = &db[i * byte_len..(i + 1) * byte_len];
3497            let distance = unsafe { neon::hamming_distance(query, vector) };
3498            *score = 1.0 - distance as f32 * inv_dim;
3499        }
3500    }
3501
3502    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3503    {
3504        for (i, score) in scores.iter_mut().enumerate() {
3505            let vector = &db[i * byte_len..(i + 1) * byte_len];
3506            let distance = hamming_distance_scalar(query, vector);
3507            *score = 1.0 - distance as f32 * inv_dim;
3508        }
3509    }
3510}
3511
3512#[cfg(test)]
3513mod tests {
3514    use super::*;
3515
3516    #[test]
3517    fn vector_simd_boundaries_reject_dimension_mismatches() {
3518        let vectors = vec![1.0f32; 6];
3519        let raw_f16 = vec![0u8; 12];
3520        let raw_u8 = vec![0u8; 6];
3521        let mut scores = vec![0.0f32; 2];
3522
3523        for invalid_query in [vec![1.0, 2.0], vec![1.0, 2.0, 3.0, 4.0]] {
3524            assert!(
3525                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3526                    batch_cosine_scores(&invalid_query, &vectors, 3, &mut scores)
3527                }))
3528                .is_err()
3529            );
3530            assert!(
3531                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3532                    batch_dot_scores_f16(&invalid_query, &raw_f16, 3, &mut scores)
3533                }))
3534                .is_err()
3535            );
3536            assert!(
3537                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3538                    batch_cosine_scores_u8(&invalid_query, &raw_u8, 3, &mut scores)
3539                }))
3540                .is_err()
3541            );
3542        }
3543    }
3544
3545    #[test]
3546    fn vector_simd_boundaries_reject_truncated_storage() {
3547        let query = [1.0f32, 2.0, 3.0];
3548        let mut scores = [0.0f32; 2];
3549
3550        assert!(
3551            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3552                batch_dot_scores(&query, &[0.0; 5], 3, &mut scores)
3553            }))
3554            .is_err()
3555        );
3556        assert!(
3557            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3558                batch_cosine_scores_f16(&query, &[0u8; 11], 3, &mut scores)
3559            }))
3560            .is_err()
3561        );
3562        assert!(
3563            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3564                dot_product_f32(&query, &query, 4)
3565            }))
3566            .is_err()
3567        );
3568    }
3569
3570    #[test]
3571    fn test_unpack_8bit() {
3572        let input: Vec<u8> = (0..128).collect();
3573        let mut output = vec![0u32; 128];
3574        unpack_8bit(&input, &mut output, 128);
3575
3576        for (i, &v) in output.iter().enumerate() {
3577            assert_eq!(v, i as u32);
3578        }
3579    }
3580
3581    #[test]
3582    fn test_unpack_16bit() {
3583        let mut input = vec![0u8; 256];
3584        for i in 0..128 {
3585            let val = (i * 100) as u16;
3586            input[i * 2] = val as u8;
3587            input[i * 2 + 1] = (val >> 8) as u8;
3588        }
3589
3590        let mut output = vec![0u32; 128];
3591        unpack_16bit(&input, &mut output, 128);
3592
3593        for (i, &v) in output.iter().enumerate() {
3594            assert_eq!(v, (i * 100) as u32);
3595        }
3596    }
3597
3598    #[test]
3599    fn test_unpack_32bit() {
3600        let mut input = vec![0u8; 512];
3601        for i in 0..128 {
3602            let val = (i * 1000) as u32;
3603            let bytes = val.to_le_bytes();
3604            input[i * 4..i * 4 + 4].copy_from_slice(&bytes);
3605        }
3606
3607        let mut output = vec![0u32; 128];
3608        unpack_32bit(&input, &mut output, 128);
3609
3610        for (i, &v) in output.iter().enumerate() {
3611            assert_eq!(v, (i * 1000) as u32);
3612        }
3613    }
3614
3615    #[test]
3616    fn test_delta_decode() {
3617        // doc_ids: [10, 15, 20, 30, 50]
3618        // gaps: [5, 5, 10, 20]
3619        // deltas (gap-1): [4, 4, 9, 19]
3620        let deltas = vec![4u32, 4, 9, 19];
3621        let mut output = vec![0u32; 5];
3622
3623        delta_decode(&mut output, &deltas, 10, 5);
3624
3625        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3626    }
3627
3628    #[test]
3629    fn test_add_one() {
3630        let mut values = vec![0u32, 1, 2, 3, 4, 5, 6, 7];
3631        add_one(&mut values, 8);
3632
3633        assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
3634    }
3635
3636    #[test]
3637    fn test_bits_needed() {
3638        assert_eq!(bits_needed(0), 0);
3639        assert_eq!(bits_needed(1), 1);
3640        assert_eq!(bits_needed(2), 2);
3641        assert_eq!(bits_needed(3), 2);
3642        assert_eq!(bits_needed(4), 3);
3643        assert_eq!(bits_needed(255), 8);
3644        assert_eq!(bits_needed(256), 9);
3645        assert_eq!(bits_needed(u32::MAX), 32);
3646    }
3647
3648    #[test]
3649    fn test_unpack_8bit_delta_decode() {
3650        // doc_ids: [10, 15, 20, 30, 50]
3651        // gaps: [5, 5, 10, 20]
3652        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3653        let input: Vec<u8> = vec![4, 4, 9, 19];
3654        let mut output = vec![0u32; 5];
3655
3656        unpack_8bit_delta_decode(&input, &mut output, 10, 5);
3657
3658        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3659    }
3660
3661    #[test]
3662    fn test_unpack_16bit_delta_decode() {
3663        // doc_ids: [100, 600, 1100, 2100, 4100]
3664        // gaps: [500, 500, 1000, 2000]
3665        // deltas (gap-1): [499, 499, 999, 1999] stored as u16
3666        let mut input = vec![0u8; 8];
3667        for (i, &delta) in [499u16, 499, 999, 1999].iter().enumerate() {
3668            input[i * 2] = delta as u8;
3669            input[i * 2 + 1] = (delta >> 8) as u8;
3670        }
3671        let mut output = vec![0u32; 5];
3672
3673        unpack_16bit_delta_decode(&input, &mut output, 100, 5);
3674
3675        assert_eq!(output, vec![100, 600, 1100, 2100, 4100]);
3676    }
3677
3678    #[test]
3679    fn test_fused_vs_separate_8bit() {
3680        // Test that fused and separate operations produce the same result
3681        let input: Vec<u8> = (0..127).collect();
3682        let first_value = 1000u32;
3683        let count = 128;
3684
3685        // Separate: unpack then delta_decode
3686        let mut unpacked = vec![0u32; 128];
3687        unpack_8bit(&input, &mut unpacked, 127);
3688        let mut separate_output = vec![0u32; 128];
3689        delta_decode(&mut separate_output, &unpacked, first_value, count);
3690
3691        // Fused
3692        let mut fused_output = vec![0u32; 128];
3693        unpack_8bit_delta_decode(&input, &mut fused_output, first_value, count);
3694
3695        assert_eq!(separate_output, fused_output);
3696    }
3697
3698    #[test]
3699    fn test_round_bit_width() {
3700        assert_eq!(round_bit_width(0), 0);
3701        assert_eq!(round_bit_width(1), 8);
3702        assert_eq!(round_bit_width(5), 8);
3703        assert_eq!(round_bit_width(8), 8);
3704        assert_eq!(round_bit_width(9), 16);
3705        assert_eq!(round_bit_width(12), 16);
3706        assert_eq!(round_bit_width(16), 16);
3707        assert_eq!(round_bit_width(17), 32);
3708        assert_eq!(round_bit_width(24), 32);
3709        assert_eq!(round_bit_width(32), 32);
3710    }
3711
3712    #[test]
3713    fn test_rounded_bitwidth_from_exact() {
3714        assert_eq!(RoundedBitWidth::from_exact(0), RoundedBitWidth::Zero);
3715        assert_eq!(RoundedBitWidth::from_exact(1), RoundedBitWidth::Bits8);
3716        assert_eq!(RoundedBitWidth::from_exact(8), RoundedBitWidth::Bits8);
3717        assert_eq!(RoundedBitWidth::from_exact(9), RoundedBitWidth::Bits16);
3718        assert_eq!(RoundedBitWidth::from_exact(16), RoundedBitWidth::Bits16);
3719        assert_eq!(RoundedBitWidth::from_exact(17), RoundedBitWidth::Bits32);
3720        assert_eq!(RoundedBitWidth::from_exact(32), RoundedBitWidth::Bits32);
3721    }
3722
3723    #[test]
3724    fn test_pack_unpack_rounded_8bit() {
3725        let values: Vec<u32> = (0..128).map(|i| i % 256).collect();
3726        let mut packed = vec![0u8; 128];
3727
3728        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits8, &mut packed);
3729        assert_eq!(bytes_written, 128);
3730
3731        let mut unpacked = vec![0u32; 128];
3732        unpack_rounded(&packed, RoundedBitWidth::Bits8, &mut unpacked, 128);
3733
3734        assert_eq!(values, unpacked);
3735    }
3736
3737    #[test]
3738    fn test_pack_unpack_rounded_16bit() {
3739        let values: Vec<u32> = (0..128).map(|i| i * 100).collect();
3740        let mut packed = vec![0u8; 256];
3741
3742        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits16, &mut packed);
3743        assert_eq!(bytes_written, 256);
3744
3745        let mut unpacked = vec![0u32; 128];
3746        unpack_rounded(&packed, RoundedBitWidth::Bits16, &mut unpacked, 128);
3747
3748        assert_eq!(values, unpacked);
3749    }
3750
3751    #[test]
3752    fn test_pack_unpack_rounded_32bit() {
3753        let values: Vec<u32> = (0..128).map(|i| i * 100000).collect();
3754        let mut packed = vec![0u8; 512];
3755
3756        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits32, &mut packed);
3757        assert_eq!(bytes_written, 512);
3758
3759        let mut unpacked = vec![0u32; 128];
3760        unpack_rounded(&packed, RoundedBitWidth::Bits32, &mut unpacked, 128);
3761
3762        assert_eq!(values, unpacked);
3763    }
3764
3765    #[test]
3766    fn test_unpack_rounded_delta_decode() {
3767        // Test 8-bit rounded delta decode
3768        // doc_ids: [10, 15, 20, 30, 50]
3769        // gaps: [5, 5, 10, 20]
3770        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3771        let input: Vec<u8> = vec![4, 4, 9, 19];
3772        let mut output = vec![0u32; 5];
3773
3774        unpack_rounded_delta_decode(&input, RoundedBitWidth::Bits8, &mut output, 10, 5);
3775
3776        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3777    }
3778
3779    #[test]
3780    fn test_unpack_rounded_delta_decode_zero() {
3781        // All zeros means gaps of 1 (consecutive doc IDs)
3782        let input: Vec<u8> = vec![];
3783        let mut output = vec![0u32; 5];
3784
3785        unpack_rounded_delta_decode(&input, RoundedBitWidth::Zero, &mut output, 100, 5);
3786
3787        assert_eq!(output, vec![100, 101, 102, 103, 104]);
3788    }
3789
3790    // ========================================================================
3791    // Sparse Vector SIMD Tests
3792    // ========================================================================
3793
3794    #[test]
3795    fn test_dequantize_uint8() {
3796        let input: Vec<u8> = vec![0, 128, 255, 64, 192];
3797        let mut output = vec![0.0f32; 5];
3798        let scale = 0.1;
3799        let min_val = 1.0;
3800
3801        dequantize_uint8(&input, &mut output, scale, min_val, 5);
3802
3803        // Expected: input[i] * scale + min_val
3804        assert!((output[0] - 1.0).abs() < 1e-6); // 0 * 0.1 + 1.0 = 1.0
3805        assert!((output[1] - 13.8).abs() < 1e-6); // 128 * 0.1 + 1.0 = 13.8
3806        assert!((output[2] - 26.5).abs() < 1e-6); // 255 * 0.1 + 1.0 = 26.5
3807        assert!((output[3] - 7.4).abs() < 1e-6); // 64 * 0.1 + 1.0 = 7.4
3808        assert!((output[4] - 20.2).abs() < 1e-6); // 192 * 0.1 + 1.0 = 20.2
3809    }
3810
3811    #[test]
3812    fn test_dequantize_uint8_large() {
3813        // Test with 128 values (full SIMD block)
3814        let input: Vec<u8> = (0..128).collect();
3815        let mut output = vec![0.0f32; 128];
3816        let scale = 2.0;
3817        let min_val = -10.0;
3818
3819        dequantize_uint8(&input, &mut output, scale, min_val, 128);
3820
3821        for (i, &out) in output.iter().enumerate().take(128) {
3822            let expected = i as f32 * scale + min_val;
3823            assert!(
3824                (out - expected).abs() < 1e-5,
3825                "Mismatch at {}: expected {}, got {}",
3826                i,
3827                expected,
3828                out
3829            );
3830        }
3831    }
3832
3833    #[test]
3834    fn test_dot_product_f32() {
3835        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
3836        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0];
3837
3838        let result = dot_product_f32(&a, &b, 5);
3839
3840        // Expected: 1*2 + 2*3 + 3*4 + 4*5 + 5*6 = 2 + 6 + 12 + 20 + 30 = 70
3841        assert!((result - 70.0).abs() < 1e-5);
3842    }
3843
3844    #[test]
3845    fn test_dot_product_f32_large() {
3846        // Test with 128 values
3847        let a: Vec<f32> = (0..128).map(|i| i as f32).collect();
3848        let b: Vec<f32> = (0..128).map(|i| (i + 1) as f32).collect();
3849
3850        let result = dot_product_f32(&a, &b, 128);
3851
3852        // Compute expected
3853        let expected: f32 = (0..128).map(|i| (i as f32) * ((i + 1) as f32)).sum();
3854        assert!(
3855            (result - expected).abs() < 1e-3,
3856            "Expected {}, got {}",
3857            expected,
3858            result
3859        );
3860    }
3861
3862    #[test]
3863    fn test_fused_dot_norm() {
3864        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
3865        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
3866        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3867
3868        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3869        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3870        assert!(
3871            (dot - expected_dot).abs() < 1e-5,
3872            "dot: expected {}, got {}",
3873            expected_dot,
3874            dot
3875        );
3876        assert!(
3877            (norm_b - expected_norm).abs() < 1e-5,
3878            "norm: expected {}, got {}",
3879            expected_norm,
3880            norm_b
3881        );
3882    }
3883
3884    #[test]
3885    fn test_fused_dot_norm_large() {
3886        let a: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01).collect();
3887        let b: Vec<f32> = (0..768).map(|i| (i as f32) * 0.02 + 0.5).collect();
3888        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3889
3890        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3891        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3892        assert!(
3893            (dot - expected_dot).abs() < 1.0,
3894            "dot: expected {}, got {}",
3895            expected_dot,
3896            dot
3897        );
3898        assert!(
3899            (norm_b - expected_norm).abs() < 1.0,
3900            "norm: expected {}, got {}",
3901            expected_norm,
3902            norm_b
3903        );
3904    }
3905
3906    #[test]
3907    fn test_batch_cosine_scores() {
3908        // 4 vectors of dim 3
3909        let query = vec![1.0f32, 0.0, 0.0];
3910        let vectors = vec![
3911            1.0, 0.0, 0.0, // identical to query
3912            0.0, 1.0, 0.0, // orthogonal
3913            -1.0, 0.0, 0.0, // opposite
3914            0.5, 0.5, 0.0, // 45 degrees
3915        ];
3916        let mut scores = vec![0f32; 4];
3917        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3918
3919        assert!((scores[0] - 1.0).abs() < 1e-5, "identical: {}", scores[0]);
3920        assert!(scores[1].abs() < 1e-5, "orthogonal: {}", scores[1]);
3921        assert!((scores[2] - (-1.0)).abs() < 1e-5, "opposite: {}", scores[2]);
3922        let expected_45 = 0.5f32 / (0.5f32.powi(2) + 0.5f32.powi(2)).sqrt();
3923        assert!(
3924            (scores[3] - expected_45).abs() < 1e-5,
3925            "45deg: expected {}, got {}",
3926            expected_45,
3927            scores[3]
3928        );
3929    }
3930
3931    #[test]
3932    fn test_batch_cosine_scores_matches_individual() {
3933        let query: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
3934        let n = 50;
3935        let dim = 128;
3936        let vectors: Vec<f32> = (0..n * dim).map(|i| ((i * 7 + 3) as f32) * 0.01).collect();
3937
3938        let mut batch_scores = vec![0f32; n];
3939        batch_cosine_scores(&query, &vectors, dim, &mut batch_scores);
3940
3941        for i in 0..n {
3942            let vec_i = &vectors[i * dim..(i + 1) * dim];
3943            let individual = cosine_similarity(&query, vec_i);
3944            assert!(
3945                (batch_scores[i] - individual).abs() < 1e-5,
3946                "vec {}: batch={}, individual={}",
3947                i,
3948                batch_scores[i],
3949                individual
3950            );
3951        }
3952    }
3953
3954    #[test]
3955    fn test_batch_cosine_scores_empty() {
3956        let query = vec![1.0f32, 2.0, 3.0];
3957        let vectors: Vec<f32> = vec![];
3958        let mut scores: Vec<f32> = vec![];
3959        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3960        assert!(scores.is_empty());
3961    }
3962
3963    #[test]
3964    fn test_batch_cosine_scores_zero_query() {
3965        let query = vec![0.0f32, 0.0, 0.0];
3966        let vectors = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
3967        let mut scores = vec![0f32; 2];
3968        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3969        assert_eq!(scores[0], 0.0);
3970        assert_eq!(scores[1], 0.0);
3971    }
3972
3973    // ================================================================
3974    // f16 conversion tests
3975    // ================================================================
3976
3977    #[test]
3978    fn test_f16_roundtrip_normal() {
3979        for &v in &[0.0f32, 1.0, -1.0, 0.5, -0.5, 0.333, 65504.0] {
3980            let h = f32_to_f16(v);
3981            let back = f16_to_f32(h);
3982            let err = (back - v).abs() / v.abs().max(1e-6);
3983            assert!(
3984                err < 0.002,
3985                "f16 roundtrip {v} → {h:#06x} → {back}, rel err {err}"
3986            );
3987        }
3988    }
3989
3990    #[test]
3991    fn test_f16_special() {
3992        // Zero
3993        assert_eq!(f16_to_f32(f32_to_f16(0.0)), 0.0);
3994        // Negative zero
3995        assert_eq!(f32_to_f16(-0.0), 0x8000);
3996        // Infinity
3997        assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
3998        // NaN
3999        assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan());
4000    }
4001
4002    #[test]
4003    fn test_f16_embedding_range() {
4004        // Typical embedding values in [-1, 1]
4005        let values: Vec<f32> = (-100..=100).map(|i| i as f32 / 100.0).collect();
4006        for &v in &values {
4007            let back = f16_to_f32(f32_to_f16(v));
4008            assert!((back - v).abs() < 0.001, "f16 error for {v}: got {back}");
4009        }
4010    }
4011
4012    // ================================================================
4013    // u8 conversion tests
4014    // ================================================================
4015
4016    #[test]
4017    fn test_u8_roundtrip() {
4018        // Boundary values
4019        assert_eq!(f32_to_u8_saturating(-1.0), 0);
4020        assert_eq!(f32_to_u8_saturating(1.0), 255);
4021        assert_eq!(f32_to_u8_saturating(0.0), 127); // ~127.5 truncated
4022
4023        // Saturation
4024        assert_eq!(f32_to_u8_saturating(-2.0), 0);
4025        assert_eq!(f32_to_u8_saturating(2.0), 255);
4026    }
4027
4028    #[test]
4029    fn test_u8_dequantize() {
4030        assert!((u8_to_f32(0) - (-1.0)).abs() < 0.01);
4031        assert!((u8_to_f32(255) - 1.0).abs() < 0.01);
4032        assert!((u8_to_f32(127) - 0.0).abs() < 0.01);
4033    }
4034
4035    // ================================================================
4036    // Batch scoring tests for quantized vectors
4037    // ================================================================
4038
4039    #[test]
4040    fn test_batch_cosine_scores_f16() {
4041        let query = vec![0.6f32, 0.8, 0.0, 0.0];
4042        let dim = 4;
4043        let vecs_f32 = vec![
4044            0.6f32, 0.8, 0.0, 0.0, // identical to query
4045            0.0, 0.0, 0.6, 0.8, // orthogonal
4046        ];
4047
4048        // Quantize to f16
4049        let mut f16_buf = vec![0u16; 8];
4050        batch_f32_to_f16(&vecs_f32, &mut f16_buf);
4051        let raw: &[u8] =
4052            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4053
4054        let mut scores = vec![0f32; 2];
4055        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4056
4057        assert!(
4058            (scores[0] - 1.0).abs() < 0.01,
4059            "identical vectors: {}",
4060            scores[0]
4061        );
4062        assert!(scores[1].abs() < 0.01, "orthogonal vectors: {}", scores[1]);
4063    }
4064
4065    #[test]
4066    fn test_batch_cosine_scores_u8() {
4067        let query = vec![0.6f32, 0.8, 0.0, 0.0];
4068        let dim = 4;
4069        let vecs_f32 = vec![
4070            0.6f32, 0.8, 0.0, 0.0, // ~identical to query
4071            -0.6, -0.8, 0.0, 0.0, // opposite
4072        ];
4073
4074        // Quantize to u8
4075        let mut u8_buf = vec![0u8; 8];
4076        batch_f32_to_u8(&vecs_f32, &mut u8_buf);
4077
4078        let mut scores = vec![0f32; 2];
4079        batch_cosine_scores_u8(&query, &u8_buf, dim, &mut scores);
4080
4081        assert!(scores[0] > 0.95, "similar vectors: {}", scores[0]);
4082        assert!(scores[1] < -0.95, "opposite vectors: {}", scores[1]);
4083    }
4084
4085    #[test]
4086    fn test_batch_cosine_scores_f16_large_dim() {
4087        // Test with typical embedding dimension
4088        let dim = 768;
4089        let query: Vec<f32> = (0..dim).map(|i| (i as f32 / dim as f32) - 0.5).collect();
4090        let vec2: Vec<f32> = query.iter().map(|x| x * 0.9 + 0.01).collect();
4091
4092        let mut all_vecs = query.clone();
4093        all_vecs.extend_from_slice(&vec2);
4094
4095        let mut f16_buf = vec![0u16; all_vecs.len()];
4096        batch_f32_to_f16(&all_vecs, &mut f16_buf);
4097        let raw: &[u8] =
4098            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4099
4100        let mut scores = vec![0f32; 2];
4101        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4102
4103        // Self-similarity should be ~1.0
4104        assert!((scores[0] - 1.0).abs() < 0.01, "self-sim: {}", scores[0]);
4105        // High similarity with scaled version
4106        assert!(scores[1] > 0.99, "scaled-sim: {}", scores[1]);
4107    }
4108
4109    // ================================================================
4110    // Hamming distance tests
4111    // ================================================================
4112
4113    #[test]
4114    fn test_hamming_distance_identical() {
4115        let a = vec![0xAA; 64];
4116        assert_eq!(hamming_distance(&a, &a), 0);
4117    }
4118
4119    #[test]
4120    fn test_hamming_distance_opposite() {
4121        let a = vec![0xFF; 32];
4122        let b = vec![0x00; 32];
4123        assert_eq!(hamming_distance(&a, &b), 256);
4124    }
4125
4126    #[test]
4127    fn test_hamming_distance_known() {
4128        // Single byte: 0b10101010 vs 0b01010101 = 8 bits differ
4129        let a = vec![0xAA];
4130        let b = vec![0x55];
4131        assert_eq!(hamming_distance(&a, &b), 8);
4132
4133        // Two bytes
4134        let a = vec![0xFF, 0x00];
4135        let b = vec![0x00, 0x00];
4136        assert_eq!(hamming_distance(&a, &b), 8);
4137    }
4138
4139    #[test]
4140    fn test_hamming_distance_single_bit() {
4141        let a = vec![0x00; 16];
4142        let mut b = vec![0x00; 16];
4143        b[7] = 0x01; // flip one bit
4144        assert_eq!(hamming_distance(&a, &b), 1);
4145    }
4146
4147    #[test]
4148    fn test_hamming_distance_empty() {
4149        let a: Vec<u8> = vec![];
4150        assert_eq!(hamming_distance(&a, &a), 0);
4151    }
4152
4153    #[test]
4154    fn test_hamming_distance_remainder_path() {
4155        // 17 bytes: not aligned to 16 (NEON) or 32 (AVX2)
4156        let a = vec![0xFF; 17];
4157        let b = vec![0x00; 17];
4158        assert_eq!(hamming_distance(&a, &b), 136); // 17 * 8
4159
4160        // 33 bytes: tests 32-byte chunk + 1 remainder for AVX2
4161        let a = vec![0xFF; 33];
4162        let b = vec![0x00; 33];
4163        assert_eq!(hamming_distance(&a, &b), 264); // 33 * 8
4164    }
4165
4166    #[test]
4167    fn test_hamming_distance_large() {
4168        // 4096 bytes = 32768 bits, all differing
4169        let a = vec![0xFF; 4096];
4170        let b = vec![0x00; 4096];
4171        assert_eq!(hamming_distance(&a, &b), 32768);
4172    }
4173
4174    #[test]
4175    fn test_hamming_distance_scalar_matches() {
4176        // Verify SIMD path matches scalar for various sizes
4177        for size in [1, 7, 8, 15, 16, 31, 32, 63, 64, 100, 128, 255, 256] {
4178            let a: Vec<u8> = (0..size).map(|i| (i * 37 + 13) as u8).collect();
4179            let b: Vec<u8> = (0..size).map(|i| (i * 53 + 7) as u8).collect();
4180            let expected = hamming_distance_scalar(&a, &b);
4181            let got = hamming_distance(&a, &b);
4182            assert_eq!(got, expected, "mismatch at size {size}");
4183        }
4184    }
4185
4186    // ================================================================
4187    // Batch Hamming scoring tests
4188    // ================================================================
4189
4190    #[test]
4191    fn test_batch_hamming_scores_identical() {
4192        let query = vec![0xAA; 16];
4193        let db = vec![0xAA; 16]; // one vector, identical
4194        let mut scores = vec![0f32; 1];
4195        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4196        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4197    }
4198
4199    #[test]
4200    fn test_batch_hamming_scores_opposite() {
4201        let query = vec![0xFF; 16];
4202        let db = vec![0x00; 16];
4203        let mut scores = vec![0f32; 1];
4204        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4205        assert!((scores[0] - 0.0).abs() < 1e-6, "opposite: {}", scores[0]);
4206    }
4207
4208    #[test]
4209    fn test_batch_hamming_scores_multiple() {
4210        let byte_len = 8;
4211        let dim_bits = 64;
4212        let query = vec![0xFF; byte_len];
4213        let mut db = Vec::new();
4214        db.extend_from_slice(&vec![0xFF; byte_len]); // identical → 1.0
4215        db.extend_from_slice(&vec![0x00; byte_len]); // opposite → 0.0
4216        db.extend_from_slice(&vec![0x0F; byte_len]); // half bits differ → 0.5
4217
4218        let mut scores = vec![0f32; 3];
4219        batch_hamming_scores(&query, &db, byte_len, dim_bits, &mut scores);
4220
4221        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4222        assert!((scores[1] - 0.0).abs() < 1e-6, "opposite: {}", scores[1]);
4223        assert!((scores[2] - 0.5).abs() < 1e-6, "half: {}", scores[2]);
4224    }
4225
4226    #[test]
4227    fn test_batch_hamming_scores_empty() {
4228        let query = vec![0xFF; 8];
4229        let db: Vec<u8> = vec![];
4230        let mut scores: Vec<f32> = vec![];
4231        batch_hamming_scores(&query, &db, 8, 64, &mut scores);
4232        assert!(scores.is_empty());
4233    }
4234
4235    #[test]
4236    fn test_batch_hamming_scores_zero_byte_len() {
4237        let query: Vec<u8> = vec![];
4238        let db: Vec<u8> = vec![];
4239        let mut scores = vec![0f32; 1];
4240        batch_hamming_scores(&query, &db, 0, 0, &mut scores);
4241        // Should return early without modifying scores
4242        assert_eq!(scores[0], 0.0);
4243    }
4244}
4245
4246// ============================================================================
4247// SIMD-accelerated linear scan for sorted u32 slices (within-block seek)
4248// ============================================================================
4249
4250/// Find index of first element >= `target` in a sorted `u32` slice.
4251///
4252/// Equivalent to `slice.partition_point(|&d| d < target)` but uses SIMD to
4253/// scan 4 elements per cycle. Faster than binary search for slices ≤ 256
4254/// elements because it avoids the data-dependency chain inherent in binary
4255/// search (~8-10 cycles/iteration vs ~1-2 cycles/iteration for SIMD scan).
4256///
4257/// Returns `slice.len()` if no element >= `target`.
4258#[inline]
4259pub fn find_first_ge_u32(slice: &[u32], target: u32) -> usize {
4260    #[cfg(target_arch = "aarch64")]
4261    {
4262        if neon::is_available() {
4263            return unsafe { find_first_ge_u32_neon(slice, target) };
4264        }
4265    }
4266
4267    #[cfg(target_arch = "x86_64")]
4268    {
4269        if sse::is_available() {
4270            return unsafe { find_first_ge_u32_sse(slice, target) };
4271        }
4272    }
4273
4274    // Scalar fallback (WASM, other architectures)
4275    slice.partition_point(|&d| d < target)
4276}
4277
4278#[cfg(target_arch = "aarch64")]
4279#[target_feature(enable = "neon")]
4280#[allow(unsafe_op_in_unsafe_fn)]
4281unsafe fn find_first_ge_u32_neon(slice: &[u32], target: u32) -> usize {
4282    use std::arch::aarch64::*;
4283
4284    let n = slice.len();
4285    let ptr = slice.as_ptr();
4286    let target_vec = vdupq_n_u32(target);
4287    // Bit positions for each lane: [1, 2, 4, 8]
4288    let bit_mask: uint32x4_t = core::mem::transmute([1u32, 2u32, 4u32, 8u32]);
4289
4290    let chunks = n / 16;
4291    let mut base = 0usize;
4292
4293    // Process 16 elements per iteration (4 × 4-wide NEON compares)
4294    for _ in 0..chunks {
4295        let v0 = vld1q_u32(ptr.add(base));
4296        let v1 = vld1q_u32(ptr.add(base + 4));
4297        let v2 = vld1q_u32(ptr.add(base + 8));
4298        let v3 = vld1q_u32(ptr.add(base + 12));
4299
4300        let c0 = vcgeq_u32(v0, target_vec);
4301        let c1 = vcgeq_u32(v1, target_vec);
4302        let c2 = vcgeq_u32(v2, target_vec);
4303        let c3 = vcgeq_u32(v3, target_vec);
4304
4305        let m0 = vaddvq_u32(vandq_u32(c0, bit_mask));
4306        if m0 != 0 {
4307            return base + m0.trailing_zeros() as usize;
4308        }
4309        let m1 = vaddvq_u32(vandq_u32(c1, bit_mask));
4310        if m1 != 0 {
4311            return base + 4 + m1.trailing_zeros() as usize;
4312        }
4313        let m2 = vaddvq_u32(vandq_u32(c2, bit_mask));
4314        if m2 != 0 {
4315            return base + 8 + m2.trailing_zeros() as usize;
4316        }
4317        let m3 = vaddvq_u32(vandq_u32(c3, bit_mask));
4318        if m3 != 0 {
4319            return base + 12 + m3.trailing_zeros() as usize;
4320        }
4321        base += 16;
4322    }
4323
4324    // Process remaining 4 elements at a time
4325    while base + 4 <= n {
4326        let vals = vld1q_u32(ptr.add(base));
4327        let cmp = vcgeq_u32(vals, target_vec);
4328        let mask = vaddvq_u32(vandq_u32(cmp, bit_mask));
4329        if mask != 0 {
4330            return base + mask.trailing_zeros() as usize;
4331        }
4332        base += 4;
4333    }
4334
4335    // Scalar remainder (0-3 elements)
4336    while base < n {
4337        if *slice.get_unchecked(base) >= target {
4338            return base;
4339        }
4340        base += 1;
4341    }
4342    n
4343}
4344
4345#[cfg(target_arch = "x86_64")]
4346#[target_feature(enable = "sse2")]
4347#[allow(unsafe_op_in_unsafe_fn)]
4348unsafe fn find_first_ge_u32_sse(slice: &[u32], target: u32) -> usize {
4349    use std::arch::x86_64::*;
4350
4351    let n = slice.len();
4352    let ptr = slice.as_ptr();
4353
4354    // For unsigned >= comparison: XOR with 0x80000000 converts to signed domain
4355    let sign_flip = _mm_set1_epi32(i32::MIN);
4356    let target_xor = _mm_xor_si128(_mm_set1_epi32(target as i32), sign_flip);
4357
4358    let chunks = n / 16;
4359    let mut base = 0usize;
4360
4361    // Process 16 elements per iteration (4 × 4-wide SSE compares)
4362    for _ in 0..chunks {
4363        let v0 = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4364        let v1 = _mm_xor_si128(
4365            _mm_loadu_si128(ptr.add(base + 4) as *const __m128i),
4366            sign_flip,
4367        );
4368        let v2 = _mm_xor_si128(
4369            _mm_loadu_si128(ptr.add(base + 8) as *const __m128i),
4370            sign_flip,
4371        );
4372        let v3 = _mm_xor_si128(
4373            _mm_loadu_si128(ptr.add(base + 12) as *const __m128i),
4374            sign_flip,
4375        );
4376
4377        // ge = eq | gt (in signed domain after XOR)
4378        let ge0 = _mm_or_si128(
4379            _mm_cmpeq_epi32(v0, target_xor),
4380            _mm_cmpgt_epi32(v0, target_xor),
4381        );
4382        let m0 = _mm_movemask_ps(_mm_castsi128_ps(ge0)) as u32;
4383        if m0 != 0 {
4384            return base + m0.trailing_zeros() as usize;
4385        }
4386
4387        let ge1 = _mm_or_si128(
4388            _mm_cmpeq_epi32(v1, target_xor),
4389            _mm_cmpgt_epi32(v1, target_xor),
4390        );
4391        let m1 = _mm_movemask_ps(_mm_castsi128_ps(ge1)) as u32;
4392        if m1 != 0 {
4393            return base + 4 + m1.trailing_zeros() as usize;
4394        }
4395
4396        let ge2 = _mm_or_si128(
4397            _mm_cmpeq_epi32(v2, target_xor),
4398            _mm_cmpgt_epi32(v2, target_xor),
4399        );
4400        let m2 = _mm_movemask_ps(_mm_castsi128_ps(ge2)) as u32;
4401        if m2 != 0 {
4402            return base + 8 + m2.trailing_zeros() as usize;
4403        }
4404
4405        let ge3 = _mm_or_si128(
4406            _mm_cmpeq_epi32(v3, target_xor),
4407            _mm_cmpgt_epi32(v3, target_xor),
4408        );
4409        let m3 = _mm_movemask_ps(_mm_castsi128_ps(ge3)) as u32;
4410        if m3 != 0 {
4411            return base + 12 + m3.trailing_zeros() as usize;
4412        }
4413        base += 16;
4414    }
4415
4416    // Process remaining 4 elements at a time
4417    while base + 4 <= n {
4418        let vals = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4419        let ge = _mm_or_si128(
4420            _mm_cmpeq_epi32(vals, target_xor),
4421            _mm_cmpgt_epi32(vals, target_xor),
4422        );
4423        let mask = _mm_movemask_ps(_mm_castsi128_ps(ge)) as u32;
4424        if mask != 0 {
4425            return base + mask.trailing_zeros() as usize;
4426        }
4427        base += 4;
4428    }
4429
4430    // Scalar remainder (0-3 elements)
4431    while base < n {
4432        if *slice.get_unchecked(base) >= target {
4433            return base;
4434        }
4435        base += 1;
4436    }
4437    n
4438}
4439
4440#[cfg(test)]
4441mod find_first_ge_tests {
4442    use super::find_first_ge_u32;
4443
4444    #[test]
4445    fn test_find_first_ge_basic() {
4446        let data: Vec<u32> = (0..128).map(|i| i * 3).collect(); // [0, 3, 6, ..., 381]
4447        assert_eq!(find_first_ge_u32(&data, 0), 0);
4448        assert_eq!(find_first_ge_u32(&data, 1), 1); // first >= 1 is 3 at idx 1
4449        assert_eq!(find_first_ge_u32(&data, 3), 1);
4450        assert_eq!(find_first_ge_u32(&data, 4), 2); // first >= 4 is 6 at idx 2
4451        assert_eq!(find_first_ge_u32(&data, 381), 127);
4452        assert_eq!(find_first_ge_u32(&data, 382), 128); // past end
4453    }
4454
4455    #[test]
4456    fn test_find_first_ge_matches_partition_point() {
4457        let data: Vec<u32> = vec![1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75];
4458        for target in 0..80 {
4459            let expected = data.partition_point(|&d| d < target);
4460            let actual = find_first_ge_u32(&data, target);
4461            assert_eq!(actual, expected, "target={}", target);
4462        }
4463    }
4464
4465    #[test]
4466    fn test_find_first_ge_small_slices() {
4467        // Empty
4468        assert_eq!(find_first_ge_u32(&[], 5), 0);
4469        // Single element
4470        assert_eq!(find_first_ge_u32(&[10], 5), 0);
4471        assert_eq!(find_first_ge_u32(&[10], 10), 0);
4472        assert_eq!(find_first_ge_u32(&[10], 11), 1);
4473        // Three elements (< SIMD width)
4474        assert_eq!(find_first_ge_u32(&[2, 4, 6], 5), 2);
4475    }
4476
4477    #[test]
4478    fn test_find_first_ge_full_block() {
4479        // Simulate a full 128-entry block
4480        let data: Vec<u32> = (100..228).collect();
4481        assert_eq!(find_first_ge_u32(&data, 100), 0);
4482        assert_eq!(find_first_ge_u32(&data, 150), 50);
4483        assert_eq!(find_first_ge_u32(&data, 227), 127);
4484        assert_eq!(find_first_ge_u32(&data, 228), 128);
4485        assert_eq!(find_first_ge_u32(&data, 99), 0);
4486    }
4487
4488    #[test]
4489    fn test_find_first_ge_u32_max() {
4490        // Test with large u32 values (unsigned correctness)
4491        let data = vec![u32::MAX - 10, u32::MAX - 5, u32::MAX - 1, u32::MAX];
4492        assert_eq!(find_first_ge_u32(&data, u32::MAX - 10), 0);
4493        assert_eq!(find_first_ge_u32(&data, u32::MAX - 7), 1);
4494        assert_eq!(find_first_ge_u32(&data, u32::MAX), 3);
4495    }
4496}