use yo_common::{Code, Error, Result};
use yo_kv::Member;
use crate::db::Handle;
#[derive(Clone)]
pub struct Sets {
pub(crate) db: Handle,
}
impl core::fmt::Debug for Sets {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Sets").finish_non_exhaustive()
}
}
impl Sets {
pub fn add(&self, key: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
self.add_many(key, &[member]).map(|n| n == 1)
}
pub fn add_many<M: AsRef<[u8]>>(&self, key: impl AsRef<[u8]>, members: &[M]) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.sadd(key.as_ref(), members.iter().map(AsRef::as_ref))
})
}
pub fn remove(&self, key: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
self.remove_many(key, &[member]).map(|n| n == 1)
}
pub fn remove_many<M: AsRef<[u8]>>(
&self,
key: impl AsRef<[u8]>,
members: &[M],
) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.srem(key.as_ref(), members.iter().map(AsRef::as_ref))
})
}
pub fn contains(&self, key: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
self.db
.run(|inner| inner.strings.sismember(key.as_ref(), member.as_ref()))
}
pub fn contains_many<M: AsRef<[u8]>>(
&self,
key: impl AsRef<[u8]>,
members: &[M],
) -> Result<Vec<bool>> {
self.db.run(|inner| {
inner
.strings
.smismember(key.as_ref(), members.iter().map(AsRef::as_ref))
})
}
pub fn len_of(&self, key: impl AsRef<[u8]>) -> Result<usize> {
self.db.run(|inner| inner.strings.scard(key.as_ref()))
}
pub fn members(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<Vec<u8>>>> {
self.db.run(|inner| {
Ok(inner
.strings
.smembers(key.as_ref())?
.map(|it| it.map(|m| m.to_vec()).collect()))
})
}
pub fn for_each(&self, key: impl AsRef<[u8]>, mut f: impl FnMut(Member<'_>)) -> Result<bool> {
self.db.run(|inner| {
inner.strings.with_set(key.as_ref(), |set| match set {
Some(set) => {
for m in set.iter() {
f(m);
}
true
}
None => false,
})
})
}
pub fn pop(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
self.db.run(|inner| inner.strings.spop(key.as_ref()))
}
pub fn pop_n(&self, key: impl AsRef<[u8]>, count: usize) -> Result<Vec<Vec<u8>>> {
self.db
.run(|inner| inner.strings.spop_n(key.as_ref(), count))
}
pub fn pick(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
self.db.run(|inner| {
inner
.strings
.srandmember(key.as_ref(), |m| m.map(|m| m.to_vec()))
})
}
pub fn pick_n(&self, key: impl AsRef<[u8]>, count: i64) -> Result<Vec<Vec<u8>>> {
self.db.run(|inner| {
let mut out = Vec::new();
inner
.strings
.srandmember_n(key.as_ref(), count, |m| out.push(m.to_vec()))?;
Ok(out)
})
}
pub fn move_member(
&self,
from: impl AsRef<[u8]>,
to: impl AsRef<[u8]>,
member: impl AsRef<[u8]>,
) -> Result<bool> {
self.db.run(|inner| {
inner
.strings
.smove(from.as_ref(), to.as_ref(), member.as_ref())
})
}
pub fn intersect<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<Vec<Vec<u8>>> {
self.collect(keys, Op::Intersect)
}
pub fn intersect_len<K: AsRef<[u8]>>(&self, keys: &[K], limit: usize) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.sintercard(keys.iter().map(AsRef::as_ref), limit)
})
}
pub fn union<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<Vec<Vec<u8>>> {
self.collect(keys, Op::Union)
}
pub fn difference<K: AsRef<[u8]>>(&self, keys: &[K]) -> Result<Vec<Vec<u8>>> {
self.collect(keys, Op::Difference)
}
pub fn intersect_into<K: AsRef<[u8]>>(
&self,
destination: impl AsRef<[u8]>,
keys: &[K],
) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.sinterstore(destination.as_ref(), keys.iter().map(AsRef::as_ref))
})
}
pub fn union_into<K: AsRef<[u8]>>(
&self,
destination: impl AsRef<[u8]>,
keys: &[K],
) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.sunionstore(destination.as_ref(), keys.iter().map(AsRef::as_ref))
})
}
pub fn difference_into<K: AsRef<[u8]>>(
&self,
destination: impl AsRef<[u8]>,
keys: &[K],
) -> Result<usize> {
self.db.run(|inner| {
inner
.strings
.sdiffstore(destination.as_ref(), keys.iter().map(AsRef::as_ref))
})
}
fn collect<K: AsRef<[u8]>>(&self, keys: &[K], op: Op) -> Result<Vec<Vec<u8>>> {
self.db.run(|inner| {
let mut out = Vec::new();
let push = |m: &[u8]| out.push(m.to_vec());
let keys = keys.iter().map(AsRef::as_ref);
match op {
Op::Intersect => inner.strings.sinter(keys, 0, push)?,
Op::Union => inner.strings.sunion(keys, push)?,
Op::Difference => inner.strings.sdiff(keys, push)?,
};
Ok(out)
})
}
}
#[derive(Clone, Copy)]
enum Op {
Intersect,
Union,
Difference,
}
#[derive(Clone)]
pub struct Set {
pub(crate) sets: Sets,
pub(crate) key: Vec<u8>,
}
impl core::fmt::Debug for Set {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("Set")
.field("key", &String::from_utf8_lossy(&self.key))
.field("len", &self.len().ok())
.finish()
}
}
impl Set {
#[must_use]
pub fn key(&self) -> &[u8] {
&self.key
}
pub fn add(&self, member: impl AsRef<[u8]>) -> Result<bool> {
self.sets.add(&self.key, member)
}
pub fn add_many<M: AsRef<[u8]>>(&self, members: &[M]) -> Result<usize> {
self.sets.add_many(&self.key, members)
}
pub fn remove(&self, member: impl AsRef<[u8]>) -> Result<bool> {
self.sets.remove(&self.key, member)
}
pub fn remove_many<M: AsRef<[u8]>>(&self, members: &[M]) -> Result<usize> {
self.sets.remove_many(&self.key, members)
}
pub fn contains(&self, member: impl AsRef<[u8]>) -> Result<bool> {
self.sets.contains(&self.key, member)
}
pub fn contains_many<M: AsRef<[u8]>>(&self, members: &[M]) -> Result<Vec<bool>> {
self.sets.contains_many(&self.key, members)
}
pub fn len(&self) -> Result<usize> {
self.sets.len_of(&self.key)
}
pub fn is_empty(&self) -> Result<bool> {
self.len().map(|n| n == 0)
}
pub fn members(&self) -> Result<Vec<Vec<u8>>> {
Ok(self.sets.members(&self.key)?.unwrap_or_default())
}
pub fn for_each(&self, f: impl FnMut(Member<'_>)) -> Result<()> {
self.sets.for_each(&self.key, f).map(|_| ())
}
pub fn pop(&self) -> Result<Option<Vec<u8>>> {
self.sets.pop(&self.key)
}
pub fn pop_n(&self, count: usize) -> Result<Vec<Vec<u8>>> {
self.sets.pop_n(&self.key, count)
}
pub fn pick(&self) -> Result<Option<Vec<u8>>> {
self.sets.pick(&self.key)
}
pub fn pick_n(&self, count: i64) -> Result<Vec<Vec<u8>>> {
self.sets.pick_n(&self.key, count)
}
pub fn move_to(&self, to: &Set, member: impl AsRef<[u8]>) -> Result<bool> {
self.same_db(to)?;
self.sets.move_member(&self.key, &to.key, member)
}
pub fn intersect(&self, others: &[&Set]) -> Result<Vec<Vec<u8>>> {
self.sets.intersect(&self.keys_with(others)?)
}
pub fn intersect_len(&self, others: &[&Set], limit: usize) -> Result<usize> {
self.sets.intersect_len(&self.keys_with(others)?, limit)
}
pub fn union(&self, others: &[&Set]) -> Result<Vec<Vec<u8>>> {
self.sets.union(&self.keys_with(others)?)
}
pub fn difference(&self, others: &[&Set]) -> Result<Vec<Vec<u8>>> {
self.sets.difference(&self.keys_with(others)?)
}
pub fn intersect_into(&self, destination: &Set, others: &[&Set]) -> Result<usize> {
self.same_db(destination)?;
self.sets
.intersect_into(&destination.key, &self.keys_with(others)?)
}
pub fn union_into(&self, destination: &Set, others: &[&Set]) -> Result<usize> {
self.same_db(destination)?;
self.sets
.union_into(&destination.key, &self.keys_with(others)?)
}
pub fn difference_into(&self, destination: &Set, others: &[&Set]) -> Result<usize> {
self.same_db(destination)?;
self.sets
.difference_into(&destination.key, &self.keys_with(others)?)
}
pub fn clear(&self) -> Result<bool> {
self.sets.db.run(|inner| Ok(inner.strings.del(&self.key)))
}
fn keys_with<'a>(&'a self, others: &[&'a Set]) -> Result<Vec<&'a [u8]>> {
let mut keys = Vec::with_capacity(others.len() + 1);
keys.push(&self.key[..]);
for other in others {
self.same_db(other)?;
keys.push(&other.key[..]);
}
Ok(keys)
}
fn same_db(&self, other: &Set) -> Result<()> {
if self.sets.db.is(&other.sets.db) {
return Ok(());
}
Err(Error::new(
Code::Invalid,
"those two sets are in different databases, and a set operation reads both of them out of one. Open both handles on the same Db",
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{MEMORY, open};
#[test]
fn the_set_commands_are_the_ones_a_redis_client_would_send() {
let db = open(MEMORY).unwrap();
let sets = db.sets();
assert!(sets.add("online", "alice").unwrap());
assert!(!sets.add("online", "alice").unwrap());
assert_eq!(sets.add_many("online", &["bob", "carol"]).unwrap(), 2);
assert!(sets.contains("online", "bob").unwrap());
assert!(!sets.contains("online", "dave").unwrap());
assert_eq!(
sets.contains_many("online", &["alice", "dave"]).unwrap(),
vec![true, false]
);
assert_eq!(sets.len_of("online").unwrap(), 3);
assert!(sets.remove("online", "alice").unwrap());
assert_eq!(sets.remove_many("online", &["bob", "dave"]).unwrap(), 1);
assert_eq!(sets.len_of("online").unwrap(), 1);
}
#[test]
fn a_set_that_is_not_there_reads_as_an_empty_one() {
let db = open(MEMORY).unwrap();
let sets = db.sets();
assert_eq!(sets.len_of("nope").unwrap(), 0);
assert!(!sets.contains("nope", "x").unwrap());
assert_eq!(sets.pop("nope").unwrap(), None);
assert_eq!(sets.pick("nope").unwrap(), None);
assert!(sets.pop_n("nope", 5).unwrap().is_empty());
assert!(sets.pick_n("nope", 5).unwrap().is_empty());
assert_eq!(sets.members("nope").unwrap(), None);
assert!(!sets.for_each("nope", |_| {}).unwrap());
}
#[test]
fn walking_a_set_of_integers_never_formats_a_digit() {
let db = open(MEMORY).unwrap();
let ids = db.set("ids");
ids.add_many(&["1", "2", "3"]).unwrap();
let mut total = 0i64;
let mut ints = 0;
ids.for_each(|m| {
if let Member::Int(n) = m {
total += n;
ints += 1;
}
})
.unwrap();
assert_eq!(total, 6);
assert_eq!(ints, 3, "stored as integers, so handed over as integers");
let mut owned = ids.members().unwrap();
owned.sort();
assert_eq!(owned, vec![b"1".to_vec(), b"2".to_vec(), b"3".to_vec()]);
}
#[test]
fn a_handle_holds_the_key_so_the_caller_does_not() {
let db = open(MEMORY).unwrap();
let online = db.set("online");
online.add("alice").unwrap();
online.add_many(&["bob", "carol"]).unwrap();
assert_eq!(online.len().unwrap(), 3);
assert!(!online.is_empty().unwrap());
assert!(online.contains("bob").unwrap());
assert_eq!(online.key(), b"online");
assert!(format!("{online:?}").contains("online"));
assert_eq!(db.sets().len_of("online").unwrap(), 3);
}
#[test]
fn a_set_that_loses_its_last_member_loses_its_key() {
let db = open(MEMORY).unwrap();
let only = db.set("only");
only.add("x").unwrap();
assert!(only.remove("x").unwrap());
assert!(only.is_empty().unwrap());
assert_eq!(db.sets().members("only").unwrap(), None);
}
#[test]
fn drawing_takes_members_out_and_picking_leaves_them() {
let db = open(MEMORY).unwrap();
let bag = db.set("bag");
bag.add_many(&["a", "b", "c", "d"]).unwrap();
assert!(bag.pick().unwrap().is_some());
assert_eq!(bag.len().unwrap(), 4, "picking leaves the set alone");
assert_eq!(bag.pick_n(3).unwrap().len(), 3);
assert_eq!(bag.len().unwrap(), 4);
assert_eq!(bag.pick_n(-9).unwrap().len(), 9);
assert_eq!(bag.len().unwrap(), 4);
assert!(bag.pop().unwrap().is_some());
assert_eq!(bag.len().unwrap(), 3, "popping takes one out");
assert_eq!(bag.pop_n(9).unwrap().len(), 3, "and never more than it has");
assert!(bag.is_empty().unwrap());
}
#[test]
fn the_three_set_operations_answer_what_they_are_named_after() {
let db = open(MEMORY).unwrap();
let a = db.set("a");
let b = db.set("b");
a.add_many(&["x", "y"]).unwrap();
b.add_many(&["y", "z"]).unwrap();
let sorted = |mut v: Vec<Vec<u8>>| {
v.sort();
v
};
assert_eq!(sorted(a.intersect(&[&b]).unwrap()), vec![b"y".to_vec()]);
assert_eq!(
sorted(a.union(&[&b]).unwrap()),
vec![b"x".to_vec(), b"y".to_vec(), b"z".to_vec()]
);
assert_eq!(sorted(a.difference(&[&b]).unwrap()), vec![b"x".to_vec()]);
assert_eq!(a.intersect_len(&[&b], 0).unwrap(), 1);
let out = db.set("out");
assert_eq!(a.union_into(&out, &[&b]).unwrap(), 3);
assert_eq!(out.len().unwrap(), 3);
assert_eq!(a.intersect_into(&out, &[&b]).unwrap(), 1);
assert_eq!(out.len().unwrap(), 1);
assert_eq!(a.difference_into(&out, &[&b]).unwrap(), 1);
assert_eq!(out.members().unwrap(), vec![b"x".to_vec()]);
}
#[test]
fn storing_an_empty_result_removes_the_destination() {
let db = open(MEMORY).unwrap();
let a = db.set("a");
let b = db.set("b");
a.add("x").unwrap();
b.add("y").unwrap();
let out = db.set("out");
out.add("stale").unwrap();
assert_eq!(a.intersect_into(&out, &[&b]).unwrap(), 0);
assert_eq!(db.sets().members("out").unwrap(), None);
}
#[test]
fn a_member_moves_between_two_sets() {
let db = open(MEMORY).unwrap();
let from = db.set("from");
let to = db.set("to");
from.add_many(&["x", "y"]).unwrap();
assert!(from.move_to(&to, "x").unwrap());
assert!(!from.contains("x").unwrap());
assert!(to.contains("x").unwrap());
assert!(!from.move_to(&to, "nope").unwrap());
assert_eq!(to.len().unwrap(), 1);
}
#[test]
fn two_databases_cannot_be_intersected_with_each_other() {
let one = open(MEMORY).unwrap();
let two = open(MEMORY).unwrap();
let a = one.set("a");
a.add_many(&["x", "y"]).unwrap();
one.set("b").add("z").unwrap();
let elsewhere = two.set("b");
elsewhere.add("x").unwrap();
let e = a.intersect(&[&elsewhere]).expect_err("different databases");
assert_eq!(e.code(), Code::Invalid);
assert!(e.message().contains("different databases"), "{e}");
assert!(a.union(&[&elsewhere]).is_err());
assert!(a.difference(&[&elsewhere]).is_err());
assert!(a.intersect_len(&[&elsewhere], 0).is_err());
assert!(a.move_to(&elsewhere, "x").is_err());
assert!(a.intersect_into(&elsewhere, &[]).is_err());
assert!(a.union_into(&elsewhere, &[]).is_err());
assert!(a.difference_into(&elsewhere, &[]).is_err());
}
#[test]
fn a_set_and_the_keyspace_are_the_same_store() {
let db = open(MEMORY).unwrap();
let tags = db.set("tags");
tags.add("rust").unwrap();
let e = db.strings().get("tags").expect_err("that is a set");
assert_eq!(e.code(), Code::WrongType);
db.strings().set("word", "nope").unwrap();
assert_eq!(db.set("word").add("x").unwrap_err().code(), Code::WrongType);
assert!(tags.clear().unwrap());
assert!(tags.is_empty().unwrap());
assert!(!tags.clear().unwrap());
}
#[test]
fn a_member_is_bytes_and_not_only_text() {
let db = open(MEMORY).unwrap();
let raw = db.set("raw");
raw.add(vec![0u8, 0xff]).unwrap();
assert!(raw.contains(b"\x00\xff").unwrap());
assert_eq!(raw.members().unwrap(), vec![vec![0u8, 0xff]]);
}
}