From 4d67564a7e840a660d9ccd4fe8e9d04a19a95bd1 Mon Sep 17 00:00:00 2001 From: Arpad Borsos Date: Sun, 30 Jul 2023 12:53:25 +0200 Subject: [PATCH] Avoid intermediate `Vec` allocations Previously, each `try_generate_hash` call would allocate a bunch of new `Vec`s just to throw them away every time. Now, all these `Vec`s are kept around in a stateful `Generator` struct and reused across multiple `try_generate_hash` calls. As generation of the PHF is probabilistic, it may need lots of trials, and reusing allocations across multiple trials may be benefitial. --- phf_generator/src/lib.rs | 180 ++++++++++++++++++++++++--------------- 1 file changed, 111 insertions(+), 69 deletions(-) diff --git a/phf_generator/src/lib.rs b/phf_generator/src/lib.rs index 7ea15866..1042b609 100644 --- a/phf_generator/src/lib.rs +++ b/phf_generator/src/lib.rs @@ -26,89 +26,131 @@ pub fn generate_hash_with_hash_fn(entries: &[T], hash_fn: F) -> HashState where F: Fn(&T, &HashKey) -> Hashes, { + let mut generator = Generator::new(entries.len()); + SmallRng::seed_from_u64(FIXED_SEED) .sample_iter(Standard) - .find_map(|key| { - let hashes: Vec<_> = entries.iter().map(|entry| hash_fn(entry, &key)).collect(); - try_generate_hash(&hashes, key) + .find(|key| { + let hashes = entries.iter().map(|entry| hash_fn(entry, key)); + generator.reset(hashes); + + generator.try_generate_hash() + }) + .map(|key| HashState { + key, + disps: generator.disps, + map: generator.map.into_iter().map(|i| i.unwrap()).collect(), }) .expect("failed to solve PHF") } -fn try_generate_hash(hashes: &[Hashes], key: HashKey) -> Option { - struct Bucket { - idx: usize, - keys: Vec, - } +struct Bucket { + idx: usize, + keys: Vec, +} - let buckets_len = (hashes.len() + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA; - let mut buckets = (0..buckets_len) - .map(|i| Bucket { - idx: i, - keys: vec![], - }) - .collect::>(); +struct Generator { + hashes: Vec, + buckets: Vec, + disps: Vec<(u32, u32)>, + map: Vec>, + try_map: Vec, +} - for (i, hash) in hashes.iter().enumerate() { - buckets[(hash.g % (buckets_len as u32)) as usize] - .keys - .push(i); +impl Generator { + fn new(table_len: usize) -> Self { + let hashes = Vec::with_capacity(table_len); + + let buckets_len = (table_len + DEFAULT_LAMBDA - 1) / DEFAULT_LAMBDA; + let buckets: Vec<_> = (0..buckets_len) + .map(|i| Bucket { + idx: i, + keys: vec![], + }) + .collect(); + let disps = vec![(0u32, 0u32); buckets_len]; + + let map = vec![None; table_len]; + let try_map = vec![0u64; table_len]; + + Self { + hashes, + buckets, + disps, + map, + try_map, + } } - // Sort descending - buckets.sort_by(|a, b| a.keys.len().cmp(&b.keys.len()).reverse()); - - let table_len = hashes.len(); - let mut map = vec![None; table_len]; - let mut disps = vec![(0u32, 0u32); buckets_len]; - - // store whether an element from the bucket being placed is - // located at a certain position, to allow for efficient overlap - // checks. It works by storing the generation in each cell and - // each new placement-attempt is a new generation, so you can tell - // if this is legitimately full by checking that the generations - // are equal. (A u64 is far too large to overflow in a reasonable - // time for current hardware.) - let mut try_map = vec![0u64; table_len]; - let mut generation = 0u64; - - // the actual values corresponding to the markers above, as - // (index, key) pairs, for adding to the main map once we've - // chosen the right disps. - let mut values_to_add = vec![]; - - 'buckets: for bucket in &buckets { - for d1 in 0..(table_len as u32) { - 'disps: for d2 in 0..(table_len as u32) { - values_to_add.clear(); - generation += 1; - - for &key in &bucket.keys { - let idx = (phf_shared::displace(hashes[key].f1, hashes[key].f2, d1, d2) - % (table_len as u32)) as usize; - if map[idx].is_some() || try_map[idx] == generation { - continue 'disps; + fn reset(&mut self, hashes: I) + where + I: Iterator, + { + self.buckets.iter_mut().for_each(|b| b.keys.clear()); + self.buckets.sort_by_key(|b| b.idx); + self.disps.iter_mut().for_each(|d| *d = (0, 0)); + self.map.iter_mut().for_each(|m| *m = None); + self.try_map.iter_mut().for_each(|m| *m = 0); + + self.hashes.clear(); + self.hashes.extend(hashes); + } + + fn try_generate_hash(&mut self) -> bool { + let buckets_len = self.buckets.len() as u32; + for (i, hash) in self.hashes.iter().enumerate() { + self.buckets[(hash.g % buckets_len) as usize].keys.push(i); + } + + // Sort descending + self.buckets + .sort_by(|a, b| a.keys.len().cmp(&b.keys.len()).reverse()); + + let table_len = self.hashes.len(); + + // store whether an element from the bucket being placed is + // located at a certain position, to allow for efficient overlap + // checks. It works by storing the generation in each cell and + // each new placement-attempt is a new generation, so you can tell + // if this is legitimately full by checking that the generations + // are equal. (A u64 is far too large to overflow in a reasonable + // time for current hardware.) + let mut generation = 0u64; + + // the actual values corresponding to the markers above, as + // (index, key) pairs, for adding to the main map once we've + // chosen the right disps. + let mut values_to_add = vec![]; + + 'buckets: for bucket in &self.buckets { + for d1 in 0..(table_len as u32) { + 'disps: for d2 in 0..(table_len as u32) { + values_to_add.clear(); + generation += 1; + + for &key in &bucket.keys { + let idx = + (phf_shared::displace(self.hashes[key].f1, self.hashes[key].f2, d1, d2) + % (table_len as u32)) as usize; + if self.map[idx].is_some() || self.try_map[idx] == generation { + continue 'disps; + } + self.try_map[idx] = generation; + values_to_add.push((idx, key)); } - try_map[idx] = generation; - values_to_add.push((idx, key)); - } - // We've picked a good set of disps - disps[bucket.idx] = (d1, d2); - for &(idx, key) in &values_to_add { - map[idx] = Some(key); + // We've picked a good set of disps + self.disps[bucket.idx] = (d1, d2); + for &(idx, key) in &values_to_add { + self.map[idx] = Some(key); + } + continue 'buckets; } - continue 'buckets; } - } - // Unable to find displacements for a bucket - return None; + // Unable to find displacements for a bucket + return false; + } + true } - - Some(HashState { - key, - disps, - map: map.into_iter().map(|i| i.unwrap()).collect(), - }) }