diff --git a/crates/storage/src/bloom.rs b/crates/storage/src/bloom.rs index f58afe92bd..953cab48e3 100644 --- a/crates/storage/src/bloom.rs +++ b/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 { - let mut block_matches = vec![]; + pub fn blocks_for_keys(&self, keys: &[Felt]) -> BTreeSet { + 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 { + (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::::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::::new()); + assert_eq!(block_matches_empty, BTreeSet::new()); } #[test] diff --git a/crates/storage/src/connection/event.rs b/crates/storage/src/connection/event.rs index fc83481918..93a909bb03 100644 --- a/crates/storage/src/connection/event.rs +++ b/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 { - 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) -> BTreeSet { + 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]) -> BTreeSet { + 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 = 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();