use crate::{
inputs::ObjectInput,
objects::{
hash::hash_object::{pick_k_random_indexes, pick_random_index},
set::set_object::SetObject,
types::object_output::ObjectOutput,
},
};
#[inline]
fn arg<'a>(input: &ObjectInput, i: usize) -> &'a [u8] {
input.parse_state.get_arg_slice_by_ref(i).as_slice()
}
pub(crate) const NO_COUNT: i32 = i32::MIN;
impl SetObject {
pub(crate) fn set_add(&mut self, input: &ObjectInput, output: &mut ObjectOutput) {
let mut added = 0_i64;
for i in 0..input.parse_state.count {
let member = arg(input, i);
if self.set.insert(member.to_vec()) {
added += 1;
self.update_size(member, true);
}
}
output.result1 = added;
}
pub(crate) fn set_members(&mut self, output: &mut ObjectOutput, resp_protocol_version: u8) {
write_set_length(output, self.set.len(), resp_protocol_version);
let mut written = 0_i64;
for item in &self.set {
output.write_bulk_string(item);
written += 1;
}
output.result1 = written;
}
pub(crate) fn set_is_member(
&mut self,
input: &ObjectInput,
output: &mut ObjectOutput,
_resp_protocol_version: u8,
) {
let member = arg(input, 0);
let is_member = self.set.contains(member);
output.write_int64(i64::from(is_member));
output.result1 = 1;
}
pub(crate) fn set_multi_is_member(
&mut self,
input: &ObjectInput,
output: &mut ObjectOutput,
_resp_protocol_version: u8,
) {
output.write_array_length(input.parse_state.count);
for i in 0..input.parse_state.count {
let member = arg(input, i);
output.write_int64(i64::from(self.set.contains(member)));
}
output.result1 = input.parse_state.count as i64;
}
pub(crate) fn set_remove(&mut self, input: &ObjectInput, output: &mut ObjectOutput) {
let mut removed = 0_i64;
for i in 0..input.parse_state.count {
let member = arg(input, i);
if self.set.remove(member) {
removed += 1;
self.update_size(member, false);
}
}
output.result1 = removed;
}
pub(crate) fn set_length(&mut self, output: &mut ObjectOutput) {
output.result1 = self.set.len() as i64;
}
pub(crate) fn set_pop(
&mut self,
input: &ObjectInput,
output: &mut ObjectOutput,
resp_protocol_version: u8,
) {
let count = input.arg1;
let mut count_done = 0_i64;
if count >= 1 {
let count_parameter = (count as usize).min(self.set.len());
let mut rng = fastrand::Rng::new();
write_set_length(output, count_parameter, resp_protocol_version);
for _ in 0..count_parameter {
if self.set.is_empty() {
break;
}
let index = rng.usize(..self.set.len());
let Some(item) = self.set.iter().nth(index).cloned() else {
break;
};
self.set.remove(&item);
self.update_size(&item, false);
output.write_bulk_string(&item);
count_done += 1;
}
count_done += i64::from(count) - count_done;
} else if count == NO_COUNT {
if !self.set.is_empty() {
let index = fastrand::usize(..self.set.len());
let item = self.set.iter().nth(index).cloned().unwrap();
self.set.remove(&item);
self.update_size(&item, false);
output.write_bulk_string(&item);
} else {
output.write_null(resp_protocol_version);
}
count_done += 1;
}
output.result1 = count_done;
}
pub(crate) fn set_random_member(
&mut self,
input: &ObjectInput,
output: &mut ObjectOutput,
resp_protocol_version: u8,
) {
let count = input.arg1;
let seed = input.arg2;
let mut count_done = 0_i64;
if count > 0 {
let count_parameter = (count as usize).min(self.set.len());
let indexes = pick_k_random_indexes(self.set.len(), count_parameter, seed, true);
write_set_length(output, count_parameter, resp_protocol_version);
for index in indexes {
let Some(element) = self.set.iter().nth(index).cloned() else {
continue;
};
output.write_bulk_string(&element);
count_done += 1;
}
count_done += i64::from(count) - count_parameter as i64;
} else if count == NO_COUNT {
if !self.set.is_empty() {
let index = pick_random_index(self.set.len(), seed);
if let Some(item) = self.set.iter().nth(index).cloned() {
output.write_bulk_string(&item);
}
} else {
output.write_null(resp_protocol_version);
}
count_done += 1;
} else {
let count_parameter = count.unsigned_abs() as usize;
let indexes = pick_k_random_indexes(self.set.len(), count_parameter, seed, false);
if !self.set.is_empty() {
output.write_array_length(count_parameter);
for index in indexes {
let Some(element) = self.set.iter().nth(index).cloned() else {
continue;
};
output.write_bulk_string(&element);
count_done += 1;
}
} else {
output.write_null(resp_protocol_version);
}
}
output.result1 = count_done;
}
}
fn write_set_length(output: &mut ObjectOutput, len: usize, resp_protocol_version: u8) {
if resp_protocol_version >= 3 {
output.payload.push(b'~');
let mut buf = itoa::Buffer::new();
output.payload.extend_from_slice(buf.format(len).as_bytes());
output.payload.extend_from_slice(b"\r\n");
} else {
output.write_array_length(len);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
arg_slice::ArgSlice,
input_header::RespInputHeader,
objects::set::set_object::SetOperation,
session_parse_state::SessionParseState,
types::{GarnetObjectType, RespInputFlags},
};
fn make_input(
op: SetOperation,
args: &[&[u8]],
arg1: i32,
arg2: i32,
) -> (ObjectInput, Vec<Vec<u8>>) {
let backing: Vec<Vec<u8>> = args.iter().map(|a| a.to_vec()).collect();
let slices: Vec<ArgSlice> = backing
.iter()
.map(|b| ArgSlice::new(b.as_ptr(), b.len()))
.collect();
let mut parse_state = SessionParseState::new();
parse_state.initialize_with_args(&slices);
let mut header = RespInputHeader::new_with_type(GarnetObjectType::Set, RespInputFlags::empty());
header.set_sub_id(op as u8);
(
ObjectInput::new_with_state(header, &mut parse_state, arg1, arg2),
backing,
)
}
fn obj_with(members: &[&str]) -> SetObject {
let mut obj = SetObject::new();
for m in members {
obj.set.insert(m.as_bytes().to_vec());
}
obj
}
#[test]
fn membership_ops() {
let mut obj = SetObject::new();
let (input, _b) = make_input(SetOperation::Sadd, &[b"a", b"b", b"a"], 0, 0);
let mut out = ObjectOutput::new();
obj.set_add(&input, &mut out);
assert_eq!(out.result1, 2);
let (input, _b) = make_input(SetOperation::Sismember, &[b"a"], 0, 0);
let mut out = ObjectOutput::new();
obj.set_is_member(&input, &mut out, 2);
assert_eq!(out.payload, b":1\r\n");
assert_eq!(out.result1, 1);
let (input, _b) = make_input(SetOperation::Sismember, &[b"nx"], 0, 0);
let mut out = ObjectOutput::new();
obj.set_is_member(&input, &mut out, 2);
assert_eq!(out.payload, b":0\r\n");
let (input, _b) = make_input(SetOperation::Smismember, &[b"a", b"nx", b"b"], 0, 0);
let mut out = ObjectOutput::new();
obj.set_multi_is_member(&input, &mut out, 2);
assert_eq!(out.payload, b"*3\r\n:1\r\n:0\r\n:1\r\n");
assert_eq!(out.result1, 3);
let (_input, _b) = make_input(SetOperation::Smembers, &[], 0, 0);
let mut out = ObjectOutput::new();
obj.set_members(&mut out, 2);
assert!(out.payload.starts_with(b"*2\r\n"));
assert_eq!(out.result1, 2);
let (_input, _b) = make_input(SetOperation::Scard, &[], 0, 0);
let mut out = ObjectOutput::new();
obj.set_length(&mut out);
assert_eq!(out.result1, 2);
let (input, _b) = make_input(SetOperation::Srem, &[b"a", b"nx"], 0, 0);
let mut out = ObjectOutput::new();
obj.set_remove(&input, &mut out);
assert_eq!(out.result1, 1);
assert_eq!(obj.set.len(), 1);
}
#[test]
fn pop_forms() {
let mut obj = obj_with(&["a", "b", "c"]);
let (input, _b) = make_input(SetOperation::Spop, &[], 2, 0);
let mut out = ObjectOutput::new();
obj.set_pop(&input, &mut out, 2);
assert!(
out.payload.starts_with(b"*2\r\n"),
"{}",
String::from_utf8_lossy(&out.payload)
);
assert_eq!(out.result1, 2);
assert_eq!(obj.set.len(), 1);
let (input, _b) = make_input(SetOperation::Spop, &[], 5, 0);
let mut out = ObjectOutput::new();
obj.set_pop(&input, &mut out, 2);
assert!(out.payload.starts_with(b"*1\r\n"));
assert_eq!(out.result1, 5); assert!(obj.set.is_empty());
let (input, _b) = make_input(SetOperation::Spop, &[], NO_COUNT, 0);
let mut out = ObjectOutput::new();
obj.set_pop(&input, &mut out, 2);
assert_eq!(out.payload, b"$-1\r\n");
assert_eq!(out.result1, 1);
let mut obj = obj_with(&["x"]);
let (input, _b) = make_input(SetOperation::Spop, &[], NO_COUNT, 0);
let mut out = ObjectOutput::new();
obj.set_pop(&input, &mut out, 2);
assert_eq!(out.payload, b"$1\r\nx\r\n");
assert!(obj.set.is_empty());
}
#[test]
fn random_member_forms() {
let mut obj = obj_with(&["a", "b", "c", "d"]);
let (input, _b) = make_input(SetOperation::Srandmember, &[], 2, 42);
let mut out = ObjectOutput::new();
obj.set_random_member(&input, &mut out, 2);
assert!(out.payload.starts_with(b"*2\r\n"));
assert_eq!(out.result1, 2); assert_eq!(obj.set.len(), 4);
let (input, _b) = make_input(SetOperation::Srandmember, &[], 6, 42);
let mut out = ObjectOutput::new();
obj.set_random_member(&input, &mut out, 2);
assert!(out.payload.starts_with(b"*4\r\n"));
assert_eq!(out.result1, 6);
let (input, _b) = make_input(SetOperation::Srandmember, &[], -3, 42);
let mut out = ObjectOutput::new();
obj.set_random_member(&input, &mut out, 2);
assert!(out.payload.starts_with(b"*3\r\n"));
assert_eq!(out.result1, 3);
let (input, _b) = make_input(SetOperation::Srandmember, &[], NO_COUNT, 1);
let mut out = ObjectOutput::new();
obj.set_random_member(&input, &mut out, 2);
assert!(out.payload.starts_with(b"$1\r\n"));
assert_eq!(out.result1, 1);
let mut empty = SetObject::new();
let (input, _b) = make_input(SetOperation::Srandmember, &[], NO_COUNT, 1);
let mut out = ObjectOutput::new();
empty.set_random_member(&input, &mut out, 2);
assert_eq!(out.payload, b"$-1\r\n");
let (input, _b) = make_input(SetOperation::Srandmember, &[], -3, 42);
let mut out = ObjectOutput::new();
empty.set_random_member(&input, &mut out, 2);
assert_eq!(out.payload, b"$-1\r\n");
}
#[test]
fn operate_dispatch_smoke() {
let mut obj = SetObject::new();
let (input, _b) = make_input(SetOperation::Sadd, &[b"m"], 0, 0);
let mut out = ObjectOutput::new();
assert!(obj.operate(&input, &mut out, 2));
assert_eq!(out.result1, 1);
let (input, _b) = make_input(SetOperation::Srem, &[b"m"], 0, 0);
let mut out = ObjectOutput::new();
obj.operate(&input, &mut out, 2);
assert!(out.has_remove_key());
let (input, _b) = make_input(SetOperation::Smove, &[b"s", b"d", b"m"], 0, 0);
let mut out = ObjectOutput::new();
obj.operate(&input, &mut out, 2);
assert_eq!(out.payload, b"-ERR unsupported operation\r\n");
let (input, _b) = make_input(SetOperation::Scard, &[], 0, 0);
let mut input = input;
input.header.data[0] = GarnetObjectType::List as u8;
let mut out = ObjectOutput::new();
obj.operate(&input, &mut out, 2);
assert!(out.has_wrong_type());
}
#[test]
fn sscan_flow() {
let mut obj = obj_with(&["one", "two", "three"]);
let (input, _b) = make_input(SetOperation::Sscan, &[b"0", b"MATCH", b"t*"], 0, 0);
let mut out = ObjectOutput::new();
assert!(obj.operate(&input, &mut out, 2));
let payload = String::from_utf8_lossy(&out.payload);
assert!(payload.starts_with("*2\r\n$1\r\n0\r\n"), "{payload}");
assert!(payload.contains("$3\r\ntwo\r\n"), "{payload}");
assert!(payload.contains("$5\r\nthree\r\n"), "{payload}");
let (input, _b) = make_input(SetOperation::Sscan, &[b"-1"], 0, 0);
let mut out = ObjectOutput::new();
obj.scan_operate(&input, &mut out);
assert_eq!(out.payload, b"-ERR invalid cursor\r\n");
}
}