perf(bloom): division-free bit indexes and inline u64 hashing

datasketches/src/bloom/sketch.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::fmt;
 use std::hash::Hash;
 use std::hash::Hasher;
 
@@ -40,7 +41,7 @@ const EMPTY_FLAG_MASK: u8 = 1 << 2;
 /// * Constant space usage
 ///
 /// See the [module level documentation](super) for more.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, PartialEq)]
 pub struct BloomFilter {
     /// Hash seed for all hash functions
     seed: u64,
@@ -50,6 +51,20 @@ pub struct BloomFilter {
     num_bits_set: u64,
     /// Bit array packed into u64 words
     bit_array: Box<[u64]>,
+    /// Reduces shifted hashes modulo the bit capacity; derived from `bit_array`.
+    reducer: CapacityReducer,
+}
+
+impl fmt::Debug for BloomFilter {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // The derived capacity multiplier is an implementation detail.
+        f.debug_struct("BloomFilter")
+            .field("seed", &self.seed)
+            .field("num_hashes", &self.num_hashes)
+            .field("num_bits_set", &self.num_bits_set)
+            .field("bit_array", &self.bit_array)
+            .finish()
+    }
 }
 
 impl BloomFilter {
@@ -589,6 +604,7 @@ impl BloomFilter {
             seed,
             num_hashes,
             num_bits_set,
+            reducer: CapacityReducer::new(bit_array.len() as u64 * 64),
             bit_array,
         })
     }
@@ -598,6 +614,7 @@ impl BloomFilter {
     /// Uses a two-hash approach:
     /// * h0 = XXHash64(item, seed)
     /// * h1 = XXHash64(item, h0)
+    #[inline]
     fn compute_hash<T: Hash>(&self, item: &T) -> (u64, u64) {
         // First hash with the configured seed
         let mut hasher = XxHash64::with_seed(self.seed);
@@ -632,8 +649,10 @@ impl BloomFilter {
     ///
     /// The right shift by 1 improves bit-distribution. The index `i` is 1-based.
     fn compute_bit_index(&self, h0: u64, h1: u64, i: u16) -> usize {
+        // The conversion to `usize` truncates the hash on 32-bit targets. Filters serialized
+        // there depend on it, so it is part of the bit layout.
         let hash = h0.wrapping_add(u64::from(i).wrapping_mul(h1)) as usize;
-        (hash >> 1) % self.capacity()
+        self.reducer.remainder((hash >> 1) as u64) as usize
     }
 
     /// Gets the value of a single bit.
@@ -650,6 +669,41 @@ impl BloomFilter {
     }
 }
 
+/// Reduces shifted hashes (below `2^63`) modulo the bit capacity without a hardware division.
+#[derive(Clone, Copy, PartialEq)]
+struct CapacityReducer {
+    capacity: u64,
+    /// `ceil(2^(63 + l) / capacity)` with `l = ceil(log2(capacity))`.
+    multiplier: u64,
+    /// `l - 1`.
+    shift: u32,
+}
+
+impl CapacityReducer {
+    /// `capacity` is a positive multiple of 64.
+    fn new(capacity: u64) -> Self {
+        // `capacity > 2^(l - 1)`, so the multiplier is below `2^64`.
+        let l = 64 - (capacity - 1).leading_zeros();
+        let multiplier = (1u128 << (63 + l)).div_ceil(u128::from(capacity)) as u64;
+        Self {
+            capacity,
+            multiplier,
+            shift: l - 1,
+        }
+    }
+
+    /// Returns `hash % capacity` for every `hash < 2^63`, using division by an invariant integer
+    /// (Granlund and Montgomery, 1994): `hash / capacity == (multiplier * hash) >> (63 + l)`,
+    /// which is one high multiply and one shift.
+    #[inline]
+    fn remainder(self, hash: u64) -> u64 {
+        debug_assert!(hash < 1 << 63);
+        let quotient =
+            ((u128::from(hash) * u128::from(self.multiplier)) >> 64) as u64 >> self.shift;
+        hash - quotient * self.capacity
+    }
+}
+
 /// Builder for creating [`BloomFilter`] instances.
 ///
 /// Provides two construction modes:
@@ -841,7 +895,105 @@ impl BloomFilterBuilder {
             seed: self.seed,
             num_hashes,
             num_bits_set: 0,
+            reducer: CapacityReducer::new(bit_array.len() as u64 * 64),
             bit_array,
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The cached-multiplier reduction must equal the `%` operator it replaces for every hash.
+    #[test]
+    fn bit_index_equals_remainder_of_shifted_hash() {
+        let mut hashes = vec![
+            0,
+            1,
+            2,
+            63,
+            64,
+            65,
+            u64::MAX,
+            u64::MAX - 1,
+            1 << 63,
+            (1 << 63) - 1,
+            0x9e37_79b9_7f4a_7c15,
+        ];
+        let mut state = 0x2545_f491_4f6c_dd1d_u64;
+        for _ in 0..200 {
+            state ^= state << 13;
+            state ^= state >> 7;
+            state ^= state << 17;
+            hashes.push(state);
+        }
+
+        for num_bits in [
+            1, 64, 65, 127, 1_000, 4_096, 4_097, 95_851, 131_072, 3_000_017,
+        ] {
+            let filter = BloomFilterBuilder::with_size(num_bits, 1).build().unwrap();
+            let capacity = filter.capacity();
+            for &h0 in &hashes {
+                for &h1 in &hashes {
+                    for i in [1, 2, 3, 7, 64, 1_000, u16::MAX] {
+                        let hash = h0.wrapping_add(u64::from(i).wrapping_mul(h1)) as usize;
+                        assert_eq!(
+                            filter.compute_bit_index(h0, h1, i),
+                            (hash >> 1) % capacity,
+                            "capacity {capacity}, h0 {h0}, h1 {h1}, i {i}"
+                        );
+                    }
+                }
+            }
+        }
+    }
+
+    /// The reducer must equal `%` for every shifted hash (below `2^63`) and every supported
+    /// capacity, including capacities too large to allocate here.
+    #[test]
+    fn capacity_reducer_equals_remainder_for_supported_capacities() {
+        let max = BloomFilterBuilder::MAX_NUM_BITS;
+        let capacities = [
+            64,
+            128,
+            192,
+            4_096,
+            4_160,
+            95_872,
+            131_072,
+            3_000_064,
+            (1 << 36) + 64,
+            (1 << 37) - 64,
+            max - 64,
+            max,
+        ];
+        let mut state = 0x2545_f491_4f6c_dd1d_u64;
+        for capacity in capacities {
+            assert_eq!(capacity % 64, 0);
+            let reducer = CapacityReducer::new(capacity);
+            let mut hashes = vec![
+                0,
+                1,
+                capacity - 1,
+                capacity,
+                capacity + 1,
+                (1 << 63) - capacity,
+                (1 << 63) - 1,
+            ];
+            for _ in 0..10_000 {
+                state ^= state << 13;
+                state ^= state >> 7;
+                state ^= state << 17;
+                hashes.push(state >> 1);
+            }
+            for hash in hashes {
+                assert_eq!(
+                    reducer.remainder(hash),
+                    hash % capacity,
+                    "capacity {capacity}, hash {hash}"
+                );
+            }
+        }
+    }
+}

datasketches/src/hash/xxhash.rs

@@ -144,7 +144,32 @@ impl Hasher for XxHash64 {
     #[inline(always)]
     fn write(&mut self, bytes: &[u8]) {
         self.total_len = self.total_len.wrapping_add(bytes.len() as u64);
+        self.write_bytes(bytes);
+    }
 
+    /// Same result as the default `write(&i.to_ne_bytes())`, with an inlined body small enough
+    /// that hashing an integer folds to a few multiplies at the call site.
+    #[inline(always)]
+    fn write_u64(&mut self, i: u64) {
+        let bytes = i.to_ne_bytes();
+        self.total_len = self.total_len.wrapping_add(8);
+        if self.buffer_len + 8 < 32 {
+            self.buffer[self.buffer_len..self.buffer_len + 8].copy_from_slice(&bytes);
+            self.buffer_len += 8;
+        } else {
+            self.write_bytes_outlined(&bytes);
+        }
+    }
+}
+
+impl XxHash64 {
+    #[inline(never)]
+    fn write_bytes_outlined(&mut self, bytes: &[u8]) {
+        self.write_bytes(bytes);
+    }
+
+    #[inline(always)]
+    fn write_bytes(&mut self, bytes: &[u8]) {
         if self.buffer_len + bytes.len() < 32 {
             self.buffer[self.buffer_len..self.buffer_len + bytes.len()].copy_from_slice(bytes);
             self.buffer_len += bytes.len();
@@ -270,4 +295,28 @@ mod tests {
         let hash2 = hasher.finish64();
         assert_eq!(hash2, hash1);
     }
+
+    #[test]
+    fn write_u64_matches_write_of_native_bytes() {
+        let input: Vec<u8> = (0..=40).collect();
+        for prefix in 0..=input.len() {
+            for value in [0u64, 1, 0x0123_4567_89ab_cdef, u64::MAX] {
+                let mut expected = XxHash64::with_seed(PRIME32);
+                expected.write(&input[..prefix]);
+                expected.write(&value.to_ne_bytes());
+                expected.write(&input[..prefix]);
+
+                let mut hasher = XxHash64::with_seed(PRIME32);
+                hasher.write(&input[..prefix]);
+                hasher.write_u64(value);
+                hasher.write(&input[..prefix]);
+
+                assert_eq!(
+                    hasher.finish64(),
+                    expected.finish64(),
+                    "prefix {prefix}, value {value:#x}"
+                );
+            }
+        }
+    }
 }