From 7c4c0376907ff174f98cb57ca1c71d1ffd46a199 Mon Sep 17 00:00:00 2001 From: Dan Draper Date: Sat, 13 Jun 2026 01:16:29 +1000 Subject: [PATCH 1/3] perf: NEON and AVX2 kernels for indicator and hash-bit masks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds primitives::simd with the two vectorisable kernels from the right- ciphertext encoder (v2 plan §3), wired into indicator_mask_xor and hash_all_into: - gt_mask_xor_256: bytewise greater-than against a broadcast value packed to a bitmask. NEON (vcgtq_u8 + and-with-bitpos + vaddv_u8 pack; the shrn trick benched slower on M1 for this exact bit layout) and AVX2 (vpcmpgtb with 0x80 bias for unsigned semantics + vpmovmskb). ~30x over the pre-PR-3 per-bit shape, ~4x over the scalar bulk form in isolation. - lsb_mask_256: stride-16 gather of AES-output LSBs. NEON via ld1q_x4 + constant-index vqtbl4q (the textbook ld4+uzp1 tree benched ~25% slower on M1); scalar elsewhere (AVX2 gains nothing without AVX-512 VBMI). Dispatch: compile-time on aarch64 (NEON is baseline), cached runtime detection on x86_64, length-generic scalar fallback everywhere else. All paths branch-free with fixed trip counts and constant tbl indices (constant-time discipline documented in the module). Equivalence pinned three ways: dispatched-vs-scalar randomized tests in the module, mask-vs-per-bit-invert quickcheck in prp.rs, and the PR 1 byte vectors end-to-end (which on this machine exercise the NEON path). Per-u64 gain on M1 Max is ~0.5% over PR 3 (PRP construction dominates at Bit8); these kernels are sized for the Bit6 scheme where per-block AES work shrinks 4x. Microarch caveat: pack/gather strategy choices are M1-tuned; re-bench on Neoverse before locking server-ARM dispatch. Part of the ORE v2 program (docs/plans/2026-06-12-ore-v2-architecture.md, PR 4). --- packages/ore-rs/src/primitives.rs | 37 +--- packages/ore-rs/src/primitives/hash.rs | 19 +- packages/ore-rs/src/primitives/prp.rs | 15 +- packages/ore-rs/src/primitives/simd.rs | 243 +++++++++++++++++++++++++ packages/ore-rs/src/scheme/width.rs | 11 +- 5 files changed, 272 insertions(+), 53 deletions(-) create mode 100644 packages/ore-rs/src/primitives/simd.rs diff --git a/packages/ore-rs/src/primitives.rs b/packages/ore-rs/src/primitives.rs index b7718290..9a0f1896 100644 --- a/packages/ore-rs/src/primitives.rs +++ b/packages/ore-rs/src/primitives.rs @@ -1,6 +1,7 @@ pub mod hash; pub mod prf; pub mod prp; +pub mod simd; use aes::cipher::{consts::U16, generic_array::GenericArray}; use aes::Block; @@ -10,36 +11,6 @@ pub type PrfKey = GenericArray; pub type HashKey = GenericArray; pub const NONCE_SIZE: usize = 16; -/// Pack one bit per element of `src` into `out`, LSB-first, eight elements per -/// byte: the bit derived from `src[j]` lands in `out[j / 8]` at position -/// `j % 8` — the bit order used by `RightBlock32::set_bit`. `bit` extracts the -/// 0/1 value of each element; `merge` folds each packed byte into the existing -/// `out` byte (assign for a fresh buffer, XOR to overlay). This is the single -/// home of the right-ciphertext bit-packing convention, shared by -/// [`Hash::hash_all_into`] and [`Prp::indicator_mask_xor`]. -/// -/// Panics unless `out.len() * 8 == src.len()`, which also guarantees -/// `chunks_exact(8)` consumes `src` with no dropped remainder. -pub(crate) fn pack_bits_lsb_first( - out: &mut [u8], - src: &[T], - bit: impl Fn(&T) -> u8, - merge: impl Fn(&mut u8, u8), -) { - assert_eq!( - out.len() * 8, - src.len(), - "pack_bits_lsb_first: out.len()*8 must equal src.len()" - ); - for (slot, chunk) in out.iter_mut().zip(src.chunks_exact(8)) { - let mut byte = 0u8; - for (i, elem) in chunk.iter().enumerate() { - byte |= (bit(elem) & 1) << i; - } - merge(slot, byte); - } -} - pub trait Prf { fn new(key: &PrfKey) -> Self; fn encrypt_all(&self, data: &mut [AesBlock]); @@ -69,9 +40,9 @@ pub trait Prp: Sized { fn invert(&self, data: T) -> PrpResult; /// XOR the indicator mask for `data` into `out`: bit `j` of the mask is - /// `1` iff `invert(j) > data`. Bit order matches `RightBlock32::set_bit` - /// (LSB-first within each byte). `out.len() * 8` must equal the - /// permutation domain. + /// `1` iff `invert(j) > data`. Bit order matches + /// `RightBitVec::set_bit` (LSB-first within each byte). `out.len() * 8` + /// must equal the permutation domain. /// /// This is the bulk form of the per-`j` `invert`-and-compare loop the /// right-ciphertext encoder needs; implementations walk their inverse diff --git a/packages/ore-rs/src/primitives/hash.rs b/packages/ore-rs/src/primitives/hash.rs index d78b506e..9321e1f3 100644 --- a/packages/ore-rs/src/primitives/hash.rs +++ b/packages/ore-rs/src/primitives/hash.rs @@ -32,15 +32,18 @@ impl Hash for Aes128Z2Hash { } fn hash_all_into(&self, data: &mut [AesBlock], out: &mut [u8]) { + debug_assert_eq!(out.len() * 8, data.len()); self.cipher.encrypt_blocks(data); - // Pack the Z2 (1-bit) outputs LSB-first, eight blocks per byte (assign - // into a fresh block); see `pack_bits_lsb_first` for the convention. - crate::primitives::pack_bits_lsb_first( - out, - data, - |block| block[0], - |slot, byte| *slot = byte, - ); + + // Pack the Z2 (1-bit) outputs LSB-first, eight blocks per byte — + // the same bit order as `RightBitVec::set_bit`. The 256-block case + // (Bit8's per-block RO output) has a vectorised gather; other sizes + // use the scalar pack. + if data.len() == 256 { + crate::primitives::simd::lsb_mask_256(data, out); + } else { + crate::primitives::simd::scalar::lsb_mask(data, out); + } } } diff --git a/packages/ore-rs/src/primitives/prp.rs b/packages/ore-rs/src/primitives/prp.rs index 3881ee36..829d68da 100644 --- a/packages/ore-rs/src/primitives/prp.rs +++ b/packages/ore-rs/src/primitives/prp.rs @@ -82,17 +82,12 @@ impl Prp for KnuthShufflePRP { } fn indicator_mask_xor(&self, data: u8, out: &mut [u8]) { + debug_assert_eq!(out.len() * 8, 256); + // `invert(j)` is `self.permutation[j]` (see `invert` above), so the - // mask is one linear pass over the table: a bytewise `> data` compare, - // XORed over the hash bits. Branch-free with a fixed trip count — the - // scalar form of a SIMD compare-and-movemask (v2 plan §3). Bit order - // is shared with `hash_all_into` via `pack_bits_lsb_first`. - crate::primitives::pack_bits_lsb_first( - out, - &self.permutation, - |&p| u8::from(p > data), - |slot, byte| *slot ^= byte, - ); + // mask is one pass over the table: a bytewise `> data` compare + // packed to bits — vectorised where the target supports it. + crate::primitives::simd::gt_mask_xor_256(&self.permutation, data, out); } } diff --git a/packages/ore-rs/src/primitives/simd.rs b/packages/ore-rs/src/primitives/simd.rs new file mode 100644 index 00000000..5db82edc --- /dev/null +++ b/packages/ore-rs/src/primitives/simd.rs @@ -0,0 +1,243 @@ +//! Vectorised kernels for the data-parallel inner loops of right-ciphertext +//! encoding, with portable scalar fallbacks (v2 plan §3). +//! +//! Two operations are covered: +//! +//! - **`gt_mask_xor_256`** — XOR the indicator mask `bit j = (table[j] > x)` +//! into a 32-byte bitvector. This is the bulk form of the per-block +//! unary encoding (a bytewise compare against a broadcast value, packed +//! to a bitmask — `cmgt`+pack on NEON, `vpcmpgtb`+`vpmovmskb` on AVX2). +//! - **`lsb_mask_256`** — pack `bit j = blocks[j][0] & 1` (a stride-16 +//! gather over AES outputs) into a 32-byte bitvector. +//! +//! Bit order everywhere is LSB-first within each byte, matching +//! `RightBitVec::set_bit`. +//! +//! # Dispatch +//! +//! - **aarch64:** NEON is baseline — compile-time dispatch, no detection. +//! - **x86_64:** AVX2 via `is_x86_feature_detected!` (cached by `std`) for +//! the indicator mask; the LSB gather stays scalar (the stride-16 gather +//! doesn't pay on AVX2 without AVX-512 VBMI, and the scalar form is +//! already branch-free and fast). +//! - **everything else:** scalar. +//! +//! # Constant-time discipline +//! +//! Every path in this module — scalar and vector — has a fixed trip count, +//! no data-dependent branches, and no data-dependent memory addressing +//! (the NEON `tbl` indices are compile-time constants). Equivalence +//! between scalar and vector paths is pinned by tests below; equivalence +//! against the per-bit reference is pinned in `prp.rs`. + +use super::AesBlock; + +/// XOR the indicator mask for `x` over `table` into `out`: +/// `out[j/8] ^= ((table[j] > x) as u8) << (j%8)`. +#[inline] +pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { + debug_assert_eq!(out.len(), 32); + + #[cfg(target_arch = "aarch64")] + // SAFETY: NEON is baseline on aarch64; `out` length asserted above. + unsafe { + neon::gt_mask_xor_256(table, x, out); + return; + } + + #[cfg(target_arch = "x86_64")] + if is_x86_feature_detected!("avx2") { + // SAFETY: AVX2 presence just checked; `out` length asserted above. + unsafe { + avx2::gt_mask_xor_256(table, x, out); + return; + } + } + + #[allow(unreachable_code)] + scalar::gt_mask_xor_256(table, x, out); +} + +/// Pack the LSB of byte 0 of each of 256 AES blocks into `out`: +/// `out[j/8] |= (blocks[j][0] & 1) << (j%8)` over zeroed positions +/// (callers pass `out` slots they own; bits are assigned, not accumulated). +#[inline] +pub(crate) fn lsb_mask_256(blocks: &[AesBlock], out: &mut [u8]) { + debug_assert_eq!(blocks.len(), 256); + debug_assert_eq!(out.len(), 32); + + #[cfg(target_arch = "aarch64")] + // SAFETY: NEON is baseline on aarch64; lengths asserted above. + unsafe { + neon::lsb_mask_256(blocks, out); + return; + } + + #[allow(unreachable_code)] + scalar::lsb_mask(blocks, out); +} + +pub(crate) mod scalar { + use super::AesBlock; + + pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { + for (slot, chunk) in out.iter_mut().zip(table.chunks_exact(8)) { + let mut byte = 0u8; + for (bit, &p) in chunk.iter().enumerate() { + byte |= u8::from(p > x) << bit; + } + *slot ^= byte; + } + } + + /// Length-generic scalar LSB pack: `blocks.len()` must be a multiple + /// of 8 and equal to `out.len() * 8`. + pub(crate) fn lsb_mask(blocks: &[AesBlock], out: &mut [u8]) { + for (slot, chunk) in out.iter_mut().zip(blocks.chunks_exact(8)) { + let mut byte = 0u8; + for (bit, block) in chunk.iter().enumerate() { + byte |= (block[0] & 1u8) << bit; + } + *slot = byte; + } + } +} + +#[cfg(target_arch = "aarch64")] +mod neon { + use super::AesBlock; + use core::arch::aarch64::*; + + /// Lane k holds `1 << (k % 8)`: the movemask bit-position vector. + const BITPOS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128]; + + /// Pack a 16-lane 0x00/0xFF compare mask into 2 bytes, LSB-first. Each + /// lane contributes a distinct power of two within its 8-lane half, so + /// the horizontal add is exactly the movemask byte. + #[inline(always)] + unsafe fn pack_addv(m: uint8x16_t) -> (u8, u8) { + let mb = vandq_u8(m, vld1q_u8(BITPOS.as_ptr())); + (vaddv_u8(vget_low_u8(mb)), vaddv_u8(vget_high_u8(mb))) + } + + #[target_feature(enable = "neon")] + pub(super) unsafe fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { + let xv = vdupq_n_u8(x); + for i in 0..16 { + let t = vld1q_u8(table.as_ptr().add(i * 16)); + let (lo, hi) = pack_addv(vcgtq_u8(t, xv)); + out[2 * i] ^= lo; + out[2 * i + 1] ^= hi; + } + } + + /// `IDX[g]` places the 4 block-start bytes of 64-byte group `g` + /// (offsets 0, 16, 32, 48) at lanes 4g..4g+3; all other lanes are 0xFF + /// (out-of-range, yielding 0 in `tbl`). + const TBL_IDX: [[u8; 16]; 4] = { + let mut idx = [[0xFFu8; 16]; 4]; + let mut g = 0; + while g < 4 { + let mut k = 0; + while k < 4 { + idx[g][4 * g + k] = (16 * k) as u8; + k += 1; + } + g += 1; + } + idx + }; + + /// Gather byte 0 of 16 consecutive 16-byte blocks (256 contiguous + /// bytes) into one vector via constant-index 4-register table lookups. + #[inline(always)] + unsafe fn gather_block0_x16(base: *const u8) -> uint8x16_t { + let r0 = vqtbl4q_u8(vld1q_u8_x4(base), vld1q_u8(TBL_IDX[0].as_ptr())); + let r1 = vqtbl4q_u8(vld1q_u8_x4(base.add(64)), vld1q_u8(TBL_IDX[1].as_ptr())); + let r2 = vqtbl4q_u8(vld1q_u8_x4(base.add(128)), vld1q_u8(TBL_IDX[2].as_ptr())); + let r3 = vqtbl4q_u8(vld1q_u8_x4(base.add(192)), vld1q_u8(TBL_IDX[3].as_ptr())); + vorrq_u8(vorrq_u8(r0, r1), vorrq_u8(r2, r3)) + } + + #[target_feature(enable = "neon")] + pub(super) unsafe fn lsb_mask_256(blocks: &[AesBlock], out: &mut [u8]) { + let p = blocks.as_ptr() as *const u8; + let one = vdupq_n_u8(1); + for i in 0..16 { + let m = vtstq_u8(gather_block0_x16(p.add(i * 256)), one); + let (lo, hi) = pack_addv(m); + out[2 * i] = lo; + out[2 * i + 1] = hi; + } + } +} + +#[cfg(target_arch = "x86_64")] +mod avx2 { + use core::arch::x86_64::*; + + /// AVX2 has only *signed* byte compares; XOR-ing 0x80 into both sides + /// maps unsigned `>` onto signed `>`. + #[target_feature(enable = "avx2")] + pub(super) unsafe fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { + let bias = _mm256_set1_epi8(i8::MIN); + let xv = _mm256_xor_si256(_mm256_set1_epi8(x as i8), bias); + for i in 0..8 { + let t = _mm256_loadu_si256(table.as_ptr().add(i * 32) as *const __m256i); + let gt = _mm256_cmpgt_epi8(_mm256_xor_si256(t, bias), xv); + let mask = _mm256_movemask_epi8(gt) as u32; + let m = mask.to_le_bytes(); + out[4 * i] ^= m[0]; + out[4 * i + 1] ^= m[1]; + out[4 * i + 2] ^= m[2]; + out[4 * i + 3] ^= m[3]; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::{Rng, SeedableRng}; + use rand_chacha::ChaCha20Rng; + + fn rand_table(rng: &mut ChaCha20Rng) -> [u8; 256] { + let mut t = [0u8; 256]; + rng.fill(&mut t[..]); + t + } + + #[test] + fn gt_mask_dispatched_matches_scalar() { + let mut rng = ChaCha20Rng::seed_from_u64(7); + for _ in 0..5_000 { + let table = rand_table(&mut rng); + let x: u8 = rng.gen(); + // Random starting contents so the XOR semantics are exercised. + let mut a = [0u8; 32]; + rng.fill(&mut a[..]); + let mut b = a; + + gt_mask_xor_256(&table, x, &mut a); + scalar::gt_mask_xor_256(&table, x, &mut b); + assert_eq!(a, b); + } + } + + #[test] + fn lsb_mask_dispatched_matches_scalar() { + let mut rng = ChaCha20Rng::seed_from_u64(11); + for _ in 0..5_000 { + let mut blocks = [AesBlock::default(); 256]; + for block in blocks.iter_mut() { + rng.fill(block.as_mut_slice()); + } + let mut a = [0u8; 32]; + let mut b = [0u8; 32]; + + lsb_mask_256(&blocks, &mut a); + scalar::lsb_mask(&blocks, &mut b); + assert_eq!(a, b); + } + } +} diff --git a/packages/ore-rs/src/scheme/width.rs b/packages/ore-rs/src/scheme/width.rs index 01f9a370..65ad39d1 100644 --- a/packages/ore-rs/src/scheme/width.rs +++ b/packages/ore-rs/src/scheme/width.rs @@ -50,16 +50,23 @@ impl AesBlockBuf for [AesBlock; 256] { /// Per-block bitvector operations on a Right ciphertext block, one bit per /// value in the block domain. pub trait RightBitVec { + /// Set bit `bit` to `value` (`0` or `1`). (Bulk encoding writes via + /// [`Self::as_mut_bytes`]; this remains for width-generic callers.) + #[allow(dead_code)] + fn set_bit(&mut self, bit: usize, value: u8); /// Read bit `bit`. (The width-generic comparator lands with the Bit6 /// scheme; the legacy comparator calls the inherent method.) #[allow(dead_code)] fn get_bit(&self, bit: usize) -> u8; - /// The raw bitvector bytes, LSB-first within each byte (bit `j` lives in - /// byte `j / 8` at position `j % 8`), for bulk mask construction. + /// The raw bitvector bytes, LSB-first within each byte (the same bit + /// order as [`Self::set_bit`]), for bulk mask construction. fn as_mut_bytes(&mut self) -> &mut [u8]; } impl RightBitVec for RightBlock32 { + fn set_bit(&mut self, bit: usize, value: u8) { + RightBlock32::set_bit(self, bit, value) + } fn get_bit(&self, bit: usize) -> u8 { RightBlock32::get_bit(self, bit) } From 893794c0b679c40ef83f34ff31e46df0008674bc Mon Sep 17 00:00:00 2001 From: Dan Draper Date: Tue, 16 Jun 2026 15:10:13 +1000 Subject: [PATCH 2/3] harden(simd): assert gather length; precise dispatch cfg (#81 self-review) Code-review follow-ups on the SIMD PR: - lsb_mask_256: promote the blocks.len()==256 / out.len()==32 debug_asserts to real assert_eq!. The NEON path gathers `blocks` via raw pointers assuming exactly 256 blocks, so a shorter slice would read out of bounds (UB) in a release build; the assert enforces the unsafe precondition at the boundary rather than relying on the single caller's length check. - Replace the blanket #[allow(unreachable_code)] on both dispatchers' scalar tails with #[cfg(not(target_arch = "aarch64"))], and drop the now-redundant aarch64 return. Paths are compile-time mutually exclusive, so a future bad early-return surfaces as a real diagnostic instead of being silently masked. - Scope scalar::gt_mask_xor_256 to cfg(any(not(aarch64), test)) so it isn't dead code on aarch64 (only reached via the non-aarch64 dispatcher or tests); no allow needed. scalar::lsb_mask stays always-compiled (hash.rs uses it directly for non-256 inputs). AVX2 pre-merge test coverage gap tracked in #87. No wire change (compat vectors byte-identical). --- packages/ore-rs/src/primitives/simd.rs | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/packages/ore-rs/src/primitives/simd.rs b/packages/ore-rs/src/primitives/simd.rs index 5db82edc..4aad8341 100644 --- a/packages/ore-rs/src/primitives/simd.rs +++ b/packages/ore-rs/src/primitives/simd.rs @@ -42,7 +42,6 @@ pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { // SAFETY: NEON is baseline on aarch64; `out` length asserted above. unsafe { neon::gt_mask_xor_256(table, x, out); - return; } #[cfg(target_arch = "x86_64")] @@ -54,7 +53,9 @@ pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { } } - #[allow(unreachable_code)] + // Scalar fallback; excluded on aarch64 (the NEON path above always handles + // it) so the dead tail needs no blanket `#[allow(unreachable_code)]`. + #[cfg(not(target_arch = "aarch64"))] scalar::gt_mask_xor_256(table, x, out); } @@ -63,23 +64,33 @@ pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { /// (callers pass `out` slots they own; bits are assigned, not accumulated). #[inline] pub(crate) fn lsb_mask_256(blocks: &[AesBlock], out: &mut [u8]) { - debug_assert_eq!(blocks.len(), 256); - debug_assert_eq!(out.len(), 32); + // Real asserts (not debug_assert): the NEON path gathers `blocks` via raw + // pointers assuming exactly 256 blocks, so a shorter slice would read out + // of bounds (UB) in a release build without this check. + assert_eq!(blocks.len(), 256); + assert_eq!(out.len(), 32); #[cfg(target_arch = "aarch64")] // SAFETY: NEON is baseline on aarch64; lengths asserted above. unsafe { neon::lsb_mask_256(blocks, out); - return; } - #[allow(unreachable_code)] + // Scalar fallback; excluded on aarch64 (the NEON path above always handles + // it) so the dead tail needs no blanket `#[allow(unreachable_code)]`. + #[cfg(not(target_arch = "aarch64"))] scalar::lsb_mask(blocks, out); } pub(crate) mod scalar { use super::AesBlock; + // Used by the non-aarch64 `gt_mask_xor_256` dispatcher and as the test + // oracle; on aarch64 the lib always reaches NEON, so outside `cfg(test)` + // this fn is unreferenced there. Compiling it exactly where it's used keeps + // it from being dead code without an `allow`. (`lsb_mask` below stays + // always compiled — `hash.rs` calls it directly for non-256 inputs.) + #[cfg(any(not(target_arch = "aarch64"), test))] pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { for (slot, chunk) in out.iter_mut().zip(table.chunks_exact(8)) { let mut byte = 0u8; From 773d717fb3ec2290028659a7d473593cc2ac7392 Mon Sep 17 00:00:00 2001 From: Dan Draper Date: Thu, 8 Oct 2026 21:46:03 +1100 Subject: [PATCH 3/3] harden(simd,hash): real length asserts; set_bit assigns Review of #81: - Removing pack_bits_lsb_first (68ec6d0) dropped its release-mode length check. hash_all_into, scalar::lsb_mask and gt_mask_xor_256 now assert their output length in release too, so a mismatch fails the same way on every backend instead of the scalar path silently truncating. - RightBitVec::set_bit promised assignment but forwarded to an OR-only setter. It had no callers outside tests, so it is removed from the trait; the inherent RightBlock32::set_bit now overwrites the bit, matching its doc. Its only production caller wrote 0/1 into zeroed blocks, where OR and assign agree, so no ciphertext byte changes. --- packages/ore-rs/src/primitives.rs | 2 +- packages/ore-rs/src/primitives/hash.rs | 10 +++++-- packages/ore-rs/src/primitives/simd.rs | 27 +++++++++++++++++-- .../ore-rs/src/scheme/bit2/block_types.rs | 26 ++++++++++++++---- packages/ore-rs/src/scheme/width.rs | 9 +------ 5 files changed, 56 insertions(+), 18 deletions(-) diff --git a/packages/ore-rs/src/primitives.rs b/packages/ore-rs/src/primitives.rs index 9a0f1896..eff30941 100644 --- a/packages/ore-rs/src/primitives.rs +++ b/packages/ore-rs/src/primitives.rs @@ -41,7 +41,7 @@ pub trait Prp: Sized { /// XOR the indicator mask for `data` into `out`: bit `j` of the mask is /// `1` iff `invert(j) > data`. Bit order matches - /// `RightBitVec::set_bit` (LSB-first within each byte). `out.len() * 8` + /// `RightBlock32::set_bit` (LSB-first within each byte). `out.len() * 8` /// must equal the permutation domain. /// /// This is the bulk form of the per-`j` `invert`-and-compare loop the diff --git a/packages/ore-rs/src/primitives/hash.rs b/packages/ore-rs/src/primitives/hash.rs index 9321e1f3..829a32a5 100644 --- a/packages/ore-rs/src/primitives/hash.rs +++ b/packages/ore-rs/src/primitives/hash.rs @@ -32,11 +32,17 @@ impl Hash for Aes128Z2Hash { } fn hash_all_into(&self, data: &mut [AesBlock], out: &mut [u8]) { - debug_assert_eq!(out.len() * 8, data.len()); + // A real assert: the scalar pack below zips `out` with `data`, so a + // mismatch would silently truncate the mask in a release build. + assert_eq!( + out.len() * 8, + data.len(), + "hash_all_into: out.len() * 8 must equal data.len()" + ); self.cipher.encrypt_blocks(data); // Pack the Z2 (1-bit) outputs LSB-first, eight blocks per byte — - // the same bit order as `RightBitVec::set_bit`. The 256-block case + // the same bit order as `RightBlock32::set_bit`. The 256-block case // (Bit8's per-block RO output) has a vectorised gather; other sizes // use the scalar pack. if data.len() == 256 { diff --git a/packages/ore-rs/src/primitives/simd.rs b/packages/ore-rs/src/primitives/simd.rs index 4aad8341..33bbfc9b 100644 --- a/packages/ore-rs/src/primitives/simd.rs +++ b/packages/ore-rs/src/primitives/simd.rs @@ -11,7 +11,7 @@ //! gather over AES outputs) into a 32-byte bitvector. //! //! Bit order everywhere is LSB-first within each byte, matching -//! `RightBitVec::set_bit`. +//! `RightBlock32::set_bit`. //! //! # Dispatch //! @@ -36,7 +36,10 @@ use super::AesBlock; /// `out[j/8] ^= ((table[j] > x) as u8) << (j%8)`. #[inline] pub(crate) fn gt_mask_xor_256(table: &[u8; 256], x: u8, out: &mut [u8]) { - debug_assert_eq!(out.len(), 32); + // A real assert, so every backend sees the same contract in release: the + // scalar path would otherwise silently truncate, and the vector paths + // would panic only on a short slice. + assert_eq!(out.len(), 32); #[cfg(target_arch = "aarch64")] // SAFETY: NEON is baseline on aarch64; `out` length asserted above. @@ -104,6 +107,11 @@ pub(crate) mod scalar { /// Length-generic scalar LSB pack: `blocks.len()` must be a multiple /// of 8 and equal to `out.len() * 8`. pub(crate) fn lsb_mask(blocks: &[AesBlock], out: &mut [u8]) { + assert_eq!( + out.len() * 8, + blocks.len(), + "lsb_mask: out.len() * 8 must equal blocks.len()" + ); for (slot, chunk) in out.iter_mut().zip(blocks.chunks_exact(8)) { let mut byte = 0u8; for (bit, block) in chunk.iter().enumerate() { @@ -235,6 +243,21 @@ mod tests { } } + #[test] + #[should_panic(expected = "lsb_mask: out.len() * 8 must equal blocks.len()")] + fn scalar_lsb_mask_rejects_a_short_output() { + let blocks = [AesBlock::default(); 64]; + let mut out = [0u8; 7]; + scalar::lsb_mask(&blocks, &mut out); + } + + #[test] + #[should_panic] + fn gt_mask_xor_256_rejects_a_short_output() { + let mut out = [0u8; 31]; + gt_mask_xor_256(&[0u8; 256], 0, &mut out); + } + #[test] fn lsb_mask_dispatched_matches_scalar() { let mut rng = ChaCha20Rng::seed_from_u64(11); diff --git a/packages/ore-rs/src/scheme/bit2/block_types.rs b/packages/ore-rs/src/scheme/bit2/block_types.rs index 8502541f..308dbc22 100644 --- a/packages/ore-rs/src/scheme/bit2/block_types.rs +++ b/packages/ore-rs/src/scheme/bit2/block_types.rs @@ -17,15 +17,15 @@ pub struct RightBlock32 { } impl RightBlock32 { - /// Set bit `bit` (in `0..256`) to `value` (which must be `0` or `1`). - /// Panics in debug if `bit >= 256`. + /// Set bit `bit` (in `0..256`) to `value` (which must be `0` or `1`), + /// overwriting whatever the bit held. Panics in debug if `bit >= 256`. #[inline] pub fn set_bit(&mut self, bit: usize, value: u8) { debug_assert!(bit < 256); let byte_index = bit / 8; - let mask = bit % 8; - let v = value << mask; - self.data[byte_index] |= v; + let shift = bit % 8; + let b = &mut self.data[byte_index]; + *b = (*b & !(1 << shift)) | ((value & 1) << shift); } /// The raw bitvector bytes, LSB-first within each byte — the bit order @@ -109,6 +109,22 @@ mod tests { assert_eq!(block.get_bit(255), 1); } + #[test] + fn set_bit_overwrites_rather_than_ors() { + let mut block = RightBlock32::default(); + block.set_bit(17, 1); + block.set_bit(17, 0); + assert_eq!(block.get_bit(17), 0); + block.set_bit(18, 1); + block.set_bit(16, 1); + block.set_bit(17, 0); + assert_eq!( + (block.get_bit(16), block.get_bit(18)), + (1, 1), + "neighbours untouched" + ); + } + #[test] fn right_default_in_place_without_new_data() { let mut right = RightBlock32::default(); diff --git a/packages/ore-rs/src/scheme/width.rs b/packages/ore-rs/src/scheme/width.rs index 65ad39d1..08013208 100644 --- a/packages/ore-rs/src/scheme/width.rs +++ b/packages/ore-rs/src/scheme/width.rs @@ -50,23 +50,16 @@ impl AesBlockBuf for [AesBlock; 256] { /// Per-block bitvector operations on a Right ciphertext block, one bit per /// value in the block domain. pub trait RightBitVec { - /// Set bit `bit` to `value` (`0` or `1`). (Bulk encoding writes via - /// [`Self::as_mut_bytes`]; this remains for width-generic callers.) - #[allow(dead_code)] - fn set_bit(&mut self, bit: usize, value: u8); /// Read bit `bit`. (The width-generic comparator lands with the Bit6 /// scheme; the legacy comparator calls the inherent method.) #[allow(dead_code)] fn get_bit(&self, bit: usize) -> u8; /// The raw bitvector bytes, LSB-first within each byte (the same bit - /// order as [`Self::set_bit`]), for bulk mask construction. + /// order as `RightBlock32::set_bit`), for bulk mask construction. fn as_mut_bytes(&mut self) -> &mut [u8]; } impl RightBitVec for RightBlock32 { - fn set_bit(&mut self, bit: usize, value: u8) { - RightBlock32::set_bit(self, bit, value) - } fn get_bit(&self, bit: usize) -> u8 { RightBlock32::get_bit(self, bit) }