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..1ed0aa5 100644 --- a/crates/bpe/src/byte_pair_encoding.rs +++ b/crates/bpe/src/byte_pair_encoding.rs @@ -63,15 +63,101 @@ fn serialize_daac( s.serialize_bytes(&daac.serialize()) } +fn read_u32_le(bytes: &[u8], offset: usize) -> u32 { + u32::from_le_bytes( + bytes[offset..offset + 4] + .try_into() + .expect("slice of length 4"), + ) +} + +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 = read_u32_le(v, 0); + 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 = read_u32_le(v, outputs_offset); + 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 = 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 = 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"); + } + } + + for (i, chunk) in v[outputs_header_end..outputs_end] + .as_chunks::() + .0 + .iter() + .enumerate() + { + 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"); + } + } + + 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 +200,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 +223,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 +488,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 +778,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, 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. @@ -712,4 +798,46 @@ 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"]).expect("valid patterns"); + 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 = 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] + .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, 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/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 9a9b0ad..6b5b0c9 100644 --- a/crates/geo_filters/src/diff_count.rs +++ b/crates/geo_filters/src/diff_count.rs @@ -4,11 +4,11 @@ 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, - mask_bit_chunks, take_ref, xor_bit_chunks, BitChunk, GeoConfig, IsBucketType, BITS_PER_BLOCK, + mask_bit_chunks, take_ref, xor_bit_chunks, BitChunk, GeoConfig, IsBucketType, Unaligned, + BITS_PER_BLOCK, }; use crate::{Count, Diff}; @@ -68,7 +68,7 @@ pub type GeoDiffCount13<'a> = GeoDiffCount<'a, GeoDiffConfig13>; 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, [Unaligned]>, lsb: BitVec<'a>, } @@ -101,7 +101,7 @@ 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); @@ -210,6 +210,7 @@ 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 { @@ -251,7 +252,7 @@ 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().unwrap_or_default(), self.lsb.bit_chunks().peekable(), @@ -272,6 +273,7 @@ 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. @@ -286,7 +288,7 @@ impl<'a, C: GeoConfig> GeoDiffCount<'a, C> { lsb.next() }; if let Some(smallest) = first { - msb.push(C::BucketType::from_usize(smallest)); + msb.push(Unaligned::from_usize(smallest)); self.lsb.resize(smallest); } else { self.lsb.resize(0); @@ -375,19 +377,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 = unsafe { - std::mem::transmute::<&[u8], &[C::BucketType]>(std::slice::from_raw_parts( - buf.as_ptr(), - msb_len, - )) - }; + 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 let msb_bytes_len = msb_len * size_of::(); Self { config: c, - msb: Cow::Borrowed(msb), + msb: Cow::Borrowed(&msb_all[..msb_len]), lsb: BitVec::from_bytes(&buf[msb_bytes_len..]), } } @@ -397,10 +394,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 = Unaligned::slice_as_bytes(&self.msb); writer.write_all(msb_bytes)?; let mut bytes_written = msb_bytes.len(); bytes_written += self.lsb.write(writer)?; @@ -570,7 +564,7 @@ pub struct GeoDiffCountBuilder> { /// (xor) into `blocks`. [`GeoDiffCountBuilder::cleanup`] reconciles the two. split: usize, numbers: Vec, - blocks: Vec, + blocks: Vec>, } impl> GeoDiffCountBuilder { @@ -587,7 +581,7 @@ impl> GeoDiffCountBuilder { expected, split, numbers: Vec::with_capacity(capacity), - blocks: vec![0; split.div_ceil(BITS_PER_BLOCK)], + blocks: vec![Unaligned::default(); split.div_ceil(BITS_PER_BLOCK)], } } @@ -602,7 +596,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), Unaligned::default()); } } @@ -657,7 +652,8 @@ 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), Unaligned::default()); let blocks = &mut self.blocks; for &bucket in &self.numbers[keep..] { let (index, bit) = bucket.into_index_and_bit(); @@ -739,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,6 +1257,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/bitvec.rs b/crates/geo_filters/src/diff_count/bitvec.rs index cf6cb93..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,7 +14,7 @@ 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, [Unaligned]>, } impl Ord for BitVec<'_> { @@ -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. @@ -49,7 +49,7 @@ impl BitVec<'_> { /// 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>, 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 +63,9 @@ impl BitVec<'_> { 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, Unaligned::default()); } self.num_bits = num_bits; self.clear_superfluous_bits(); @@ -93,7 +95,7 @@ impl BitVec<'_> { } /// The raw 64-bit blocks backing the vector, from least to most significant. - pub fn blocks(&self) -> &[u64] { + pub fn blocks(&self) -> &[Unaligned] { self.blocks.deref() } @@ -134,7 +136,7 @@ impl BitVec<'_> { .enumerate() .rev() .filter(|(_, block)| **block != 0) - .map(|(index, block)| BitChunk::new(index, *block)) + .map(|(index, block)| BitChunk::new(index, block.get())) } // Returns an owned BitVec with static lifetime that can be used for cases when we want @@ -175,34 +177,31 @@ 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 {}", + 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 {}", buf.len() ); - let blocks = unsafe { - std::mem::transmute::<&[u8], &[u64]>(std::slice::from_raw_parts( - buf.as_ptr(), - buf.len() / BYTES_PER_BLOCK, - )) - }; - let blocks = Cow::Borrowed(blocks); - Self { num_bits, blocks } + let num_bits = buf.len() * 8 - unoccupied_bits; + Self { + num_bits, + blocks: Cow::Borrowed(blocks), + } } #[cfg(target_endian = "little")] @@ -213,10 +212,7 @@ impl BitVec<'_> { // 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 = Unaligned::slice_as_bytes(&self.blocks); writer.write_all(block_bytes)?; Ok(block_bytes.len() + 1) } @@ -240,7 +236,7 @@ impl Index for BitVec<'_> { /// Obtained via [`BitVec::toggler`]. pub(crate) struct BitToggler<'a> { num_bits: usize, - blocks: &'a mut [u64], + blocks: &'a mut [Unaligned], } impl BitToggler<'_> { @@ -260,7 +256,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 = Unaligned::new(if value { u64::MAX } else { u64::MIN }); + let blocks = vec![fill; num_blocks]; BitVec { blocks: Cow::from(blocks), num_bits, @@ -358,4 +355,33 @@ 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!(matches!(decoded.blocks, Cow::Borrowed(_))); + 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]); + } } diff --git a/crates/geo_filters/src/diff_count/sim_hash.rs b/crates/geo_filters/src/diff_count/sim_hash.rs index 5c92aa6..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(self.msb[n]) + Some(self.msb[n].get()) } else { n -= self.msb.len(); self.lsb diff --git a/crates/hash-sorted-map/src/group_ops.rs b/crates/hash-sorted-map/src/group_ops.rs index cc14813..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 { @@ -56,12 +65,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 +103,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 +141,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..b8c6c20 100644 --- a/crates/hash-sorted-map/src/hash_sorted_map.rs +++ b/crates/hash-sorted-map/src/hash_sorted_map.rs @@ -214,22 +214,25 @@ 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() }; + 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() { @@ -238,23 +241,19 @@ 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. + 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); + self.groups[gj].ctrl[sj] = cur_ctrl; } - // 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. - // 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); } } } @@ -1043,4 +1042,94 @@ 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"); + } + } + + 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(u32, Arc); + impl Drop for DropTrackedVal { + fn drop(&mut self) { + self.1.fetch_add(1, AtomicOrdering::SeqCst); + } + } + + let mut map = HashSortedMap::with_capacity_and_hasher(1, KeyHasher); + for id in [0, 1, 2] { + map.insert( + DropTrackedKey { + id, + drops: Arc::clone(&key_drops), + }, + DropTrackedVal(id * 10, Arc::clone(&val_drops)), + ); + } + + let result = catch_unwind(AssertUnwindSafe(|| { + map.sort_by_hash(); + })); + + assert!(result.is_err()); + // 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); + } }