use yo_common::num::parse_i64;
use yo_common::{Code, Error, Result};
use yo_kv::{Foreign, Keyspace};
use yo_sketch::cuckoo::{
Cuckoo, HEADER, Insert, MAX_BUCKET_SIZE, MAX_CAPACITY, MAX_EXPANSION, MAX_ITERATIONS,
};
use super::args::{self, Args};
use super::table::Spec;
use crate::reply::Out;
const DEFAULT_CAPACITY: u64 = 1024;
const DEFAULT_BUCKET_SIZE: u16 = 2;
const DEFAULT_ITERATIONS: u16 = 20;
const DEFAULT_EXPANSION: u16 = 1;
const NOT_FOUND: &str = "not found";
const ITEM_EXISTS: &str = "item exists";
const DEL_NOT_FOUND: &[u8] = b"Not found";
const NO_FILTER: &[u8] = b"Cuckoo filter was not found";
const BAD_CAPACITY: &[u8] = b"Bad capacity";
const RESERVE_RANGE: &[u8] = b"Capacity must be in the range [2 * BUCKETSIZE, 1073741824]";
const INSERT_RANGE: &[u8] = b"Capacity must be in the range [cf-bucket-size * 2, 1073741824]";
const BAD_BUCKET: &[u8] = b"Couldn't parse BUCKETSIZE";
const BUCKET_RANGE: &[u8] = b"BUCKETSIZE: value must be in the range [1, 255]";
const BAD_ITERATIONS: &[u8] = b"Couldn't parse MAXITERATIONS";
const ITERATION_RANGE: &[u8] = b"MAXITERATIONS: value must be in the range [1, 65535]";
const BAD_EXPANSION: &[u8] = b"Couldn't parse EXPANSION";
const EXPANSION_RANGE: &[u8] = b"EXPANSION: value must be in the range [0, 32768]";
const FULL: &[u8] = b"Filter is full";
const MAX_EXPANSIONS: &[u8] = b"Maximum expansions reached";
const UNKNOWN_ARG: &[u8] = b"Unknown argument received";
const BAD_POSITION: &[u8] = b"Invalid position";
const BAD_HEADER: &[u8] = b"Invalid header";
const NO_CREATE: &[u8] = b"Couldn't create filter!";
const BAD_CHUNK: &[u8] = b"Couldn't load chunk!";
const WRONG_KIND: &str = "Operation against a key holding the wrong kind of value";
#[derive(Debug)]
pub(super) struct CuckooBody {
c: Cuckoo,
}
impl Foreign for CuckooBody {
fn type_name(&self) -> &'static str {
"MBbloomCF"
}
fn encoding(&self) -> &'static str {
"raw"
}
fn memory_bytes(&self) -> usize {
self.c.memory_bytes()
}
fn is_empty(&self) -> bool {
false
}
}
pub(super) fn execute(db: &mut Keyspace, spec: &Spec, args: Args<'_>, out: &mut Out) -> Result<()> {
match spec.name {
"cf.reserve" => reserve(db, args, out),
"cf.add" => add(db, args, out, false),
"cf.addnx" => add(db, args, out, true),
"cf.insert" => insert(db, args, out, false),
"cf.insertnx" => insert(db, args, out, true),
"cf.exists" => exists(db, args, out),
"cf.mexists" => mexists(db, args, out),
"cf.count" => count(db, args, out),
"cf.del" => del(db, args, out),
"cf.scandump" => scandump(db, args, out),
"cf.loadchunk" => loadchunk(db, args, out),
"cf.info" => info(db, args, out),
"cf.debug" => debug(db, args, out),
"cf.compact" => compact(db, args, out),
other => unreachable!("{other} is not a cuckoo filter command"),
}
}
fn reserve(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
if args.len().is_multiple_of(2) {
return Err(args::wrong_arity("cf.reserve"));
}
let Some(capacity) = parse_i64(args.get(2)) else {
out.error(BAD_CAPACITY);
return Ok(());
};
let (bucket_size, max_iterations, expansion) = match geometry(args) {
Ok(g) => g,
Err(msg) => {
out.error(msg);
return Ok(());
}
};
if capacity < 2 * i64::from(bucket_size) || capacity > MAX_CAPACITY {
out.error(RESERVE_RANGE);
return Ok(());
}
let key = args.get(1);
if write(db, key)?.is_some() {
return Err(cf(ITEM_EXISTS));
}
put(
db,
key,
Cuckoo::new(capacity as u64, bucket_size, max_iterations, expansion),
);
out.ok();
Ok(())
}
fn geometry(args: Args<'_>) -> core::result::Result<(u16, u16, u16), &'static [u8]> {
let iterations = setting(
args,
b"maxiterations",
(1, MAX_ITERATIONS),
BAD_ITERATIONS,
ITERATION_RANGE,
)?;
let bucket_size = setting(
args,
b"bucketsize",
(1, MAX_BUCKET_SIZE),
BAD_BUCKET,
BUCKET_RANGE,
)?;
let expansion = setting(
args,
b"expansion",
(0, MAX_EXPANSION),
BAD_EXPANSION,
EXPANSION_RANGE,
)?;
Ok((
bucket_size.unwrap_or(i64::from(DEFAULT_BUCKET_SIZE)) as u16,
iterations.unwrap_or(i64::from(DEFAULT_ITERATIONS)) as u16,
expansion.unwrap_or(i64::from(DEFAULT_EXPANSION)) as u16,
))
}
fn setting(
args: Args<'_>,
name: &[u8],
bounds: (i64, i64),
bad: &'static [u8],
range: &'static [u8],
) -> core::result::Result<Option<i64>, &'static [u8]> {
let Some(value) = pair(args, name) else {
return Ok(None);
};
match parse_i64(value) {
None => Err(bad),
Some(n) if n < bounds.0 || n > bounds.1 => Err(range),
Some(n) => Ok(Some(n)),
}
}
fn pair<'a>(args: Args<'a>, name: &[u8]) -> Option<&'a [u8]> {
(3..args.len())
.step_by(2)
.find(|&i| args::is(args.get(i), name))
.map(|i| args.get(i + 1))
}
fn add(db: &mut Keyspace, args: Args<'_>, out: &mut Out, unique: bool) -> Result<()> {
let body = open(db, args.get(1))?;
let item = args.get(2);
let done = match unique {
true => body.c.insert_unique(item),
false => body.c.insert(item),
};
match done {
Insert::Yes => out.bool(true),
Insert::Exists => out.bool(false),
Insert::Full => out.error(FULL),
Insert::MaxFilters => out.error(MAX_EXPANSIONS),
}
Ok(())
}
fn insert(db: &mut Keyspace, args: Args<'_>, out: &mut Out, unique: bool) -> Result<()> {
let name = match unique {
true => "cf.insertnx",
false => "cf.insert",
};
let mut capacity = DEFAULT_CAPACITY;
let mut create = true;
let mut items = None;
let mut i = 2;
while i < args.len() {
match word(args.get(i)) {
Some(Opt::Items) => {
items = Some(i + 1);
break;
}
Some(Opt::NoCreate) => {
create = false;
i += 1;
}
Some(Opt::Capacity) => {
let Some(value) = args.opt(i + 1) else {
return Err(args::wrong_arity(name));
};
let Some(n) = parse_i64(value) else {
out.error(BAD_CAPACITY);
return Ok(());
};
if !(2 * i64::from(DEFAULT_BUCKET_SIZE)..=MAX_CAPACITY).contains(&n) {
out.error(INSERT_RANGE);
return Ok(());
}
capacity = n as u64;
i += 2;
}
None => {
out.error(UNKNOWN_ARG);
return Ok(());
}
}
}
let Some(first) = items.filter(|&at| at < args.len()) else {
return Err(args::wrong_arity(name));
};
let key = args.get(1);
let body = match write(db, key)? {
Some(body) => body,
None if create => {
put(
db,
key,
Cuckoo::new(
capacity,
DEFAULT_BUCKET_SIZE,
DEFAULT_ITERATIONS,
DEFAULT_EXPANSION,
),
);
write(db, key)?.expect("the filter was just created")
}
None => return Err(cf(NOT_FOUND)),
};
out.array(args.len() - first);
for i in first..args.len() {
let item = args.get(i);
let done = match unique {
true => body.c.insert_unique(item),
false => body.c.insert(item),
};
match (unique, done) {
(true, Insert::Yes) => out.int(1),
(true, Insert::Exists) => out.int(0),
(true, _) => out.int(-1),
(false, Insert::Yes) => out.bool(true),
(false, _) if out.proto().is_resp3() => out.bool(false),
(false, _) => out.int(-1),
}
}
Ok(())
}
enum Opt {
Capacity,
NoCreate,
Items,
}
fn word(arg: &[u8]) -> Option<Opt> {
match arg.first().copied().unwrap_or(0).to_ascii_uppercase() {
b'C' => Some(Opt::Capacity),
b'I' => Some(Opt::Items),
b'N' => Some(Opt::NoCreate),
_ => None,
}
}
fn exists(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let found = peek(db, args.get(1)).is_some_and(|b| b.c.contains(args.get(2)));
out.bool(found);
Ok(())
}
fn mexists(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let body = peek(db, args.get(1));
out.array(args.len() - 2);
for i in 2..args.len() {
out.bool(body.is_some_and(|b| b.c.contains(args.get(i))));
}
Ok(())
}
fn count(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let n = peek(db, args.get(1)).map_or(0, |b| b.c.count(args.get(2)));
out.uint(n);
Ok(())
}
fn del(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let Some(body) = peek_mut(db, args.get(1)) else {
out.error(DEL_NOT_FOUND);
return Ok(());
};
let gone = body.c.remove(args.get(2));
out.bool(gone);
Ok(())
}
fn scandump(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let Some(body) = read(db, args.get(1))? else {
return Err(cf(NOT_FOUND));
};
let Some(pos) = parse_i64(args.get(2)).filter(|&n| n >= 0) else {
out.error(BAD_POSITION);
return Ok(());
};
out.array(2);
let header;
let (next, data) = if body.c.is_empty() {
(0, &[][..])
} else if pos == 0 {
header = body.c.header();
(1, &header[..])
} else {
body.c.chunk(pos)
};
out.int(next);
if next == 0 {
out.nil();
} else {
out.bulk(data);
}
Ok(())
}
fn loadchunk(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let Some(pos) = parse_i64(args.get(2)).filter(|&n| n != 0) else {
out.error(BAD_POSITION);
return Ok(());
};
let key = args.get(1);
let data = args.get(3);
let Some(body) = write(db, key)? else {
if pos != 1 {
return Err(cf(NOT_FOUND));
}
if data.len() != HEADER {
out.error(BAD_HEADER);
return Ok(());
}
let Some(c) = Cuckoo::from_header(data) else {
out.error(NO_CREATE);
return Ok(());
};
put(db, key, c);
out.ok();
return Ok(());
};
if pos == 1 {
return Err(cf(ITEM_EXISTS));
}
if body.c.load(pos, data) {
out.ok();
} else {
out.error(BAD_CHUNK);
}
Ok(())
}
fn info(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let Some(body) = read(db, args.get(1))? else {
return Err(cf(NOT_FOUND));
};
let c = &body.c;
out.map(8);
out.simple(b"Size");
out.uint(c.reported_size());
out.simple(b"Number of buckets");
out.uint(c.buckets());
out.simple(b"Number of filters");
out.uint(c.filters() as u64);
out.simple(b"Number of items inserted");
out.uint(c.len());
out.simple(b"Number of items deleted");
out.uint(c.deleted());
out.simple(b"Bucket size");
out.uint(u64::from(c.bucket_size()));
out.simple(b"Expansion rate");
out.uint(u64::from(c.expansion()));
out.simple(b"Max iterations");
out.uint(u64::from(c.max_iterations()));
Ok(())
}
fn debug(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
let Some(body) = read(db, args.get(1))? else {
return Err(cf(NOT_FOUND));
};
let c = &body.c;
let line = yo_alloc::allow(|| {
format!(
"bktsize:{} buckets:{} items:{} deletes:{} filters:{} max_iterations:{} expansion:{}",
c.bucket_size(),
c.buckets(),
c.len(),
c.deleted(),
c.filters(),
c.max_iterations(),
c.expansion(),
)
});
out.bulk(line.as_bytes());
Ok(())
}
fn compact(db: &mut Keyspace, args: Args<'_>, out: &mut Out) -> Result<()> {
if args.len() != 2 {
return Err(args::wrong_arity("cf.compact"));
}
let Some(body) = peek_mut(db, args.get(1)) else {
out.error(NO_FILTER);
return Ok(());
};
body.c.compact(true);
out.ok();
Ok(())
}
fn cf(msg: &'static str) -> Error {
Error::new(Code::Invalid, msg)
}
fn put(db: &mut Keyspace, key: &[u8], c: Cuckoo) {
db.put_foreign(key, Box::new(CuckooBody { c }));
}
fn open<'d>(db: &'d mut Keyspace, key: &[u8]) -> Result<&'d mut CuckooBody> {
if write(db, key)?.is_none() {
put(
db,
key,
Cuckoo::new(
DEFAULT_CAPACITY,
DEFAULT_BUCKET_SIZE,
DEFAULT_ITERATIONS,
DEFAULT_EXPANSION,
),
);
}
Ok(write(db, key)?.expect("the filter is there either way"))
}
fn write<'d>(db: &'d mut Keyspace, key: &[u8]) -> Result<Option<&'d mut CuckooBody>> {
match db.foreign_mut(key)? {
Some(body) => match body.downcast_mut::<CuckooBody>() {
Some(body) => Ok(Some(body)),
None => Err(Error::new(Code::WrongType, WRONG_KIND)),
},
None => Ok(None),
}
}
fn read<'d>(db: &'d mut Keyspace, key: &[u8]) -> Result<Option<&'d CuckooBody>> {
match db.foreign(key)? {
Some(body) => match body.downcast_ref::<CuckooBody>() {
Some(body) => Ok(Some(body)),
None => Err(Error::new(Code::WrongType, WRONG_KIND)),
},
None => Ok(None),
}
}
fn peek<'d>(db: &'d mut Keyspace, key: &[u8]) -> Option<&'d CuckooBody> {
db.foreign(key)
.ok()
.flatten()
.and_then(<dyn Foreign>::downcast_ref::<CuckooBody>)
}
fn peek_mut<'d>(db: &'d mut Keyspace, key: &[u8]) -> Option<&'d mut CuckooBody> {
db.foreign_mut(key)
.ok()
.flatten()
.and_then(<dyn Foreign>::downcast_mut::<CuckooBody>)
}