From 6d0f3f1f42076198f7f169c5e2ef47415bbf1839 Mon Sep 17 00:00:00 2001 From: Kristoffer Sahlin Date: Wed, 30 Sep 2026 13:45:37 +0200 Subject: [PATCH] Use galloping search in IndexEntry::get_count When at least 8 entries remained in the bucket, get_count binary-searched the entire rest of the bucket. Callers usually need only small counts (is_too_frequent calls it only when the count is at most the filter cutoff), but for repetitive seeds the rest of the bucket can contain hundreds of thousands of entries, so most binary search probes were cache misses far away from the entry. get_count now gallops from the entry: it probes offsets 1, 2, 4, ... until an entry no longer matches and then binary-searches between the last two probes. This needs O(log count) steps and only touches entries close to this one. The result is identical; a new test compares get_count with the previous implementation. Speedup on top of index prefetching (#XXX) on an Apple M4 Max, -x, 1 thread, median of 5 runs, measured with a prototype of this change: - CHM13, 150 nt reads: finding hits -7%, total mapping time -4% - CHM13, 10 kb reads: finding hits -3%, total mapping time -1% - maize, 150 nt reads: finding hits -7%, total mapping time -4% - maize, 10 kb reads: finding hits -3%, total mapping time -1% - chrY, 150 nt reads: finding hits -4%, total mapping time -1% - chrY, 10 kb reads: finding hits -3%, total mapping time -0.4% --- src/index.rs | 91 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 89 insertions(+), 2 deletions(-) diff --git a/src/index.rs b/src/index.rs index cd0f7f82..df3ec7d6 100644 --- a/src/index.rs +++ b/src/index.rs @@ -253,6 +253,17 @@ impl<'a> IndexEntry<'a> { self.get_count(self.strobemer_index.parameters.randstrobe.main_hash_mask) } + /// Return the number of consecutive entries in this bucket, starting at + /// this one, whose hash masked by `hash_mask` equals the masked hash of + /// this entry. + /// + /// Callers typically need only small counts, but the rest of the bucket + /// can be very large for repetitive seeds. A galloping (exponential) + /// search is therefore used, which needs O(log count) steps and only + /// touches entries close to this one. + /// + /// This requires that masked hashes are non-decreasing within a bucket, + /// which holds because all lookup masks keep a contiguous run of top bits. pub fn get_count(&self, hash_mask: u64) -> usize { let position = self.position; const MAX_LINEAR_SEARCH: usize = 8; @@ -273,8 +284,20 @@ impl<'a> IndexEntry<'a> { } count } else { - let bucket = &self.strobemer_index.randstrobes[position..position_end]; - custom_partition_point(bucket, |h| h.hash() & hash_mask <= masked_key) + // Probe offsets 1, 2, 4, ... until an entry does not match or the + // end of the bucket is reached, then binary search between the + // last two probes + let rest = &self.strobemer_index.randstrobes[position..position_end]; + let matches = |h: &RefRandstrobe| h.hash() & hash_mask == masked_key; + let mut lo = 0; + let mut hi = 1; + while hi < rest.len() && matches(&rest[hi]) { + lo = hi; + hi *= 2; + } + // rest[lo] matches and rest[hi] (if it exists) does not + let hi = hi.min(rest.len()); + lo + 1 + custom_partition_point(&rest[lo + 1..hi], matches) } } @@ -605,4 +628,68 @@ mod tests { assert!(rev_pos_forward.is_some()); } } + + /// get_count as it was implemented before galloping search was introduced + fn get_count_reference(index: &StrobemerIndex, position: usize, hash_mask: u64) -> usize { + let key = index.randstrobes[position].hash(); + let masked_key = key & hash_mask; + let top_n = (key >> (64 - index.bits)) as usize; + let rest = &index.randstrobes[position..index.bucket_starts[top_n + 1]]; + if rest.len() < 8 { + rest.iter() + .take_while(|h| h.hash() & hash_mask == masked_key) + .count() + } else { + custom_partition_point(rest, |h| h.hash() & hash_mask <= masked_key) + } + } + + #[test] + fn get_count_matches_reference() { + let parameters = SeedingParameters::new(150); + let rp = parameters.randstrobe.clone(); + let bits = 6; + let mut rng = fastrand::Rng::with_seed(1); + let mut randstrobes = vec![]; + let mut bucket_starts = vec![0]; + for bucket in 0..1u64 << bits { + // Main hash, orientation and auxiliary hash are each drawn from at + // most n values, creating runs for all masks (n = 1 fills the + // bucket with a single key). Every combination of bucket size and + // n occurs at least once. + let size = [0, 1, 2, 7, 8, 9, 17, 100, 3000][bucket as usize % 9]; + let n = [1, 2, 4, 1 << 16][bucket as usize / 9 % 4]; + for _ in 0..size { + let hash = (bucket << (64 - bits)) + | (rng.u64(..n) << (rp.partial_orientation_pos + 1)) + | (rng.u64(..n.min(2)) << rp.partial_orientation_pos) + | (rng.u64(..n) << STROBE2_OFFSET_BITS); + randstrobes.push(RefRandstrobe::new(hash, rng.usize(..1000), rng.u8(..))); + } + bucket_starts.push(randstrobes.len()); + } + randstrobes.sort(); + let index = StrobemerIndex::new( + parameters, + bits, + 1000, + randstrobes, + bucket_starts, + ContigStarts::default(), + ); + + for hash_mask in [ + REF_RANDSTROBE_HASH_MASK, + rp.forward_main_hash_mask, + rp.main_hash_mask, + ] { + for position in 0..index.len() { + assert_eq!( + index.entry(position).get_count(hash_mask), + get_count_reference(&index, position, hash_mask), + "position {position}, hash_mask {hash_mask:#x}" + ); + } + } + } }