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    for i in 0..n {
3474        let vec = &db[i * byte_len..(i + 1) * byte_len];
3475        let dist = hamming_distance(query, vec);
3476        scores[i] = 1.0 - dist as f32 * inv_dim;
3477    }
3478}
3479
3480#[cfg(test)]
3481mod tests {
3482    use super::*;
3483
3484    #[test]
3485    fn vector_simd_boundaries_reject_dimension_mismatches() {
3486        let vectors = vec![1.0f32; 6];
3487        let raw_f16 = vec![0u8; 12];
3488        let raw_u8 = vec![0u8; 6];
3489        let mut scores = vec![0.0f32; 2];
3490
3491        for invalid_query in [vec![1.0, 2.0], vec![1.0, 2.0, 3.0, 4.0]] {
3492            assert!(
3493                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3494                    batch_cosine_scores(&invalid_query, &vectors, 3, &mut scores)
3495                }))
3496                .is_err()
3497            );
3498            assert!(
3499                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3500                    batch_dot_scores_f16(&invalid_query, &raw_f16, 3, &mut scores)
3501                }))
3502                .is_err()
3503            );
3504            assert!(
3505                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3506                    batch_cosine_scores_u8(&invalid_query, &raw_u8, 3, &mut scores)
3507                }))
3508                .is_err()
3509            );
3510        }
3511    }
3512
3513    #[test]
3514    fn vector_simd_boundaries_reject_truncated_storage() {
3515        let query = [1.0f32, 2.0, 3.0];
3516        let mut scores = [0.0f32; 2];
3517
3518        assert!(
3519            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3520                batch_dot_scores(&query, &[0.0; 5], 3, &mut scores)
3521            }))
3522            .is_err()
3523        );
3524        assert!(
3525            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3526                batch_cosine_scores_f16(&query, &[0u8; 11], 3, &mut scores)
3527            }))
3528            .is_err()
3529        );
3530        assert!(
3531            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3532                dot_product_f32(&query, &query, 4)
3533            }))
3534            .is_err()
3535        );
3536    }
3537
3538    #[test]
3539    fn test_unpack_8bit() {
3540        let input: Vec<u8> = (0..128).collect();
3541        let mut output = vec![0u32; 128];
3542        unpack_8bit(&input, &mut output, 128);
3543
3544        for (i, &v) in output.iter().enumerate() {
3545            assert_eq!(v, i as u32);
3546        }
3547    }
3548
3549    #[test]
3550    fn test_unpack_16bit() {
3551        let mut input = vec![0u8; 256];
3552        for i in 0..128 {
3553            let val = (i * 100) as u16;
3554            input[i * 2] = val as u8;
3555            input[i * 2 + 1] = (val >> 8) as u8;
3556        }
3557
3558        let mut output = vec![0u32; 128];
3559        unpack_16bit(&input, &mut output, 128);
3560
3561        for (i, &v) in output.iter().enumerate() {
3562            assert_eq!(v, (i * 100) as u32);
3563        }
3564    }
3565
3566    #[test]
3567    fn test_unpack_32bit() {
3568        let mut input = vec![0u8; 512];
3569        for i in 0..128 {
3570            let val = (i * 1000) as u32;
3571            let bytes = val.to_le_bytes();
3572            input[i * 4..i * 4 + 4].copy_from_slice(&bytes);
3573        }
3574
3575        let mut output = vec![0u32; 128];
3576        unpack_32bit(&input, &mut output, 128);
3577
3578        for (i, &v) in output.iter().enumerate() {
3579            assert_eq!(v, (i * 1000) as u32);
3580        }
3581    }
3582
3583    #[test]
3584    fn test_delta_decode() {
3585        // doc_ids: [10, 15, 20, 30, 50]
3586        // gaps: [5, 5, 10, 20]
3587        // deltas (gap-1): [4, 4, 9, 19]
3588        let deltas = vec![4u32, 4, 9, 19];
3589        let mut output = vec![0u32; 5];
3590
3591        delta_decode(&mut output, &deltas, 10, 5);
3592
3593        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3594    }
3595
3596    #[test]
3597    fn test_add_one() {
3598        let mut values = vec![0u32, 1, 2, 3, 4, 5, 6, 7];
3599        add_one(&mut values, 8);
3600
3601        assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]);
3602    }
3603
3604    #[test]
3605    fn test_bits_needed() {
3606        assert_eq!(bits_needed(0), 0);
3607        assert_eq!(bits_needed(1), 1);
3608        assert_eq!(bits_needed(2), 2);
3609        assert_eq!(bits_needed(3), 2);
3610        assert_eq!(bits_needed(4), 3);
3611        assert_eq!(bits_needed(255), 8);
3612        assert_eq!(bits_needed(256), 9);
3613        assert_eq!(bits_needed(u32::MAX), 32);
3614    }
3615
3616    #[test]
3617    fn test_unpack_8bit_delta_decode() {
3618        // doc_ids: [10, 15, 20, 30, 50]
3619        // gaps: [5, 5, 10, 20]
3620        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3621        let input: Vec<u8> = vec![4, 4, 9, 19];
3622        let mut output = vec![0u32; 5];
3623
3624        unpack_8bit_delta_decode(&input, &mut output, 10, 5);
3625
3626        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3627    }
3628
3629    #[test]
3630    fn test_unpack_16bit_delta_decode() {
3631        // doc_ids: [100, 600, 1100, 2100, 4100]
3632        // gaps: [500, 500, 1000, 2000]
3633        // deltas (gap-1): [499, 499, 999, 1999] stored as u16
3634        let mut input = vec![0u8; 8];
3635        for (i, &delta) in [499u16, 499, 999, 1999].iter().enumerate() {
3636            input[i * 2] = delta as u8;
3637            input[i * 2 + 1] = (delta >> 8) as u8;
3638        }
3639        let mut output = vec![0u32; 5];
3640
3641        unpack_16bit_delta_decode(&input, &mut output, 100, 5);
3642
3643        assert_eq!(output, vec![100, 600, 1100, 2100, 4100]);
3644    }
3645
3646    #[test]
3647    fn test_fused_vs_separate_8bit() {
3648        // Test that fused and separate operations produce the same result
3649        let input: Vec<u8> = (0..127).collect();
3650        let first_value = 1000u32;
3651        let count = 128;
3652
3653        // Separate: unpack then delta_decode
3654        let mut unpacked = vec![0u32; 128];
3655        unpack_8bit(&input, &mut unpacked, 127);
3656        let mut separate_output = vec![0u32; 128];
3657        delta_decode(&mut separate_output, &unpacked, first_value, count);
3658
3659        // Fused
3660        let mut fused_output = vec![0u32; 128];
3661        unpack_8bit_delta_decode(&input, &mut fused_output, first_value, count);
3662
3663        assert_eq!(separate_output, fused_output);
3664    }
3665
3666    #[test]
3667    fn test_round_bit_width() {
3668        assert_eq!(round_bit_width(0), 0);
3669        assert_eq!(round_bit_width(1), 8);
3670        assert_eq!(round_bit_width(5), 8);
3671        assert_eq!(round_bit_width(8), 8);
3672        assert_eq!(round_bit_width(9), 16);
3673        assert_eq!(round_bit_width(12), 16);
3674        assert_eq!(round_bit_width(16), 16);
3675        assert_eq!(round_bit_width(17), 32);
3676        assert_eq!(round_bit_width(24), 32);
3677        assert_eq!(round_bit_width(32), 32);
3678    }
3679
3680    #[test]
3681    fn test_rounded_bitwidth_from_exact() {
3682        assert_eq!(RoundedBitWidth::from_exact(0), RoundedBitWidth::Zero);
3683        assert_eq!(RoundedBitWidth::from_exact(1), RoundedBitWidth::Bits8);
3684        assert_eq!(RoundedBitWidth::from_exact(8), RoundedBitWidth::Bits8);
3685        assert_eq!(RoundedBitWidth::from_exact(9), RoundedBitWidth::Bits16);
3686        assert_eq!(RoundedBitWidth::from_exact(16), RoundedBitWidth::Bits16);
3687        assert_eq!(RoundedBitWidth::from_exact(17), RoundedBitWidth::Bits32);
3688        assert_eq!(RoundedBitWidth::from_exact(32), RoundedBitWidth::Bits32);
3689    }
3690
3691    #[test]
3692    fn test_pack_unpack_rounded_8bit() {
3693        let values: Vec<u32> = (0..128).map(|i| i % 256).collect();
3694        let mut packed = vec![0u8; 128];
3695
3696        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits8, &mut packed);
3697        assert_eq!(bytes_written, 128);
3698
3699        let mut unpacked = vec![0u32; 128];
3700        unpack_rounded(&packed, RoundedBitWidth::Bits8, &mut unpacked, 128);
3701
3702        assert_eq!(values, unpacked);
3703    }
3704
3705    #[test]
3706    fn test_pack_unpack_rounded_16bit() {
3707        let values: Vec<u32> = (0..128).map(|i| i * 100).collect();
3708        let mut packed = vec![0u8; 256];
3709
3710        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits16, &mut packed);
3711        assert_eq!(bytes_written, 256);
3712
3713        let mut unpacked = vec![0u32; 128];
3714        unpack_rounded(&packed, RoundedBitWidth::Bits16, &mut unpacked, 128);
3715
3716        assert_eq!(values, unpacked);
3717    }
3718
3719    #[test]
3720    fn test_pack_unpack_rounded_32bit() {
3721        let values: Vec<u32> = (0..128).map(|i| i * 100000).collect();
3722        let mut packed = vec![0u8; 512];
3723
3724        let bytes_written = pack_rounded(&values, RoundedBitWidth::Bits32, &mut packed);
3725        assert_eq!(bytes_written, 512);
3726
3727        let mut unpacked = vec![0u32; 128];
3728        unpack_rounded(&packed, RoundedBitWidth::Bits32, &mut unpacked, 128);
3729
3730        assert_eq!(values, unpacked);
3731    }
3732
3733    #[test]
3734    fn test_unpack_rounded_delta_decode() {
3735        // Test 8-bit rounded delta decode
3736        // doc_ids: [10, 15, 20, 30, 50]
3737        // gaps: [5, 5, 10, 20]
3738        // deltas (gap-1): [4, 4, 9, 19] stored as u8
3739        let input: Vec<u8> = vec![4, 4, 9, 19];
3740        let mut output = vec![0u32; 5];
3741
3742        unpack_rounded_delta_decode(&input, RoundedBitWidth::Bits8, &mut output, 10, 5);
3743
3744        assert_eq!(output, vec![10, 15, 20, 30, 50]);
3745    }
3746
3747    #[test]
3748    fn test_unpack_rounded_delta_decode_zero() {
3749        // All zeros means gaps of 1 (consecutive doc IDs)
3750        let input: Vec<u8> = vec![];
3751        let mut output = vec![0u32; 5];
3752
3753        unpack_rounded_delta_decode(&input, RoundedBitWidth::Zero, &mut output, 100, 5);
3754
3755        assert_eq!(output, vec![100, 101, 102, 103, 104]);
3756    }
3757
3758    // ========================================================================
3759    // Sparse Vector SIMD Tests
3760    // ========================================================================
3761
3762    #[test]
3763    fn test_dequantize_uint8() {
3764        let input: Vec<u8> = vec![0, 128, 255, 64, 192];
3765        let mut output = vec![0.0f32; 5];
3766        let scale = 0.1;
3767        let min_val = 1.0;
3768
3769        dequantize_uint8(&input, &mut output, scale, min_val, 5);
3770
3771        // Expected: input[i] * scale + min_val
3772        assert!((output[0] - 1.0).abs() < 1e-6); // 0 * 0.1 + 1.0 = 1.0
3773        assert!((output[1] - 13.8).abs() < 1e-6); // 128 * 0.1 + 1.0 = 13.8
3774        assert!((output[2] - 26.5).abs() < 1e-6); // 255 * 0.1 + 1.0 = 26.5
3775        assert!((output[3] - 7.4).abs() < 1e-6); // 64 * 0.1 + 1.0 = 7.4
3776        assert!((output[4] - 20.2).abs() < 1e-6); // 192 * 0.1 + 1.0 = 20.2
3777    }
3778
3779    #[test]
3780    fn test_dequantize_uint8_large() {
3781        // Test with 128 values (full SIMD block)
3782        let input: Vec<u8> = (0..128).collect();
3783        let mut output = vec![0.0f32; 128];
3784        let scale = 2.0;
3785        let min_val = -10.0;
3786
3787        dequantize_uint8(&input, &mut output, scale, min_val, 128);
3788
3789        for (i, &out) in output.iter().enumerate().take(128) {
3790            let expected = i as f32 * scale + min_val;
3791            assert!(
3792                (out - expected).abs() < 1e-5,
3793                "Mismatch at {}: expected {}, got {}",
3794                i,
3795                expected,
3796                out
3797            );
3798        }
3799    }
3800
3801    #[test]
3802    fn test_dot_product_f32() {
3803        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0];
3804        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0];
3805
3806        let result = dot_product_f32(&a, &b, 5);
3807
3808        // Expected: 1*2 + 2*3 + 3*4 + 4*5 + 5*6 = 2 + 6 + 12 + 20 + 30 = 70
3809        assert!((result - 70.0).abs() < 1e-5);
3810    }
3811
3812    #[test]
3813    fn test_dot_product_f32_large() {
3814        // Test with 128 values
3815        let a: Vec<f32> = (0..128).map(|i| i as f32).collect();
3816        let b: Vec<f32> = (0..128).map(|i| (i + 1) as f32).collect();
3817
3818        let result = dot_product_f32(&a, &b, 128);
3819
3820        // Compute expected
3821        let expected: f32 = (0..128).map(|i| (i as f32) * ((i + 1) as f32)).sum();
3822        assert!(
3823            (result - expected).abs() < 1e-3,
3824            "Expected {}, got {}",
3825            expected,
3826            result
3827        );
3828    }
3829
3830    #[test]
3831    fn test_fused_dot_norm() {
3832        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
3833        let b = vec![2.0f32, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
3834        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3835
3836        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3837        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3838        assert!(
3839            (dot - expected_dot).abs() < 1e-5,
3840            "dot: expected {}, got {}",
3841            expected_dot,
3842            dot
3843        );
3844        assert!(
3845            (norm_b - expected_norm).abs() < 1e-5,
3846            "norm: expected {}, got {}",
3847            expected_norm,
3848            norm_b
3849        );
3850    }
3851
3852    #[test]
3853    fn test_fused_dot_norm_large() {
3854        let a: Vec<f32> = (0..768).map(|i| (i as f32) * 0.01).collect();
3855        let b: Vec<f32> = (0..768).map(|i| (i as f32) * 0.02 + 0.5).collect();
3856        let (dot, norm_b) = fused_dot_norm(&a, &b, a.len());
3857
3858        let expected_dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
3859        let expected_norm: f32 = b.iter().map(|x| x * x).sum();
3860        assert!(
3861            (dot - expected_dot).abs() < 1.0,
3862            "dot: expected {}, got {}",
3863            expected_dot,
3864            dot
3865        );
3866        assert!(
3867            (norm_b - expected_norm).abs() < 1.0,
3868            "norm: expected {}, got {}",
3869            expected_norm,
3870            norm_b
3871        );
3872    }
3873
3874    #[test]
3875    fn test_batch_cosine_scores() {
3876        // 4 vectors of dim 3
3877        let query = vec![1.0f32, 0.0, 0.0];
3878        let vectors = vec![
3879            1.0, 0.0, 0.0, // identical to query
3880            0.0, 1.0, 0.0, // orthogonal
3881            -1.0, 0.0, 0.0, // opposite
3882            0.5, 0.5, 0.0, // 45 degrees
3883        ];
3884        let mut scores = vec![0f32; 4];
3885        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3886
3887        assert!((scores[0] - 1.0).abs() < 1e-5, "identical: {}", scores[0]);
3888        assert!(scores[1].abs() < 1e-5, "orthogonal: {}", scores[1]);
3889        assert!((scores[2] - (-1.0)).abs() < 1e-5, "opposite: {}", scores[2]);
3890        let expected_45 = 0.5f32 / (0.5f32.powi(2) + 0.5f32.powi(2)).sqrt();
3891        assert!(
3892            (scores[3] - expected_45).abs() < 1e-5,
3893            "45deg: expected {}, got {}",
3894            expected_45,
3895            scores[3]
3896        );
3897    }
3898
3899    #[test]
3900    fn test_batch_cosine_scores_matches_individual() {
3901        let query: Vec<f32> = (0..128).map(|i| (i as f32) * 0.1).collect();
3902        let n = 50;
3903        let dim = 128;
3904        let vectors: Vec<f32> = (0..n * dim).map(|i| ((i * 7 + 3) as f32) * 0.01).collect();
3905
3906        let mut batch_scores = vec![0f32; n];
3907        batch_cosine_scores(&query, &vectors, dim, &mut batch_scores);
3908
3909        for i in 0..n {
3910            let vec_i = &vectors[i * dim..(i + 1) * dim];
3911            let individual = cosine_similarity(&query, vec_i);
3912            assert!(
3913                (batch_scores[i] - individual).abs() < 1e-5,
3914                "vec {}: batch={}, individual={}",
3915                i,
3916                batch_scores[i],
3917                individual
3918            );
3919        }
3920    }
3921
3922    #[test]
3923    fn test_batch_cosine_scores_empty() {
3924        let query = vec![1.0f32, 2.0, 3.0];
3925        let vectors: Vec<f32> = vec![];
3926        let mut scores: Vec<f32> = vec![];
3927        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3928        assert!(scores.is_empty());
3929    }
3930
3931    #[test]
3932    fn test_batch_cosine_scores_zero_query() {
3933        let query = vec![0.0f32, 0.0, 0.0];
3934        let vectors = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
3935        let mut scores = vec![0f32; 2];
3936        batch_cosine_scores(&query, &vectors, 3, &mut scores);
3937        assert_eq!(scores[0], 0.0);
3938        assert_eq!(scores[1], 0.0);
3939    }
3940
3941    // ================================================================
3942    // f16 conversion tests
3943    // ================================================================
3944
3945    #[test]
3946    fn test_f16_roundtrip_normal() {
3947        for &v in &[0.0f32, 1.0, -1.0, 0.5, -0.5, 0.333, 65504.0] {
3948            let h = f32_to_f16(v);
3949            let back = f16_to_f32(h);
3950            let err = (back - v).abs() / v.abs().max(1e-6);
3951            assert!(
3952                err < 0.002,
3953                "f16 roundtrip {v} → {h:#06x} → {back}, rel err {err}"
3954            );
3955        }
3956    }
3957
3958    #[test]
3959    fn test_f16_special() {
3960        // Zero
3961        assert_eq!(f16_to_f32(f32_to_f16(0.0)), 0.0);
3962        // Negative zero
3963        assert_eq!(f32_to_f16(-0.0), 0x8000);
3964        // Infinity
3965        assert!(f16_to_f32(f32_to_f16(f32::INFINITY)).is_infinite());
3966        // NaN
3967        assert!(f16_to_f32(f32_to_f16(f32::NAN)).is_nan());
3968    }
3969
3970    #[test]
3971    fn test_f16_embedding_range() {
3972        // Typical embedding values in [-1, 1]
3973        let values: Vec<f32> = (-100..=100).map(|i| i as f32 / 100.0).collect();
3974        for &v in &values {
3975            let back = f16_to_f32(f32_to_f16(v));
3976            assert!((back - v).abs() < 0.001, "f16 error for {v}: got {back}");
3977        }
3978    }
3979
3980    // ================================================================
3981    // u8 conversion tests
3982    // ================================================================
3983
3984    #[test]
3985    fn test_u8_roundtrip() {
3986        // Boundary values
3987        assert_eq!(f32_to_u8_saturating(-1.0), 0);
3988        assert_eq!(f32_to_u8_saturating(1.0), 255);
3989        assert_eq!(f32_to_u8_saturating(0.0), 127); // ~127.5 truncated
3990
3991        // Saturation
3992        assert_eq!(f32_to_u8_saturating(-2.0), 0);
3993        assert_eq!(f32_to_u8_saturating(2.0), 255);
3994    }
3995
3996    #[test]
3997    fn test_u8_dequantize() {
3998        assert!((u8_to_f32(0) - (-1.0)).abs() < 0.01);
3999        assert!((u8_to_f32(255) - 1.0).abs() < 0.01);
4000        assert!((u8_to_f32(127) - 0.0).abs() < 0.01);
4001    }
4002
4003    // ================================================================
4004    // Batch scoring tests for quantized vectors
4005    // ================================================================
4006
4007    #[test]
4008    fn test_batch_cosine_scores_f16() {
4009        let query = vec![0.6f32, 0.8, 0.0, 0.0];
4010        let dim = 4;
4011        let vecs_f32 = vec![
4012            0.6f32, 0.8, 0.0, 0.0, // identical to query
4013            0.0, 0.0, 0.6, 0.8, // orthogonal
4014        ];
4015
4016        // Quantize to f16
4017        let mut f16_buf = vec![0u16; 8];
4018        batch_f32_to_f16(&vecs_f32, &mut f16_buf);
4019        let raw: &[u8] =
4020            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4021
4022        let mut scores = vec![0f32; 2];
4023        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4024
4025        assert!(
4026            (scores[0] - 1.0).abs() < 0.01,
4027            "identical vectors: {}",
4028            scores[0]
4029        );
4030        assert!(scores[1].abs() < 0.01, "orthogonal vectors: {}", scores[1]);
4031    }
4032
4033    #[test]
4034    fn test_batch_cosine_scores_u8() {
4035        let query = vec![0.6f32, 0.8, 0.0, 0.0];
4036        let dim = 4;
4037        let vecs_f32 = vec![
4038            0.6f32, 0.8, 0.0, 0.0, // ~identical to query
4039            -0.6, -0.8, 0.0, 0.0, // opposite
4040        ];
4041
4042        // Quantize to u8
4043        let mut u8_buf = vec![0u8; 8];
4044        batch_f32_to_u8(&vecs_f32, &mut u8_buf);
4045
4046        let mut scores = vec![0f32; 2];
4047        batch_cosine_scores_u8(&query, &u8_buf, dim, &mut scores);
4048
4049        assert!(scores[0] > 0.95, "similar vectors: {}", scores[0]);
4050        assert!(scores[1] < -0.95, "opposite vectors: {}", scores[1]);
4051    }
4052
4053    #[test]
4054    fn test_batch_cosine_scores_f16_large_dim() {
4055        // Test with typical embedding dimension
4056        let dim = 768;
4057        let query: Vec<f32> = (0..dim).map(|i| (i as f32 / dim as f32) - 0.5).collect();
4058        let vec2: Vec<f32> = query.iter().map(|x| x * 0.9 + 0.01).collect();
4059
4060        let mut all_vecs = query.clone();
4061        all_vecs.extend_from_slice(&vec2);
4062
4063        let mut f16_buf = vec![0u16; all_vecs.len()];
4064        batch_f32_to_f16(&all_vecs, &mut f16_buf);
4065        let raw: &[u8] =
4066            unsafe { std::slice::from_raw_parts(f16_buf.as_ptr() as *const u8, f16_buf.len() * 2) };
4067
4068        let mut scores = vec![0f32; 2];
4069        batch_cosine_scores_f16(&query, raw, dim, &mut scores);
4070
4071        // Self-similarity should be ~1.0
4072        assert!((scores[0] - 1.0).abs() < 0.01, "self-sim: {}", scores[0]);
4073        // High similarity with scaled version
4074        assert!(scores[1] > 0.99, "scaled-sim: {}", scores[1]);
4075    }
4076
4077    // ================================================================
4078    // Hamming distance tests
4079    // ================================================================
4080
4081    #[test]
4082    fn test_hamming_distance_identical() {
4083        let a = vec![0xAA; 64];
4084        assert_eq!(hamming_distance(&a, &a), 0);
4085    }
4086
4087    #[test]
4088    fn test_hamming_distance_opposite() {
4089        let a = vec![0xFF; 32];
4090        let b = vec![0x00; 32];
4091        assert_eq!(hamming_distance(&a, &b), 256);
4092    }
4093
4094    #[test]
4095    fn test_hamming_distance_known() {
4096        // Single byte: 0b10101010 vs 0b01010101 = 8 bits differ
4097        let a = vec![0xAA];
4098        let b = vec![0x55];
4099        assert_eq!(hamming_distance(&a, &b), 8);
4100
4101        // Two bytes
4102        let a = vec![0xFF, 0x00];
4103        let b = vec![0x00, 0x00];
4104        assert_eq!(hamming_distance(&a, &b), 8);
4105    }
4106
4107    #[test]
4108    fn test_hamming_distance_single_bit() {
4109        let a = vec![0x00; 16];
4110        let mut b = vec![0x00; 16];
4111        b[7] = 0x01; // flip one bit
4112        assert_eq!(hamming_distance(&a, &b), 1);
4113    }
4114
4115    #[test]
4116    fn test_hamming_distance_empty() {
4117        let a: Vec<u8> = vec![];
4118        assert_eq!(hamming_distance(&a, &a), 0);
4119    }
4120
4121    #[test]
4122    fn test_hamming_distance_remainder_path() {
4123        // 17 bytes: not aligned to 16 (NEON) or 32 (AVX2)
4124        let a = vec![0xFF; 17];
4125        let b = vec![0x00; 17];
4126        assert_eq!(hamming_distance(&a, &b), 136); // 17 * 8
4127
4128        // 33 bytes: tests 32-byte chunk + 1 remainder for AVX2
4129        let a = vec![0xFF; 33];
4130        let b = vec![0x00; 33];
4131        assert_eq!(hamming_distance(&a, &b), 264); // 33 * 8
4132    }
4133
4134    #[test]
4135    fn test_hamming_distance_large() {
4136        // 4096 bytes = 32768 bits, all differing
4137        let a = vec![0xFF; 4096];
4138        let b = vec![0x00; 4096];
4139        assert_eq!(hamming_distance(&a, &b), 32768);
4140    }
4141
4142    #[test]
4143    fn test_hamming_distance_scalar_matches() {
4144        // Verify SIMD path matches scalar for various sizes
4145        for size in [1, 7, 8, 15, 16, 31, 32, 63, 64, 100, 128, 255, 256] {
4146            let a: Vec<u8> = (0..size).map(|i| (i * 37 + 13) as u8).collect();
4147            let b: Vec<u8> = (0..size).map(|i| (i * 53 + 7) as u8).collect();
4148            let expected = hamming_distance_scalar(&a, &b);
4149            let got = hamming_distance(&a, &b);
4150            assert_eq!(got, expected, "mismatch at size {size}");
4151        }
4152    }
4153
4154    // ================================================================
4155    // Batch Hamming scoring tests
4156    // ================================================================
4157
4158    #[test]
4159    fn test_batch_hamming_scores_identical() {
4160        let query = vec![0xAA; 16];
4161        let db = vec![0xAA; 16]; // one vector, identical
4162        let mut scores = vec![0f32; 1];
4163        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4164        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4165    }
4166
4167    #[test]
4168    fn test_batch_hamming_scores_opposite() {
4169        let query = vec![0xFF; 16];
4170        let db = vec![0x00; 16];
4171        let mut scores = vec![0f32; 1];
4172        batch_hamming_scores(&query, &db, 16, 128, &mut scores);
4173        assert!((scores[0] - 0.0).abs() < 1e-6, "opposite: {}", scores[0]);
4174    }
4175
4176    #[test]
4177    fn test_batch_hamming_scores_multiple() {
4178        let byte_len = 8;
4179        let dim_bits = 64;
4180        let query = vec![0xFF; byte_len];
4181        let mut db = Vec::new();
4182        db.extend_from_slice(&vec![0xFF; byte_len]); // identical → 1.0
4183        db.extend_from_slice(&vec![0x00; byte_len]); // opposite → 0.0
4184        db.extend_from_slice(&vec![0x0F; byte_len]); // half bits differ → 0.5
4185
4186        let mut scores = vec![0f32; 3];
4187        batch_hamming_scores(&query, &db, byte_len, dim_bits, &mut scores);
4188
4189        assert!((scores[0] - 1.0).abs() < 1e-6, "identical: {}", scores[0]);
4190        assert!((scores[1] - 0.0).abs() < 1e-6, "opposite: {}", scores[1]);
4191        assert!((scores[2] - 0.5).abs() < 1e-6, "half: {}", scores[2]);
4192    }
4193
4194    #[test]
4195    fn test_batch_hamming_scores_empty() {
4196        let query = vec![0xFF; 8];
4197        let db: Vec<u8> = vec![];
4198        let mut scores: Vec<f32> = vec![];
4199        batch_hamming_scores(&query, &db, 8, 64, &mut scores);
4200        assert!(scores.is_empty());
4201    }
4202
4203    #[test]
4204    fn test_batch_hamming_scores_zero_byte_len() {
4205        let query: Vec<u8> = vec![];
4206        let db: Vec<u8> = vec![];
4207        let mut scores = vec![0f32; 1];
4208        batch_hamming_scores(&query, &db, 0, 0, &mut scores);
4209        // Should return early without modifying scores
4210        assert_eq!(scores[0], 0.0);
4211    }
4212}
4213
4214// ============================================================================
4215// SIMD-accelerated linear scan for sorted u32 slices (within-block seek)
4216// ============================================================================
4217
4218/// Find index of first element >= `target` in a sorted `u32` slice.
4219///
4220/// Equivalent to `slice.partition_point(|&d| d < target)` but uses SIMD to
4221/// scan 4 elements per cycle. Faster than binary search for slices ≤ 256
4222/// elements because it avoids the data-dependency chain inherent in binary
4223/// search (~8-10 cycles/iteration vs ~1-2 cycles/iteration for SIMD scan).
4224///
4225/// Returns `slice.len()` if no element >= `target`.
4226#[inline]
4227pub fn find_first_ge_u32(slice: &[u32], target: u32) -> usize {
4228    #[cfg(target_arch = "aarch64")]
4229    {
4230        if neon::is_available() {
4231            return unsafe { find_first_ge_u32_neon(slice, target) };
4232        }
4233    }
4234
4235    #[cfg(target_arch = "x86_64")]
4236    {
4237        if sse::is_available() {
4238            return unsafe { find_first_ge_u32_sse(slice, target) };
4239        }
4240    }
4241
4242    // Scalar fallback (WASM, other architectures)
4243    slice.partition_point(|&d| d < target)
4244}
4245
4246#[cfg(target_arch = "aarch64")]
4247#[target_feature(enable = "neon")]
4248#[allow(unsafe_op_in_unsafe_fn)]
4249unsafe fn find_first_ge_u32_neon(slice: &[u32], target: u32) -> usize {
4250    use std::arch::aarch64::*;
4251
4252    let n = slice.len();
4253    let ptr = slice.as_ptr();
4254    let target_vec = vdupq_n_u32(target);
4255    // Bit positions for each lane: [1, 2, 4, 8]
4256    let bit_mask: uint32x4_t = core::mem::transmute([1u32, 2u32, 4u32, 8u32]);
4257
4258    let chunks = n / 16;
4259    let mut base = 0usize;
4260
4261    // Process 16 elements per iteration (4 × 4-wide NEON compares)
4262    for _ in 0..chunks {
4263        let v0 = vld1q_u32(ptr.add(base));
4264        let v1 = vld1q_u32(ptr.add(base + 4));
4265        let v2 = vld1q_u32(ptr.add(base + 8));
4266        let v3 = vld1q_u32(ptr.add(base + 12));
4267
4268        let c0 = vcgeq_u32(v0, target_vec);
4269        let c1 = vcgeq_u32(v1, target_vec);
4270        let c2 = vcgeq_u32(v2, target_vec);
4271        let c3 = vcgeq_u32(v3, target_vec);
4272
4273        let m0 = vaddvq_u32(vandq_u32(c0, bit_mask));
4274        if m0 != 0 {
4275            return base + m0.trailing_zeros() as usize;
4276        }
4277        let m1 = vaddvq_u32(vandq_u32(c1, bit_mask));
4278        if m1 != 0 {
4279            return base + 4 + m1.trailing_zeros() as usize;
4280        }
4281        let m2 = vaddvq_u32(vandq_u32(c2, bit_mask));
4282        if m2 != 0 {
4283            return base + 8 + m2.trailing_zeros() as usize;
4284        }
4285        let m3 = vaddvq_u32(vandq_u32(c3, bit_mask));
4286        if m3 != 0 {
4287            return base + 12 + m3.trailing_zeros() as usize;
4288        }
4289        base += 16;
4290    }
4291
4292    // Process remaining 4 elements at a time
4293    while base + 4 <= n {
4294        let vals = vld1q_u32(ptr.add(base));
4295        let cmp = vcgeq_u32(vals, target_vec);
4296        let mask = vaddvq_u32(vandq_u32(cmp, bit_mask));
4297        if mask != 0 {
4298            return base + mask.trailing_zeros() as usize;
4299        }
4300        base += 4;
4301    }
4302
4303    // Scalar remainder (0-3 elements)
4304    while base < n {
4305        if *slice.get_unchecked(base) >= target {
4306            return base;
4307        }
4308        base += 1;
4309    }
4310    n
4311}
4312
4313#[cfg(target_arch = "x86_64")]
4314#[target_feature(enable = "sse2")]
4315#[allow(unsafe_op_in_unsafe_fn)]
4316unsafe fn find_first_ge_u32_sse(slice: &[u32], target: u32) -> usize {
4317    use std::arch::x86_64::*;
4318
4319    let n = slice.len();
4320    let ptr = slice.as_ptr();
4321
4322    // For unsigned >= comparison: XOR with 0x80000000 converts to signed domain
4323    let sign_flip = _mm_set1_epi32(i32::MIN);
4324    let target_xor = _mm_xor_si128(_mm_set1_epi32(target as i32), sign_flip);
4325
4326    let chunks = n / 16;
4327    let mut base = 0usize;
4328
4329    // Process 16 elements per iteration (4 × 4-wide SSE compares)
4330    for _ in 0..chunks {
4331        let v0 = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4332        let v1 = _mm_xor_si128(
4333            _mm_loadu_si128(ptr.add(base + 4) as *const __m128i),
4334            sign_flip,
4335        );
4336        let v2 = _mm_xor_si128(
4337            _mm_loadu_si128(ptr.add(base + 8) as *const __m128i),
4338            sign_flip,
4339        );
4340        let v3 = _mm_xor_si128(
4341            _mm_loadu_si128(ptr.add(base + 12) as *const __m128i),
4342            sign_flip,
4343        );
4344
4345        // ge = eq | gt (in signed domain after XOR)
4346        let ge0 = _mm_or_si128(
4347            _mm_cmpeq_epi32(v0, target_xor),
4348            _mm_cmpgt_epi32(v0, target_xor),
4349        );
4350        let m0 = _mm_movemask_ps(_mm_castsi128_ps(ge0)) as u32;
4351        if m0 != 0 {
4352            return base + m0.trailing_zeros() as usize;
4353        }
4354
4355        let ge1 = _mm_or_si128(
4356            _mm_cmpeq_epi32(v1, target_xor),
4357            _mm_cmpgt_epi32(v1, target_xor),
4358        );
4359        let m1 = _mm_movemask_ps(_mm_castsi128_ps(ge1)) as u32;
4360        if m1 != 0 {
4361            return base + 4 + m1.trailing_zeros() as usize;
4362        }
4363
4364        let ge2 = _mm_or_si128(
4365            _mm_cmpeq_epi32(v2, target_xor),
4366            _mm_cmpgt_epi32(v2, target_xor),
4367        );
4368        let m2 = _mm_movemask_ps(_mm_castsi128_ps(ge2)) as u32;
4369        if m2 != 0 {
4370            return base + 8 + m2.trailing_zeros() as usize;
4371        }
4372
4373        let ge3 = _mm_or_si128(
4374            _mm_cmpeq_epi32(v3, target_xor),
4375            _mm_cmpgt_epi32(v3, target_xor),
4376        );
4377        let m3 = _mm_movemask_ps(_mm_castsi128_ps(ge3)) as u32;
4378        if m3 != 0 {
4379            return base + 12 + m3.trailing_zeros() as usize;
4380        }
4381        base += 16;
4382    }
4383
4384    // Process remaining 4 elements at a time
4385    while base + 4 <= n {
4386        let vals = _mm_xor_si128(_mm_loadu_si128(ptr.add(base) as *const __m128i), sign_flip);
4387        let ge = _mm_or_si128(
4388            _mm_cmpeq_epi32(vals, target_xor),
4389            _mm_cmpgt_epi32(vals, target_xor),
4390        );
4391        let mask = _mm_movemask_ps(_mm_castsi128_ps(ge)) as u32;
4392        if mask != 0 {
4393            return base + mask.trailing_zeros() as usize;
4394        }
4395        base += 4;
4396    }
4397
4398    // Scalar remainder (0-3 elements)
4399    while base < n {
4400        if *slice.get_unchecked(base) >= target {
4401            return base;
4402        }
4403        base += 1;
4404    }
4405    n
4406}
4407
4408#[cfg(test)]
4409mod find_first_ge_tests {
4410    use super::find_first_ge_u32;
4411
4412    #[test]
4413    fn test_find_first_ge_basic() {
4414        let data: Vec<u32> = (0..128).map(|i| i * 3).collect(); // [0, 3, 6, ..., 381]
4415        assert_eq!(find_first_ge_u32(&data, 0), 0);
4416        assert_eq!(find_first_ge_u32(&data, 1), 1); // first >= 1 is 3 at idx 1
4417        assert_eq!(find_first_ge_u32(&data, 3), 1);
4418        assert_eq!(find_first_ge_u32(&data, 4), 2); // first >= 4 is 6 at idx 2
4419        assert_eq!(find_first_ge_u32(&data, 381), 127);
4420        assert_eq!(find_first_ge_u32(&data, 382), 128); // past end
4421    }
4422
4423    #[test]
4424    fn test_find_first_ge_matches_partition_point() {
4425        let data: Vec<u32> = vec![1, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 75];
4426        for target in 0..80 {
4427            let expected = data.partition_point(|&d| d < target);
4428            let actual = find_first_ge_u32(&data, target);
4429            assert_eq!(actual, expected, "target={}", target);
4430        }
4431    }
4432
4433    #[test]
4434    fn test_find_first_ge_small_slices() {
4435        // Empty
4436        assert_eq!(find_first_ge_u32(&[], 5), 0);
4437        // Single element
4438        assert_eq!(find_first_ge_u32(&[10], 5), 0);
4439        assert_eq!(find_first_ge_u32(&[10], 10), 0);
4440        assert_eq!(find_first_ge_u32(&[10], 11), 1);
4441        // Three elements (< SIMD width)
4442        assert_eq!(find_first_ge_u32(&[2, 4, 6], 5), 2);
4443    }
4444
4445    #[test]
4446    fn test_find_first_ge_full_block() {
4447        // Simulate a full 128-entry block
4448        let data: Vec<u32> = (100..228).collect();
4449        assert_eq!(find_first_ge_u32(&data, 100), 0);
4450        assert_eq!(find_first_ge_u32(&data, 150), 50);
4451        assert_eq!(find_first_ge_u32(&data, 227), 127);
4452        assert_eq!(find_first_ge_u32(&data, 228), 128);
4453        assert_eq!(find_first_ge_u32(&data, 99), 0);
4454    }
4455
4456    #[test]
4457    fn test_find_first_ge_u32_max() {
4458        // Test with large u32 values (unsigned correctness)
4459        let data = vec![u32::MAX - 10, u32::MAX - 5, u32::MAX - 1, u32::MAX];
4460        assert_eq!(find_first_ge_u32(&data, u32::MAX - 10), 0);
4461        assert_eq!(find_first_ge_u32(&data, u32::MAX - 7), 1);
4462        assert_eq!(find_first_ge_u32(&data, u32::MAX), 3);
4463    }
4464}