AND event Bloom filter output instead of OR

crates/storage/src/bloom.rs

@@ -60,6 +60,8 @@
 //! specific set of keys without having to load and check each individual bloom
 //! filter.
 
+use std::collections::BTreeSet;
+
 use bloomfilter::Bloom;
 use pathfinder_common::BlockNumber;
 use pathfinder_crypto::Felt;
@@ -165,8 +167,12 @@ impl AggregateBloom {
 
     /// Returns a set of [block numbers](BlockNumber) for which the given keys
     /// are present in the aggregate.
-    pub fn blocks_for_keys(&self, keys: &[Felt]) -> Vec<BlockNumber> {
-        let mut block_matches = vec![];
+    pub fn blocks_for_keys(&self, keys: &[Felt]) -> BTreeSet<BlockNumber> {
+        if keys.is_empty() {
+            return self.all_blocks();
+        }
+
+        let mut block_matches = BTreeSet::new();
 
         for k in keys {
             let mut row_to_check = vec![u8::MAX; Self::BLOCK_RANGE_BYTES as usize];
@@ -186,7 +192,7 @@ impl AggregateBloom {
                 for i in 0..8 {
                     if byte & (1 << i) != 0 {
                         let match_number = self.from_block + col_idx as u64 * 8 + i as u64;
-                        block_matches.push(match_number);
+                        block_matches.insert(match_number);
                     }
                 }
             }
@@ -195,6 +201,12 @@ impl AggregateBloom {
         block_matches
     }
 
+    pub(super) fn all_blocks(&self) -> BTreeSet<BlockNumber> {
+        (self.from_block.get()..=self.to_block.get())
+            .map(BlockNumber::new_or_panic)
+            .collect()
+    }
+
     fn bitmap_at(&self, row: usize, col: usize) -> u8 {
         let idx = row * Self::BLOCK_RANGE_BYTES as usize + col;
         self.bitmap[idx]
@@ -315,7 +327,8 @@ mod tests {
         aggregate_bloom_filter.add_bloom(&bloom, from_block);
 
         let block_matches = aggregate_bloom_filter.blocks_for_keys(&[KEY]);
-        assert_eq!(block_matches, vec![from_block]);
+        let expected = BTreeSet::from_iter(vec![from_block]);
+        assert_eq!(block_matches, expected);
     }
 
     #[test]
@@ -330,7 +343,8 @@ mod tests {
         aggregate_bloom_filter.add_bloom(&bloom, from_block + 1);
 
         let block_matches = aggregate_bloom_filter.blocks_for_keys(&[KEY]);
-        assert_eq!(block_matches, vec![from_block, from_block + 1]);
+        let expected = BTreeSet::from_iter(vec![from_block, from_block + 1]);
+        assert_eq!(block_matches, expected);
     }
 
     #[test]
@@ -346,7 +360,7 @@ mod tests {
         aggregate_bloom_filter.add_bloom(&bloom, from_block + 1);
 
         let block_matches_empty = aggregate_bloom_filter.blocks_for_keys(&[KEY_NOT_IN_FILTER]);
-        assert_eq!(block_matches_empty, Vec::<BlockNumber>::new());
+        assert_eq!(block_matches_empty, BTreeSet::new());
     }
 
     #[test]
@@ -369,12 +383,11 @@ mod tests {
         decompressed.add_bloom(&bloom, from_block + 2);
 
         let block_matches = decompressed.blocks_for_keys(&[KEY]);
+        let expected = BTreeSet::from_iter(vec![from_block, from_block + 1, from_block + 2]);
+        assert_eq!(block_matches, expected,);
+
         let block_matches_empty = decompressed.blocks_for_keys(&[KEY_NOT_IN_FILTER]);
-        assert_eq!(
-            block_matches,
-            vec![from_block, from_block + 1, from_block + 2]
-        );
-        assert_eq!(block_matches_empty, Vec::<BlockNumber>::new());
+        assert_eq!(block_matches_empty, BTreeSet::new());
     }
 
     #[test]

crates/storage/src/connection/event.rs

@@ -499,21 +499,28 @@ impl Transaction<'_> {
 impl AggregateBloom {
     /// Returns the block numbers that match the given constraints.
     pub fn check(&self, constraints: &EventConstraints) -> BTreeSet<BlockNumber> {
-        if constraints.contract_address.is_none()
-            && (constraints.keys.iter().flatten().count() == 0)
-        {
-            // No constraints, return filter's entire block range.
-            return (self.from_block.get()..=self.to_block.get())
-                .map(BlockNumber::new_or_panic)
-                .collect();
+        let addr_blocks = self.check_address(constraints.contract_address);
+        let keys_blocks = self.check_keys(&constraints.keys);
+
+        addr_blocks.intersection(&keys_blocks).cloned().collect()
+    }
+
+    fn check_address(&self, address: Option<ContractAddress>) -> BTreeSet<BlockNumber> {
+        match address {
+            Some(addr) => self.blocks_for_keys(&[addr.0]),
+            None => self.all_blocks(),
         }
+    }
 
-        let mut blocks: BTreeSet<_> = constraints
-            .keys
-            .iter()
+    fn check_keys(&self, keys: &[Vec<EventKey>]) -> BTreeSet<BlockNumber> {
+        if keys.is_empty() {
+            return self.all_blocks();
+        }
+
+        keys.iter()
             .enumerate()
-            .flat_map(|(idx, keys)| {
-                let keys: Vec<_> = keys
+            .map(|(idx, key_group)| {
+                let indexed_keys: Vec<_> = key_group
                     .iter()
                     .map(|key| {
                         let mut key_with_idx = key.0;
@@ -522,15 +529,15 @@ impl AggregateBloom {
                     })
                     .collect();
 
-                self.blocks_for_keys(&keys)
+                self.blocks_for_keys(&indexed_keys)
             })
-            .collect();
-
-        if let Some(contract_address) = constraints.contract_address {
-            blocks.extend(self.blocks_for_keys(&[contract_address.0]));
-        }
-
-        blocks
+            .reduce(|intersection, blocks_for_key| {
+                intersection
+                    .intersection(&blocks_for_key)
+                    .cloned()
+                    .collect()
+            })
+            .unwrap_or_default()
     }
 }
 
@@ -699,6 +706,130 @@ mod tests {
     static MAX_BLOOM_FILTERS_TO_LOAD: LazyLock<NonZeroUsize> =
         LazyLock::new(|| NonZeroUsize::new(3).unwrap());
 
+    mod event_bloom {
+        use pretty_assertions_sorted::assert_eq;
+
+        use super::*;
+
+        #[test]
+        fn matching_contraints() {
+            let mut aggregate = AggregateBloom::new(BlockNumber::GENESIS);
+
+            let mut filter = BloomFilter::new();
+            filter.set_keys(&[event_key!("0xdeadbeef")]);
+            filter.set_address(&contract_address!("0x1234"));
+
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS);
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS + 1);
+            let constraints = EventConstraints {
+                from_block: None,
+                to_block: None,
+                contract_address: Some(contract_address!("0x1234")),
+                keys: vec![vec![event_key!("0xdeadbeef")]],
+                page_size: 1024,
+                offset: 0,
+            };
+
+            assert_eq!(
+                aggregate.check(&constraints),
+                BTreeSet::from_iter(vec![BlockNumber::GENESIS, BlockNumber::GENESIS + 1])
+            );
+        }
+
+        #[test]
+        fn correct_key_wrong_address() {
+            let mut aggregate = AggregateBloom::new(BlockNumber::GENESIS);
+
+            let mut filter = BloomFilter::new();
+            filter.set_keys(&[event_key!("0xdeadbeef")]);
+            filter.set_address(&contract_address!("0x1234"));
+
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS);
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS + 1);
+            let constraints = EventConstraints {
+                from_block: None,
+                to_block: None,
+                contract_address: Some(contract_address!("0x4321")),
+                keys: vec![vec![event_key!("0xdeadbeef")]],
+                page_size: 1024,
+                offset: 0,
+            };
+
+            assert_eq!(aggregate.check(&constraints), BTreeSet::new());
+        }
+
+        #[test]
+        fn correct_address_wrong_key() {
+            let mut aggregate = AggregateBloom::new(BlockNumber::GENESIS);
+
+            let mut filter = BloomFilter::new();
+            filter.set_keys(&[event_key!("0xdeadbeef")]);
+            filter.set_address(&contract_address!("0x1234"));
+
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS);
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS + 1);
+            let constraints = EventConstraints {
+                from_block: None,
+                to_block: None,
+                contract_address: Some(contract_address!("0x1234")),
+                keys: vec![vec![event_key!("0xfeebdaed"), event_key!("0x4321")]],
+                page_size: 1024,
+                offset: 0,
+            };
+
+            assert_eq!(aggregate.check(&constraints), BTreeSet::new());
+        }
+
+        #[test]
+        fn wrong_and_correct_key() {
+            let mut aggregate = AggregateBloom::new(BlockNumber::GENESIS);
+
+            let mut filter = BloomFilter::new();
+            filter.set_address(&contract_address!("0x1234"));
+            filter.set_keys(&[event_key!("0xdeadbeef")]);
+
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS);
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS + 1);
+            let constraints = EventConstraints {
+                from_block: None,
+                to_block: None,
+                contract_address: None,
+                keys: vec![
+                    // Key present in both blocks as the first key.
+                    vec![event_key!("0xdeadbeef")],
+                    // Key that does not exist in any block.
+                    vec![event_key!("0xbeefdead")],
+                ],
+                page_size: 1024,
+                offset: 0,
+            };
+
+            assert_eq!(aggregate.check(&constraints), BTreeSet::new());
+        }
+
+        #[test]
+        fn no_constraints() {
+            let mut aggregate = AggregateBloom::new(BlockNumber::GENESIS);
+
+            let mut filter = BloomFilter::new();
+            filter.set_keys(&[event_key!("0xdeadbeef")]);
+            filter.set_address(&contract_address!("0x1234"));
+
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS);
+            aggregate.add_bloom(&filter, BlockNumber::GENESIS + 1);
+            let constraints = EventConstraints {
+                from_block: None,
+                to_block: None,
+                contract_address: None,
+                keys: vec![],
+                page_size: 1024,
+                offset: 0,
+            };
+
+            assert_eq!(aggregate.check(&constraints), aggregate.all_blocks());
+        }
+    }
+
     #[test_log::test(test)]
     fn get_events_with_fully_specified_filter() {
         let (storage, test_data) = test_utils::setup_test_storage();