chore: use AMQFilter trait and decode StandardBloomFilter from AMQFilterBuilder

Author: k-jingyang

benches/bloom.rs

@@ -14,7 +14,7 @@ fn standard_filter_construction(c: &mut Criterion) {
 }
 
 fn standard_filter_contains(c: &mut Criterion) {
-    use lsm_tree::segment::filter::standard_bloom::Builder;
+    use lsm_tree::segment::filter::{standard_bloom::Builder, AMQFilter};
 
     let keys = (0..100_000u128)
         .map(|x| x.to_be_bytes().to_vec())

src/segment/filter/mod.rs

@@ -6,10 +6,13 @@ pub mod bit_array;
 pub mod blocked_bloom;
 pub mod standard_bloom;
 
-const CACHE_LINE_BYTES: usize = 64;
+use crate::{coding::DecodeError, file::MAGIC_BYTES};
+use byteorder::ReadBytesExt;
+use std::io::Read;
+
+use standard_bloom::{Builder as StandardBloomFilterBuilder, StandardBloomFilter};
 
-use blocked_bloom::Builder as BlockedBloomFilterBuilder;
-use standard_bloom::Builder as StandardBloomFilterBuilder;
+const CACHE_LINE_BYTES: usize = 64;
 
 #[derive(Copy, Clone, Debug)]
 pub enum BloomConstructionPolicy {
@@ -42,3 +45,56 @@ impl BloomConstructionPolicy {
         }
     }
 }
+
+enum FilterType {
+    StandardBloom = 0,
+    BlockedBloom = 1,
+}
+
+impl TryFrom<u8> for FilterType {
+    type Error = ();
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(Self::StandardBloom),
+            1 => Ok(Self::BlockedBloom),
+            _ => Err(()),
+        }
+    }
+}
+
+pub trait AMQFilter: Sync + Send {
+    fn bytes(&self) -> &[u8];
+    fn len(&self) -> usize;
+    fn contains(&self, item: &[u8]) -> bool;
+    fn contains_hash(&self, hash: (u64, u64)) -> bool;
+}
+
+pub struct AMQFilterBuilder {}
+
+impl AMQFilterBuilder {
+    pub fn decode_from<R: Read>(reader: &mut R) -> Result<Box<dyn AMQFilter + Sync>, DecodeError> {
+        // Check header
+        let mut magic = [0u8; MAGIC_BYTES.len()];
+        reader.read_exact(&mut magic)?;
+
+        if magic != MAGIC_BYTES {
+            return Err(DecodeError::InvalidHeader("BloomFilter"));
+        }
+
+        let filter_type = reader.read_u8()?;
+        let filter_type = FilterType::try_from(filter_type)
+            .map_err(|_| DecodeError::InvalidHeader("FilterType"))?;
+
+        match filter_type {
+            FilterType::StandardBloom => StandardBloomFilter::decode_from(reader)
+                .map(Self::wrap_filter)
+                .map_err(|e| DecodeError::from(e)),
+            // TODO: Implement
+            FilterType::BlockedBloom => Err(DecodeError::InvalidHeader("BlockedBloom")),
+        }
+    }
+
+    fn wrap_filter<T: 'static + AMQFilter + Sync>(filter: T) -> Box<dyn AMQFilter + Sync> {
+        Box::new(filter)
+    }
+}

src/segment/filter/standard_bloom/mod.rs

@@ -1,4 +1,4 @@
-use super::bit_array::BitArrayReader;
+use super::{bit_array::BitArrayReader, AMQFilter};
 use crate::{
     coding::{Decode, DecodeError, Encode, EncodeError},
     file::MAGIC_BYTES,
@@ -30,6 +30,49 @@ pub struct StandardBloomFilter {
 }
 
 // TODO: change encode/decode to be Filter enum
+impl AMQFilter for StandardBloomFilter {
+    /// Size of bloom filter in bytes.
+    #[must_use]
+    fn len(&self) -> usize {
+        self.inner.bytes().len()
+    }
+
+    /// Returns the raw bytes of the filter.
+    fn bytes(&self) -> &[u8] {
+        self.inner.bytes()
+    }
+
+    /// Returns `true` if the item may be contained.
+    ///
+    /// Will never have a false negative.
+    #[must_use]
+    fn contains(&self, key: &[u8]) -> bool {
+        self.contains_hash(Self::get_hash(key))
+    }
+
+    /// Returns `true` if the hash may be contained.
+    ///
+    /// Will never have a false negative.
+    #[must_use]
+    fn contains_hash(&self, hash: CompositeHash) -> bool {
+        let (mut h1, mut h2) = hash;
+
+        for i in 1..=(self.k as u64) {
+            let idx = h1 % (self.m as u64);
+
+            // NOTE: should be in bounds because of modulo
+            #[allow(clippy::expect_used)]
+            if !self.has_bit(idx as usize) {
+                return false;
+            }
+
+            h1 = h1.wrapping_add(h2);
+            h2 = h2.wrapping_mul(i);
+        }
+
+        true
+    }
+}
 
 impl Encode for StandardBloomFilter {
     fn encode_into<W: Write>(&self, writer: &mut W) -> Result<(), EncodeError> {
@@ -50,20 +93,10 @@ impl Encode for StandardBloomFilter {
     }
 }
 
-impl Decode for StandardBloomFilter {
-    fn decode_from<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
-        // Check header
-        let mut magic = [0u8; MAGIC_BYTES.len()];
-        reader.read_exact(&mut magic)?;
-
-        if magic != MAGIC_BYTES {
-            return Err(DecodeError::InvalidHeader("BloomFilter"));
-        }
-
-        // NOTE: Filter type (unused)
-        let filter_type = reader.read_u8()?;
-        assert_eq!(0, filter_type, "Invalid filter type");
-
+#[allow(clippy::len_without_is_empty)]
+impl StandardBloomFilter {
+    // To be used by AMQFilter after magic bytes and filter type have been read and parsed
+    pub(super) fn decode_from<R: Read>(reader: &mut R) -> Result<Self, DecodeError> {
         // NOTE: Hash type (unused)
         let hash_type = reader.read_u8()?;
         assert_eq!(0, hash_type, "Invalid bloom hash type");
@@ -76,15 +109,6 @@ impl Decode for StandardBloomFilter {
 
         Ok(Self::from_raw(m, k, bytes.into()))
     }
-}
-
-#[allow(clippy::len_without_is_empty)]
-impl StandardBloomFilter {
-    /// Size of bloom filter in bytes.
-    #[must_use]
-    pub fn len(&self) -> usize {
-        self.inner.bytes().len()
-    }
 
     fn from_raw(m: usize, k: usize, slice: crate::Slice) -> Self {
         Self {
@@ -94,37 +118,6 @@ impl StandardBloomFilter {
         }
     }
 
-    /// Returns `true` if the hash may be contained.
-    ///
-    /// Will never have a false negative.
-    #[must_use]
-    pub fn contains_hash(&self, hash: CompositeHash) -> bool {
-        let (mut h1, mut h2) = hash;
-
-        for i in 1..=(self.k as u64) {
-            let idx = h1 % (self.m as u64);
-
-            // NOTE: should be in bounds because of modulo
-            #[allow(clippy::expect_used)]
-            if !self.has_bit(idx as usize) {
-                return false;
-            }
-
-            h1 = h1.wrapping_add(h2);
-            h2 = h2.wrapping_mul(i);
-        }
-
-        true
-    }
-
-    /// Returns `true` if the item may be contained.
-    ///
-    /// Will never have a false negative.
-    #[must_use]
-    pub fn contains(&self, key: &[u8]) -> bool {
-        self.contains_hash(Self::get_hash(key))
-    }
-
     /// Returns `true` if the bit at `idx` is `1`.
     fn has_bit(&self, idx: usize) -> bool {
         self.inner.get(idx)
@@ -138,6 +131,8 @@ impl StandardBloomFilter {
 
 #[cfg(test)]
 mod tests {
+    use crate::segment::filter::AMQFilterBuilder;
+
     use super::*;
     use std::fs::File;
     use test_log::test;
@@ -174,9 +169,9 @@ mod tests {
         drop(file);
 
         let mut file = File::open(&path)?;
-        let filter_copy = StandardBloomFilter::decode_from(&mut file)?;
+        let filter_copy = AMQFilterBuilder::decode_from(&mut file)?;
 
-        assert_eq!(filter.inner, filter_copy.inner);
+        assert_eq!(filter.inner.bytes(), filter_copy.bytes());
 
         for key in keys {
             assert!(filter.contains(&**key));

src/segment/inner.rs

@@ -3,8 +3,7 @@
 // (found in the LICENSE-* files in the repository)
 
 use super::{
-    block_index::NewBlockIndexImpl, filter::standard_bloom::StandardBloomFilter, meta::ParsedMeta,
-    trailer::Trailer,
+    block_index::NewBlockIndexImpl, filter::AMQFilter, meta::ParsedMeta, trailer::Trailer,
 };
 use crate::{
     cache::Cache, descriptor_table::DescriptorTable, tree::inner::TreeId, GlobalSegmentId,
@@ -39,7 +38,7 @@ pub struct Inner {
     pub cache: Arc<Cache>,
 
     /// Pinned AMQ filter
-    pub pinned_filter: Option<StandardBloomFilter>,
+    pub pinned_filter: Option<Box<dyn AMQFilter + Sync>>,
 
     // /// Pinned filter
     // #[doc(hidden)]

src/segment/mod.rs

@@ -27,7 +27,7 @@ use crate::{
     cache::Cache, descriptor_table::DescriptorTable, InternalValue, SeqNo, TreeId, UserKey,
 };
 use block_index::{NewBlockIndex, NewBlockIndexImpl, NewFullBlockIndex};
-use filter::standard_bloom::{CompositeHash, StandardBloomFilter};
+use filter::{standard_bloom::CompositeHash, AMQFilterBuilder};
 use inner::Inner;
 use meta::ParsedMeta;
 use std::{
@@ -341,8 +341,6 @@ impl Segment {
         let pinned_filter = trailer
             .filter
             .map(|filter_ptr| {
-                use crate::coding::Decode;
-
                 log::debug!("Reading filter block for pinning, with filter_ptr={filter_ptr:?}");
 
                 let block = Block::from_file(
@@ -353,7 +351,7 @@ impl Segment {
                 )?;
 
                 let mut reader = &block.data[..];
-                StandardBloomFilter::decode_from(&mut reader).map_err(Into::<crate::Error>::into)
+                AMQFilterBuilder::decode_from(&mut reader).map_err(Into::<crate::Error>::into)
             })
             .transpose()?;