From 13e6c6f46ec2e5ce2ff6ae8bcf1a13a418fcfe1b Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 09:58:54 +0100 Subject: [PATCH 01/12] Avoid unaligned slice casts in GeoDiffCount and BitVec::from_bytes Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/geo_filters/src/diff_count.rs | 20 ++++--- crates/geo_filters/src/diff_count/bitvec.rs | 62 ++++++++++++++++----- 2 files changed, 60 insertions(+), 22 deletions(-) diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index 9a9b0ad..dc8c7cd 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -376,18 +376,24 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { } // The number of most significant bits stores in the MSB sparse repr let msb_len = (buf.len() / size_of::()).min(c.max_msb_len()); - let msb = unsafe { - std::mem::transmute::<&[u8], &[C::BucketType]>(std::slice::from_raw_parts( - buf.as_ptr(), - msb_len, - )) - }; // The number of bytes representing the MSB - this is how many bytes we need to // skip over to reach the LSB let msb_bytes_len = msb_len * size_of::(); + let msb_bytes = &buf[..msb_bytes_len]; + let msb = match unsafe { msb_bytes.align_to::() } { + ([], aligned, []) => Cow::Borrowed(aligned), + _ => { + let mut vec = Vec::with_capacity(msb_len); + let ptr = msb_bytes.as_ptr() as *const C::BucketType; + for i in 0..msb_len { + vec.push(unsafe { ptr.add(i).read_unaligned() }); + } + Cow::Owned(vec) + } + }; Self { config: c, - msb: Cow::Borrowed(msb), + msb, lsb: BitVec::from_bytes(&buf[msb_bytes_len..]), } } diff --git a/crates/geo_filters/src/diff_count/bitvec.rs b/crates/geo_filters/src/diff_count/bitvec.rs index cf6cb93..4793d31 100644 --- a/crates/geo_filters/src/diff_count/bitvec.rs +++ b/crates/geo_filters/src/diff_count/bitvec.rs @@ -32,7 +32,7 @@ impl PartialOrd for BitVec<'_> { } } -impl BitVec<'_> { +impl<'a> BitVec<'a> { /// Takes an iterator of `BitChunk` items as input and returns the corresponding `BitVec`. /// The order of `BitChunk`s doesn't matter for this function and `BitChunk` may be hitting /// the same block. In this case, the function will simply xor them together. @@ -175,33 +175,37 @@ impl BitVec<'_> { } #[cfg(target_endian = "little")] - pub fn from_bytes(mut buf: &[u8]) -> Self { + pub fn from_bytes(mut buf: &'a [u8]) -> Self { if buf.is_empty() { return Self::default(); } // The first byte of the serialized BitVec is used to indicate how many // of the bits in the left-most u64 block are *unoccupied*. // See [`BitVec::write`] implementation for how this is done. + let unoccupied_bits = buf[0] as usize; assert!( - buf[0] < 64, + unoccupied_bits < 64, "Number of unoccupied bits should be <64, got {}", - buf[0] + unoccupied_bits ); - let num_bits = (buf.len() - 1) * 8 - buf[0] as usize; buf = &buf[1..]; - assert_eq!( - buf.len() % BYTES_PER_BLOCK, - 0, - "buffer should be a multiple of 8 bytes, got {}", + assert!( + !buf.is_empty() && buf.len().is_multiple_of(BYTES_PER_BLOCK), + "buffer should be a non-empty multiple of 8 bytes, got {}", buf.len() ); - let blocks = unsafe { - std::mem::transmute::<&[u8], &[u64]>(std::slice::from_raw_parts( - buf.as_ptr(), - buf.len() / BYTES_PER_BLOCK, - )) + let num_bits = buf.len() * 8 - unoccupied_bits; + let blocks = match unsafe { buf.align_to::() } { + ([], aligned, []) => Cow::Borrowed(aligned), + _ => Cow::Owned( + buf.as_chunks::() + .0 + .iter() + .copied() + .map(u64::from_le_bytes) + .collect(), + ), }; - let blocks = Cow::Borrowed(blocks); Self { num_bits, blocks } } @@ -358,4 +362,32 @@ mod tests { let bitvec = bitvec_from_element(128, false); bitvec.bit_range(&(0..65)); } + + #[test] + #[cfg(target_endian = "little")] + fn test_bitvec_from_bytes_unaligned() { + let mut original = BitVec::default(); + original.resize(130); + original.toggle(0); + original.toggle(63); + original.toggle(64); + original.toggle(129); + + let mut raw = Vec::new(); + original.write(&mut raw).unwrap(); + + let mut padded = vec![0u8; raw.len() + 8]; + for offset in 0..8 { + padded[offset..offset + raw.len()].copy_from_slice(&raw); + let decoded = BitVec::from_bytes(&padded[offset..offset + raw.len()]); + assert_eq!(decoded, original); + } + } + + #[test] + #[cfg(target_endian = "little")] + #[should_panic(expected = "buffer should be a non-empty multiple of 8 bytes")] + fn test_bitvec_from_bytes_header_only_panics() { + let _ = BitVec::from_bytes(&[1]); + } } From c96b17c0c3ad551d8120a155202db88eb622c009 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 09:58:54 +0100 Subject: [PATCH 02/12] Fix panic-unwind double-drop in sort_by_hash and big-endian group_ops Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/hash-sorted-map/src/group_ops.rs | 15 +-- crates/hash-sorted-map/src/hash_sorted_map.rs | 98 +++++++++++++++++-- 2 files changed, 98 insertions(+), 15 deletions(-) diff --git a/crates/hash-sorted-map/src/group_ops.rs b/crates/hash-sorted-map/src/group_ops.rs index cc14813..e936ce2 100644 --- a/crates/hash-sorted-map/src/group_ops.rs +++ b/crates/hash-sorted-map/src/group_ops.rs @@ -56,12 +56,12 @@ mod arch { /// Number of trailing occupied (non-zero) bytes in the ctrl array. #[inline(always)] pub fn count_occupied(ctrl: &[u8; GROUP_SIZE]) -> usize { - let word = u128::from_ne_bytes(*ctrl); + let word = u128::from_le_bytes(*ctrl); GROUP_SIZE - (word.leading_zeros() / 8) as usize } } -#[cfg(target_arch = "aarch64")] +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] mod arch { use core::arch::aarch64 as neon; @@ -94,18 +94,21 @@ mod arch { /// Number of trailing occupied (non-zero) bytes in the ctrl array. #[inline(always)] pub fn count_occupied(ctrl: &[u8; GROUP_SIZE]) -> usize { - let word = u64::from_ne_bytes(*ctrl); + let word = u64::from_le_bytes(*ctrl); GROUP_SIZE - (word.leading_zeros() / 8) as usize } } -#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] +#[cfg(not(any( + target_arch = "x86_64", + all(target_arch = "aarch64", target_endian = "little") +)))] mod arch { use super::{Mask, GROUP_SIZE}; #[inline(always)] pub fn match_tag(ctrl: &[u8; GROUP_SIZE], tag: u8) -> Mask { - let word = u64::from_ne_bytes(*ctrl); + let word = u64::from_le_bytes(*ctrl); let broadcast = 0x0101010101010101u64 * (tag as u64); let xor = word ^ broadcast; (xor.wrapping_sub(0x0101010101010101)) & !xor & 0x8080808080808080 @@ -129,7 +132,7 @@ mod arch { /// Number of trailing occupied (non-zero) bytes in the ctrl array. #[inline(always)] pub fn count_occupied(ctrl: &[u8; GROUP_SIZE]) -> usize { - let word = u64::from_ne_bytes(*ctrl); + let word = u64::from_le_bytes(*ctrl); GROUP_SIZE - (word.leading_zeros() / 8) as usize } } diff --git a/crates/hash-sorted-map/src/hash_sorted_map.rs b/crates/hash-sorted-map/src/hash_sorted_map.rs index 6608349..088eb0e 100644 --- a/crates/hash-sorted-map/src/hash_sorted_map.rs +++ b/crates/hash-sorted-map/src/hash_sorted_map.rs @@ -214,22 +214,24 @@ impl> HashSortedMap { let n = hashes.len(); // Insertion sort by (hash, key). for i in 1..n { - // Extract element at position i. let cur_hash = hashes[i]; let (gi, si) = chain_slot(&chain, i); - let cur_key = unsafe { self.groups[gi].keys[si].assume_init_read() }; - let cur_val = unsafe { self.groups[gi].values[si].assume_init_read() }; // Find insertion point via linear scan backward. + // Borrow `cur_key` in place during comparisons so that if `Ord::cmp` + // panics, all slots remain valid and no double-drop can occur on unwind. + let cur_key = unsafe { self.groups[gi].keys[si].assume_init_ref() }; let mut j = i; while j > 0 { let (gj, sj) = chain_slot(&chain, j - 1); let prev_key = unsafe { self.groups[gj].keys[sj].assume_init_ref() }; - if (hashes[j - 1], prev_key) <= (cur_hash, &cur_key) { + if (hashes[j - 1], prev_key) <= (cur_hash, cur_key) { break; } j -= 1; } if j < i { + let cur_key = unsafe { self.groups[gi].keys[si].assume_init_read() }; + let cur_val = unsafe { self.groups[gi].values[si].assume_init_read() }; // Shift positions j..i up by one. hashes.copy_within(j..i, j + 1); for pos in (j..i).rev() { @@ -242,12 +244,12 @@ impl> HashSortedMap { self.groups[dst_g].values[dst_s] = v; } } + // Insert at position j. + hashes[j] = cur_hash; + let (gj, sj) = chain_slot(&chain, j); + self.groups[gj].keys[sj] = MaybeUninit::new(cur_key); + self.groups[gj].values[sj] = MaybeUninit::new(cur_val); } - // Insert at position j (or write back to i if already in place). - hashes[j] = cur_hash; - let (gj, sj) = chain_slot(&chain, j); - self.groups[gj].keys[sj] = MaybeUninit::new(cur_key); - self.groups[gj].values[sj] = MaybeUninit::new(cur_val); } // Rebuild ctrl/tag bytes from the sorted hashes so that // get/insert/entry still work after sorting. @@ -1043,4 +1045,82 @@ mod tests { first = false; } } + + #[test] + fn sort_by_hash_panic_in_ord_does_not_double_drop() { + use std::cmp::Ordering; + use std::panic::{catch_unwind, AssertUnwindSafe}; + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::Arc; + + struct DropTrackedKey { + id: u32, + drops: Arc, + } + + impl Drop for DropTrackedKey { + fn drop(&mut self) { + self.drops.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + impl PartialEq for DropTrackedKey { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } + } + + impl Eq for DropTrackedKey {} + + impl PartialOrd for DropTrackedKey { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } + } + + impl Ord for DropTrackedKey { + fn cmp(&self, _other: &Self) -> Ordering { + panic!("comparison panic during sort_by_hash"); + } + } + + impl SortingHash for FixedState { + fn hash(&self, _key: &DropTrackedKey) -> u64 { + self.0 + } + } + + let key_drops = Arc::new(AtomicUsize::new(0)); + let val_drops = Arc::new(AtomicUsize::new(0)); + + struct DropTrackedVal(Arc); + impl Drop for DropTrackedVal { + fn drop(&mut self) { + self.0.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + let result = catch_unwind(AssertUnwindSafe(|| { + let mut map = HashSortedMap::with_capacity_and_hasher(1, FixedState(42)); + map.insert( + DropTrackedKey { + id: 1, + drops: Arc::clone(&key_drops), + }, + DropTrackedVal(Arc::clone(&val_drops)), + ); + map.insert( + DropTrackedKey { + id: 2, + drops: Arc::clone(&key_drops), + }, + DropTrackedVal(Arc::clone(&val_drops)), + ); + map.sort_by_hash(); + })); + + assert!(result.is_err()); + assert_eq!(key_drops.load(AtomicOrdering::SeqCst), 2); + assert_eq!(val_drops.load(AtomicOrdering::SeqCst), 2); + } } From 53cec1ff2c178020f1e53b0701f1c44def08eb51 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 09:58:54 +0100 Subject: [PATCH 03/12] Validate DAAC deserialization and check token bounds in bpe Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/bpe/src/appendable_encoder.rs | 7 +- crates/bpe/src/byte_pair_encoding.rs | 143 +++++++++++++++++++++++--- crates/bpe/src/prependable_encoder.rs | 7 +- 3 files changed, 137 insertions(+), 20 deletions(-) diff --git a/crates/bpe/src/appendable_encoder.rs b/crates/bpe/src/appendable_encoder.rs index 3dc13db..9d70e56 100644 --- a/crates/bpe/src/appendable_encoder.rs +++ b/crates/bpe/src/appendable_encoder.rs @@ -56,11 +56,10 @@ impl<'a> AppendableEncoder<'a> { }); break; } else { - let prev_token = - unsafe { self.states.get_unchecked(new_range.start - 1).last_token }; + let prev = &self.states[new_range.start - 1]; + let prev_token = prev.last_token; if self.bpe.is_valid_token_pair(prev_token, new_token) { - let prev_count = - unsafe { self.states.get_unchecked(new_range.start - 1).count }; + let prev_count = prev.count; self.states.push(State { state, last_token: new_token, diff --git a/crates/bpe/src/byte_pair_encoding.rs b/crates/bpe/src/byte_pair_encoding.rs index 298532d..3e8c574 100644 --- a/crates/bpe/src/byte_pair_encoding.rs +++ b/crates/bpe/src/byte_pair_encoding.rs @@ -63,15 +63,93 @@ fn serialize_daac( s.serialize_bytes(&daac.serialize()) } +fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { + const STATE_SIZE: usize = 12; + const OUTPUT_SIZE: usize = 12; + const MAX_OUTPUTS: u32 = 0x00ff_ffff; + + if v.len() < 4 { + return Err("truncated DoubleArrayAhoCorasick states length"); + } + let states_len = u32::from_le_bytes(v[0..4].try_into().unwrap()); + if states_len < 256 || !states_len.is_multiple_of(256) { + return Err("invalid DoubleArrayAhoCorasick states length"); + } + let states_bytes_len = (states_len as usize) + .checked_mul(STATE_SIZE) + .ok_or("DoubleArrayAhoCorasick states byte length overflow")?; + let outputs_offset = 4usize + .checked_add(states_bytes_len) + .ok_or("DoubleArrayAhoCorasick states offset overflow")?; + let outputs_header_end = outputs_offset + .checked_add(4) + .ok_or("DoubleArrayAhoCorasick outputs header overflow")?; + if v.len() < outputs_header_end { + return Err("truncated DoubleArrayAhoCorasick states data"); + } + + let outputs_len = u32::from_le_bytes(v[outputs_offset..outputs_header_end].try_into().unwrap()); + if outputs_len > MAX_OUTPUTS { + return Err("invalid DoubleArrayAhoCorasick outputs length"); + } + let outputs_bytes_len = (outputs_len as usize) + .checked_mul(OUTPUT_SIZE) + .ok_or("DoubleArrayAhoCorasick outputs byte length overflow")?; + let outputs_end = outputs_header_end + .checked_add(outputs_bytes_len) + .ok_or("DoubleArrayAhoCorasick outputs end overflow")?; + let total_len = outputs_end + .checked_add(5) + .ok_or("DoubleArrayAhoCorasick total length overflow")?; + if v.len() != total_len { + return Err("unexpected DoubleArrayAhoCorasick byte length"); + } + + let match_kind = v[outputs_end]; + if match_kind > 2 { + return Err("invalid DoubleArrayAhoCorasick match kind"); + } + let num_states = u32::from_le_bytes(v[outputs_end + 1..total_len].try_into().unwrap()); + if num_states > states_len { + return Err("invalid DoubleArrayAhoCorasick num_states"); + } + + for chunk in v[4..outputs_offset].as_chunks::().0 { + let base = u32::from_le_bytes(chunk[0..4].try_into().unwrap()); + let fail = u32::from_le_bytes(chunk[4..8].try_into().unwrap()); + let opos_ch = u32::from_le_bytes(chunk[8..12].try_into().unwrap()); + let output_pos = opos_ch >> 8; + if base >= states_len || fail >= states_len || output_pos > outputs_len { + return Err("DoubleArrayAhoCorasick state out of bounds"); + } + } + + for (i, chunk) in v[outputs_header_end..outputs_end] + .as_chunks::() + .0 + .iter() + .enumerate() + { + let length = u32::from_le_bytes(chunk[4..8].try_into().unwrap()); + let parent = u32::from_le_bytes(chunk[8..12].try_into().unwrap()); + if length == 0 || (parent as usize) > i { + return Err("DoubleArrayAhoCorasick output out of bounds"); + } + } + + Ok(()) +} + struct DaacVisitor; impl Visitor<'_> for DaacVisitor { type Value = DoubleArrayAhoCorasick; - fn expecting(&self, _formatter: &mut std::fmt::Formatter) -> std::fmt::Result { - Err(std::fmt::Error) + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("a serialized DoubleArrayAhoCorasick automaton") } fn visit_bytes(self, v: &[u8]) -> Result { + validate_daac_bytes(v).map_err(E::custom)?; Ok(unsafe { DoubleArrayAhoCorasick::deserialize_unchecked(v).0 }) } } @@ -114,20 +192,20 @@ fn is_valid_token_pair( // Reverse the merge operation from BPE. if token1 > token2 { limit = token1; - token1 = unsafe { split_table.get_unchecked(token1 as usize).1 }; + token1 = split_table[token1 as usize].1; if token1 == limit { limit = token2 + 1; - token2 = unsafe { split_table.get_unchecked(token2 as usize).0 }; + token2 = split_table[token2 as usize].0; if token2 + 1 == limit { return true; } } } else { limit = token2 + 1; - token2 = unsafe { split_table.get_unchecked(token2 as usize).0 }; + token2 = split_table[token2 as usize].0; if token2 + 1 == limit { limit = token1; - token1 = unsafe { split_table.get_unchecked(token1 as usize).1 }; + token1 = split_table[token1 as usize].1; if token1 == limit { return true; } @@ -137,10 +215,8 @@ fn is_valid_token_pair( } fn token_range(token_starts: &[u32], token_id: u32) -> Range { - unsafe { - *token_starts.get_unchecked(token_id as usize) as usize - ..*token_starts.get_unchecked(token_id as usize + 1) as usize - } + let idx = token_id as usize; + token_starts[idx] as usize..token_starts[idx + 1] as usize } fn token_bytes<'a>(all_tokens: &'a [u8], token_starts: &[u32], token_id: u32) -> &'a [u8] { @@ -404,7 +480,7 @@ impl BytePairEncoding { last_token.push(new_token); break; } else { - let prev_token = unsafe { *last_token.get_unchecked(new_range.start - 1) }; + let prev_token = last_token[new_range.start - 1]; if self.is_valid_token_pair(prev_token, new_token) { last_token.push(new_token); break; @@ -694,7 +770,9 @@ pub fn create_test_bytes(bpe: &BytePairEncoding, min_bytes: usize) -> Vec { #[cfg(test)] mod tests { - use super::hash_bytes; + use aneubeck_daachorse::DoubleArrayAhoCorasick; + + use super::{hash_bytes, validate_daac_bytes, BytePairEncoding}; /// The hash must not depend on the target. It keys the serialized `bytes_hash_to_token`, /// and a hardcoded hash factor is only collision-free for one particular mapping. @@ -712,4 +790,45 @@ mod tests { // Tiktoken vocabularies contain arbitrary bytes, not just UTF-8. assert_eq!(hash_bytes(&[0, 0xff, 0x80], FACTOR), 3276605559); } + + #[test] + #[should_panic] + fn token_bytes_out_of_bounds_panics() { + let tokens: Vec> = (0..=255u8).map(|b| vec![b]).collect(); + let bpe = BytePairEncoding::from_dictionary(tokens, None); + let _ = bpe.token_bytes(999); + } + + #[test] + fn validate_daac_bytes_checks_invariants() { + let pma = DoubleArrayAhoCorasick::::new(vec!["a", "ab", "bab"]).unwrap(); + let valid = pma.serialize(); + assert!(validate_daac_bytes(&valid).is_ok()); + + // Truncated input + assert!(validate_daac_bytes(&valid[..valid.len() - 1]).is_err()); + + // Non-multiple-of-256 states_len + let mut bad_states_len = valid.clone(); + bad_states_len[0..4].copy_from_slice(&255u32.to_le_bytes()); + assert!(validate_daac_bytes(&bad_states_len).is_err()); + + // Out-of-bounds state `base` + let mut bad_base = valid.clone(); + bad_base[4..8].copy_from_slice(&9999u32.to_le_bytes()); + assert!(validate_daac_bytes(&bad_base).is_err()); + + // Out-of-bounds state `fail` + let mut bad_fail = valid.clone(); + bad_fail[8..12].copy_from_slice(&9999u32.to_le_bytes()); + assert!(validate_daac_bytes(&bad_fail).is_err()); + + // Out-of-bounds output `parent` (self-referential 1-based parent == 1 at index 0) + let states_len = u32::from_le_bytes(valid[0..4].try_into().unwrap()) as usize; + let first_output_parent_offset = 4 + states_len * 12 + 4 + 8; + let mut bad_parent = valid.clone(); + bad_parent[first_output_parent_offset..first_output_parent_offset + 4] + .copy_from_slice(&1u32.to_le_bytes()); + assert!(validate_daac_bytes(&bad_parent).is_err()); + } } diff --git a/crates/bpe/src/prependable_encoder.rs b/crates/bpe/src/prependable_encoder.rs index 264c6a5..72b3e9c 100644 --- a/crates/bpe/src/prependable_encoder.rs +++ b/crates/bpe/src/prependable_encoder.rs @@ -56,11 +56,10 @@ impl<'a> PrependableEncoder<'a> { }); break; } else { - let next_token = - unsafe { self.states.get_unchecked(new_range.start - 1).prev_token }; + let prev = &self.states[new_range.start - 1]; + let next_token = prev.prev_token; if self.bpe.is_valid_token_pair(new_token, next_token) { - let prev_count = - unsafe { self.states.get_unchecked(new_range.start - 1).count }; + let prev_count = prev.count; self.states.push(State { state, prev_token: new_token, From 5c983a0fe1c2d825a37ad42491b02e81befc3b7c Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 10:08:23 +0100 Subject: [PATCH 04/12] Avoid unwrap() in validate_daac_bytes to satisfy clippy::unwrap_used Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/bpe/src/byte_pair_encoding.rs | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/crates/bpe/src/byte_pair_encoding.rs b/crates/bpe/src/byte_pair_encoding.rs index 3e8c574..604d8fd 100644 --- a/crates/bpe/src/byte_pair_encoding.rs +++ b/crates/bpe/src/byte_pair_encoding.rs @@ -63,6 +63,15 @@ fn serialize_daac( s.serialize_bytes(&daac.serialize()) } +fn read_u32_le(bytes: &[u8], offset: usize) -> u32 { + u32::from_le_bytes([ + bytes[offset], + bytes[offset + 1], + bytes[offset + 2], + bytes[offset + 3], + ]) +} + fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { const STATE_SIZE: usize = 12; const OUTPUT_SIZE: usize = 12; @@ -71,7 +80,7 @@ fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { if v.len() < 4 { return Err("truncated DoubleArrayAhoCorasick states length"); } - let states_len = u32::from_le_bytes(v[0..4].try_into().unwrap()); + let states_len = read_u32_le(v, 0); if states_len < 256 || !states_len.is_multiple_of(256) { return Err("invalid DoubleArrayAhoCorasick states length"); } @@ -88,7 +97,7 @@ fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { return Err("truncated DoubleArrayAhoCorasick states data"); } - let outputs_len = u32::from_le_bytes(v[outputs_offset..outputs_header_end].try_into().unwrap()); + let outputs_len = read_u32_le(v, outputs_offset); if outputs_len > MAX_OUTPUTS { return Err("invalid DoubleArrayAhoCorasick outputs length"); } @@ -109,15 +118,15 @@ fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { if match_kind > 2 { return Err("invalid DoubleArrayAhoCorasick match kind"); } - let num_states = u32::from_le_bytes(v[outputs_end + 1..total_len].try_into().unwrap()); + let num_states = read_u32_le(v, outputs_end + 1); if num_states > states_len { return Err("invalid DoubleArrayAhoCorasick num_states"); } for chunk in v[4..outputs_offset].as_chunks::().0 { - let base = u32::from_le_bytes(chunk[0..4].try_into().unwrap()); - let fail = u32::from_le_bytes(chunk[4..8].try_into().unwrap()); - let opos_ch = u32::from_le_bytes(chunk[8..12].try_into().unwrap()); + let base = read_u32_le(chunk, 0); + let fail = read_u32_le(chunk, 4); + let opos_ch = read_u32_le(chunk, 8); let output_pos = opos_ch >> 8; if base >= states_len || fail >= states_len || output_pos > outputs_len { return Err("DoubleArrayAhoCorasick state out of bounds"); @@ -130,8 +139,8 @@ fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { .iter() .enumerate() { - let length = u32::from_le_bytes(chunk[4..8].try_into().unwrap()); - let parent = u32::from_le_bytes(chunk[8..12].try_into().unwrap()); + let length = read_u32_le(chunk, 4); + let parent = read_u32_le(chunk, 8); if length == 0 || (parent as usize) > i { return Err("DoubleArrayAhoCorasick output out of bounds"); } From ce9ac6626dce650c7455db56df0344df791d43c5 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 10:36:49 +0100 Subject: [PATCH 05/12] Use zero-copy unaligned block reads in BitVec Store BitVec blocks as Cow<'a, [[u8; BYTES_PER_BLOCK]]> and read/write u64 values via u64::from_le_bytes/to_le_bytes so BitVec::from_bytes always borrows directly from the input slice without copying or unsafe code. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/geo_filters/src/diff_count.rs | 22 ++++-- crates/geo_filters/src/diff_count/bitvec.rs | 83 ++++++++++++--------- crates/geo_filters/src/diff_count/metric.rs | 3 +- 3 files changed, 62 insertions(+), 46 deletions(-) diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index dc8c7cd..8fe57a7 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -9,6 +9,7 @@ use std::ops::Deref as _; use crate::config::{ count_ones_from_bitchunks, count_ones_from_msb_and_lsb, iter_bit_chunks, iter_ones, mask_bit_chunks, take_ref, xor_bit_chunks, BitChunk, GeoConfig, IsBucketType, BITS_PER_BLOCK, + BYTES_PER_BLOCK, }; use crate::{Count, Diff}; @@ -576,7 +577,7 @@ pub struct GeoDiffCountBuilder> { /// (xor) into `blocks`. [`GeoDiffCountBuilder::cleanup`] reconciles the two. split: usize, numbers: Vec, - blocks: Vec, + blocks: Vec<[u8; BYTES_PER_BLOCK]>, } impl> GeoDiffCountBuilder { @@ -593,7 +594,7 @@ impl> GeoDiffCountBuilder { expected, split, numbers: Vec::with_capacity(capacity), - blocks: vec![0; split.div_ceil(BITS_PER_BLOCK)], + blocks: vec![[0; BYTES_PER_BLOCK]; split.div_ceil(BITS_PER_BLOCK)], } } @@ -608,7 +609,8 @@ impl> GeoDiffCountBuilder { let new_split = estimate_split_bucket(&self.config, self.expected); if new_split > self.split { self.split = new_split; - self.blocks.resize(new_split.div_ceil(BITS_PER_BLOCK), 0); + self.blocks + .resize(new_split.div_ceil(BITS_PER_BLOCK), [0; BYTES_PER_BLOCK]); } } @@ -633,7 +635,8 @@ impl> GeoDiffCountBuilder { if (next - read) % 2 == 1 { if bucket < split { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] ^= bit.into_block(); + blocks[index] = + (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); } else { numbers[write] = bucket; write += 1; @@ -663,11 +666,12 @@ impl> GeoDiffCountBuilder { new_split += bits_per_level; keep = self.numbers.partition_point(|&b| b >= new_split); } - self.blocks.resize(new_split.div_ceil(BITS_PER_BLOCK), 0); + self.blocks + .resize(new_split.div_ceil(BITS_PER_BLOCK), [0; BYTES_PER_BLOCK]); let blocks = &mut self.blocks; for &bucket in &self.numbers[keep..] { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] ^= bit.into_block(); + blocks[index] = (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); } self.numbers.truncate(keep); self.split = new_split; @@ -688,7 +692,8 @@ impl> GeoDiffCountBuilder { } else { // `bucket < split`, so the block index is always in range; toggling cancels repeats. let (index, bit) = bucket.into_index_and_bit(); - self.blocks[index] ^= bit.into_block(); + self.blocks[index] = + (u64::from_le_bytes(self.blocks[index]) ^ bit.into_block()).to_le_bytes(); } } @@ -724,7 +729,8 @@ impl> GeoDiffCountBuilder { } } else { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] ^= bit.into_block(); + blocks[index] = + (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); } } filled diff --git a/crates/geo_filters/src/diff_count/bitvec.rs b/crates/geo_filters/src/diff_count/bitvec.rs index 4793d31..efebfa6 100644 --- a/crates/geo_filters/src/diff_count/bitvec.rs +++ b/crates/geo_filters/src/diff_count/bitvec.rs @@ -14,13 +14,18 @@ use crate::config::{BitChunk, BYTES_PER_BLOCK}; #[derive(Clone, Default, Debug, Eq, PartialEq)] pub(crate) struct BitVec<'a> { num_bits: usize, - blocks: Cow<'a, [u64]>, + blocks: Cow<'a, [[u8; BYTES_PER_BLOCK]]>, } impl Ord for BitVec<'_> { fn cmp(&self, other: &Self) -> std::cmp::Ordering { match self.num_bits.cmp(&other.num_bits) { - Ordering::Equal => self.blocks.iter().rev().cmp(other.blocks.iter().rev()), + Ordering::Equal => self + .blocks + .iter() + .rev() + .map(|&b| u64::from_le_bytes(b)) + .cmp(other.blocks.iter().rev().map(|&b| u64::from_le_bytes(b))), ord => ord, } } @@ -41,7 +46,7 @@ impl<'a> BitVec<'a> { result.resize(num_bits); let blocks = result.blocks.to_mut(); for BitChunk { index, block } in chunks { - blocks[index] ^= block; + blocks[index] = (u64::from_le_bytes(blocks[index]) ^ block).to_le_bytes(); } result.clear_superfluous_bits(); result @@ -49,7 +54,7 @@ impl<'a> BitVec<'a> { /// Wraps raw `blocks` covering `[0, num_bits)` into a `BitVec`. The number of blocks must match /// `num_bits`, and any bits at or above `num_bits` are cleared. - pub fn from_blocks(blocks: Vec, num_bits: usize) -> BitVec<'static> { + pub fn from_blocks(blocks: Vec<[u8; BYTES_PER_BLOCK]>, num_bits: usize) -> BitVec<'static> { debug_assert_eq!(blocks.len(), num_bits.div_ceil(BITS_PER_BLOCK)); let mut result = BitVec { num_bits, @@ -63,7 +68,9 @@ impl<'a> BitVec<'a> { pub fn resize(&mut self, num_bits: usize) { let num_blocks = num_bits.div_ceil(BITS_PER_BLOCK); if num_blocks != self.blocks.len() { - self.blocks.to_mut().resize(num_blocks, 0); + self.blocks + .to_mut() + .resize(num_blocks, [0; BYTES_PER_BLOCK]); } self.num_bits = num_bits; self.clear_superfluous_bits(); @@ -76,7 +83,9 @@ impl<'a> BitVec<'a> { .to_mut() .last_mut() .into_iter() - .for_each(|block| *block &= bit.into_mask()); + .for_each(|block| { + *block = (u64::from_le_bytes(*block) & bit.into_mask()).to_le_bytes(); + }); } } @@ -88,15 +97,20 @@ impl<'a> BitVec<'a> { pub fn count_ones(&self) -> usize { self.blocks .iter() - .map(|block| block.count_ones() as usize) + .map(|&block| u64::from_le_bytes(block).count_ones() as usize) .sum() } /// The raw 64-bit blocks backing the vector, from least to most significant. - pub fn blocks(&self) -> &[u64] { + pub fn blocks(&self) -> &[[u8; BYTES_PER_BLOCK]] { self.blocks.deref() } + #[inline] + pub fn block(&self, index: usize) -> u64 { + u64::from_le_bytes(self.blocks[index]) + } + pub fn is_empty(&self) -> bool { self.num_bits() == 0 } @@ -105,14 +119,15 @@ impl<'a> BitVec<'a> { pub fn test_bit(&self, index: usize) -> bool { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - (self.blocks[block_idx] & bit_idx.into_block()) != 0 + (self.block(block_idx) & bit_idx.into_block()) != 0 } /// Toggles the bit specified by the provided zero-based bit position. pub fn toggle(&mut self, index: usize) { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - self.blocks.to_mut()[block_idx] ^= bit_idx.into_block(); + let block = &mut self.blocks.to_mut()[block_idx]; + *block = (u64::from_le_bytes(*block) ^ bit_idx.into_block()).to_le_bytes(); } /// Returns a [`BitToggler`] that toggles many bits without re-resolving the `Cow` on every @@ -131,10 +146,11 @@ impl<'a> BitVec<'a> { pub fn bit_chunks(&self) -> impl Iterator + '_ { self.blocks .iter() + .map(|&block| u64::from_le_bytes(block)) .enumerate() .rev() - .filter(|(_, block)| **block != 0) - .map(|(index, block)| BitChunk::new(index, *block)) + .filter(|(_, block)| *block != 0) + .map(|(index, block)| BitChunk::new(index, block)) } // Returns an owned BitVec with static lifetime that can be used for cases when we want @@ -161,10 +177,10 @@ impl<'a> BitVec<'a> { if block_start >= self.blocks.len() { return 0; } - let mut result = self.blocks[block_start] >> start_idx; + let mut result = self.block(block_start) >> start_idx; let block_end = (range.end - 1) / BITS_PER_BLOCK; if block_end < self.blocks.len() && block_end != block_start { - result |= self.blocks[block_end] << (BITS_PER_BLOCK - start_idx); + result |= self.block(block_end) << (BITS_PER_BLOCK - start_idx); } result & (u64::MAX >> (BITS_PER_BLOCK - range.len())) } @@ -189,24 +205,17 @@ impl<'a> BitVec<'a> { unoccupied_bits ); buf = &buf[1..]; + let (blocks, remainder) = buf.as_chunks::(); assert!( - !buf.is_empty() && buf.len().is_multiple_of(BYTES_PER_BLOCK), + !blocks.is_empty() && remainder.is_empty(), "buffer should be a non-empty multiple of 8 bytes, got {}", buf.len() ); let num_bits = buf.len() * 8 - unoccupied_bits; - let blocks = match unsafe { buf.align_to::() } { - ([], aligned, []) => Cow::Borrowed(aligned), - _ => Cow::Owned( - buf.as_chunks::() - .0 - .iter() - .copied() - .map(u64::from_le_bytes) - .collect(), - ), - }; - Self { num_bits, blocks } + Self { + num_bits, + blocks: Cow::Borrowed(blocks), + } } #[cfg(target_endian = "little")] @@ -217,10 +226,7 @@ impl<'a> BitVec<'a> { // First serialize the number of unoccupied bits in the last block as one byte. let unoccupied_bits = 63 - ((self.num_bits - 1) % 64) as u8; writer.write_all(&[unoccupied_bits])?; - let blocks = self.blocks.deref(); - let block_bytes = unsafe { - std::slice::from_raw_parts(blocks.as_ptr() as *const u8, blocks.len() * BYTES_PER_BLOCK) - }; + let block_bytes = self.blocks.as_flattened(); writer.write_all(block_bytes)?; Ok(block_bytes.len() + 1) } @@ -233,7 +239,7 @@ impl Index for BitVec<'_> { fn index(&self, index: usize) -> &Self::Output { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - match (self.blocks[block_idx] & bit_idx.into_block()) != 0 { + match (self.block(block_idx) & bit_idx.into_block()) != 0 { true => &true, false => &false, } @@ -244,7 +250,7 @@ impl Index for BitVec<'_> { /// Obtained via [`BitVec::toggler`]. pub(crate) struct BitToggler<'a> { num_bits: usize, - blocks: &'a mut [u64], + blocks: &'a mut [[u8; BYTES_PER_BLOCK]], } impl BitToggler<'_> { @@ -253,7 +259,8 @@ impl BitToggler<'_> { pub fn toggle(&mut self, index: usize) { debug_assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - self.blocks[block_idx] ^= bit_idx.into_block(); + let block = &mut self.blocks[block_idx]; + *block = (u64::from_le_bytes(*block) ^ bit_idx.into_block()).to_le_bytes(); } } @@ -264,7 +271,8 @@ mod tests { /// Constructs a BitVec instance with num_bits and all bits initialized to the provided default value. fn bitvec_from_element<'a>(num_bits: usize, value: bool) -> BitVec<'a> { let num_blocks = num_bits.div_ceil(BITS_PER_BLOCK) + 1; - let blocks = vec![if value { u64::MAX } else { u64::MIN }; num_blocks]; + let fill = if value { u64::MAX } else { u64::MIN }.to_le_bytes(); + let blocks = vec![fill; num_blocks]; BitVec { blocks: Cow::from(blocks), num_bits, @@ -306,7 +314,7 @@ mod tests { assert!(bitvec.test_bit(14)); assert!(bitvec.test_bit(48)); assert!(bitvec.test_bit(72)); - let last = bitvec.blocks.last().expect("there must be a block"); + let last = u64::from_le_bytes(*bitvec.blocks.last().expect("there must be a block")); assert!(last & 120usize.into_bit().into_block() == 0); } @@ -314,7 +322,7 @@ mod tests { { assert!(bitvec.test_bit(14)); assert!(bitvec.test_bit(48)); - let last = bitvec.blocks.last().expect("there must be a block"); + let last = u64::from_le_bytes(*bitvec.blocks.last().expect("there must be a block")); assert!(last & 72usize.into_bit().into_block() == 0); } @@ -380,6 +388,7 @@ mod tests { for offset in 0..8 { padded[offset..offset + raw.len()].copy_from_slice(&raw); let decoded = BitVec::from_bytes(&padded[offset..offset + raw.len()]); + assert!(matches!(decoded.blocks, Cow::Borrowed(_))); assert_eq!(decoded, original); } } diff --git a/crates/geo_filters/src/diff_count/metric.rs b/crates/geo_filters/src/diff_count/metric.rs index d618826..4262652 100644 --- a/crates/geo_filters/src/diff_count/metric.rs +++ b/crates/geo_filters/src/diff_count/metric.rs @@ -184,7 +184,8 @@ impl + Default> MetricSpace for GeoDiffMetric<'_, C> { let mut count = 0; let (a_blocks, b_blocks) = (a.lsb.blocks(), b.lsb.blocks()); for i in 0..common_full { - count += (a_blocks[i] ^ b_blocks[i]).count_ones() as usize; + count += (u64::from_le_bytes(a_blocks[i]) ^ u64::from_le_bytes(b_blocks[i])) + .count_ones() as usize; if count >= bound { return OnesMetric::infinite(); } From 43773ddf73ef8c663faba69a53b0336d78407d38 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 12:12:58 +0100 Subject: [PATCH 06/12] Use zero-copy unaligned bucket reads for GeoDiffCount::msb Store GeoDiffCount::msb as Cow<[::Bytes]> and decode/encode bucket values via IsBucketType::from_le_bytes and to_le_bytes so GeoDiffCount::from_bytes_with_config and write are zero-copy without requiring BucketType alignment or unsafe code. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/geo_filters/src/config/buckets.rs | 114 ++++++--------- crates/geo_filters/src/diff_count.rs | 137 +++++++++++------- crates/geo_filters/src/diff_count/sim_hash.rs | 6 +- 3 files changed, 137 insertions(+), 120 deletions(-) diff --git a/crates/geo_filters/src/config/buckets.rs b/crates/geo_filters/src/config/buckets.rs index 25154a4..1d4b8a5 100644 --- a/crates/geo_filters/src/config/buckets.rs +++ b/crates/geo_filters/src/config/buckets.rs @@ -12,9 +12,14 @@ where Self: Add + Rem, { const BITS: u32; + type Bytes: Copy + Debug + Default + Eq + PartialEq + Send + Sync + 'static; fn from_usize(v: usize) -> Self; fn into_usize(self) -> usize; + fn from_le_bytes(bytes: Self::Bytes) -> Self; + fn to_le_bytes(self) -> Self::Bytes; + fn slice_from_bytes(buf: &[u8]) -> (&[Self::Bytes], &[u8]); + fn slice_as_bytes(slice: &[Self::Bytes]) -> &[u8]; /// Returns a block with the bit in this bit position set to one. #[inline] @@ -75,75 +80,50 @@ where } } -impl IsBucketType for u8 { - const BITS: u32 = Self::BITS; - - #[inline] - fn from_usize(v: usize) -> Self { - v as Self - } - - #[inline] - fn into_usize(self) -> usize { - self as usize - } +macro_rules! impl_is_bucket_type { + ($t:ty) => { + impl IsBucketType for $t { + const BITS: u32 = Self::BITS; + type Bytes = [u8; std::mem::size_of::<$t>()]; + + #[inline] + fn from_usize(v: usize) -> Self { + v as Self + } + + #[inline] + fn into_usize(self) -> usize { + self as usize + } + + #[inline] + fn from_le_bytes(bytes: Self::Bytes) -> Self { + Self::from_le_bytes(bytes) + } + + #[inline] + fn to_le_bytes(self) -> Self::Bytes { + Self::to_le_bytes(self) + } + + #[inline] + fn slice_from_bytes(buf: &[u8]) -> (&[Self::Bytes], &[u8]) { + buf.as_chunks::<{ std::mem::size_of::<$t>() }>() + } + + #[inline] + fn slice_as_bytes(slice: &[Self::Bytes]) -> &[u8] { + slice.as_flattened() + } + } + }; } -impl IsBucketType for u16 { - const BITS: u32 = Self::BITS; - - #[inline] - fn from_usize(v: usize) -> Self { - v as Self - } - - #[inline] - fn into_usize(self) -> usize { - self as usize - } -} - -impl IsBucketType for u32 { - const BITS: u32 = Self::BITS; - - #[inline] - fn from_usize(v: usize) -> Self { - v as Self - } - - #[inline] - fn into_usize(self) -> usize { - self as usize - } -} - -impl IsBucketType for u64 { - const BITS: u32 = Self::BITS; - - #[inline] - fn from_usize(v: usize) -> Self { - v as Self - } - - #[inline] - fn into_usize(self) -> usize { - self as usize - } -} - -impl IsBucketType for usize { - const BITS: u32 = Self::BITS; - - #[inline] - fn from_usize(v: usize) -> Self { - v - } - - #[inline] - fn into_usize(self) -> usize { - self - } -} +impl_is_bucket_type!(u8); +impl_is_bucket_type!(u16); +impl_is_bucket_type!(u32); +impl_is_bucket_type!(u64); +impl_is_bucket_type!(usize); /// Computes the bits required to store bucket positions for 64-bit hashes given that (1 << B) /// buckets cover half the hash space. diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index 8fe57a7..d7261ad 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -4,7 +4,6 @@ use std::borrow::Cow; use std::cmp::Ordering; use std::hash::BuildHasher as _; use std::mem::{size_of, size_of_val}; -use std::ops::Deref as _; use crate::config::{ count_ones_from_bitchunks, count_ones_from_msb_and_lsb, iter_bit_chunks, iter_ones, @@ -65,14 +64,49 @@ pub type GeoDiffCount13<'a> = GeoDiffCount<'a, GeoDiffConfig13>; /// with an error of +/-1200 items using the GeoDiffCount. A HyperLogLog data structure with /// 1500 bytes achieves a precision of ~0.022 and could estimate the 10k items with an error /// of +/-22k items in the best case which is 20 times worse despite using 5 times more space! -#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord)] +#[derive(Clone, Default, PartialEq, Eq)] pub struct GeoDiffCount<'a, C: GeoConfig> { config: C, /// The bit positions are stored from largest to smallest. - msb: Cow<'a, [C::BucketType]>, + msb: Cow<'a, [::Bytes]>, lsb: BitVec<'a>, } +impl + Ord> Ord for GeoDiffCount<'_, C> { + fn cmp(&self, other: &Self) -> Ordering { + self.config + .cmp(&other.config) + .then_with(|| { + self.msb + .iter() + .copied() + .map(C::BucketType::from_le_bytes) + .cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes)) + }) + .then_with(|| self.lsb.cmp(&other.lsb)) + } +} + +impl + PartialOrd> PartialOrd for GeoDiffCount<'_, C> { + fn partial_cmp(&self, other: &Self) -> Option { + match self.config.partial_cmp(&other.config)? { + Ordering::Equal => {} + ord => return Some(ord), + } + match self + .msb + .iter() + .copied() + .map(C::BucketType::from_le_bytes) + .partial_cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes))? + { + Ordering::Equal => {} + ord => return Some(ord), + } + self.lsb.partial_cmp(&other.lsb) + } +} + impl> std::fmt::Debug for GeoDiffCount<'_, C> { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!( @@ -105,12 +139,12 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { let mut ones = iter_ones::(chunks.peekable()); let mut msb = Vec::default(); take_ref(&mut ones, config.max_msb_len() - 1).for_each(|bucket| { - msb.push(bucket); + msb.push(bucket.to_le_bytes()); }); let smallest_msb = ones .next() .inspect(|bucket| { - msb.push(*bucket); + msb.push(bucket.to_le_bytes()); }) .unwrap_or_default(); let lsb = BitVec::from_bit_chunks(ones.into_bitchunks(), smallest_msb.into_usize()); @@ -181,7 +215,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { let msb = self .msb .iter() - .map(|&p| p.into_usize()) + .map(|&p| C::BucketType::from_le_bytes(p).into_usize()) .filter(move |&pos| (mask >> (pos % mask_size)) & 1 != 0) .map(|pos| pos as u64); let lsb = iter_ones::( @@ -214,7 +248,9 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { if index.into_usize() < self.lsb.num_bits() { self.lsb.test_bit(index.into_usize()) } else { - self.msb.binary_search_by(|k| index.cmp(k)).is_ok() + self.msb + .binary_search_by(|&k| index.cmp(&C::BucketType::from_le_bytes(k))) + .is_ok() } } @@ -242,7 +278,9 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// most to least significant, overlapping `BitChunk`s got merged, and empty ones filtered. fn bit_chunks(&self) -> impl Iterator + '_ { iter_bit_chunks( - self.msb.iter().map(|b| b.into_usize()), + self.msb + .iter() + .map(|&b| C::BucketType::from_le_bytes(b).into_usize()), self.lsb.bit_chunks(), ) } @@ -254,7 +292,11 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { fn estimate_size(&self) -> f32 { let (total, lowest_bucket) = count_ones_from_msb_and_lsb::( self.msb.len(), - self.msb.last().copied().unwrap_or_default(), + self.msb + .last() + .copied() + .map(C::BucketType::from_le_bytes) + .unwrap_or_default(), self.lsb.bit_chunks().peekable(), self.config.max_bytes(), self.config.max_msb_len(), @@ -279,7 +321,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { self.lsb.toggle(bucket.into_usize()); } else { let msb = self.msb.to_mut(); - match msb.binary_search_by(|k| bucket.cmp(k)) { + match msb.binary_search_by(|&k| bucket.cmp(&C::BucketType::from_le_bytes(k))) { Ok(idx) => { msb.remove(idx); let first = { @@ -287,33 +329,33 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { lsb.next() }; if let Some(smallest) = first { - msb.push(C::BucketType::from_usize(smallest)); + msb.push(C::BucketType::from_usize(smallest).to_le_bytes()); self.lsb.resize(smallest); } else { self.lsb.resize(0); }; } Err(idx) => { - msb.insert(idx, bucket); + msb.insert(idx, bucket.to_le_bytes()); if msb.len() > self.config.max_msb_len() { // We have too many values in the MSB sparse index vector, // let's move the smalles MSB value into the LSB bit vector - let smallest = msb - .pop() - .expect("we should have at least one element!") - .into_usize(); - let new_smallest = msb - .last() - .expect("should have at least one element") - .into_usize(); + let smallest = C::BucketType::from_le_bytes( + msb.pop().expect("we should have at least one element!"), + ) + .into_usize(); + let new_smallest = C::BucketType::from_le_bytes( + *msb.last().expect("should have at least one element"), + ) + .into_usize(); // ensure LSB bit vector has the space for `smallest` self.lsb.resize(new_smallest); self.lsb.toggle(smallest); } else if msb.len() == self.config.max_msb_len() { - let smallest = msb - .last() - .expect("should have at least one element") - .into_usize(); + let smallest = C::BucketType::from_le_bytes( + *msb.last().expect("should have at least one element"), + ) + .into_usize(); self.lsb.resize(smallest); } } @@ -350,7 +392,11 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { ); debug_assert!( self.lsb.num_bits() == 0 - || self.msb.last().map_or(0, |b| b.into_usize()) == self.lsb.num_bits(), + || self + .msb + .last() + .map_or(0, |&b| C::BucketType::from_le_bytes(b).into_usize()) + == self.lsb.num_bits(), "msb and lsb non-contiguous: msb = {:?} |lsb| = {:?}", self.msb, self.lsb.num_bits(), @@ -376,25 +422,14 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { return Self::new(c); } // The number of most significant bits stores in the MSB sparse repr - let msb_len = (buf.len() / size_of::()).min(c.max_msb_len()); + let (msb_all, _) = C::BucketType::slice_from_bytes(buf); + let msb_len = msb_all.len().min(c.max_msb_len()); // The number of bytes representing the MSB - this is how many bytes we need to // skip over to reach the LSB let msb_bytes_len = msb_len * size_of::(); - let msb_bytes = &buf[..msb_bytes_len]; - let msb = match unsafe { msb_bytes.align_to::() } { - ([], aligned, []) => Cow::Borrowed(aligned), - _ => { - let mut vec = Vec::with_capacity(msb_len); - let ptr = msb_bytes.as_ptr() as *const C::BucketType; - for i in 0..msb_len { - vec.push(unsafe { ptr.add(i).read_unaligned() }); - } - Cow::Owned(vec) - } - }; Self { config: c, - msb, + msb: Cow::Borrowed(&msb_all[..msb_len]), lsb: BitVec::from_bytes(&buf[msb_bytes_len..]), } } @@ -404,10 +439,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { if self.msb.is_empty() { return Ok(0); } - let msb_buckets = self.msb.deref(); - let msb_bytes = unsafe { - std::slice::from_raw_parts(msb_buckets.as_ptr() as *const u8, size_of_val(msb_buckets)) - }; + let msb_bytes = C::BucketType::slice_as_bytes(&self.msb); writer.write_all(msb_bytes)?; let mut bytes_written = msb_bytes.len(); bytes_written += self.lsb.write(writer)?; @@ -530,12 +562,12 @@ fn split_into_msb( mut buckets: impl Iterator, lsb: &mut BitVec<'_>, max_msb_len: usize, -) -> Vec { - let mut msb: Vec = Vec::with_capacity(max_msb_len); - msb.extend(buckets.by_ref().take(max_msb_len)); +) -> Vec { + let mut msb: Vec = Vec::with_capacity(max_msb_len); + msb.extend(buckets.by_ref().take(max_msb_len).map(T::to_le_bytes)); if msb.len() == max_msb_len { // The msb is full: its smallest entry is the new boundary, the rest folds into the bits. - let smallest = msb[max_msb_len - 1].into_usize(); + let smallest = T::from_le_bytes(msb[max_msb_len - 1]).into_usize(); lsb.resize(smallest); let mut toggler = lsb.toggler(); for bucket in buckets { @@ -544,11 +576,12 @@ fn split_into_msb( } else { // Refill the msb from the highest bits, then truncate the bits to the new boundary. let need = max_msb_len - msb.len(); - let pulled: Vec = iter_ones::(lsb.bit_chunks().peekable()) + let pulled: Vec = iter_ones::(lsb.bit_chunks().peekable()) .take(need) + .map(T::to_le_bytes) .collect(); let smallest = if pulled.len() == need { - pulled[need - 1].into_usize() + T::from_le_bytes(pulled[need - 1]).into_usize() } else { 0 }; @@ -1228,7 +1261,10 @@ mod tests { a.xor_bit(7); a.xor_bit(17); a.xor_bit(11); - assert_eq!(vec![17, 11, 7], a.msb.iter().copied().collect_vec()); + assert_eq!( + vec![17, 11, 7], + a.msb.iter().copied().map(u16::from_le_bytes).collect_vec() + ); } #[test] @@ -1271,6 +1307,7 @@ mod tests { before.config.clone(), &writer[pad_amount..], ); + assert!(matches!(after.msb, Cow::Borrowed(_))); assert_eq!(before, after); } diff --git a/crates/geo_filters/src/diff_count/sim_hash.rs b/crates/geo_filters/src/diff_count/sim_hash.rs index 5c92aa6..f49cbe5 100644 --- a/crates/geo_filters/src/diff_count/sim_hash.rs +++ b/crates/geo_filters/src/diff_count/sim_hash.rs @@ -179,7 +179,7 @@ impl> GeoDiffCount<'_, C> { /// The zero-based index of the desired one bit is returned. fn nth_most_significant_one(&self, mut n: usize) -> Option { if n < self.msb.len() { - Some(self.msb[n]) + Some(C::BucketType::from_le_bytes(self.msb[n])) } else { n -= self.msb.len(); self.lsb @@ -194,8 +194,8 @@ impl> GeoDiffCount<'_, C> { fn bit_range(&self, range: &Range) -> u64 { let mut result = self.lsb.bit_range(range); if self.lsb.num_bits() < range.end { - for bit in self.msb.iter() { - let bit = bit.into_usize(); + for &bit in self.msb.iter() { + let bit = C::BucketType::from_le_bytes(bit).into_usize(); if range.contains(&bit) { result ^= (bit - range.start).into_block(); } From 4be98478c999cd9d9500160a44c6dc79358e3e1b Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 12:15:47 +0100 Subject: [PATCH 07/12] Use slice try_into in read_u32_le Convert bytes[offset..offset + 4] into [u8; 4] via try_into().expect(...) instead of indexing 4 elements individually. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/bpe/src/byte_pair_encoding.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/crates/bpe/src/byte_pair_encoding.rs b/crates/bpe/src/byte_pair_encoding.rs index 604d8fd..32568a9 100644 --- a/crates/bpe/src/byte_pair_encoding.rs +++ b/crates/bpe/src/byte_pair_encoding.rs @@ -64,12 +64,11 @@ fn serialize_daac( } fn read_u32_le(bytes: &[u8], offset: usize) -> u32 { - u32::from_le_bytes([ - bytes[offset], - bytes[offset + 1], - bytes[offset + 2], - bytes[offset + 3], - ]) + u32::from_le_bytes( + bytes[offset..offset + 4] + .try_into() + .expect("slice of length 4"), + ) } fn validate_daac_bytes(v: &[u8]) -> Result<(), &'static str> { From 1d0b6b158669c9353249bd29ede9083a124bffad Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 12:35:42 +0100 Subject: [PATCH 08/12] Nest PartialOrd match branches in GeoDiffCount Move msb and lsb comparisons into the Ordering::Equal branches of the preceding match expressions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/geo_filters/src/diff_count.rs | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index d7261ad..00b7beb 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -90,20 +90,18 @@ impl + Ord> Ord for GeoDiffCount<'_, C> { impl + PartialOrd> PartialOrd for GeoDiffCount<'_, C> { fn partial_cmp(&self, other: &Self) -> Option { match self.config.partial_cmp(&other.config)? { - Ordering::Equal => {} - ord => return Some(ord), - } - match self - .msb - .iter() - .copied() - .map(C::BucketType::from_le_bytes) - .partial_cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes))? - { - Ordering::Equal => {} - ord => return Some(ord), + Ordering::Equal => match self + .msb + .iter() + .copied() + .map(C::BucketType::from_le_bytes) + .partial_cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes))? + { + Ordering::Equal => self.lsb.partial_cmp(&other.lsb), + ord => Some(ord), + }, + ord => Some(ord), } - self.lsb.partial_cmp(&other.lsb) } } From 8e12238410e6a13ee0cde6c23a993a508936e5f3 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 13:00:03 +0100 Subject: [PATCH 09/12] Refactor geo_filters unaligned storage into Unaligned wrapper Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/geo_filters/src/config.rs | 2 + crates/geo_filters/src/config/buckets.rs | 114 +++++---- crates/geo_filters/src/config/unaligned.rs | 232 ++++++++++++++++++ crates/geo_filters/src/diff_count.rs | 148 ++++------- crates/geo_filters/src/diff_count/bitvec.rs | 61 ++--- crates/geo_filters/src/diff_count/metric.rs | 3 +- crates/geo_filters/src/diff_count/sim_hash.rs | 6 +- 7 files changed, 378 insertions(+), 188 deletions(-) create mode 100644 crates/geo_filters/src/config/unaligned.rs diff --git a/crates/geo_filters/src/config.rs b/crates/geo_filters/src/config.rs index adefa29..2ef00dd 100644 --- a/crates/geo_filters/src/config.rs +++ b/crates/geo_filters/src/config.rs @@ -8,12 +8,14 @@ mod bitchunks; mod buckets; mod estimation; mod lookup; +mod unaligned; pub(crate) use bitchunks::*; pub(crate) use buckets::*; pub(crate) use estimation::*; pub(crate) use lookup::*; use once_cell::sync::Lazy; +pub(crate) use unaligned::*; /// Trait that configures geo filters and determines counting precision. /// diff --git a/crates/geo_filters/src/config/buckets.rs b/crates/geo_filters/src/config/buckets.rs index 1d4b8a5..25154a4 100644 --- a/crates/geo_filters/src/config/buckets.rs +++ b/crates/geo_filters/src/config/buckets.rs @@ -12,14 +12,9 @@ where Self: Add + Rem, { const BITS: u32; - type Bytes: Copy + Debug + Default + Eq + PartialEq + Send + Sync + 'static; fn from_usize(v: usize) -> Self; fn into_usize(self) -> usize; - fn from_le_bytes(bytes: Self::Bytes) -> Self; - fn to_le_bytes(self) -> Self::Bytes; - fn slice_from_bytes(buf: &[u8]) -> (&[Self::Bytes], &[u8]); - fn slice_as_bytes(slice: &[Self::Bytes]) -> &[u8]; /// Returns a block with the bit in this bit position set to one. #[inline] @@ -80,50 +75,75 @@ where } } -macro_rules! impl_is_bucket_type { - ($t:ty) => { - impl IsBucketType for $t { - const BITS: u32 = Self::BITS; - type Bytes = [u8; std::mem::size_of::<$t>()]; - - #[inline] - fn from_usize(v: usize) -> Self { - v as Self - } - - #[inline] - fn into_usize(self) -> usize { - self as usize - } - - #[inline] - fn from_le_bytes(bytes: Self::Bytes) -> Self { - Self::from_le_bytes(bytes) - } - - #[inline] - fn to_le_bytes(self) -> Self::Bytes { - Self::to_le_bytes(self) - } - - #[inline] - fn slice_from_bytes(buf: &[u8]) -> (&[Self::Bytes], &[u8]) { - buf.as_chunks::<{ std::mem::size_of::<$t>() }>() - } - - #[inline] - fn slice_as_bytes(slice: &[Self::Bytes]) -> &[u8] { - slice.as_flattened() - } - } - }; +impl IsBucketType for u8 { + const BITS: u32 = Self::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + v as Self + } + + #[inline] + fn into_usize(self) -> usize { + self as usize + } } -impl_is_bucket_type!(u8); -impl_is_bucket_type!(u16); -impl_is_bucket_type!(u32); -impl_is_bucket_type!(u64); -impl_is_bucket_type!(usize); +impl IsBucketType for u16 { + const BITS: u32 = Self::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + v as Self + } + + #[inline] + fn into_usize(self) -> usize { + self as usize + } +} + +impl IsBucketType for u32 { + const BITS: u32 = Self::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + v as Self + } + + #[inline] + fn into_usize(self) -> usize { + self as usize + } +} + +impl IsBucketType for u64 { + const BITS: u32 = Self::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + v as Self + } + + #[inline] + fn into_usize(self) -> usize { + self as usize + } +} + +impl IsBucketType for usize { + const BITS: u32 = Self::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + v + } + + #[inline] + fn into_usize(self) -> usize { + self + } +} /// Computes the bits required to store bucket positions for 64-bit hashes given that (1 << B) /// buckets cover half the hash space. diff --git a/crates/geo_filters/src/config/unaligned.rs b/crates/geo_filters/src/config/unaligned.rs new file mode 100644 index 0000000..51b5755 --- /dev/null +++ b/crates/geo_filters/src/config/unaligned.rs @@ -0,0 +1,232 @@ +use std::cmp::Ordering; +use std::fmt::{self, Debug, Formatter}; +use std::mem::{align_of, size_of, size_of_val}; +use std::ops::{Add, BitAnd, BitAndAssign, BitXor, BitXorAssign, Rem, Shl, Shr}; + +use crate::config::IsBucketType; + +/// A wrapper that stores a value of type `T` with 1-byte alignment (`#[repr(C, packed)]`), +/// allowing zero-copy borrowing of unaligned byte slices while supporting value comparisons +/// and bitwise/bucket operations. +#[derive(Copy, Default)] +#[repr(C, packed)] +pub(crate) struct Unaligned(T); + +impl Clone for Unaligned { + #[inline] + fn clone(&self) -> Self { + *self + } +} + +impl Unaligned { + #[inline] + pub const fn new(value: T) -> Self { + Self(value) + } + + #[inline] + pub const fn get(self) -> T { + self.0 + } +} + +impl Unaligned { + #[inline] + pub fn slice_from_bytes(bytes: &[u8]) -> (&[Self], &[u8]) { + const { + assert!(align_of::() == 1); + assert!(size_of::() == size_of::()); + } + let elem_size = size_of::(); + let len = bytes.len() / elem_size; + let (head, tail) = bytes.split_at(len * elem_size); + // SAFETY: + // - `Unaligned` is `#[repr(C, packed)]` with alignment 1 and size `size_of::()`. + // - `T: IsBucketType` is only implemented for primitive unsigned integers (`u8`, `u16`, + // `u32`, `u64`, `usize`), for which all initialized byte patterns are valid. + // - `head` has exact byte length `len * size_of::()`. + let slice = unsafe { std::slice::from_raw_parts(head.as_ptr().cast::(), len) }; + (slice, tail) + } + + #[inline] + pub fn slice_as_bytes(slice: &[Self]) -> &[u8] { + // SAFETY: `Unaligned` is `#[repr(C, packed)]` wrapping a primitive unsigned integer `T` + // with no padding, so all bytes in `slice` are initialized. + unsafe { std::slice::from_raw_parts(slice.as_ptr().cast::(), size_of_val(slice)) } + } +} + +impl From for Unaligned { + #[inline] + fn from(value: T) -> Self { + Self::new(value) + } +} + +impl Debug for Unaligned { + #[inline] + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + Debug::fmt(&self.get(), f) + } +} + +impl PartialEq for Unaligned { + #[inline] + fn eq(&self, other: &Self) -> bool { + self.get() == other.get() + } +} + +impl PartialEq for Unaligned { + #[inline] + fn eq(&self, other: &T) -> bool { + self.get() == *other + } +} + +macro_rules! impl_partial_eq_for_primitive { + ($($t:ty),* $(,)?) => { + $( + impl PartialEq> for $t { + #[inline] + fn eq(&self, other: &Unaligned<$t>) -> bool { + *self == other.get() + } + } + )* + }; +} + +impl_partial_eq_for_primitive!(u8, u16, u32, u64, usize); + +impl Eq for Unaligned {} + +impl PartialOrd for Unaligned { + #[inline] + fn partial_cmp(&self, other: &Self) -> Option { + self.get().partial_cmp(&other.get()) + } +} + +impl PartialOrd for Unaligned { + #[inline] + fn partial_cmp(&self, other: &T) -> Option { + self.get().partial_cmp(other) + } +} + +impl Ord for Unaligned { + #[inline] + fn cmp(&self, other: &Self) -> Ordering { + self.get().cmp(&other.get()) + } +} + +impl> Add for Unaligned { + type Output = Self; + + #[inline] + fn add(self, rhs: Self) -> Self::Output { + Self::new(self.get() + rhs.get()) + } +} + +impl> Rem for Unaligned { + type Output = Self; + + #[inline] + fn rem(self, rhs: Self) -> Self::Output { + Self::new(self.get() % rhs.get()) + } +} + +impl> BitAnd for Unaligned { + type Output = T; + + #[inline] + fn bitand(self, rhs: T) -> Self::Output { + self.get() & rhs + } +} + +impl> BitAnd for &Unaligned { + type Output = T; + + #[inline] + fn bitand(self, rhs: T) -> Self::Output { + self.get() & rhs + } +} + +impl> BitAndAssign for Unaligned { + #[inline] + fn bitand_assign(&mut self, rhs: T) { + self.0 = self.get() & rhs; + } +} + +impl> BitXor for Unaligned { + type Output = T; + + #[inline] + fn bitxor(self, rhs: Self) -> Self::Output { + self.get() ^ rhs.get() + } +} + +impl> BitXor for Unaligned { + type Output = T; + + #[inline] + fn bitxor(self, rhs: T) -> Self::Output { + self.get() ^ rhs + } +} + +impl> BitXorAssign for Unaligned { + #[inline] + fn bitxor_assign(&mut self, rhs: T) { + self.0 = self.get() ^ rhs; + } +} + +impl> Shl for Unaligned { + type Output = T; + + #[inline] + fn shl(self, rhs: usize) -> Self::Output { + self.get() << rhs + } +} + +impl> Shr for Unaligned { + type Output = T; + + #[inline] + fn shr(self, rhs: usize) -> Self::Output { + self.get() >> rhs + } +} + +impl Unaligned { + #[inline] + pub const fn count_ones(self) -> u32 { + self.get().count_ones() + } +} + +impl IsBucketType for Unaligned { + const BITS: u32 = T::BITS; + + #[inline] + fn from_usize(v: usize) -> Self { + Self::new(T::from_usize(v)) + } + + #[inline] + fn into_usize(self) -> usize { + self.get().into_usize() + } +} diff --git a/crates/geo_filters/src/diff_count.rs b/crates/geo_filters/src/diff_count.rs index 00b7beb..6b5b0c9 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -7,8 +7,8 @@ use std::mem::{size_of, size_of_val}; use crate::config::{ count_ones_from_bitchunks, count_ones_from_msb_and_lsb, iter_bit_chunks, iter_ones, - mask_bit_chunks, take_ref, xor_bit_chunks, BitChunk, GeoConfig, IsBucketType, BITS_PER_BLOCK, - BYTES_PER_BLOCK, + mask_bit_chunks, take_ref, xor_bit_chunks, BitChunk, GeoConfig, IsBucketType, Unaligned, + BITS_PER_BLOCK, }; use crate::{Count, Diff}; @@ -64,47 +64,14 @@ pub type GeoDiffCount13<'a> = GeoDiffCount<'a, GeoDiffConfig13>; /// with an error of +/-1200 items using the GeoDiffCount. A HyperLogLog data structure with /// 1500 bytes achieves a precision of ~0.022 and could estimate the 10k items with an error /// of +/-22k items in the best case which is 20 times worse despite using 5 times more space! -#[derive(Clone, Default, PartialEq, Eq)] +#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord)] pub struct GeoDiffCount<'a, C: GeoConfig> { config: C, /// The bit positions are stored from largest to smallest. - msb: Cow<'a, [::Bytes]>, + msb: Cow<'a, [Unaligned]>, lsb: BitVec<'a>, } -impl + Ord> Ord for GeoDiffCount<'_, C> { - fn cmp(&self, other: &Self) -> Ordering { - self.config - .cmp(&other.config) - .then_with(|| { - self.msb - .iter() - .copied() - .map(C::BucketType::from_le_bytes) - .cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes)) - }) - .then_with(|| self.lsb.cmp(&other.lsb)) - } -} - -impl + PartialOrd> PartialOrd for GeoDiffCount<'_, C> { - fn partial_cmp(&self, other: &Self) -> Option { - match self.config.partial_cmp(&other.config)? { - Ordering::Equal => match self - .msb - .iter() - .copied() - .map(C::BucketType::from_le_bytes) - .partial_cmp(other.msb.iter().copied().map(C::BucketType::from_le_bytes))? - { - Ordering::Equal => self.lsb.partial_cmp(&other.lsb), - ord => Some(ord), - }, - ord => Some(ord), - } - } -} - impl> std::fmt::Debug for GeoDiffCount<'_, C> { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!( @@ -134,15 +101,15 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// Note: we need a peekable iterator, such that we can extract the most significant bits without /// having to construct another iterator with the remaining `BitChunk`s. fn from_bit_chunks>(config: C, chunks: I) -> Self { - let mut ones = iter_ones::(chunks.peekable()); + let mut ones = iter_ones::, _>(chunks.peekable()); let mut msb = Vec::default(); take_ref(&mut ones, config.max_msb_len() - 1).for_each(|bucket| { - msb.push(bucket.to_le_bytes()); + msb.push(bucket); }); let smallest_msb = ones .next() .inspect(|bucket| { - msb.push(bucket.to_le_bytes()); + msb.push(*bucket); }) .unwrap_or_default(); let lsb = BitVec::from_bit_chunks(ones.into_bitchunks(), smallest_msb.into_usize()); @@ -213,7 +180,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { let msb = self .msb .iter() - .map(|&p| C::BucketType::from_le_bytes(p).into_usize()) + .map(|&p| p.into_usize()) .filter(move |&pos| (mask >> (pos % mask_size)) & 1 != 0) .map(|pos| pos as u64); let lsb = iter_ones::( @@ -243,12 +210,11 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { } fn test_bit(&self, index: C::BucketType) -> bool { + let index = Unaligned::new(index); if index.into_usize() < self.lsb.num_bits() { self.lsb.test_bit(index.into_usize()) } else { - self.msb - .binary_search_by(|&k| index.cmp(&C::BucketType::from_le_bytes(k))) - .is_ok() + self.msb.binary_search_by(|k| index.cmp(k)).is_ok() } } @@ -276,9 +242,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// most to least significant, overlapping `BitChunk`s got merged, and empty ones filtered. fn bit_chunks(&self) -> impl Iterator + '_ { iter_bit_chunks( - self.msb - .iter() - .map(|&b| C::BucketType::from_le_bytes(b).into_usize()), + self.msb.iter().map(|b| b.into_usize()), self.lsb.bit_chunks(), ) } @@ -288,13 +252,9 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// With ones = 128, a relative precision of ~0.12 is achieved. /// Using numbers larger than 128 won't increase the precision! fn estimate_size(&self) -> f32 { - let (total, lowest_bucket) = count_ones_from_msb_and_lsb::( + let (total, lowest_bucket) = count_ones_from_msb_and_lsb::>( self.msb.len(), - self.msb - .last() - .copied() - .map(C::BucketType::from_le_bytes) - .unwrap_or_default(), + self.msb.last().copied().unwrap_or_default(), self.lsb.bit_chunks().peekable(), self.config.max_bytes(), self.config.max_msb_len(), @@ -313,13 +273,14 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { /// executions has a constant complexity, but one that is rather high compared to the fast path. In total /// that makes the cost of the else case negligible. fn xor_bit(&mut self, bucket: C::BucketType) { + let bucket = Unaligned::new(bucket); if bucket.into_usize() < self.lsb.num_bits() { // The bit being toggled is within our LSB bit vector // so toggle it directly. self.lsb.toggle(bucket.into_usize()); } else { let msb = self.msb.to_mut(); - match msb.binary_search_by(|&k| bucket.cmp(&C::BucketType::from_le_bytes(k))) { + match msb.binary_search_by(|k| bucket.cmp(k)) { Ok(idx) => { msb.remove(idx); let first = { @@ -327,33 +288,33 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { lsb.next() }; if let Some(smallest) = first { - msb.push(C::BucketType::from_usize(smallest).to_le_bytes()); + msb.push(Unaligned::from_usize(smallest)); self.lsb.resize(smallest); } else { self.lsb.resize(0); }; } Err(idx) => { - msb.insert(idx, bucket.to_le_bytes()); + msb.insert(idx, bucket); if msb.len() > self.config.max_msb_len() { // We have too many values in the MSB sparse index vector, // let's move the smalles MSB value into the LSB bit vector - let smallest = C::BucketType::from_le_bytes( - msb.pop().expect("we should have at least one element!"), - ) - .into_usize(); - let new_smallest = C::BucketType::from_le_bytes( - *msb.last().expect("should have at least one element"), - ) - .into_usize(); + let smallest = msb + .pop() + .expect("we should have at least one element!") + .into_usize(); + let new_smallest = msb + .last() + .expect("should have at least one element") + .into_usize(); // ensure LSB bit vector has the space for `smallest` self.lsb.resize(new_smallest); self.lsb.toggle(smallest); } else if msb.len() == self.config.max_msb_len() { - let smallest = C::BucketType::from_le_bytes( - *msb.last().expect("should have at least one element"), - ) - .into_usize(); + let smallest = msb + .last() + .expect("should have at least one element") + .into_usize(); self.lsb.resize(smallest); } } @@ -390,11 +351,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { ); debug_assert!( self.lsb.num_bits() == 0 - || self - .msb - .last() - .map_or(0, |&b| C::BucketType::from_le_bytes(b).into_usize()) - == self.lsb.num_bits(), + || self.msb.last().map_or(0, |b| b.into_usize()) == self.lsb.num_bits(), "msb and lsb non-contiguous: msb = {:?} |lsb| = {:?}", self.msb, self.lsb.num_bits(), @@ -420,7 +377,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { return Self::new(c); } // The number of most significant bits stores in the MSB sparse repr - let (msb_all, _) = C::BucketType::slice_from_bytes(buf); + let (msb_all, _) = Unaligned::::slice_from_bytes(buf); let msb_len = msb_all.len().min(c.max_msb_len()); // The number of bytes representing the MSB - this is how many bytes we need to // skip over to reach the LSB @@ -437,7 +394,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { if self.msb.is_empty() { return Ok(0); } - let msb_bytes = C::BucketType::slice_as_bytes(&self.msb); + let msb_bytes = Unaligned::slice_as_bytes(&self.msb); writer.write_all(msb_bytes)?; let mut bytes_written = msb_bytes.len(); bytes_written += self.lsb.write(writer)?; @@ -560,12 +517,12 @@ fn split_into_msb( mut buckets: impl Iterator, lsb: &mut BitVec<'_>, max_msb_len: usize, -) -> Vec { - let mut msb: Vec = Vec::with_capacity(max_msb_len); - msb.extend(buckets.by_ref().take(max_msb_len).map(T::to_le_bytes)); +) -> Vec { + let mut msb: Vec = Vec::with_capacity(max_msb_len); + msb.extend(buckets.by_ref().take(max_msb_len)); if msb.len() == max_msb_len { // The msb is full: its smallest entry is the new boundary, the rest folds into the bits. - let smallest = T::from_le_bytes(msb[max_msb_len - 1]).into_usize(); + let smallest = msb[max_msb_len - 1].into_usize(); lsb.resize(smallest); let mut toggler = lsb.toggler(); for bucket in buckets { @@ -574,12 +531,11 @@ fn split_into_msb( } else { // Refill the msb from the highest bits, then truncate the bits to the new boundary. let need = max_msb_len - msb.len(); - let pulled: Vec = iter_ones::(lsb.bit_chunks().peekable()) + let pulled: Vec = iter_ones::(lsb.bit_chunks().peekable()) .take(need) - .map(T::to_le_bytes) .collect(); let smallest = if pulled.len() == need { - T::from_le_bytes(pulled[need - 1]).into_usize() + pulled[need - 1].into_usize() } else { 0 }; @@ -608,7 +564,7 @@ pub struct GeoDiffCountBuilder> { /// (xor) into `blocks`. [`GeoDiffCountBuilder::cleanup`] reconciles the two. split: usize, numbers: Vec, - blocks: Vec<[u8; BYTES_PER_BLOCK]>, + blocks: Vec>, } impl> GeoDiffCountBuilder { @@ -625,7 +581,7 @@ impl> GeoDiffCountBuilder { expected, split, numbers: Vec::with_capacity(capacity), - blocks: vec![[0; BYTES_PER_BLOCK]; split.div_ceil(BITS_PER_BLOCK)], + blocks: vec![Unaligned::default(); split.div_ceil(BITS_PER_BLOCK)], } } @@ -641,7 +597,7 @@ impl> GeoDiffCountBuilder { if new_split > self.split { self.split = new_split; self.blocks - .resize(new_split.div_ceil(BITS_PER_BLOCK), [0; BYTES_PER_BLOCK]); + .resize(new_split.div_ceil(BITS_PER_BLOCK), Unaligned::default()); } } @@ -666,8 +622,7 @@ impl> GeoDiffCountBuilder { if (next - read) % 2 == 1 { if bucket < split { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] = - (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); + blocks[index] ^= bit.into_block(); } else { numbers[write] = bucket; write += 1; @@ -698,11 +653,11 @@ impl> GeoDiffCountBuilder { keep = self.numbers.partition_point(|&b| b >= new_split); } self.blocks - .resize(new_split.div_ceil(BITS_PER_BLOCK), [0; BYTES_PER_BLOCK]); + .resize(new_split.div_ceil(BITS_PER_BLOCK), Unaligned::default()); let blocks = &mut self.blocks; for &bucket in &self.numbers[keep..] { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] = (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); + blocks[index] ^= bit.into_block(); } self.numbers.truncate(keep); self.split = new_split; @@ -723,8 +678,7 @@ impl> GeoDiffCountBuilder { } else { // `bucket < split`, so the block index is always in range; toggling cancels repeats. let (index, bit) = bucket.into_index_and_bit(); - self.blocks[index] = - (u64::from_le_bytes(self.blocks[index]) ^ bit.into_block()).to_le_bytes(); + self.blocks[index] ^= bit.into_block(); } } @@ -760,8 +714,7 @@ impl> GeoDiffCountBuilder { } } else { let (index, bit) = bucket.into_index_and_bit(); - blocks[index] = - (u64::from_le_bytes(blocks[index]) ^ bit.into_block()).to_le_bytes(); + blocks[index] ^= bit.into_block(); } } filled @@ -782,7 +735,9 @@ impl> GeoDiffCountBuilder { self.cleanup(); let mut lsb = BitVec::from_blocks(self.blocks, self.split); let msb = split_into_msb( - self.numbers.iter().map(|&b| C::BucketType::from_usize(b)), + self.numbers + .iter() + .map(|&b| Unaligned::::from_usize(b)), &mut lsb, max_msb_len, ); @@ -1259,10 +1214,7 @@ mod tests { a.xor_bit(7); a.xor_bit(17); a.xor_bit(11); - assert_eq!( - vec![17, 11, 7], - a.msb.iter().copied().map(u16::from_le_bytes).collect_vec() - ); + assert_eq!(vec![17, 11, 7], a.msb.iter().copied().collect_vec()); } #[test] diff --git a/crates/geo_filters/src/diff_count/bitvec.rs b/crates/geo_filters/src/diff_count/bitvec.rs index efebfa6..e58c647 100644 --- a/crates/geo_filters/src/diff_count/bitvec.rs +++ b/crates/geo_filters/src/diff_count/bitvec.rs @@ -5,7 +5,7 @@ use std::ops::{Deref as _, Index, Range}; use crate::config::IsBucketType; use crate::config::BITS_PER_BLOCK; -use crate::config::{BitChunk, BYTES_PER_BLOCK}; +use crate::config::{BitChunk, Unaligned}; /// A bit vector where every bit occupies exactly one bit (in contrast to `Vec` where each /// bit consumes 1 byte). It only implements the minimum number of operations that we need for our @@ -14,18 +14,13 @@ use crate::config::{BitChunk, BYTES_PER_BLOCK}; #[derive(Clone, Default, Debug, Eq, PartialEq)] pub(crate) struct BitVec<'a> { num_bits: usize, - blocks: Cow<'a, [[u8; BYTES_PER_BLOCK]]>, + blocks: Cow<'a, [Unaligned]>, } impl Ord for BitVec<'_> { fn cmp(&self, other: &Self) -> std::cmp::Ordering { match self.num_bits.cmp(&other.num_bits) { - Ordering::Equal => self - .blocks - .iter() - .rev() - .map(|&b| u64::from_le_bytes(b)) - .cmp(other.blocks.iter().rev().map(|&b| u64::from_le_bytes(b))), + Ordering::Equal => self.blocks.iter().rev().cmp(other.blocks.iter().rev()), ord => ord, } } @@ -46,7 +41,7 @@ impl<'a> BitVec<'a> { result.resize(num_bits); let blocks = result.blocks.to_mut(); for BitChunk { index, block } in chunks { - blocks[index] = (u64::from_le_bytes(blocks[index]) ^ block).to_le_bytes(); + blocks[index] ^= block; } result.clear_superfluous_bits(); result @@ -54,7 +49,7 @@ impl<'a> BitVec<'a> { /// Wraps raw `blocks` covering `[0, num_bits)` into a `BitVec`. The number of blocks must match /// `num_bits`, and any bits at or above `num_bits` are cleared. - pub fn from_blocks(blocks: Vec<[u8; BYTES_PER_BLOCK]>, num_bits: usize) -> BitVec<'static> { + pub fn from_blocks(blocks: Vec>, num_bits: usize) -> BitVec<'static> { debug_assert_eq!(blocks.len(), num_bits.div_ceil(BITS_PER_BLOCK)); let mut result = BitVec { num_bits, @@ -70,7 +65,7 @@ impl<'a> BitVec<'a> { if num_blocks != self.blocks.len() { self.blocks .to_mut() - .resize(num_blocks, [0; BYTES_PER_BLOCK]); + .resize(num_blocks, Unaligned::default()); } self.num_bits = num_bits; self.clear_superfluous_bits(); @@ -83,9 +78,7 @@ impl<'a> BitVec<'a> { .to_mut() .last_mut() .into_iter() - .for_each(|block| { - *block = (u64::from_le_bytes(*block) & bit.into_mask()).to_le_bytes(); - }); + .for_each(|block| *block &= bit.into_mask()); } } @@ -97,20 +90,15 @@ impl<'a> BitVec<'a> { pub fn count_ones(&self) -> usize { self.blocks .iter() - .map(|&block| u64::from_le_bytes(block).count_ones() as usize) + .map(|block| block.count_ones() as usize) .sum() } /// The raw 64-bit blocks backing the vector, from least to most significant. - pub fn blocks(&self) -> &[[u8; BYTES_PER_BLOCK]] { + pub fn blocks(&self) -> &[Unaligned] { self.blocks.deref() } - #[inline] - pub fn block(&self, index: usize) -> u64 { - u64::from_le_bytes(self.blocks[index]) - } - pub fn is_empty(&self) -> bool { self.num_bits() == 0 } @@ -119,15 +107,14 @@ impl<'a> BitVec<'a> { pub fn test_bit(&self, index: usize) -> bool { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - (self.block(block_idx) & bit_idx.into_block()) != 0 + (self.blocks[block_idx] & bit_idx.into_block()) != 0 } /// Toggles the bit specified by the provided zero-based bit position. pub fn toggle(&mut self, index: usize) { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - let block = &mut self.blocks.to_mut()[block_idx]; - *block = (u64::from_le_bytes(*block) ^ bit_idx.into_block()).to_le_bytes(); + self.blocks.to_mut()[block_idx] ^= bit_idx.into_block(); } /// Returns a [`BitToggler`] that toggles many bits without re-resolving the `Cow` on every @@ -146,11 +133,10 @@ impl<'a> BitVec<'a> { pub fn bit_chunks(&self) -> impl Iterator + '_ { self.blocks .iter() - .map(|&block| u64::from_le_bytes(block)) .enumerate() .rev() - .filter(|(_, block)| *block != 0) - .map(|(index, block)| BitChunk::new(index, block)) + .filter(|(_, block)| **block != 0) + .map(|(index, block)| BitChunk::new(index, block.get())) } // Returns an owned BitVec with static lifetime that can be used for cases when we want @@ -177,10 +163,10 @@ impl<'a> BitVec<'a> { if block_start >= self.blocks.len() { return 0; } - let mut result = self.block(block_start) >> start_idx; + let mut result = self.blocks[block_start] >> start_idx; let block_end = (range.end - 1) / BITS_PER_BLOCK; if block_end < self.blocks.len() && block_end != block_start { - result |= self.block(block_end) << (BITS_PER_BLOCK - start_idx); + result |= self.blocks[block_end] << (BITS_PER_BLOCK - start_idx); } result & (u64::MAX >> (BITS_PER_BLOCK - range.len())) } @@ -205,7 +191,7 @@ impl<'a> BitVec<'a> { unoccupied_bits ); buf = &buf[1..]; - let (blocks, remainder) = buf.as_chunks::(); + let (blocks, remainder) = Unaligned::::slice_from_bytes(buf); assert!( !blocks.is_empty() && remainder.is_empty(), "buffer should be a non-empty multiple of 8 bytes, got {}", @@ -226,7 +212,7 @@ impl<'a> BitVec<'a> { // First serialize the number of unoccupied bits in the last block as one byte. let unoccupied_bits = 63 - ((self.num_bits - 1) % 64) as u8; writer.write_all(&[unoccupied_bits])?; - let block_bytes = self.blocks.as_flattened(); + let block_bytes = Unaligned::slice_as_bytes(&self.blocks); writer.write_all(block_bytes)?; Ok(block_bytes.len() + 1) } @@ -239,7 +225,7 @@ impl Index for BitVec<'_> { fn index(&self, index: usize) -> &Self::Output { assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - match (self.block(block_idx) & bit_idx.into_block()) != 0 { + match (self.blocks[block_idx] & bit_idx.into_block()) != 0 { true => &true, false => &false, } @@ -250,7 +236,7 @@ impl Index for BitVec<'_> { /// Obtained via [`BitVec::toggler`]. pub(crate) struct BitToggler<'a> { num_bits: usize, - blocks: &'a mut [[u8; BYTES_PER_BLOCK]], + blocks: &'a mut [Unaligned], } impl BitToggler<'_> { @@ -259,8 +245,7 @@ impl BitToggler<'_> { pub fn toggle(&mut self, index: usize) { debug_assert!(index < self.num_bits); let (block_idx, bit_idx) = index.into_index_and_bit(); - let block = &mut self.blocks[block_idx]; - *block = (u64::from_le_bytes(*block) ^ bit_idx.into_block()).to_le_bytes(); + self.blocks[block_idx] ^= bit_idx.into_block(); } } @@ -271,7 +256,7 @@ mod tests { /// Constructs a BitVec instance with num_bits and all bits initialized to the provided default value. fn bitvec_from_element<'a>(num_bits: usize, value: bool) -> BitVec<'a> { let num_blocks = num_bits.div_ceil(BITS_PER_BLOCK) + 1; - let fill = if value { u64::MAX } else { u64::MIN }.to_le_bytes(); + let fill = Unaligned::new(if value { u64::MAX } else { u64::MIN }); let blocks = vec![fill; num_blocks]; BitVec { blocks: Cow::from(blocks), @@ -314,7 +299,7 @@ mod tests { assert!(bitvec.test_bit(14)); assert!(bitvec.test_bit(48)); assert!(bitvec.test_bit(72)); - let last = u64::from_le_bytes(*bitvec.blocks.last().expect("there must be a block")); + let last = bitvec.blocks.last().expect("there must be a block"); assert!(last & 120usize.into_bit().into_block() == 0); } @@ -322,7 +307,7 @@ mod tests { { assert!(bitvec.test_bit(14)); assert!(bitvec.test_bit(48)); - let last = u64::from_le_bytes(*bitvec.blocks.last().expect("there must be a block")); + let last = bitvec.blocks.last().expect("there must be a block"); assert!(last & 72usize.into_bit().into_block() == 0); } diff --git a/crates/geo_filters/src/diff_count/metric.rs b/crates/geo_filters/src/diff_count/metric.rs index 4262652..d618826 100644 --- a/crates/geo_filters/src/diff_count/metric.rs +++ b/crates/geo_filters/src/diff_count/metric.rs @@ -184,8 +184,7 @@ impl + Default> MetricSpace for GeoDiffMetric<'_, C> { let mut count = 0; let (a_blocks, b_blocks) = (a.lsb.blocks(), b.lsb.blocks()); for i in 0..common_full { - count += (u64::from_le_bytes(a_blocks[i]) ^ u64::from_le_bytes(b_blocks[i])) - .count_ones() as usize; + count += (a_blocks[i] ^ b_blocks[i]).count_ones() as usize; if count >= bound { return OnesMetric::infinite(); } diff --git a/crates/geo_filters/src/diff_count/sim_hash.rs b/crates/geo_filters/src/diff_count/sim_hash.rs index f49cbe5..28eb69e 100644 --- a/crates/geo_filters/src/diff_count/sim_hash.rs +++ b/crates/geo_filters/src/diff_count/sim_hash.rs @@ -179,7 +179,7 @@ impl> GeoDiffCount<'_, C> { /// The zero-based index of the desired one bit is returned. fn nth_most_significant_one(&self, mut n: usize) -> Option { if n < self.msb.len() { - Some(C::BucketType::from_le_bytes(self.msb[n])) + Some(self.msb[n].get()) } else { n -= self.msb.len(); self.lsb @@ -194,8 +194,8 @@ impl> GeoDiffCount<'_, C> { fn bit_range(&self, range: &Range) -> u64 { let mut result = self.lsb.bit_range(range); if self.lsb.num_bits() < range.end { - for &bit in self.msb.iter() { - let bit = C::BucketType::from_le_bytes(bit).into_usize(); + for bit in self.msb.iter() { + let bit = bit.into_usize(); if range.contains(&bit) { result ^= (bit - range.start).into_block(); } From 4234986e9dffcaaa512dfbf5acef7ebec5898648 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 14:26:56 +0100 Subject: [PATCH 10/12] Update ctrl tags immediately during sort_by_hash insertion sort Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/hash-sorted-map/src/hash_sorted_map.rs | 59 +++++++++++-------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/crates/hash-sorted-map/src/hash_sorted_map.rs b/crates/hash-sorted-map/src/hash_sorted_map.rs index 088eb0e..b8c6c20 100644 --- a/crates/hash-sorted-map/src/hash_sorted_map.rs +++ b/crates/hash-sorted-map/src/hash_sorted_map.rs @@ -232,6 +232,7 @@ impl> HashSortedMap { if j < i { let cur_key = unsafe { self.groups[gi].keys[si].assume_init_read() }; let cur_val = unsafe { self.groups[gi].values[si].assume_init_read() }; + let cur_ctrl = self.groups[gi].ctrl[si]; // Shift positions j..i up by one. hashes.copy_within(j..i, j + 1); for pos in (j..i).rev() { @@ -240,8 +241,10 @@ impl> HashSortedMap { unsafe { let k = std::ptr::read(&self.groups[src_g].keys[src_s]); let v = std::ptr::read(&self.groups[src_g].values[src_s]); + let c = self.groups[src_g].ctrl[src_s]; self.groups[dst_g].keys[dst_s] = k; self.groups[dst_g].values[dst_s] = v; + self.groups[dst_g].ctrl[dst_s] = c; } } // Insert at position j. @@ -249,15 +252,9 @@ impl> HashSortedMap { let (gj, sj) = chain_slot(&chain, j); self.groups[gj].keys[sj] = MaybeUninit::new(cur_key); self.groups[gj].values[sj] = MaybeUninit::new(cur_val); + self.groups[gj].ctrl[sj] = cur_ctrl; } } - // Rebuild ctrl/tag bytes from the sorted hashes so that - // get/insert/entry still work after sorting. - // This adds a small performance penalty of maybe 6%. - for (pos, &h) in hashes.iter().enumerate() { - let (gi, si) = chain_slot(&chain, pos); - self.groups[gi].ctrl[si] = tag(h); - } } } } @@ -1084,43 +1081,55 @@ mod tests { } } - impl SortingHash for FixedState { - fn hash(&self, _key: &DropTrackedKey) -> u64 { - self.0 + struct KeyHasher; + impl SortingHash for KeyHasher { + fn hash(&self, key: &DropTrackedKey) -> u64 { + // Key 0 gets hash 2; keys 1 and 2 get hash 1 (all map to group 0 when n_bits == 0). + if key.id == 0 { + 2 + } else { + 1 + } } } let key_drops = Arc::new(AtomicUsize::new(0)); let val_drops = Arc::new(AtomicUsize::new(0)); - struct DropTrackedVal(Arc); + struct DropTrackedVal(u32, Arc); impl Drop for DropTrackedVal { fn drop(&mut self) { - self.0.fetch_add(1, AtomicOrdering::SeqCst); + self.1.fetch_add(1, AtomicOrdering::SeqCst); } } - let result = catch_unwind(AssertUnwindSafe(|| { - let mut map = HashSortedMap::with_capacity_and_hasher(1, FixedState(42)); - map.insert( - DropTrackedKey { - id: 1, - drops: Arc::clone(&key_drops), - }, - DropTrackedVal(Arc::clone(&val_drops)), - ); + let mut map = HashSortedMap::with_capacity_and_hasher(1, KeyHasher); + for id in [0, 1, 2] { map.insert( DropTrackedKey { - id: 2, + id, drops: Arc::clone(&key_drops), }, - DropTrackedVal(Arc::clone(&val_drops)), + DropTrackedVal(id * 10, Arc::clone(&val_drops)), ); + } + + let result = catch_unwind(AssertUnwindSafe(|| { map.sort_by_hash(); })); assert!(result.is_err()); - assert_eq!(key_drops.load(AtomicOrdering::SeqCst), 2); - assert_eq!(val_drops.load(AtomicOrdering::SeqCst), 2); + // Even after a panic following a completed insertion-sort move (hashes [2, 1, 1]), + // ctrl tags stay in sync with moved keys. + for id in [0, 1, 2] { + let probe = DropTrackedKey { + id, + drops: Arc::new(AtomicUsize::new(0)), + }; + assert_eq!(map.get(&probe).map(|v| v.0), Some(id * 10)); + } + drop(map); + assert_eq!(key_drops.load(AtomicOrdering::SeqCst), 3); + assert_eq!(val_drops.load(AtomicOrdering::SeqCst), 3); } } From 28bebc7a9df373bf911df465df8c7b6542c31868 Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 14:35:50 +0100 Subject: [PATCH 11/12] Document why group_ops uses from_le_bytes instead of from_ne_bytes Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/hash-sorted-map/src/group_ops.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/crates/hash-sorted-map/src/group_ops.rs b/crates/hash-sorted-map/src/group_ops.rs index e936ce2..3bda3b8 100644 --- a/crates/hash-sorted-map/src/group_ops.rs +++ b/crates/hash-sorted-map/src/group_ops.rs @@ -19,6 +19,15 @@ pub type Mask = u32; pub type Mask = u64; // ── SIMD group operations ─────────────────────────────────────────────────── +// +// Note on byte order: `from_le_bytes` is used (rather than `from_ne_bytes`) when +// packing `ctrl` into an integer so that `ctrl[0]` always occupies the least- +// significant byte and `ctrl[GROUP_SIZE - 1]` occupies the most-significant byte. +// `lowest` relies on `trailing_zeros()` to map the first matching byte to slot +// index `0..GROUP_SIZE`, and `count_occupied` relies on `leading_zeros()` to +// count empty trailing slots at the high-index end of `ctrl`. On little-endian +// targets `from_le_bytes` compiles to a single unaligned load identical to +// `from_ne_bytes`. #[cfg(target_arch = "x86_64")] mod arch { From 1b5785cc4526a3e7aa43a4a0699542660c28cd8f Mon Sep 17 00:00:00 2001 From: Alexander Neubeck Date: Wed, 4 Nov 2026 15:39:32 +0100 Subject: [PATCH 12/12] Replace unwrap() calls in validate_daac_bytes_checks_invariants test Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- crates/bpe/src/byte_pair_encoding.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/bpe/src/byte_pair_encoding.rs b/crates/bpe/src/byte_pair_encoding.rs index 32568a9..1ed0aa5 100644 --- a/crates/bpe/src/byte_pair_encoding.rs +++ b/crates/bpe/src/byte_pair_encoding.rs @@ -780,7 +780,7 @@ pub fn create_test_bytes(bpe: &BytePairEncoding, min_bytes: usize) -> Vec { mod tests { use aneubeck_daachorse::DoubleArrayAhoCorasick; - use super::{hash_bytes, validate_daac_bytes, BytePairEncoding}; + use super::{hash_bytes, read_u32_le, validate_daac_bytes, BytePairEncoding}; /// The hash must not depend on the target. It keys the serialized `bytes_hash_to_token`, /// and a hardcoded hash factor is only collision-free for one particular mapping. @@ -809,7 +809,8 @@ mod tests { #[test] fn validate_daac_bytes_checks_invariants() { - let pma = DoubleArrayAhoCorasick::::new(vec!["a", "ab", "bab"]).unwrap(); + let pma = + DoubleArrayAhoCorasick::::new(vec!["a", "ab", "bab"]).expect("valid patterns"); let valid = pma.serialize(); assert!(validate_daac_bytes(&valid).is_ok()); @@ -832,7 +833,7 @@ mod tests { assert!(validate_daac_bytes(&bad_fail).is_err()); // Out-of-bounds output `parent` (self-referential 1-based parent == 1 at index 0) - let states_len = u32::from_le_bytes(valid[0..4].try_into().unwrap()) as usize; + let states_len = read_u32_le(&valid, 0) as usize; let first_output_parent_offset = 4 + states_len * 12 + 4 + 8; let mut bad_parent = valid.clone(); bad_parent[first_output_parent_offset..first_output_parent_offset + 4]