introduce a more generic extend function

Author: aneubeck

crates/geo_filters/evaluation/performance.rs

@@ -133,8 +133,8 @@ fn criterion_benchmark(c: &mut Criterion) {
     }
 
     // Compare building a diff filter from a precomputed slice of hashes one by one
-    // (`push_hash`) versus in a single batch (`from_hashes`). The hashes are precomputed so that
-    // only the construction cost is measured.
+    // (`push_hash`) versus in a single batch (`extend_by_hashes` on an empty filter). The hashes
+    // are precomputed so that only the construction cost is measured.
     for size in [1000usize, 10000, 100000, 1000000] {
         let mut group = c.benchmark_group(format!("construct:{size}"));
         let build_hasher = UnstableDefaultBuildHasher::default();
@@ -149,12 +149,11 @@ fn criterion_benchmark(c: &mut Criterion) {
                 black_box(&gc);
             })
         });
-        group.bench_function("geo_diff_count_7_from_hashes", |b| {
+        group.bench_function("geo_diff_count_7_extend", |b| {
             b.iter(|| {
-                black_box(GeoDiffCount7::from_hashes(
-                    Default::default(),
-                    hashes.iter().copied(),
-                ));
+                let mut gc = GeoDiffCount7::default();
+                gc.extend_by_hashes(hashes.iter().copied());
+                black_box(&gc);
             })
         });
         group.bench_function("geo_diff_count_13_push", |b| {
@@ -166,15 +165,60 @@ fn criterion_benchmark(c: &mut Criterion) {
                 black_box(&gc);
             })
         });
-        group.bench_function("geo_diff_count_13_from_hashes", |b| {
+        group.bench_function("geo_diff_count_13_extend", |b| {
             b.iter(|| {
-                black_box(GeoDiffCount13::from_hashes(
-                    Default::default(),
-                    hashes.iter().copied(),
-                ));
+                let mut gc = GeoDiffCount13::default();
+                gc.extend_by_hashes(hashes.iter().copied());
+                black_box(&gc);
             })
         });
     }
+
+    // Insert a batch of hashes into an existing filter: one-by-one `push_hash` vs the batched
+    // `extend_by_hashes`. The existing filter is cloned in the (untimed) batched setup so that only
+    // the insertion cost is measured, not the clone.
+    for (existing_n, batch_n) in [
+        (10000usize, 10000usize),
+        (100000, 10000),
+        (100000, 100000),
+        (1000, 100000),
+    ] {
+        let mut group = c.benchmark_group(format!("insert_batch:{existing_n}+{batch_n}"));
+        let build_hasher = UnstableDefaultBuildHasher::default();
+        let batch: Vec<u64> = (existing_n..existing_n + batch_n)
+            .map(|i| build_hasher.hash_one(i))
+            .collect();
+        let base = {
+            let mut gc = GeoDiffCount13::default();
+            for i in 0..existing_n {
+                gc.push_hash(build_hasher.hash_one(i));
+            }
+            gc
+        };
+
+        group.bench_function("push", |b| {
+            b.iter_batched(
+                || base.clone(),
+                |mut gc| {
+                    for &hash in &batch {
+                        gc.push_hash(hash);
+                    }
+                    gc
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        });
+        group.bench_function("extend_by_hashes", |b| {
+            b.iter_batched(
+                || base.clone(),
+                |mut gc| {
+                    gc.extend_by_hashes(batch.iter().copied());
+                    gc
+                },
+                criterion::BatchSize::SmallInput,
+            )
+        });
+    }
 }
 
 criterion_group!(benches, criterion_benchmark);

crates/geo_filters/src/config/bitchunks.rs

@@ -20,6 +20,11 @@ impl BitChunk {
     }
 }
 
+/// Merges a descending stream of distinct one-bit positions (`leading`) with a descending stream
+/// of `BitChunk`s (`trailing`) into a single descending `BitChunk` stream. All leading positions
+/// must be more significant than all trailing bits, except that the least-significant leading block
+/// may overlap the most-significant trailing block (the two are or-ed). Leading positions must be
+/// distinct.
 pub(crate) fn iter_bit_chunks(
     leading: impl Iterator<Item = usize>,
     trailing: impl Iterator<Item = BitChunk>,
@@ -55,8 +60,7 @@ impl<I: Iterator<Item = BitChunk>, J: Iterator<Item = usize>> Iterator for BitCh
                     _ => break,
                 }
             }
-            // All leading bits were consumed, test whether it can be merged with
-            // trailing bits.
+            // All leading bits were consumed, test whether it can be merged with trailing bits.
             match self.trailing.peek() {
                 Some(BitChunk {
                     index: other_index,
@@ -78,6 +82,39 @@ impl<I: Iterator<Item = BitChunk>, J: Iterator<Item = usize>> Iterator for BitCh
     }
 }
 
+/// Turns a descending stream of one-bit positions into a descending `BitChunk` stream containing
+/// each position's parity. Unlike [`iter_bit_chunks`], positions may repeat: a position occurring
+/// an even number of times cancels out (xor), and a block that cancels out entirely is skipped.
+pub(crate) fn parity_bit_positions(
+    positions: impl Iterator<Item = usize>,
+) -> impl Iterator<Item = BitChunk> {
+    ParityBitPositions(positions.peekable())
+}
+
+struct ParityBitPositions<I: Iterator<Item = usize>>(Peekable<I>);
+
+impl<I: Iterator<Item = usize>> Iterator for ParityBitPositions<I> {
+    type Item = BitChunk;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let (index, bit) = self.0.next()?.into_index_and_bit();
+            let mut block: u64 = bit.into_block();
+            while let Some(&other) = self.0.peek() {
+                if other.into_index() != index {
+                    break;
+                }
+                block ^= other.into_bit().into_block();
+                self.0.next();
+            }
+            // The block is empty only if its positions cancelled out entirely; skip it.
+            if block != 0 {
+                return Some(BitChunk::new(index, block));
+            }
+        }
+    }
+}
+
 /// Combine-s two `BitChunk` iterators using the given operator. Both iterators have to output
 /// `BitChunk`s from most to least significant and must not report duplicates!
 struct BinOpBitChunksIterator<
@@ -314,7 +351,7 @@ impl<T: IsBucketType, I: Iterator<Item = BitChunk>> Iterator for BitChunksOnes<T
 mod tests {
     use itertools::Itertools;
 
-    use super::{iter_ones, BitChunk};
+    use super::{iter_bit_chunks, iter_ones, parity_bit_positions, BitChunk};
 
     #[test]
     fn test_iter_ones() {
@@ -338,4 +375,40 @@ mod tests {
             iter_ones::<usize, _>(chunks.into_iter().peekable()).collect_vec()
         );
     }
+
+    #[test]
+    fn test_iter_bit_chunks() {
+        // Distinct leading bits merge within a block (via or) and merge with the trailing block at
+        // the boundary index.
+        let chunks = iter_bit_chunks(
+            vec![70, 67, 5].into_iter(),
+            vec![BitChunk::new(0, 1 << 2)].into_iter(),
+        )
+        .collect_vec();
+        assert_eq!(
+            chunks,
+            vec![
+                BitChunk::new(1, (1 << 6) | (1 << 3)), // 70, 67
+                BitChunk::new(0, (1 << 5) | (1 << 2)), // 5 (leading) and bit 2 (trailing)
+            ]
+        );
+    }
+
+    #[test]
+    fn test_parity_bit_positions() {
+        // Equal positions cancel (xor): 70 twice cancels, 67 stays; 5 three times stays, 3 once.
+        let chunks = parity_bit_positions(vec![70, 70, 67, 5, 5, 5, 3].into_iter()).collect_vec();
+        assert_eq!(
+            chunks,
+            vec![
+                BitChunk::new(1, 1 << 3),              // 67
+                BitChunk::new(0, (1 << 5) | (1 << 3)), // 5, 3
+            ]
+        );
+
+        // A block whose positions cancel out entirely is skipped.
+        assert!(parity_bit_positions(vec![5, 5].into_iter())
+            .collect_vec()
+            .is_empty());
+    }
 }