1use std::hash::{Hash, Hasher};
11use std::marker::PhantomData;
12use std::mem::size_of_val;
13
14use num::{Integer, PrimInt, Unsigned};
15use wyhash::WyHash;
16
17use crate::mphf::MphfError::*;
18use crate::rank::{RankedBits, RankedBitsAccess};
19
20#[derive(Default)]
28#[cfg_attr(feature = "rkyv_derive", derive(rkyv::Archive, rkyv::Deserialize, rkyv::Serialize))]
29#[cfg_attr(feature = "rkyv_derive", archive_attr(derive(rkyv::CheckBytes)))]
30#[cfg_attr(feature = "serde", derive(serde::Serialize))]
31#[cfg_attr(
32 feature = "serde",
33 serde(bound(serialize = "ST: serde::Serialize", deserialize = "ST: serde::Deserialize<'de>"))
34)]
35pub struct Mphf<const B: usize = 32, const S: usize = 8, ST: PrimInt + Unsigned = u8, H: Hasher + Default = WyHash> {
36 ranked_bits: RankedBits,
38 level_groups: Box<[u32]>,
40 group_seeds: Box<[ST]>,
42 #[cfg_attr(feature = "serde", serde(skip))]
44 _phantom_hasher: PhantomData<H>,
45}
46
47#[cfg(feature = "serde")]
48#[derive(serde::Deserialize)]
49struct MphfUnchecked<const B: usize = 32, const S: usize = 8, ST: PrimInt + Unsigned = u8> {
50 ranked_bits: RankedBits,
51 level_groups: Box<[u32]>,
52 group_seeds: Box<[ST]>,
53}
54
55#[cfg(feature = "serde")]
56impl<'de, const B: usize, const S: usize, ST, H> serde::Deserialize<'de> for Mphf<B, S, ST, H>
57where
58 ST: serde::Deserialize<'de> + PrimInt + Unsigned,
59 H: Hasher + Default,
60{
61 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
62 where
63 D: serde::Deserializer<'de>,
64 {
65 use serde::de::Error;
66 let this = MphfUnchecked::<B, S, ST>::deserialize(deserializer)?;
67
68 if !Self::is_valid_seed_type() {
69 return Err(Error::custom(format!(
70 "ST (u{}) should be able to store (1 << S) - 1",
71 std::mem::size_of::<ST>() * 8
72 )));
73 }
74
75 if this.level_groups.len() > MAX_LEVELS {
76 return Err(Error::custom(format!("MAX_LEVELS {MAX_LEVELS} exceeded")));
77 }
78
79 if !this.level_groups.iter().all(|&x| x > 0) {
80 return Err(Error::custom("no level_groups should be empty"));
81 }
82
83 let level_groups_sum = this
84 .level_groups
85 .iter()
86 .try_fold(0usize, |acc, &x| acc.checked_add(x as usize))
87 .ok_or_else(|| Error::custom("level_groups sum overflowed"))?;
88 if level_groups_sum != this.group_seeds.len() {
89 return Err(Error::custom(
90 "sum of all level_groups should equal length of group_seeds",
91 ));
92 }
93
94 let needed_bits = level_groups_sum
95 .checked_mul(Self::B)
96 .ok_or_else(|| Error::custom("needed_bits overflowed"))?;
97 let ranked_bit_len = this.ranked_bits.bits.len() * 64;
98 if ranked_bit_len < needed_bits {
99 return Err(Error::custom(
100 "count of ranked bits should not be less than the sum of level_groups * B",
101 ));
102 }
103
104 Ok(Self {
105 ranked_bits: this.ranked_bits,
106 level_groups: this.level_groups,
107 group_seeds: this.group_seeds,
108 _phantom_hasher: PhantomData,
109 })
110 }
111}
112
113const MAX_LEVELS: usize = 64;
115
116#[derive(Debug)]
118pub enum MphfError {
119 MaxLevelsExceeded,
121 InvalidSeedType,
123 InvalidGammaParameter,
125}
126
127pub const DEFAULT_GAMMA: f32 = 2.0;
129
130#[cfg(feature = "serde")]
131#[derive(Debug, PartialEq, Eq)]
132pub(crate) enum ValidateKeyResult {
133 InvalidKeyCount,
134 IncorrectKeyOrder,
135}
136
137#[cfg(feature = "serde")]
138#[derive(Debug, PartialEq, Eq)]
139pub(crate) enum ValidateValueResult {
140 KeyValueLenMismatch,
141 InvalidValueIndex,
142}
143
144impl<const B: usize, const S: usize, ST: PrimInt + Unsigned, H: Hasher + Default> Mphf<B, S, ST, H> {
145 const B: usize = {
147 assert!(B >= 1 && B <= 64);
148 B
149 };
150 const S: usize = {
152 assert!(S <= 16);
153 S
154 };
155 const ST_SIZE: usize = {
158 let sz = std::mem::size_of::<ST>();
159 assert!((sz * 8) >= Self::S);
160 assert!(sz <= 4);
161 sz
162 };
163
164 const fn is_valid_seed_type() -> bool {
167 Self::ST_SIZE * 8 >= Self::S
169 }
170
171 #[cfg(feature = "serde")]
173 pub(crate) fn validate_keys<K>(&self, keys: &[K]) -> Result<(), ValidateKeyResult>
174 where
175 K: Hash + Sized,
176 {
177 if keys.len() != self.ranked_bits.count_ones() {
179 return Err(ValidateKeyResult::InvalidKeyCount);
180 }
181
182 if keys.iter().enumerate().any(|(i, k)| self.get(k) != Some(i)) {
184 return Err(ValidateKeyResult::IncorrectKeyOrder);
185 }
186
187 Ok(())
188 }
189
190 #[cfg(feature = "serde")]
192 pub(crate) fn validate_values<K, V>(
193 &self,
194 keys: &[K],
195 values_indices: &[usize],
196 values_dict: &[V],
197 ) -> Result<(), ValidateValueResult>
198 where
199 K: Hash + Sized,
200 {
201 if keys.len() != values_indices.len() {
202 return Err(ValidateValueResult::KeyValueLenMismatch);
203 }
204
205 if values_indices.iter().any(|&i| i >= values_dict.len()) {
206 return Err(ValidateValueResult::InvalidValueIndex);
207 }
208
209 Ok(())
210 }
211
212 pub fn from_slice<K: Hash>(keys: &[K], gamma: f32) -> Result<Self, MphfError> {
214 if gamma < 1.0 {
215 return Err(InvalidGammaParameter);
216 }
217
218 if !Self::is_valid_seed_type() {
219 return Err(InvalidSeedType);
220 }
221
222 let mut hashes: Vec<u64> = keys.iter().map(|key| hash_key::<H, _>(key)).collect();
223 let mut group_bits = vec![];
224 let mut group_seeds = vec![];
225 let mut level_groups = vec![];
226
227 while !hashes.is_empty() {
228 let level = level_groups.len() as u32;
229 let (level_group_bits, level_group_seeds) = Self::build_level(level, &mut hashes, gamma);
230
231 group_bits.extend_from_slice(&level_group_bits);
232 group_seeds.extend_from_slice(&level_group_seeds);
233 level_groups.push(level_group_seeds.len() as u32);
234
235 if level_groups.len() == MAX_LEVELS && !hashes.is_empty() {
236 return Err(MaxLevelsExceeded);
237 }
238 }
239
240 Ok(Mphf {
241 ranked_bits: RankedBits::new(group_bits.into_boxed_slice()),
242 level_groups: level_groups.into_boxed_slice(),
243 group_seeds: group_seeds.into_boxed_slice(),
244 _phantom_hasher: PhantomData,
245 })
246 }
247
248 fn build_level(level: u32, hashes: &mut Vec<u64>, gamma: f32) -> (Vec<u64>, Vec<ST>) {
250 let level_size = ((hashes.len() as f32) * gamma).ceil() as usize;
252 let (groups, segments) = Self::level_size_groups_segments(level_size);
253 let max_group_seed = 1 << S;
254
255 let mut group_bits = vec![0u64; 3 * segments + 3];
261 let mut best_group_seeds = vec![ST::zero(); groups];
262
263 for group_seed in 0..max_group_seed {
265 Self::update_group_bits_with_seed(
266 level,
267 groups,
268 group_seed,
269 hashes,
270 &mut group_bits,
271 &mut best_group_seeds,
272 );
273 }
274
275 let best_group_bits: Vec<u64> = group_bits[..group_bits.len() - 3]
277 .as_chunks::<3>()
278 .0
279 .iter()
280 .map(|group_bits| group_bits[2])
281 .collect();
282
283 hashes.retain(|&hash| {
285 let level_hash = hash_with_seed(hash, level);
286 let group_idx = fastmod32(level_hash as u32, groups as u32);
287 let group_seed = best_group_seeds[group_idx].to_u32().unwrap();
288 let bit_idx = bit_index_for_seed::<B>(level_hash, group_seed, group_idx);
289 *unsafe { best_group_bits.get_unchecked(bit_idx / 64) } & (1 << (bit_idx % 64)) == 0
291 });
292
293 (best_group_bits, best_group_seeds)
294 }
295
296 #[inline]
298 fn level_size_groups_segments(size: usize) -> (usize, usize) {
299 let lcm_value = Self::B.lcm(&64);
301
302 let adjusted_size = size.div_ceil(lcm_value) * lcm_value;
304
305 (adjusted_size / Self::B, adjusted_size / 64)
306 }
307
308 #[inline]
310 fn update_group_bits_with_seed(
311 level: u32,
312 groups: usize,
313 group_seed: u32,
314 hashes: &[u64],
315 group_bits: &mut [u64],
316 best_group_seeds: &mut [ST],
317 ) {
318 let group_bits_len = group_bits.len();
320 for bits in group_bits[..group_bits_len - 3].as_chunks_mut::<3>().0 {
321 bits[0] = 0;
322 bits[1] = 0;
323 }
324
325 for &hash in hashes {
327 let level_hash = hash_with_seed(hash, level);
328 let group_idx = fastmod32(level_hash as u32, groups as u32);
329 let bit_idx = bit_index_for_seed::<B>(level_hash, group_seed, group_idx);
330 let mask = 1 << (bit_idx % 64);
331 let idx = (bit_idx / 64) * 3;
332
333 let bits = unsafe { group_bits.get_unchecked_mut(idx..idx + 2) };
335
336 bits[1] |= bits[0] & mask;
337 bits[0] |= mask;
338 }
339
340 for bits in group_bits.as_chunks_mut::<3>().0 {
342 bits[0] &= !bits[1];
343 }
344
345 for (group_idx, best_group_seed) in best_group_seeds.iter_mut().enumerate() {
347 let bit_idx = group_idx * Self::B;
348 let bit_pos = bit_idx % 64;
349 let idx = (bit_idx / 64) * 3;
350
351 let bits = unsafe { group_bits.get_unchecked_mut(idx..idx + 6) };
353
354 let bits_1 = Self::B.min(64 - bit_pos);
355 let bits_2 = Self::B - bits_1;
356 let mask_1 = u64::MAX >> (64 - bits_1);
357 let mask_2 = (1 << bits_2) - 1;
358
359 let new_bits_1 = (bits[0] >> bit_pos) & mask_1;
360 let new_bits_2 = bits[3] & mask_2;
361 let new_ones = new_bits_1.count_ones() + new_bits_2.count_ones();
362
363 let best_bits_1 = (bits[2] >> bit_pos) & mask_1;
364 let best_bits_2 = bits[5] & mask_2;
365 let best_ones = best_bits_1.count_ones() + best_bits_2.count_ones();
366
367 if new_ones > best_ones {
368 bits[2] &= !(mask_1 << bit_pos);
369 bits[2] |= new_bits_1 << bit_pos;
370
371 bits[5] &= !mask_2;
372 bits[5] |= new_bits_2;
373
374 *best_group_seed = ST::from(group_seed).unwrap();
375 }
376 }
377 }
378
379 #[inline]
382 pub fn get<K: Hash + ?Sized>(&self, key: &K) -> Option<usize> {
383 Self::get_impl(key, &self.level_groups, &self.group_seeds, &self.ranked_bits)
384 }
385
386 #[inline]
389 fn get_impl<K: Hash + ?Sized>(
390 key: &K,
391 level_groups: &[u32],
392 group_seeds: &[ST],
393 ranked_bits: &impl RankedBitsAccess,
394 ) -> Option<usize> {
395 let mut groups_before = 0;
396 for (level, &groups) in level_groups.iter().enumerate() {
397 let level_hash = hash_with_seed(hash_key::<H, _>(key), level as u32);
398 let group_idx = groups_before + fastmod32(level_hash as u32, groups);
399 let group_seed = unsafe { group_seeds.get_unchecked(group_idx).to_u32().unwrap() };
401 let bit_idx = bit_index_for_seed::<B>(level_hash, group_seed, group_idx);
402 if let Some(rank) = ranked_bits.rank(bit_idx) {
403 return Some(rank);
404 }
405 groups_before += groups as usize;
406 }
407
408 None
409 }
410
411 pub fn size(&self) -> usize {
413 size_of_val(self)
414 + size_of_val(self.level_groups.as_ref())
415 + size_of_val(self.group_seeds.as_ref())
416 + self.ranked_bits.size()
417 }
418}
419
420#[inline]
422fn hash_key<H: Hasher + Default, T: Hash + ?Sized>(key: &T) -> u64 {
423 let mut hasher = H::default();
424 key.hash(&mut hasher);
425 hasher.finish()
426}
427
428#[inline]
430fn bit_index_for_seed<const B: usize>(hash: u64, group_seed: u32, groups_before: usize) -> usize {
431 let mut x = (hash as u32) ^ group_seed;
433
434 x = (x ^ (x >> 16)).wrapping_mul(0x85ebca6b);
436 x = (x ^ (x >> 13)).wrapping_mul(0xc2b2ae35);
437 x ^= x >> 16;
438
439 groups_before * B + fastmod32(x, B as u32)
440}
441
442#[inline]
444fn hash_with_seed(hash: u64, seed: u32) -> u64 {
445 let x = ((hash as u128) ^ (seed as u128)).wrapping_mul(0x5851f42d4c957f2d);
446 ((x & 0xFFFFFFFFFFFFFFFF) as u64) ^ ((x >> 64) as u64)
447}
448
449#[inline]
452fn fastmod32(x: u32, n: u32) -> usize {
453 (((x as u64) * (n as u64)) >> 32) as usize
454}
455
456#[cfg(feature = "rkyv_derive")]
458impl<const B: usize, const S: usize, ST, H> ArchivedMphf<B, S, ST, H>
459where
460 ST: PrimInt + Unsigned + rkyv::Archive<Archived = ST>,
461 H: Hasher + Default,
462{
463 #[inline]
464 pub fn get<K: Hash + ?Sized>(&self, key: &K) -> Option<usize> {
465 Mphf::<B, S, ST, H>::get_impl(key, &self.level_groups, &self.group_seeds, &self.ranked_bits)
466 }
467}
468
469#[cfg(test)]
470mod tests {
471 use super::*;
472 use paste::paste;
473 use std::collections::HashSet;
474 use test_case::test_case;
475
476 #[cfg(feature = "serde")]
480 pub(crate) fn decode_mutated<T: serde::de::DeserializeOwned>(
481 bytes: &[u8],
482 f: impl FnOnce(&mut rmpv::Value),
483 ) -> Result<T, rmp_serde::decode::Error> {
484 let mut value = rmpv::decode::value::read_value(&mut &bytes[..]).unwrap();
485 f(&mut value);
486
487 let mut out = Vec::new();
488 rmpv::encode::write_value(&mut out, &value).unwrap();
489 rmp_serde::from_slice(&out)
490 }
491
492 fn test_mphfs_impl<const B: usize, const S: usize>(n: usize, gamma: f32) -> String {
494 let keys = (0..n as u64).collect::<Vec<u64>>();
495 let mphf = Mphf::<B, S>::from_slice(&keys, gamma).expect("failed to create mphf");
496
497 let mut set = HashSet::with_capacity(n);
499 for key in &keys {
500 let idx = mphf.get(key).unwrap();
501 assert!(idx < n, "idx = {} n = {}", idx, n);
502 if !set.insert(idx) {
503 panic!("duplicate idx = {} for key {}", idx, key);
504 }
505 }
506 assert_eq!(set.len(), n);
507
508 let mut avg_levels = 0f32;
510 let total_groups: u32 = mphf.level_groups.iter().sum();
511 for (i, &groups) in mphf.level_groups.iter().enumerate() {
512 avg_levels += ((i + 1) as f32 * groups as f32) / (total_groups as f32);
513 }
514 let bits = mphf.size() as f32 * (8.0 / n as f32);
515
516 format!(
517 "bits: {:.2} total_levels: {} avg_levels: {:.2}",
518 bits,
519 mphf.level_groups.len(),
520 avg_levels
521 )
522 }
523
524 macro_rules! generate_tests {
526 ($(($b:expr, $s:expr, $n: expr, $gamma:expr, $expected:expr)),* $(,)?) => {
527 $(
528 paste! {
529 #[test_case($n, $gamma => $expected)]
530 fn [<test_mphfs_ $b _ $s _ $n _ $gamma>](n: usize, gamma_scaled: usize) -> String {
531 let gamma = (gamma_scaled as f32) / 100.0;
532 test_mphfs_impl::<$b, $s>(n, gamma)
533 }
534 }
535 )*
536 };
537 }
538
539 generate_tests!(
541 (1, 8, 10000, 100, "bits: 26.64 total_levels: 42 avg_levels: 4.34"),
542 (2, 8, 10000, 100, "bits: 9.00 total_levels: 8 avg_levels: 1.76"),
543 (4, 8, 10000, 100, "bits: 4.39 total_levels: 6 avg_levels: 1.42"),
544 (7, 8, 10000, 100, "bits: 3.12 total_levels: 4 avg_levels: 1.39"),
545 (8, 8, 10000, 100, "bits: 2.80 total_levels: 6 avg_levels: 1.34"),
546 (15, 8, 10000, 100, "bits: 2.50 total_levels: 4 avg_levels: 1.50"),
547 (16, 8, 10000, 100, "bits: 2.30 total_levels: 6 avg_levels: 1.43"),
548 (23, 8, 10000, 100, "bits: 2.53 total_levels: 4 avg_levels: 1.67"),
549 (24, 8, 10000, 100, "bits: 2.25 total_levels: 6 avg_levels: 1.57"),
550 (31, 8, 10000, 100, "bits: 2.40 total_levels: 3 avg_levels: 1.44"),
551 (32, 8, 10000, 100, "bits: 2.20 total_levels: 7 avg_levels: 1.63"),
552 (33, 8, 10000, 100, "bits: 2.52 total_levels: 4 avg_levels: 1.78"),
553 (48, 8, 10000, 100, "bits: 2.25 total_levels: 7 avg_levels: 1.78"),
554 (53, 8, 10000, 100, "bits: 2.90 total_levels: 4 avg_levels: 2.00"),
555 (61, 8, 10000, 100, "bits: 2.82 total_levels: 4 avg_levels: 2.00"),
556 (63, 8, 10000, 100, "bits: 2.89 total_levels: 4 avg_levels: 2.00"),
557 (64, 8, 10000, 100, "bits: 2.25 total_levels: 8 avg_levels: 1.84"),
558 (32, 7, 10000, 100, "bits: 2.29 total_levels: 7 avg_levels: 1.70"),
559 (32, 5, 10000, 100, "bits: 2.47 total_levels: 8 avg_levels: 1.84"),
560 (32, 4, 10000, 100, "bits: 2.58 total_levels: 9 avg_levels: 1.92"),
561 (32, 3, 10000, 100, "bits: 2.75 total_levels: 10 avg_levels: 2.05"),
562 (32, 1, 10000, 100, "bits: 3.22 total_levels: 11 avg_levels: 2.39"),
563 (32, 0, 10000, 100, "bits: 3.65 total_levels: 14 avg_levels: 2.73"),
564 (32, 8, 100000, 100, "bits: 2.11 total_levels: 10 avg_levels: 1.64"),
565 (32, 8, 100000, 200, "bits: 2.73 total_levels: 4 avg_levels: 1.06"),
566 (32, 6, 100000, 200, "bits: 2.84 total_levels: 5 avg_levels: 1.11"),
567 );
568
569 #[cfg(feature = "rkyv_derive")]
570 #[test]
571 fn test_rkyv() {
572 let n = 10000;
573 let keys = (0..n as u64).collect::<Vec<u64>>();
574 let mphf = Mphf::<32, 4>::from_slice(&keys, DEFAULT_GAMMA).expect("failed to create mphf");
575 let rkyv_bytes = rkyv::to_bytes::<_, 1024>(&mphf).unwrap();
576
577 let rkyv_mphf = rkyv::check_archived_root::<Mphf<32, 4>>(&rkyv_bytes).unwrap();
578
579 let mut set = HashSet::with_capacity(n);
581 for key in &keys {
582 let idx = mphf.get(key).unwrap();
583 let rkyv_idx = rkyv_mphf.get(key).unwrap();
584
585 assert_eq!(idx, rkyv_idx);
586 assert!(idx < n, "idx = {} n = {}", idx, n);
587 if !set.insert(idx) {
588 panic!("duplicate idx = {} for key {}", idx, key);
589 }
590 }
591 assert_eq!(set.len(), n);
592 }
593
594 #[cfg(feature = "serde")]
595 #[test]
596 fn test_serde() {
597 let n = 10000;
598 let keys = (0..n as u64).collect::<Vec<u64>>();
599 let mphf = Mphf::<32, 4>::from_slice(&keys, DEFAULT_GAMMA).expect("failed to create mphf");
600
601 let bytes = rmp_serde::to_vec(&mphf).unwrap();
604 let de: Mphf<32, 4> = rmp_serde::from_slice(&bytes).unwrap();
605
606 let mut set = HashSet::with_capacity(n);
609 for key in &keys {
610 let idx = mphf.get(key).unwrap();
611 let de_idx = de.get(key).unwrap();
612
613 assert_eq!(idx, de_idx);
614 assert!(idx < n, "idx = {} n = {}", idx, n);
615 assert!(set.insert(idx), "duplicate idx = {} for key {}", idx, key);
616 }
617 assert_eq!(set.len(), n);
618 }
619
620 #[cfg(feature = "serde")]
621 #[test]
622 fn test_serde_rejects_invalid_structures() {
623 use rmpv::Value;
624
625 let keys = (0..1000u64).collect::<Vec<_>>();
626 let mphf = Mphf::<32, 4>::from_slice(&keys, DEFAULT_GAMMA).unwrap();
627 let bytes = rmp_serde::to_vec(&mphf).unwrap();
628
629 assert!(decode_mutated::<Mphf<32, 4>>(&bytes, |v| {
631 if let Value::Array(f) = v {
632 if let Value::Array(seeds) = &mut f[2] {
633 seeds.pop();
634 }
635 }
636 })
637 .is_err());
638
639 assert!(decode_mutated::<Mphf<32, 4>>(&bytes, |v| {
641 if let Value::Array(f) = v {
642 if let Value::Binary(bits) = &mut f[0] {
643 let n = bits.len();
644 bits.truncate(n - 8);
645 }
646 }
647 })
648 .is_err());
649 }
650}