pub mod conf;
pub mod meta;
pub use conf::{
BfInsert, BfReserve, BloomFilterAddResult, BloomFilterInfo, BloomFilterInsertOptions,
BloomInfoType, CfInsert, CfReserve, CuckooFilterInfo, CuckooFilterInsertOptions,
DEFAULT_BF_CAPACITY, DEFAULT_BF_ERROR_RATE, DEFAULT_BF_EXPANSION, DEFAULT_CF_BUCKET_SIZE,
DEFAULT_CF_CAPACITY, DEFAULT_CF_EXPANSION, DEFAULT_CF_MAX_ITERATIONS, DEFAULT_CF_PAGE_SIZE,
MAX_CF_EXPANSION,
};
pub use meta::{BloomChainMeta, CuckooChainMeta};
use core::hint::unreachable_unchecked;
use memchr::memchr;
use rapidhash::{HashMapExt, HashSetExt, RapidHashMap as HashMap, RapidHashSet as HashSet};
use std::str;
use crate::db::WeDb;
use crate::error::{Error, Result};
use crate::key_composer::KeyComposer;
use crate::meta::current_now_ms;
use fjall::Keyspace;
use rapidhash::v3::rapidhash_v3;
#[inline]
fn rapid_hash(bytes: &[u8]) -> u64 {
rapidhash_v3(bytes)
}
pub struct BlockSplitBloomFilter;
impl BlockSplitBloomFilter {
pub const MIN_BYTES: usize = 32;
pub const MAX_BYTES: usize = 128 * 1024 * 1024; pub const BYTES_PER_FILTER_BLOCK: usize = 32;
pub const BITS_SET_PER_BLOCK: usize = 8;
pub const SALT: [u32; 8] = [
0x47b6_137b,
0x4497_4d91,
0x8824_ad5b,
0xa2b7_289d,
0x7054_95c7,
0x2df1_424b,
0x9efc_4947,
0x5c6b_fb31,
];
#[inline]
pub fn hash(data: &[u8]) -> u64 {
rapid_hash(data)
}
pub fn optimal_num_of_bits(ndv: u32, fpp: f64) -> usize {
let fpp_clamped = fpp.clamp(1e-15, 0.99999);
let m = -8.0 * (ndv as f64) / (1.0 - fpp_clamped.powf(1.0 / 8.0)).ln();
let max_bits = Self::MAX_BYTES * 8;
let mut num_bits = if m.is_nan() || m < 0.0 || m > (max_bits as f64) {
max_bits
} else {
m as usize
};
let min_bits = Self::MIN_BYTES * 8;
if num_bits < min_bits {
num_bits = min_bits;
}
if !num_bits.is_power_of_two() {
num_bits = num_bits.next_power_of_two();
}
if num_bits > max_bits {
num_bits = max_bits;
}
num_bits
}
#[inline]
pub fn optimal_num_of_bytes(ndv: u32, fpp: f64) -> usize {
Self::optimal_num_of_bits(ndv, fpp) >> 3
}
#[inline]
pub fn find_hash(data: &[u8], hash: u64) -> bool {
if data.len() < Self::BYTES_PER_FILTER_BLOCK {
return false;
}
let num_blocks = (data.len() / Self::BYTES_PER_FILTER_BLOCK) as u64;
if num_blocks == 0 {
return false;
}
let bucket_index = (((hash >> 32).wrapping_mul(num_blocks)) >> 32) as usize;
let key = hash as u32;
let block_start = bucket_index * Self::BYTES_PER_FILTER_BLOCK;
if block_start + Self::BYTES_PER_FILTER_BLOCK > data.len() {
return false;
}
let block =
unsafe { data.get_unchecked(block_start..block_start + Self::BYTES_PER_FILTER_BLOCK) };
let (chunks, _) = block.as_chunks::<4>();
for (&salt, chunk) in Self::SALT.iter().zip(chunks) {
let bit_shift = (key.wrapping_mul(salt)) >> 27;
let mask = 1u32 << bit_shift;
let word = u32::from_le_bytes(*chunk);
if (word & mask) == 0 {
return false;
}
}
true
}
#[inline]
pub fn insert_hash(data: &mut [u8], hash: u64) -> bool {
if data.len() < Self::BYTES_PER_FILTER_BLOCK {
return false;
}
let num_blocks = (data.len() / Self::BYTES_PER_FILTER_BLOCK) as u64;
if num_blocks == 0 {
return false;
}
let bucket_index = (((hash >> 32).wrapping_mul(num_blocks)) >> 32) as usize;
let key = hash as u32;
let block_start = bucket_index * Self::BYTES_PER_FILTER_BLOCK;
if block_start + Self::BYTES_PER_FILTER_BLOCK > data.len() {
return false;
}
let block = unsafe {
data.get_unchecked_mut(block_start..block_start + Self::BYTES_PER_FILTER_BLOCK)
};
let mut modified = false;
let (chunks, _) = block.as_chunks_mut::<4>();
for (&salt, chunk) in Self::SALT.iter().zip(chunks) {
let bit_shift = (key.wrapping_mul(salt)) >> 27;
let mask = 1u32 << bit_shift;
let word = u32::from_le_bytes(*chunk);
if (word & mask) == 0 {
*chunk = (word | mask).to_le_bytes();
modified = true;
}
}
modified
}
}
enum BloomFilterBuffer {
Clean(fjall::Slice),
Dirty(Vec<u8>),
}
impl BloomFilterBuffer {
#[inline]
fn as_slice(&self) -> &[u8] {
match self {
Self::Clean(s) => s.as_ref(),
Self::Dirty(v) => v.as_slice(),
}
}
#[inline]
fn ensure_mut(&mut self, expected_bytes: usize) -> &mut Vec<u8> {
if let Self::Clean(s) = self {
let mut vec = s.to_vec();
if vec.len() < expected_bytes {
vec.resize(expected_bytes, 0);
}
*self = Self::Dirty(vec);
}
match self {
Self::Dirty(v) => {
if v.len() < expected_bytes {
v.resize(expected_bytes, 0);
}
v
}
_ => unsafe { unreachable_unchecked() },
}
}
}
pub struct CuckooFilterHelper;
impl CuckooFilterHelper {
pub const LOAD_FACTOR: f64 = 0.955;
pub const ALT_HASH_MULTIPLIER: u64 = 0x5bd1_e995;
pub const FINGERPRINT_MODULUS: u64 = 255;
pub const DEFAULT_PAGE_SIZE: u32 = DEFAULT_CF_PAGE_SIZE;
pub const DEFAULT_CAPACITY: u64 = DEFAULT_CF_CAPACITY;
pub const DEFAULT_BUCKET_SIZE: u8 = DEFAULT_CF_BUCKET_SIZE;
pub const DEFAULT_MAX_ITERATIONS: u16 = DEFAULT_CF_MAX_ITERATIONS;
pub const DEFAULT_EXPANSION: u16 = DEFAULT_CF_EXPANSION;
pub const MAX_EXPANSION: u16 = MAX_CF_EXPANSION;
#[inline]
pub fn hash(data: &[u8]) -> u64 {
rapid_hash(data)
}
#[inline]
pub fn generate_fingerprint(hash: u64) -> u8 {
((hash % Self::FINGERPRINT_MODULUS) + 1) as u8
}
#[inline]
pub fn get_alt_hash(fingerprint: u8, hash: u64) -> u64 {
hash ^ ((fingerprint as u64).wrapping_mul(Self::ALT_HASH_MULTIPLIER))
}
#[inline]
pub fn get_alt_bucket_index(bucket_idx: u32, fingerprint: u8, num_buckets: u32) -> u32 {
let alt_hash = Self::get_alt_hash(fingerprint, bucket_idx as u64);
(alt_hash % (num_buckets as u64)) as u32
}
#[inline]
pub fn normalize_expansion(expansion: u16) -> u16 {
if expansion <= 1 {
expansion
} else {
(expansion as u32).next_power_of_two().min(32768) as u16
}
}
pub fn calculate_required_buckets(capacity: u64, bucket_size: u8) -> Result<u32> {
if bucket_size == 0 {
return Err(Error::invalid_data("bucket_size must be larger than 0"));
}
let max_supported_capacity =
((1u64 << 31) as f64 * (bucket_size as f64) * Self::LOAD_FACTOR) as u64;
if capacity > max_supported_capacity {
return Err(Error::invalid_data("capacity is too large"));
}
let exact_buckets = (capacity as f64) / (bucket_size as f64) / Self::LOAD_FACTOR;
let mut req_buckets = exact_buckets.ceil() as u64;
if req_buckets == 0 {
req_buckets = 1;
}
let num_buckets = req_buckets.next_power_of_two();
if num_buckets > (1u64 << 31) {
return Err(Error::invalid_data("capacity is too large"));
}
Ok(num_buckets as u32)
}
}
#[derive(Debug, Clone, Copy)]
struct PageSlotUndo {
filter_idx: u16,
page_idx: u32,
byte_offset: usize,
old_fp: u8,
}
enum CuckooPageData {
Clean(fjall::Slice),
Dirty(Vec<u8>),
Empty(usize),
}
impl CuckooPageData {
#[inline]
fn get_subslice(&self, offset: usize, len: usize) -> Option<&[u8]> {
match self {
Self::Clean(s) => s.get(offset..offset + len),
Self::Dirty(v) => v.get(offset..offset + len),
Self::Empty(_) => None,
}
}
#[inline]
fn get_byte(&self, offset: usize) -> u8 {
match self {
Self::Clean(s) => s.get(offset).copied().unwrap_or(0),
Self::Dirty(v) => v.get(offset).copied().unwrap_or(0),
Self::Empty(_) => 0,
}
}
#[inline]
fn ensure_mut(&mut self, expected_size: usize) -> &mut Vec<u8> {
match self {
Self::Dirty(v) => {
if v.len() < expected_size {
v.resize(expected_size, 0);
}
v
}
Self::Clean(s) => {
let mut vec = s.to_vec();
if vec.len() < expected_size {
vec.resize(expected_size, 0);
}
*self = Self::Dirty(vec);
if let Self::Dirty(v) = self {
v
} else {
unsafe { unreachable_unchecked() }
}
}
Self::Empty(exp) => {
let size = (*exp).max(expected_size);
*self = Self::Dirty(vec![0u8; size]);
if let Self::Dirty(v) = self {
v
} else {
unsafe { unreachable_unchecked() }
}
}
}
}
}
struct CuckooPageCache<'a> {
data: &'a Keyspace,
kc: KeyComposer<'a>,
key: &'a str,
bucket_size: u8,
buckets_per_page: u32,
pages: HashMap<(u16, u32), CuckooPageData>,
}
impl<'a> CuckooPageCache<'a> {
#[inline]
fn new(
data: &'a Keyspace,
kc: KeyComposer<'a>,
key: &'a str,
bucket_size: u8,
page_size: u32,
) -> Self {
let buckets_per_page = (page_size / bucket_size as u32).max(1);
Self {
data,
kc,
key,
bucket_size,
buckets_per_page,
pages: HashMap::new(),
}
}
#[inline]
fn get_bucket_location(&self, bucket_idx: u32) -> (u32, usize) {
let page_idx = bucket_idx / self.buckets_per_page;
let offset = ((bucket_idx % self.buckets_per_page) as usize) * (self.bucket_size as usize);
(page_idx, offset)
}
#[inline]
fn get_expected_page_size(&self, num_buckets: u32, page_idx: u32) -> usize {
let first_bucket = page_idx * self.buckets_per_page;
let page_bucket_count = self
.buckets_per_page
.min(num_buckets.saturating_sub(first_bucket));
(page_bucket_count as usize) * (self.bucket_size as usize)
}
fn get_page_entry_mut(
&mut self,
filter_idx: u16,
num_buckets: u32,
page_idx: u32,
) -> Result<&mut CuckooPageData> {
use std::collections::hash_map::Entry;
let expected_size = self.get_expected_page_size(num_buckets, page_idx);
let key = (filter_idx, page_idx);
match self.pages.entry(key) {
Entry::Occupied(entry) => Ok(entry.into_mut()),
Entry::Vacant(entry) => {
let page_key = self.kc.cf_page(self.key, filter_idx, page_idx);
let page_data = match self.data.get(page_key.as_bytes())? {
Some(slice) => CuckooPageData::Clean(slice),
None => CuckooPageData::Empty(expected_size),
};
Ok(entry.insert(page_data))
}
}
}
#[inline]
fn try_insert_in_bucket(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
fp: u8,
) -> Result<bool> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let bs = self.bucket_size as usize;
let expected_size = self.get_expected_page_size(num_buckets, page_idx);
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
match entry {
CuckooPageData::Empty(_) => {
let page_vec = entry.ensure_mut(expected_size);
page_vec[offset] = fp;
Ok(true)
}
CuckooPageData::Clean(s) => {
if let Some(sub) = s.get(offset..offset + bs) {
if let Some(pos) = memchr(0, sub) {
let page_vec = entry.ensure_mut(expected_size);
page_vec[offset + pos] = fp;
Ok(true)
} else {
Ok(false)
}
} else {
let page_vec = entry.ensure_mut(expected_size);
page_vec[offset] = fp;
Ok(true)
}
}
CuckooPageData::Dirty(v) => {
if offset + bs > v.len() {
v.resize(offset + bs, 0);
}
if let Some(pos) = memchr(0, &v[offset..offset + bs]) {
v[offset + pos] = fp;
Ok(true)
} else {
Ok(false)
}
}
}
}
#[inline]
fn try_insert_in_bucket_recorded(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
fp: u8,
undo_log: &mut Vec<PageSlotUndo>,
) -> Result<bool> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let bs = self.bucket_size as usize;
let expected_size = self.get_expected_page_size(num_buckets, page_idx);
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
match entry {
CuckooPageData::Empty(_) => {
let page_vec = entry.ensure_mut(expected_size);
let byte_offset = offset;
undo_log.push(PageSlotUndo {
filter_idx,
page_idx,
byte_offset,
old_fp: 0,
});
page_vec[byte_offset] = fp;
Ok(true)
}
CuckooPageData::Clean(s) => {
if let Some(sub) = s.get(offset..offset + bs) {
if let Some(pos) = memchr(0, sub) {
let page_vec = entry.ensure_mut(expected_size);
let byte_offset = offset + pos;
undo_log.push(PageSlotUndo {
filter_idx,
page_idx,
byte_offset,
old_fp: 0,
});
page_vec[byte_offset] = fp;
Ok(true)
} else {
Ok(false)
}
} else {
let page_vec = entry.ensure_mut(expected_size);
let byte_offset = offset;
undo_log.push(PageSlotUndo {
filter_idx,
page_idx,
byte_offset,
old_fp: 0,
});
page_vec[byte_offset] = fp;
Ok(true)
}
}
CuckooPageData::Dirty(v) => {
if offset + bs > v.len() {
v.resize(offset + bs, 0);
}
if let Some(pos) = memchr(0, &v[offset..offset + bs]) {
let byte_offset = offset + pos;
undo_log.push(PageSlotUndo {
filter_idx,
page_idx,
byte_offset,
old_fp: 0,
});
v[byte_offset] = fp;
Ok(true)
} else {
Ok(false)
}
}
}
}
#[inline]
fn get_bucket_slot(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
slot: usize,
) -> Result<u8> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
Ok(entry.get_byte(offset + slot))
}
#[inline]
fn set_bucket_slot_recorded(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
slot: usize,
fp: u8,
undo_log: &mut Vec<PageSlotUndo>,
) -> Result<()> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let expected_size = self.get_expected_page_size(num_buckets, page_idx);
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
let byte_offset = offset + slot;
let old_fp = entry.get_byte(byte_offset);
let page_vec = entry.ensure_mut(expected_size);
undo_log.push(PageSlotUndo {
filter_idx,
page_idx,
byte_offset,
old_fp,
});
page_vec[byte_offset] = fp;
Ok(())
}
#[inline]
fn apply_undo_log(&mut self, undo_log: &[PageSlotUndo]) {
for undo in undo_log.iter().rev() {
if let Some(CuckooPageData::Dirty(vec)) =
self.pages.get_mut(&(undo.filter_idx, undo.page_idx))
&& undo.byte_offset < vec.len()
{
vec[undo.byte_offset] = undo.old_fp;
}
}
}
#[inline]
fn contains_in_bucket(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
fp: u8,
) -> Result<bool> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let bs = self.bucket_size as usize;
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
if let Some(sub) = entry.get_subslice(offset, bs) {
Ok(memchr(fp, sub).is_some())
} else {
Ok(false)
}
}
#[inline]
fn count_in_bucket(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
fp: u8,
) -> Result<usize> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let bs = self.bucket_size as usize;
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
if let Some(mut slice) = entry.get_subslice(offset, bs) {
let mut count = 0;
while let Some(pos) = memchr(fp, slice) {
count += 1;
slice = &slice[pos + 1..];
}
Ok(count)
} else {
Ok(0)
}
}
#[inline]
fn delete_from_bucket(
&mut self,
filter_idx: u16,
num_buckets: u32,
bucket_idx: u32,
fp: u8,
) -> Result<bool> {
let (page_idx, offset) = self.get_bucket_location(bucket_idx);
let bs = self.bucket_size as usize;
let expected_size = self.get_expected_page_size(num_buckets, page_idx);
let entry = self.get_page_entry_mut(filter_idx, num_buckets, page_idx)?;
match entry {
CuckooPageData::Empty(_) => Ok(false),
CuckooPageData::Clean(s) => {
if let Some(sub) = s.get(offset..offset + bs) {
if let Some(pos) = memchr(fp, sub) {
let page_vec = entry.ensure_mut(expected_size);
page_vec[offset + pos] = 0;
Ok(true)
} else {
Ok(false)
}
} else {
Ok(false)
}
}
CuckooPageData::Dirty(v) => {
if let Some(sub) = v.get(offset..offset + bs) {
if let Some(pos) = memchr(fp, sub) {
v[offset + pos] = 0;
Ok(true)
} else {
Ok(false)
}
} else {
Ok(false)
}
}
}
}
fn commit_dirty_to_batch(&self, batch: &mut fjall::OwnedWriteBatch) {
for (&(filter_idx, page_idx), page_data) in &self.pages {
if let CuckooPageData::Dirty(vec) = page_data {
let page_key = self.kc.cf_page(self.key, filter_idx, page_idx);
batch.insert(self.data, page_key.as_bytes(), vec.as_slice());
}
}
}
}
fn try_cuckoo_kickout(
page_cache: &mut CuckooPageCache<'_>,
filter_idx: u16,
num_buckets: u32,
bucket_size: u8,
max_iterations: u16,
hash: u64,
fp: u8,
) -> Result<bool> {
let mut undo_log = Vec::with_capacity(max_iterations as usize * 2);
let mut cur_i = (hash % (num_buckets as u64)) as u32;
let mut cur_fp = fp;
let mut victim_slot = 0;
for _ in 0..max_iterations {
let old_fp = match page_cache.get_bucket_slot(filter_idx, num_buckets, cur_i, victim_slot) {
Ok(v) => v,
Err(e) => {
page_cache.apply_undo_log(&undo_log);
return Err(e);
}
};
if let Err(e) = page_cache.set_bucket_slot_recorded(
filter_idx,
num_buckets,
cur_i,
victim_slot,
cur_fp,
&mut undo_log,
) {
page_cache.apply_undo_log(&undo_log);
return Err(e);
}
cur_fp = old_fp;
if cur_fp == 0 {
return Ok(true);
}
let alt_bucket_idx = CuckooFilterHelper::get_alt_bucket_index(cur_i, cur_fp, num_buckets);
let inserted_in_alt = match page_cache.try_insert_in_bucket_recorded(
filter_idx,
num_buckets,
alt_bucket_idx,
cur_fp,
&mut undo_log,
) {
Ok(v) => v,
Err(e) => {
page_cache.apply_undo_log(&undo_log);
return Err(e);
}
};
if inserted_in_alt {
return Ok(true);
}
cur_i = alt_bucket_idx;
victim_slot = (victim_slot + 1) % (bucket_size as usize);
}
page_cache.apply_undo_log(&undo_log);
Ok(false)
}
impl WeDb {
#[inline]
fn get_bf_meta_checked(&self, meta_k: &str, now_ms: u64) -> Result<Option<BloomChainMeta>> {
match self.meta.get(meta_k.as_bytes())? {
Some(m_bytes) => {
let meta = BloomChainMeta::decode(&m_bytes)
.ok_or_else(|| Error::invalid_data("corrupted bloom metadata"))?;
if meta.is_expired(now_ms) {
Ok(None)
} else {
Ok(Some(meta))
}
}
None => Ok(None),
}
}
#[inline]
fn get_cf_meta_checked(&self, meta_k: &str, now_ms: u64) -> Result<Option<CuckooChainMeta>> {
match self.meta.get(meta_k.as_bytes())? {
Some(bytes) => {
let meta = CuckooChainMeta::decode(&bytes)
.ok_or_else(|| Error::invalid_data("corrupted cuckoo metadata"))?;
if meta.is_expired(now_ms) {
Ok(None)
} else {
Ok(Some(meta))
}
}
None => Ok(None),
}
}
pub fn bf_reserve<K: AsRef<[u8]>>(
&self,
key: K,
error_rate: f64,
capacity: u32,
expansion: u16,
) -> Result<()> {
if capacity == 0 {
return Err(Error::invalid_data("capacity must be larger than 0"));
}
if error_rate <= 0.0 || error_rate >= 1.0 {
return Err(Error::invalid_data("error_rate must be between 0 and 1"));
}
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.bf_meta(k_str);
let now_ms = current_now_ms();
if self.get_bf_meta_checked(&meta_k, now_ms)?.is_some() {
return Err(Error::invalid_data("ERR item exists"));
}
let bloom_bytes = BlockSplitBloomFilter::optimal_num_of_bytes(capacity, error_rate) as u32;
let meta = BloomChainMeta::new(capacity, error_rate, expansion, 0, 0, bloom_bytes);
let mut batch = self.db.batch();
batch.insert(&self.meta, meta_k.as_bytes(), meta.encode());
let item_k = kc.bf_item(k_str, 0);
let block_data = vec![0u8; bloom_bytes as usize];
batch.insert(&self.data, item_k.as_bytes(), block_data.as_slice());
batch.commit()?;
Ok(())
}
pub fn bf_add<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let res = self.bf_insert(key, &[item], &BloomFilterInsertOptions::default())?;
match res.first() {
Some(BloomFilterAddResult::Ok) => Ok(true),
Some(BloomFilterAddResult::Exist) => Ok(false),
Some(BloomFilterAddResult::Full) => {
Err(Error::invalid_data("ERR nonscaling filter is full"))
}
None => Ok(false),
}
}
pub fn bf_madd<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
) -> Result<Vec<bool>> {
let res = self.bf_insert(key, items, &BloomFilterInsertOptions::default())?;
Ok(res
.into_iter()
.map(|r| matches!(r, BloomFilterAddResult::Ok))
.collect())
}
pub fn bf_insert<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
opt: &BloomFilterInsertOptions,
) -> Result<Vec<BloomFilterAddResult>> {
if items.is_empty() {
return Ok(Vec::new());
}
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.bf_meta(k_str);
let now_ms = current_now_ms();
let mut meta = match self.get_bf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => {
if !opt.auto_create {
return Err(Error::invalid_data("ERR not found"));
}
if opt.capacity == 0 {
return Err(Error::invalid_data("capacity must be larger than 0"));
}
if opt.error_rate <= 0.0 || opt.error_rate >= 1.0 {
return Err(Error::invalid_data("error_rate must be between 0 and 1"));
}
let bloom_bytes =
BlockSplitBloomFilter::optimal_num_of_bytes(opt.capacity, opt.error_rate)
as u32;
BloomChainMeta::new(
opt.capacity,
opt.error_rate,
opt.expansion,
0,
0,
bloom_bytes,
)
}
};
let mut sub_filters = Vec::with_capacity(meta.n_filters as usize);
let mut dirty_filters = HashSet::new();
for i in 0..meta.n_filters {
let sub_k = kc.bf_item(k_str, i);
let cap = meta.sub_filter_capacity(i);
let expected_bytes = BlockSplitBloomFilter::optimal_num_of_bytes(cap, meta.error_rate);
let buf = match self.data.get(sub_k.as_bytes())? {
Some(slice) => {
if slice.len() < expected_bytes {
let mut vec = slice.to_vec();
vec.resize(expected_bytes, 0);
BloomFilterBuffer::Dirty(vec)
} else {
BloomFilterBuffer::Clean(slice)
}
}
None => BloomFilterBuffer::Dirty(vec![0u8; expected_bytes]),
};
sub_filters.push(buf);
}
let origin_size = meta.base.size;
let mut results = Vec::with_capacity(items.len());
for item in items {
let h = BlockSplitBloomFilter::hash(item.as_ref());
let exists = sub_filters
.iter()
.rev()
.any(|buf| BlockSplitBloomFilter::find_hash(buf.as_slice(), h));
if exists {
results.push(BloomFilterAddResult::Exist);
} else {
if meta.base.size + 1 > meta.get_capacity() as u64 {
if meta.is_scaling() && meta.n_filters < u16::MAX {
let new_filter_idx = meta.n_filters;
let new_cap = meta.sub_filter_capacity(new_filter_idx);
let new_bytes =
BlockSplitBloomFilter::optimal_num_of_bytes(new_cap, meta.error_rate);
meta.n_filters += 1;
meta.bloom_bytes += new_bytes as u32;
sub_filters.push(BloomFilterBuffer::Dirty(vec![0u8; new_bytes]));
dirty_filters.insert(new_filter_idx);
} else {
results.push(BloomFilterAddResult::Full);
continue;
}
}
let last_idx = (meta.n_filters - 1) as usize;
let cap = meta.sub_filter_capacity(last_idx as u16);
let expected_bytes =
BlockSplitBloomFilter::optimal_num_of_bytes(cap, meta.error_rate);
let filter_buf = sub_filters[last_idx].ensure_mut(expected_bytes);
BlockSplitBloomFilter::insert_hash(filter_buf, h);
dirty_filters.insert(last_idx as u16);
meta.base.size += 1;
results.push(BloomFilterAddResult::Ok);
}
}
if meta.base.size != origin_size || !dirty_filters.is_empty() {
let mut batch = self.db.batch();
batch.insert(&self.meta, meta_k.as_bytes(), meta.encode());
for idx in dirty_filters {
let sub_k = kc.bf_item(k_str, idx);
batch.insert(
&self.data,
sub_k.as_bytes(),
sub_filters[idx as usize].as_slice(),
);
}
batch.commit()?;
}
Ok(results)
}
pub fn bf_exists<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let res = self.bf_mexists(key, &[item])?;
Ok(res.first().copied().unwrap_or(false))
}
pub fn bf_mexists<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
) -> Result<Vec<bool>> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.bf_meta(k_str);
let now_ms = current_now_ms();
let meta = match self.get_bf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => return Ok(vec![false; items.len()]),
};
let mut blocks = Vec::with_capacity(meta.n_filters as usize);
for idx in 0..meta.n_filters {
let item_k = kc.bf_item(k_str, idx);
if let Some(b) = self.data.get(item_k.as_bytes())? {
blocks.push(b);
}
}
let mut results = Vec::with_capacity(items.len());
for it in items {
let h = BlockSplitBloomFilter::hash(it.as_ref());
let found = blocks
.iter()
.rev()
.any(|blk| BlockSplitBloomFilter::find_hash(blk.as_ref(), h));
results.push(found);
}
Ok(results)
}
pub fn bf_info<K: AsRef<[u8]>>(&self, key: K) -> Result<BloomFilterInfo> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.bf_meta(k_str);
let now_ms = current_now_ms();
let meta = self
.get_bf_meta_checked(&meta_k, now_ms)?
.ok_or_else(|| Error::invalid_data("ERR not found"))?;
Ok(BloomFilterInfo {
capacity: meta.get_capacity(),
bloom_bytes: meta.bloom_bytes,
n_filters: meta.n_filters,
size: meta.base.size,
expansion: meta.expansion,
})
}
pub fn bf_info_by_type<K: AsRef<[u8]>>(
&self,
key: K,
info_type: BloomInfoType,
) -> Result<Option<u64>> {
let info = self.bf_info(key)?;
match info_type {
BloomInfoType::All => Ok(None),
BloomInfoType::Capacity => Ok(Some(info.capacity as u64)),
BloomInfoType::Size => Ok(Some(info.bloom_bytes as u64)),
BloomInfoType::Filters => Ok(Some(info.n_filters as u64)),
BloomInfoType::Items => Ok(Some(info.size)),
BloomInfoType::Expansion => {
if info.expansion == 0 {
Ok(None)
} else {
Ok(Some(info.expansion as u64))
}
}
}
}
pub fn bf_card<K: AsRef<[u8]>>(&self, key: K) -> Result<u64> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.bf_meta(k_str);
let now_ms = current_now_ms();
Ok(self
.get_bf_meta_checked(&meta_k, now_ms)?
.map(|m| m.base.size)
.unwrap_or(0))
}
pub fn cf_reserve_ext<K: AsRef<[u8]>>(
&self,
key: K,
capacity: u64,
bucket_size: u8,
max_iterations: u16,
expansion: u16,
page_size: u32,
) -> Result<()> {
if capacity == 0 {
return Err(Error::invalid_data("capacity must be larger than 0"));
}
if capacity < 2 {
return Err(Error::invalid_data("capacity must be at least 2"));
}
if bucket_size == 0 {
return Err(Error::invalid_data("bucket_size must be between 1 and 255"));
}
if max_iterations == 0 {
return Err(Error::invalid_data("max_iterations must be larger than 0"));
}
if page_size == 0 {
return Err(Error::invalid_data("page_size must be larger than 0"));
}
if page_size < bucket_size as u32 {
return Err(Error::invalid_data(
"page_size must be at least bucket_size",
));
}
if expansion > CuckooFilterHelper::MAX_EXPANSION {
return Err(Error::invalid_data("expansion must be between 0 and 32768"));
}
CuckooFilterHelper::calculate_required_buckets(capacity, bucket_size)?;
let normalized_expansion = CuckooFilterHelper::normalize_expansion(expansion);
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
if self.get_cf_meta_checked(&meta_k, now_ms)?.is_some() {
return Err(Error::invalid_data("ERR item exists"));
}
let meta = CuckooChainMeta::new(
capacity,
bucket_size,
max_iterations,
normalized_expansion,
page_size,
0,
0,
);
self.meta.insert(meta_k.as_bytes(), meta.encode())?;
Ok(())
}
#[inline]
pub fn cf_reserve<K: AsRef<[u8]>>(
&self,
key: K,
capacity: u64,
bucket_size: u8,
max_iterations: u16,
expansion: u16,
) -> Result<()> {
self.cf_reserve_ext(
key,
capacity,
bucket_size,
max_iterations,
expansion,
CuckooFilterHelper::DEFAULT_PAGE_SIZE,
)
}
pub fn cf_add<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let res = self.cf_insert(key, &[item], &CuckooFilterInsertOptions::default())?;
Ok(res.first().copied().unwrap_or(false))
}
pub fn cf_addnx<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let opt = CuckooFilterInsertOptions {
nx: true,
..Default::default()
};
let res = self.cf_insert(key, &[item], &opt)?;
Ok(res.first().copied().unwrap_or(false))
}
pub fn cf_insert<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
opt: &CuckooFilterInsertOptions,
) -> Result<Vec<bool>> {
if items.is_empty() {
return Ok(Vec::new());
}
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
let mut meta = match self.get_cf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => {
if !opt.auto_create {
return Err(Error::invalid_data("ERR not found"));
}
if opt.capacity < 2 {
return Err(Error::invalid_data("capacity must be at least 2"));
}
if opt.bucket_size == 0 {
return Err(Error::invalid_data("bucket_size must be between 1 and 255"));
}
if opt.max_iterations == 0 {
return Err(Error::invalid_data("max_iterations must be larger than 0"));
}
if opt.page_size == 0 {
return Err(Error::invalid_data("page_size must be larger than 0"));
}
if opt.page_size < opt.bucket_size as u32 {
return Err(Error::invalid_data(
"page_size must be at least bucket_size",
));
}
if opt.expansion > CuckooFilterHelper::MAX_EXPANSION {
return Err(Error::invalid_data("expansion must be between 0 and 32768"));
}
CuckooFilterHelper::calculate_required_buckets(opt.capacity, opt.bucket_size)?;
CuckooChainMeta::new(
opt.capacity,
opt.bucket_size,
opt.max_iterations,
opt.expansion,
opt.page_size,
0,
0,
)
}
};
let mut page_cache =
CuckooPageCache::new(&self.data, kc, k_str, meta.bucket_size, meta.page_size);
let mut results = Vec::with_capacity(items.len());
for item in items {
let h = CuckooFilterHelper::hash(item.as_ref());
let fp = CuckooFilterHelper::generate_fingerprint(h);
if opt.nx {
let mut exists = false;
for filter_idx in (0..meta.n_filters).rev() {
let num_buckets = meta.sub_filter_num_buckets(filter_idx)?;
let b1 = (h % (num_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, num_buckets);
if page_cache.contains_in_bucket(filter_idx, num_buckets, b1, fp)?
|| (b1 != b2
&& page_cache.contains_in_bucket(filter_idx, num_buckets, b2, fp)?)
{
exists = true;
break;
}
}
if exists {
results.push(false);
continue;
}
}
let mut inserted = false;
for filter_idx in (0..meta.n_filters).rev() {
let num_buckets = meta.sub_filter_num_buckets(filter_idx)?;
let b1 = (h % (num_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, num_buckets);
if page_cache.try_insert_in_bucket(filter_idx, num_buckets, b1, fp)? {
inserted = true;
break;
}
if b1 != b2 && page_cache.try_insert_in_bucket(filter_idx, num_buckets, b2, fp)? {
inserted = true;
break;
}
}
if !inserted {
let last_filter_idx = meta.n_filters - 1;
let num_buckets = meta.sub_filter_num_buckets(last_filter_idx)?;
inserted = try_cuckoo_kickout(
&mut page_cache,
last_filter_idx,
num_buckets,
meta.bucket_size,
meta.max_iterations,
h,
fp,
)?;
}
if !inserted && meta.is_scaling() && meta.n_filters < u16::MAX {
let new_filter_idx = meta.n_filters;
if let Ok(new_buckets) = meta.sub_filter_num_buckets(new_filter_idx) {
let b1 = (h % (new_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, new_buckets);
let ok =
page_cache.try_insert_in_bucket(new_filter_idx, new_buckets, b1, fp)?
|| (b1 != b2
&& page_cache.try_insert_in_bucket(
new_filter_idx,
new_buckets,
b2,
fp,
)?);
if ok {
meta.n_filters += 1;
inserted = true;
}
}
}
if !inserted {
return Err(Error::invalid_data("ERR filter is full"));
}
meta.base.size += 1;
results.push(true);
}
let mut batch = self.db.batch();
batch.insert(&self.meta, meta_k.as_bytes(), meta.encode());
page_cache.commit_dirty_to_batch(&mut batch);
batch.commit()?;
Ok(results)
}
pub fn cf_insertnx<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
opt: &CuckooFilterInsertOptions,
) -> Result<Vec<bool>> {
let mut nx_opt = opt.clone();
nx_opt.nx = true;
self.cf_insert(key, items, &nx_opt)
}
pub fn cf_exists<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let res = self.cf_mexists(key, &[item])?;
Ok(res.first().copied().unwrap_or(false))
}
pub fn cf_mexists<K: AsRef<[u8]>, I: AsRef<[u8]>>(
&self,
key: K,
items: &[I],
) -> Result<Vec<bool>> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
let meta = match self.get_cf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => return Ok(vec![false; items.len()]),
};
let mut page_cache =
CuckooPageCache::new(&self.data, kc, k_str, meta.bucket_size, meta.page_size);
let mut results = Vec::with_capacity(items.len());
for item in items {
let h = CuckooFilterHelper::hash(item.as_ref());
let fp = CuckooFilterHelper::generate_fingerprint(h);
let mut exists = false;
for filter_idx in (0..meta.n_filters).rev() {
let num_buckets = meta.sub_filter_num_buckets(filter_idx)?;
let b1 = (h % (num_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, num_buckets);
if page_cache.contains_in_bucket(filter_idx, num_buckets, b1, fp)?
|| (b1 != b2
&& page_cache.contains_in_bucket(filter_idx, num_buckets, b2, fp)?)
{
exists = true;
break;
}
}
results.push(exists);
}
Ok(results)
}
pub fn cf_del<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<bool> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
let mut meta = match self.get_cf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => return Ok(false),
};
let h = CuckooFilterHelper::hash(item.as_ref());
let fp = CuckooFilterHelper::generate_fingerprint(h);
let mut page_cache =
CuckooPageCache::new(&self.data, kc, k_str, meta.bucket_size, meta.page_size);
let mut deleted = false;
for filter_idx in (0..meta.n_filters).rev() {
let num_buckets = meta.sub_filter_num_buckets(filter_idx)?;
let b1 = (h % (num_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, num_buckets);
if page_cache.delete_from_bucket(filter_idx, num_buckets, b1, fp)? {
deleted = true;
break;
}
if b1 != b2 && page_cache.delete_from_bucket(filter_idx, num_buckets, b2, fp)? {
deleted = true;
break;
}
}
if deleted {
meta.base.size = meta.base.size.saturating_sub(1);
meta.num_deleted_items = meta.num_deleted_items.saturating_add(1);
let mut batch = self.db.batch();
batch.insert(&self.meta, meta_k.as_bytes(), meta.encode());
page_cache.commit_dirty_to_batch(&mut batch);
batch.commit()?;
Ok(true)
} else {
Ok(false)
}
}
pub fn cf_count<K: AsRef<[u8]>, I: AsRef<[u8]>>(&self, key: K, item: I) -> Result<u64> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
let meta = match self.get_cf_meta_checked(&meta_k, now_ms)? {
Some(m) => m,
None => return Ok(0),
};
let h = CuckooFilterHelper::hash(item.as_ref());
let fp = CuckooFilterHelper::generate_fingerprint(h);
let mut page_cache =
CuckooPageCache::new(&self.data, kc, k_str, meta.bucket_size, meta.page_size);
let mut total_count = 0u64;
for filter_idx in 0..meta.n_filters {
let num_buckets = meta.sub_filter_num_buckets(filter_idx)?;
let b1 = (h % (num_buckets as u64)) as u32;
let b2 = CuckooFilterHelper::get_alt_bucket_index(b1, fp, num_buckets);
total_count += page_cache.count_in_bucket(filter_idx, num_buckets, b1, fp)? as u64;
if b1 != b2 {
total_count += page_cache.count_in_bucket(filter_idx, num_buckets, b2, fp)? as u64;
}
}
Ok(total_count)
}
pub fn cf_info<K: AsRef<[u8]>>(&self, key: K) -> Result<CuckooFilterInfo> {
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key.as_ref()).unwrap_or("");
let meta_k = kc.cf_meta(k_str);
let now_ms = current_now_ms();
let meta = self
.get_cf_meta_checked(&meta_k, now_ms)?
.ok_or_else(|| Error::invalid_data("ERR not found"))?;
let mut total_buckets = 0u64;
for i in 0..meta.n_filters {
if let Ok(nb) = meta.sub_filter_num_buckets(i) {
total_buckets += nb as u64;
}
}
Ok(CuckooFilterInfo {
size: meta.base.size,
num_buckets: total_buckets,
num_filters: meta.n_filters,
num_items_inserted: meta.base.size + meta.num_deleted_items,
num_items_deleted: meta.num_deleted_items,
bucket_size: meta.bucket_size,
expansion: meta.expansion,
max_iterations: meta.max_iterations,
})
}
}