Skip to main content

simd_popcnt/
lib.rs

1//! # simd-popcnt
2//!
3//! Count the number of 1 bits (bit population count, a.k.a. Hamming weight) in
4//! a byte slice as quickly as possible using specialized CPU instructions:
5//! POPCNT, AVX2 and AVX512 on x86/x86-64, and NEON and SVE on AArch64.
6//!
7//! This is a Rust port of the C/C++ [`libpopcnt.h`](https://github.com/kimwalisch/libpopcnt)
8//! header-only library by Kim Walisch.
9//!
10//! ## Usage
11//!
12//! ```
13//! let bytes = [0xFFu8; 16];
14//! assert_eq!(simd_popcnt::popcnt(&bytes), 128);
15//! ```
16//!
17//! ## Performance
18//!
19//! For the fastest possible code, compile with `RUSTFLAGS="-C target-cpu=native"`.
20//! This lets the crate select the best SIMD path at compile time with zero
21//! runtime dispatch overhead. Otherwise the best available instruction set is
22//! detected once at runtime and cached.
23
24// Enable SVE intrinsics only when build.rs confirmed they compile on this rustc.
25#![cfg_attr(simd_popcnt_have_sve, feature(stdarch_aarch64_sve))]
26
27#[cfg(target_arch = "aarch64")]
28use core::arch::aarch64::*;
29#[cfg(target_arch = "x86")]
30use core::arch::x86::*;
31#[cfg(target_arch = "x86_64")]
32use core::arch::x86_64::*;
33#[cfg(all(target_arch = "aarch64", simd_popcnt_have_sve))]
34use std::arch::is_aarch64_feature_detected;
35
36/// Counts the number of one bits (population count) in `bytes`.
37///
38/// Dispatches to the fastest implementation for the running CPU: SIMD where
39/// available, a scalar fallback otherwise.
40///
41/// To count the bits in a slice of a wider integer type (`&[u64]`, `&[u32]`, …),
42/// use the [`PopcntExt::popcnt`] method rather than converting to bytes by hand.
43///
44/// # Examples
45///
46/// ```
47/// assert_eq!(simd_popcnt::popcnt(&[]), 0);
48/// assert_eq!(simd_popcnt::popcnt(&[0xFF, 0x0F]), 12);
49/// ```
50#[must_use]
51#[inline]
52pub fn popcnt(bytes: &[u8]) -> u64 {
53    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
54    {
55        popcnt_x86(bytes)
56    }
57
58    #[cfg(target_arch = "aarch64")]
59    {
60        popcnt_aarch64(bytes)
61    }
62
63    #[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
64    {
65        popcnt_scalar(bytes)
66    }
67}
68
69// ────────────────────────────────────────────────────────────────────────────
70// Ergonomic extension trait for integer slices
71// ────────────────────────────────────────────────────────────────────────────
72
73/// Adds a [`popcnt`](PopcntExt::popcnt) method to slices of the built-in integer
74/// types, counting their bits without a manual byte cast. Implemented for slices,
75/// arrays and `Vec`s of `u8`/`u16`/`u32`/`u64`/`u128`/`usize` and their signed
76/// counterparts; bring it into scope with `use simd_popcnt::PopcntExt;`.
77///
78/// ```
79/// use simd_popcnt::PopcntExt;
80///
81/// let words: &[u64] = &[u64::MAX, 0x0F0F_0F0F_0F0F_0F0F];
82/// assert_eq!(words.popcnt(), 64 + 32);
83/// assert_eq!(vec![1u32, 2, 3].popcnt(), 4);
84/// ```
85pub trait PopcntExt {
86    /// Count the total number of 1 bits across all elements of the slice.
87    #[must_use]
88    fn popcnt(&self) -> u64;
89}
90
91/// Implement [`PopcntExt`] for `[$t]` by viewing the slice as bytes and
92/// delegating to [`popcnt`]. Population count is byte-order independent, so this
93/// is correct on both little- and big-endian targets.
94macro_rules! impl_popcnt_ext {
95    ($($t:ty),+ $(,)?) => {$(
96        impl PopcntExt for [$t] {
97            #[inline]
98            fn popcnt(&self) -> u64 {
99                // SAFETY: `$t` is a plain integer (no padding, every bit pattern
100                // valid) and `u8` is always 1-aligned, so the slice is a valid
101                // `&[u8]` of `size_of_val` bytes.
102                let bytes = unsafe {
103                    core::slice::from_raw_parts(
104                        self.as_ptr().cast::<u8>(),
105                        core::mem::size_of_val(self),
106                    )
107                };
108                popcnt(bytes)
109            }
110        }
111    )+};
112}
113
114impl_popcnt_ext!(
115    u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
116);
117
118// ────────────────────────────────────────────────────────────────────────────
119// Portable scalar fallbacks (available on every architecture)
120// ────────────────────────────────────────────────────────────────────────────
121
122/// Packs the trailing `rem.len()` (0..=7) bytes into a zero-padded `u64` for
123/// counting. Uses native byte order (the popcount is order-independent), which
124/// avoids a byte swap on big-endian targets.
125#[inline]
126fn tail_u64(rem: &[u8]) -> u64 {
127    let mut buf = [0u8; 8];
128    buf[..rem.len()].copy_from_slice(rem);
129    u64::from_ne_bytes(buf)
130}
131
132/// Scalar population count loop, summing `count_ones()` over 8-byte chunks.
133/// `count_ones()` is inlined in release and lowers to the target's hardware
134/// popcount where available (x86 POPCNT, PowerPC popcntd, WebAssembly
135/// `i64.popcnt`, …), otherwise to an inline bit-twiddling sequence — never a
136/// library call. Shared with the POPCNT-`target_feature` variant below.
137macro_rules! popcnt_scalar_loop {
138    ($bytes:expr) => {{
139        let mut cnt = 0u64;
140        let (chunks, rem) = $bytes.as_chunks::<8>();
141        for chunk in chunks {
142            cnt += u64::from_ne_bytes(*chunk).count_ones() as u64;
143        }
144        if !rem.is_empty() {
145            cnt += tail_u64(rem).count_ones() as u64;
146        }
147        cnt
148    }};
149}
150
151/// Portable scalar population count via [`u64::count_ones`].
152#[allow(dead_code)]
153fn popcnt_scalar(bytes: &[u8]) -> u64 {
154    popcnt_scalar_loop!(bytes)
155}
156
157// ════════════════════════════════════════════════════════════════════════════
158// x86 / x86-64
159// ════════════════════════════════════════════════════════════════════════════
160
161#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
162fn popcnt_x86(bytes: &[u8]) -> u64 {
163    // Compile-time AVX512 path (e.g. with `-C target-cpu=native`).
164    #[cfg(target_feature = "avx512vpopcntdq")]
165    {
166        // AVX512 isn't worth its setup cost for tiny arrays.
167        if bytes.len() >= 40 {
168            unsafe { popcnt_avx512(bytes) }
169        } else {
170            popcnt_scalar_static(bytes)
171        }
172    }
173
174    // Compile-time AVX2 path.
175    #[cfg(all(target_feature = "avx2", not(target_feature = "avx512vpopcntdq")))]
176    {
177        let mut cnt = 0u64;
178        let mut rest = bytes;
179        // AVX2 only wins for arrays >= 512 bytes.
180        if bytes.len() >= 512 {
181            let n = bytes.len() / 32 * 32;
182            cnt += unsafe { popcnt_avx2(&bytes[..n]) };
183            rest = &bytes[n..];
184        }
185        cnt + popcnt_scalar_static(rest)
186    }
187
188    // No SIMD enabled at compile time: detect at runtime.
189    #[cfg(not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq")))]
190    {
191        popcnt_x86_runtime(bytes)
192    }
193}
194
195/// Scalar count for the compile-time SIMD paths' small arrays and tails:
196/// hardware POPCNT when statically enabled, otherwise the integer fallback.
197#[cfg(all(
198    any(target_arch = "x86", target_arch = "x86_64"),
199    any(target_feature = "avx2", target_feature = "avx512vpopcntdq")
200))]
201#[inline]
202fn popcnt_scalar_static(bytes: &[u8]) -> u64 {
203    #[cfg(target_feature = "popcnt")]
204    {
205        // SAFETY: `popcnt` is statically enabled for the whole crate.
206        unsafe { popcnt_scalar_hw(bytes) }
207    }
208    #[cfg(not(target_feature = "popcnt"))]
209    {
210        popcnt_scalar(bytes)
211    }
212}
213
214/// Cached runtime check for AVX-512F + AVX-512BW + AVX-512VPOPCNTDQ support.
215/// Reads three `is_x86_feature_detected!` results once and stores the combined
216/// outcome so subsequent calls only load a single atomic.
217#[cfg(all(
218    any(target_arch = "x86", target_arch = "x86_64"),
219    not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq"))
220))]
221fn has_avx512() -> bool {
222    use core::sync::atomic::{AtomicI32, Ordering};
223    static HAS_AVX512: AtomicI32 = AtomicI32::new(-1);
224    let cached = HAS_AVX512.load(Ordering::Relaxed);
225    if cached != -1 {
226        return cached != 0;
227    }
228    let v = (is_x86_feature_detected!("avx512f")
229        && is_x86_feature_detected!("avx512bw")
230        && is_x86_feature_detected!("avx512vpopcntdq")) as i32;
231    HAS_AVX512.store(v, Ordering::Relaxed);
232    v != 0
233}
234
235/// Runtime dispatch using cached CPU feature detection. Only compiled when no
236/// SIMD feature is statically enabled (otherwise the compile-time paths run).
237#[cfg(all(
238    any(target_arch = "x86", target_arch = "x86_64"),
239    not(any(target_feature = "avx2", target_feature = "avx512vpopcntdq"))
240))]
241fn popcnt_x86_runtime(bytes: &[u8]) -> u64 {
242    // AVX512: not worth its setup cost below ~40 bytes, handles any length.
243    if bytes.len() >= 40 && has_avx512() {
244        return unsafe { popcnt_avx512(bytes) };
245    }
246
247    let mut cnt = 0u64;
248    let mut rest = bytes;
249
250    // AVX2 only wins for arrays >= 512 bytes.
251    if bytes.len() >= 512 && is_x86_feature_detected!("avx2") {
252        let n = bytes.len() / 32 * 32;
253        cnt += unsafe { popcnt_avx2(&bytes[..n]) };
254        rest = &bytes[n..];
255    }
256
257    // Scalar tail (or the whole array if AVX2 didn't fire). Dispatching on
258    // POPCNT is essential: outside a `#[target_feature(enable = "popcnt")]`
259    // function, `count_ones()` compiles to a software fallback even on
260    // POPCNT-capable CPUs.
261    cnt += if is_x86_feature_detected!("popcnt") {
262        unsafe { popcnt_scalar_hw(rest) }
263    } else {
264        popcnt_scalar(rest)
265    };
266
267    cnt
268}
269
270/// Scalar population count via the hardware POPCNT instruction. The
271/// `#[target_feature(enable = "popcnt")]` attribute is what lets `count_ones()`
272/// lower to a single `popcnt`; only call it once POPCNT support is confirmed.
273#[allow(dead_code)]
274#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
275#[target_feature(enable = "popcnt")]
276fn popcnt_scalar_hw(bytes: &[u8]) -> u64 {
277    popcnt_scalar_loop!(bytes) // count_ones() lowers to popcntq here
278}
279
280// ── AVX2 ────────────────────────────────────────────────────────────────────
281
282/// Carry-save adder: returns the `(carry, sum)` bit-planes of `a + b + c`,
283/// computed across all lanes in parallel.
284#[cfg(all(
285    any(target_arch = "x86", target_arch = "x86_64"),
286    not(target_feature = "avx512vpopcntdq")
287))]
288#[target_feature(enable = "avx2")]
289#[inline]
290fn csa256(a: __m256i, b: __m256i, c: __m256i) -> (__m256i, __m256i) {
291    let u = _mm256_xor_si256(a, b);
292    let h = _mm256_or_si256(_mm256_and_si256(a, b), _mm256_and_si256(u, c));
293    let l = _mm256_xor_si256(u, c);
294    (h, l)
295}
296
297/// Per-byte population count of a 256-bit vector using the nibble lookup, then
298/// horizontal sum of each 8-byte lane via `_mm256_sad_epu8` (result in 4 u64s).
299#[cfg(all(
300    any(target_arch = "x86", target_arch = "x86_64"),
301    not(target_feature = "avx512vpopcntdq")
302))]
303#[target_feature(enable = "avx2")]
304#[inline]
305fn popcnt256(v: __m256i) -> __m256i {
306    let lookup1 = _mm256_setr_epi8(
307        4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8, 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7,
308        7, 8,
309    );
310    let lookup2 = _mm256_setr_epi8(
311        4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1, 1, 0, 4, 3, 3, 2, 3, 2, 2, 1, 3, 2, 2, 1, 2, 1,
312        1, 0,
313    );
314    let low_mask = _mm256_set1_epi8(0x0f);
315    let lo = _mm256_and_si256(v, low_mask);
316    let hi = _mm256_and_si256(_mm256_srli_epi16(v, 4), low_mask);
317    let popcnt1 = _mm256_shuffle_epi8(lookup1, lo);
318    let popcnt2 = _mm256_shuffle_epi8(lookup2, hi);
319    _mm256_sad_epu8(popcnt1, popcnt2)
320}
321
322/// AVX2 Harley-Seal population count (4th iteration), from "Faster Population
323/// Counts using AVX2 Instructions" by Lemire, Kurz and Muła (2016),
324/// <https://arxiv.org/abs/1611.07612>.
325///
326/// `bytes.len()` must be a multiple of 32.
327#[cfg(all(
328    any(target_arch = "x86", target_arch = "x86_64"),
329    not(target_feature = "avx512vpopcntdq")
330))]
331#[target_feature(enable = "avx2")]
332// Hand-aligned: keep the 16-way CSA tree readable.
333#[rustfmt::skip]
334fn popcnt_avx2(bytes: &[u8]) -> u64 {
335    let zero = _mm256_setzero_si256();
336    let mut cnt = zero;
337    let mut ones = zero;
338    let mut twos = zero;
339    let mut fours = zero;
340    let mut eights = zero;
341    let mut twos_a;
342    let mut twos_b;
343    let mut fours_a;
344    let mut fours_b;
345    let mut eights_a;
346    let mut eights_b;
347    let mut sixteens;
348
349    // 16 vectors (512 bytes) per iteration.
350    let (blocks, tail) = bytes.as_chunks::<512>();
351    for chunk in blocks {
352        let p = chunk.as_ptr() as *const __m256i;
353        // SAFETY: `chunk` is 512 bytes, so all 16 loads (32 bytes each) are in bounds.
354        unsafe {
355            (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(0)), _mm256_loadu_si256(p.add(1)));
356            (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(2)), _mm256_loadu_si256(p.add(3)));
357            (fours_a, twos) = csa256(twos, twos_a, twos_b);
358            (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(4)), _mm256_loadu_si256(p.add(5)));
359            (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(6)), _mm256_loadu_si256(p.add(7)));
360            (fours_b, twos) = csa256(twos, twos_a, twos_b);
361            (eights_a, fours) = csa256(fours, fours_a, fours_b);
362            (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(8)), _mm256_loadu_si256(p.add(9)));
363            (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(10)), _mm256_loadu_si256(p.add(11)));
364            (fours_a, twos) = csa256(twos, twos_a, twos_b);
365            (twos_a, ones) = csa256(ones, _mm256_loadu_si256(p.add(12)), _mm256_loadu_si256(p.add(13)));
366            (twos_b, ones) = csa256(ones, _mm256_loadu_si256(p.add(14)), _mm256_loadu_si256(p.add(15)));
367            (fours_b, twos) = csa256(twos, twos_a, twos_b);
368            (eights_b, fours) = csa256(fours, fours_a, fours_b);
369            (sixteens, eights) = csa256(eights, eights_a, eights_b);
370            cnt = _mm256_add_epi64(cnt, popcnt256(sixteens));
371        }
372    }
373
374    cnt = _mm256_slli_epi64(cnt, 4);
375    cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(eights), 3));
376    cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(fours), 2));
377    cnt = _mm256_add_epi64(cnt, _mm256_slli_epi64(popcnt256(twos), 1));
378    cnt = _mm256_add_epi64(cnt, popcnt256(ones));
379
380    // Remaining whole 32-byte vectors.
381    let (vecs, _) = tail.as_chunks::<32>();
382    for chunk in vecs {
383        let v = unsafe { _mm256_loadu_si256(chunk.as_ptr() as *const __m256i) };
384        cnt = _mm256_add_epi64(cnt, popcnt256(v));
385    }
386
387    // Sum the four 64-bit lanes.
388    // SAFETY: `__m256i` and `[u64; 4]` are both 32 bytes with no invalid bit patterns.
389    let lanes: [u64; 4] = unsafe { core::mem::transmute(cnt) };
390    lanes[0] + lanes[1] + lanes[2] + lanes[3]
391}
392
393// ── AVX512 ──────────────────────────────────────────────────────────────────
394
395/// AVX512-VPOPCNTDQ population count, handling any length: a 4×-unrolled
396/// 256-byte loop, then a 64-byte loop, then a masked load for the final
397/// 1..=63 bytes.
398#[cfg(all(
399    any(target_arch = "x86", target_arch = "x86_64"),
400    any(not(target_feature = "avx2"), target_feature = "avx512vpopcntdq")
401))]
402#[target_feature(enable = "avx512f,avx512bw,avx512vpopcntdq")]
403fn popcnt_avx512(bytes: &[u8]) -> u64 {
404    let mut cnt = _mm512_setzero_si512();
405
406    // 4× unrolled 64-byte loop (256 bytes per iteration).
407    let (blocks, tail256) = bytes.as_chunks::<256>();
408    for chunk in blocks {
409        let p = chunk.as_ptr();
410        // SAFETY: `chunk` is 256 bytes, so the four 64-byte loads are in bounds.
411        unsafe {
412            let v0 = _mm512_loadu_si512(p.add(0) as *const _);
413            let v1 = _mm512_loadu_si512(p.add(64) as *const _);
414            let v2 = _mm512_loadu_si512(p.add(128) as *const _);
415            let v3 = _mm512_loadu_si512(p.add(192) as *const _);
416            cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v0));
417            cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v1));
418            cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v2));
419            cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v3));
420        }
421    }
422
423    // Remaining complete 64-byte blocks.
424    let (vecs, tail64) = tail256.as_chunks::<64>();
425    for chunk in vecs {
426        let v = unsafe { _mm512_loadu_si512(chunk.as_ptr() as *const _) };
427        cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v));
428    }
429
430    // Masked load for the final 1..=63 bytes.
431    if !tail64.is_empty() {
432        let len = tail64.len();
433        // Mask covering the final `len` bytes.
434        let mask = (u64::MAX >> (64 - len)) as __mmask64;
435        // SAFETY: the mask selects only the `len` valid bytes; masked-off lanes
436        // are not accessed.
437        unsafe {
438            let v = _mm512_maskz_loadu_epi8(mask, tail64.as_ptr() as *const _);
439            cnt = _mm512_add_epi64(cnt, _mm512_popcnt_epi64(v));
440        }
441    }
442
443    _mm512_reduce_add_epi64(cnt) as u64
444}
445
446// ════════════════════════════════════════════════════════════════════════════
447// AArch64
448// ════════════════════════════════════════════════════════════════════════════
449
450#[cfg(target_arch = "aarch64")]
451fn popcnt_aarch64(bytes: &[u8]) -> u64 {
452    // Compile-time SVE path.
453    #[cfg(all(target_feature = "sve", simd_popcnt_have_sve))]
454    {
455        unsafe { popcnt_arm_sve(bytes) }
456    }
457
458    // NEON baseline; `popcnt_neon` dispatches to SVE at runtime when available.
459    #[cfg(not(all(target_feature = "sve", simd_popcnt_have_sve)))]
460    {
461        popcnt_neon(bytes)
462    }
463}
464
465#[cfg(target_arch = "aarch64")]
466#[inline]
467fn vpadalq(sum: uint64x2_t, t: uint8x16_t) -> uint64x2_t {
468    unsafe { vpadalq_u32(sum, vpaddlq_u16(vpaddlq_u8(t))) }
469}
470
471#[cfg(target_arch = "aarch64")]
472fn popcnt_neon(bytes: &[u8]) -> u64 {
473    // Runtime SVE dispatch (present only when the build probe enabled SVE).
474    #[cfg(simd_popcnt_have_sve)]
475    if is_aarch64_feature_detected!("sve") {
476        return unsafe { popcnt_arm_sve(bytes) };
477    }
478
479    // NEON path.
480    const CHUNK: usize = 64;
481    let mut cnt = 0u64;
482    let iters = bytes.len() / CHUNK;
483    let ptr = bytes.as_ptr();
484
485    if iters > 0 {
486        // SAFETY: `iters = len / 64`, so every `vld4q_u8` at `i * 64` (i < iters)
487        // reads 64 in-bounds bytes; the final store targets a local array.
488        unsafe {
489            let mut sum = vdupq_n_u64(0);
490            let zero = vdupq_n_u8(0);
491            let mut i = 0usize;
492
493            while i < iters {
494                let mut t0 = zero;
495                let mut t1 = zero;
496                let mut t2 = zero;
497                let mut t3 = zero;
498
499                // Accumulate at most 31 chunks before draining into `sum`:
500                // 31 × 8 bits = 248 ≤ 255 guarantees no u8 lane overflow.
501                let limit = (i + 31).min(iters);
502                while i < limit {
503                    let input = vld4q_u8(ptr.add(i * CHUNK));
504                    t0 = vaddq_u8(t0, vcntq_u8(input.0));
505                    t1 = vaddq_u8(t1, vcntq_u8(input.1));
506                    t2 = vaddq_u8(t2, vcntq_u8(input.2));
507                    t3 = vaddq_u8(t3, vcntq_u8(input.3));
508                    i += 1;
509                }
510
511                sum = vpadalq(sum, t0);
512                sum = vpadalq(sum, t1);
513                sum = vpadalq(sum, t2);
514                sum = vpadalq(sum, t3);
515            }
516
517            let mut tmp = [0u64; 2];
518            vst1q_u64(tmp.as_mut_ptr(), sum);
519            cnt += tmp[0] + tmp[1];
520        }
521    }
522
523    // Scalar tail. On AArch64 `count_ones()` always lowers to NEON `cnt`, so no
524    // POPCNT runtime check is needed here.
525    let rest = &bytes[iters * CHUNK..];
526    cnt += popcnt_scalar_loop!(rest);
527    cnt
528}
529
530// ── ARM SVE ─────────────────────────────────────────────────────────────────
531
532/// SVE population count: a 4×-unrolled main loop over full vectors, then a
533/// predicated tail loop that needs no separate scalar remainder.
534#[cfg(all(target_arch = "aarch64", simd_popcnt_have_sve))]
535#[target_feature(enable = "sve")]
536fn popcnt_arm_sve(bytes: &[u8]) -> u64 {
537    // SAFETY: the loop bound keeps each full load within `len`; the tail loop's
538    // predicate masks off any lanes past the end.
539    unsafe {
540        let mut i = 0usize;
541        let mut vcnt = svdup_n_u64(0);
542        let vl = svcntb() as usize; // SVE vector length in bytes (hardware-defined)
543        let ptr = bytes.as_ptr();
544        let len = bytes.len();
545
546        // 4× unrolled full-predicate loop.
547        while i + vl * 4 <= len {
548            let v0 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i)));
549            let v1 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl)));
550            let v2 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 2)));
551            let v3 = svreinterpret_u64_u8(svld1_u8(svptrue_b8(), ptr.add(i + vl * 3)));
552            vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v0));
553            vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v1));
554            vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v2));
555            vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v3));
556            i += vl * 4;
557        }
558
559        // Predicated tail: the load zero-fills inactive lanes, so no separate
560        // scalar remainder is needed.
561        let mut pg = svwhilelt_b8_u64(i as u64, len as u64);
562        while svptest_any(svptrue_b8(), pg) {
563            let v = svreinterpret_u64_u8(svld1_u8(pg, ptr.add(i)));
564            vcnt = svadd_u64_x(svptrue_b64(), vcnt, svcnt_u64_x(svptrue_b64(), v));
565            i += vl;
566            pg = svwhilelt_b8_u64(i as u64, len as u64);
567        }
568
569        svaddv_u64(svptrue_b64(), vcnt)
570    }
571}
572
573// ════════════════════════════════════════════════════════════════════════════
574// Tests
575// ════════════════════════════════════════════════════════════════════════════
576
577#[cfg(test)]
578mod tests {
579    use super::*;
580
581    /// Reference implementation: count bits one byte at a time.
582    fn reference(bytes: &[u8]) -> u64 {
583        bytes.iter().map(|b| b.count_ones() as u64).sum()
584    }
585
586    /// Independent integer-only popcount oracle (does not use `count_ones`),
587    /// so the sweep cross-checks the crate against a different algorithm.
588    fn popcnt64_bitwise(x: u64) -> u64 {
589        const M1: u64 = 0x5555555555555555;
590        const M2: u64 = 0x3333333333333333;
591        const M4: u64 = 0x0F0F0F0F0F0F0F0F;
592        const H01: u64 = 0x0101010101010101;
593        let x = x - ((x >> 1) & M1);
594        let x = (x & M2) + ((x >> 2) & M2);
595        let x = (x + (x >> 4)) & M4;
596        x.wrapping_mul(H01) >> 56
597    }
598
599    #[test]
600    fn empty() {
601        assert_eq!(popcnt(&[]), 0);
602    }
603
604    #[test]
605    fn all_ones() {
606        for &size in &[
607            0, 1, 7, 8, 31, 32, 39, 40, 63, 64, 255, 256, 511, 512, 4095, 4096, 65537,
608        ] {
609            let bytes = vec![0xFFu8; size];
610            assert_eq!(popcnt(&bytes), size as u64 * 8, "size={size}");
611        }
612    }
613
614    #[test]
615    fn all_zeros() {
616        let bytes = vec![0u8; 65536];
617        assert_eq!(popcnt(&bytes), 0);
618    }
619
620    #[test]
621    fn single_bits() {
622        for bit in 0u64..64 {
623            let val = 1u64 << bit;
624            assert_eq!(popcnt(&val.to_le_bytes()), 1, "bit={bit}");
625        }
626    }
627
628    /// `PopcntExt::popcnt` on each integer width must equal the per-element
629    /// `count_ones()` sum (an oracle independent of the byte reinterpretation).
630    #[test]
631    fn ext_trait_widths() {
632        let u8s: &[u8] = &[0xFF, 0x0F, 0x00, 0xAB, 0x01];
633        assert_eq!(
634            u8s.popcnt(),
635            u8s.iter().map(|x| x.count_ones() as u64).sum()
636        );
637
638        let u16s: &[u16] = &[0xFFFF, 0x0F0F, 0x1234, 0];
639        assert_eq!(
640            u16s.popcnt(),
641            u16s.iter().map(|x| x.count_ones() as u64).sum()
642        );
643
644        let u32s: &[u32] = &[u32::MAX, 0, 0x8000_0001];
645        assert_eq!(
646            u32s.popcnt(),
647            u32s.iter().map(|x| x.count_ones() as u64).sum()
648        );
649
650        let u64s: &[u64] = &[u64::MAX, 0x0F0F_0F0F_0F0F_0F0F, 0];
651        assert_eq!(
652            u64s.popcnt(),
653            u64s.iter().map(|x| x.count_ones() as u64).sum()
654        );
655
656        // Signed types and arrays resolve through the same impls (the doc
657        // example covers `Vec`).
658        let i32s = [-1i32, 0, 1, i32::MIN];
659        assert_eq!(
660            i32s.popcnt(),
661            i32s.iter().map(|x| x.count_ones() as u64).sum()
662        );
663        assert_eq!([u128::MAX, 0].popcnt(), 128);
664    }
665
666    /// Sweep every boundary-relevant size against the byte-wise reference using
667    /// a deterministic pseudo-random fill (xorshift). Covers tail handling,
668    /// the AVX2/AVX512 thresholds and multiple Harley-Seal outer iterations.
669    #[test]
670    fn pseudorandom_all_sizes() {
671        let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
672        let mut next = || {
673            state ^= state << 13;
674            state ^= state >> 7;
675            state ^= state << 17;
676            state
677        };
678
679        // Largest size + largest offset exercised below, plus margin. 4695
680        // bytes spans several 512-byte Harley-Seal iterations.
681        const MAX_SIZE: usize = 4695;
682        const MAX_OFF: usize = 7;
683        let mut bytes = vec![0u8; MAX_SIZE + MAX_OFF + 1];
684        for b in bytes.iter_mut() {
685            *b = (next() & 0xFF) as u8;
686        }
687
688        // Every size from 0 up through the AVX2/AVX512 active range, plus a few
689        // larger ones, exercised at multiple start offsets so alignment varies.
690        let sizes =
691            (0usize..=600).chain([1023, 1024, 1025, 2048, 4095, 4096, 4097, 4608, MAX_SIZE]);
692        for size in sizes {
693            for &off in &[0usize, 1, 3, MAX_OFF] {
694                let slice = &bytes[off..off + size];
695                assert_eq!(popcnt(slice), reference(slice), "size={size} off={off}");
696            }
697        }
698    }
699
700    /// Verify `popcnt()` of every suffix `bytes[i..]` against an independent
701    /// byte-wise reference, covering every length and a range of start
702    /// alignments in one sweep.
703    ///
704    /// Size defaults to 20_000 to keep `cargo test` fast — the sweep is O(n²) in
705    /// the work `popcnt` performs. Override with `SIMD_POPCNT_TEST_SIZE` for a
706    /// heavier run, e.g. `SIMD_POPCNT_TEST_SIZE=100000 cargo test --release suffix_sweep`.
707    #[test]
708    fn suffix_sweep() {
709        let size = std::env::var("SIMD_POPCNT_TEST_SIZE")
710            .ok()
711            .and_then(|s| s.parse::<usize>().ok())
712            .unwrap_or(20_000);
713
714        // All-ones array.
715        let ones = vec![0xFFu8; size];
716        check_all_suffixes(&ones);
717
718        // Deterministic pseudo-random array (fixed seed → reproducible failures).
719        let mut state: u64 = 0x2545_F491_4F6C_DD1D;
720        let mut bytes = vec![0u8; size];
721        for b in bytes.iter_mut() {
722            state ^= state << 13;
723            state ^= state >> 7;
724            state ^= state << 17;
725            *b = state as u8;
726        }
727        check_all_suffixes(&bytes);
728    }
729
730    /// Assert `popcnt(&bytes[i..])` for every `i` against an O(1) prefix-sum
731    /// reference, so only `popcnt` itself does O(n) work per suffix.
732    fn check_all_suffixes(bytes: &[u8]) {
733        let total: u64 = bytes.iter().map(|&b| popcnt64_bitwise(b as u64)).sum();
734        let mut prefix = 0u64; // popcount of bytes[..i]
735        for (i, &byte) in bytes.iter().enumerate() {
736            assert_eq!(popcnt(&bytes[i..]), total - prefix, "suffix at offset {i}");
737            prefix += popcnt64_bitwise(byte as u64);
738        }
739        // Empty suffix.
740        assert_eq!(popcnt(&bytes[bytes.len()..]), 0);
741    }
742}