use crate::elem::Elements;
use crate::intset::Intset;
use crate::listpack::{self, Listpack};
use crate::scan::Cursor;
use yo_common::num::{DIGITS_MAX, i64_digits, i64_len, parse_i64};
pub type Member<'a> = listpack::Entry<'a>;
#[derive(Debug, Clone, Copy)]
pub struct Needle<'a> {
bytes: &'a [u8],
int: Option<i64>,
hash: u64,
}
impl<'a> Needle<'a> {
#[must_use]
pub fn new(bytes: &'a [u8]) -> Needle<'a> {
Needle {
bytes,
int: parse_i64(bytes),
hash: Elements::<()>::hash_of(bytes),
}
}
#[must_use]
pub fn of(member: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> Needle<'a> {
match member {
Member::Str(s) => Needle::new(s),
Member::Int(n) => {
let bytes = i64_digits(digits, n);
Needle {
bytes,
int: Some(n),
hash: Elements::<()>::hash_of(bytes),
}
}
}
}
#[must_use]
pub const fn bytes(&self) -> &'a [u8] {
self.bytes
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Limits {
pub max_intset_entries: usize,
pub max_listpack_entries: usize,
pub max_listpack_value: usize,
}
impl Limits {
pub const DEFAULT: Limits = Limits {
max_intset_entries: 512,
max_listpack_entries: 128,
max_listpack_value: 64,
};
}
impl Default for Limits {
fn default() -> Limits {
Limits::DEFAULT
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Encoding {
Intset,
Listpack,
Hashtable,
}
impl Encoding {
#[inline]
pub const fn name(self) -> &'static str {
match self {
Encoding::Intset => "intset",
Encoding::Listpack => "listpack",
Encoding::Hashtable => "hashtable",
}
}
}
#[derive(Debug, Clone)]
enum Body {
Ints(Intset),
Packed(Listpack),
Table(Elements<()>),
}
#[derive(Debug, Clone)]
pub struct Set {
body: Body,
}
impl Set {
#[must_use]
pub fn new() -> Set {
Set {
body: Body::Ints(Intset::new()),
}
}
#[must_use]
pub fn with_hint(first: &[u8], hint: usize, limits: &Limits) -> Set {
if parse_i64(first).is_some() && hint <= limits.max_intset_entries {
Set {
body: Body::Ints(Intset::with_capacity(hint)),
}
} else if hint <= limits.max_listpack_entries {
Set {
body: Body::Packed(Listpack::new()),
}
} else {
Set {
body: Body::Table(Elements::with_capacity(hint)),
}
}
}
#[inline]
#[must_use]
pub const fn encoding(&self) -> Encoding {
match self.body {
Body::Ints(_) => Encoding::Intset,
Body::Packed(_) => Encoding::Listpack,
Body::Table(_) => Encoding::Hashtable,
}
}
#[inline]
pub fn len(&self) -> usize {
match &self.body {
Body::Ints(s) => s.len(),
Body::Packed(lp) => lp.len(),
Body::Table(t) => t.len(),
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub fn contains(&self, member: &[u8]) -> bool {
match &self.body {
Body::Ints(s) => parse_i64(member).is_some_and(|v| s.contains(v)),
Body::Packed(lp) => lp.find(member, 1).is_some(),
Body::Table(t) => t.contains(member),
}
}
#[must_use]
#[inline]
pub fn has(&self, needle: &Needle<'_>) -> bool {
match &self.body {
Body::Ints(s) => needle.int.is_some_and(|v| s.contains(v)),
Body::Packed(lp) => lp.find_parsed(needle.bytes, needle.int, 1).is_some(),
Body::Table(t) => t.contains_hashed(needle.hash, needle.bytes),
}
}
#[must_use]
pub fn at(&self, index: usize) -> Option<Member<'_>> {
match &self.body {
Body::Ints(s) => s.get(index).map(Member::Int),
Body::Packed(lp) => lp.get(index),
Body::Table(t) => t.at(index).map(|(name, _)| Member::Str(name)),
}
}
pub fn iter(&self) -> impl Iterator<Item = Member<'_>> {
(0..self.len()).map(|i| self.at(i).expect("index is under the length"))
}
pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
where
F: FnMut(Member<'_>),
{
match &self.body {
Body::Table(t) => t.scan(cursor, count, |name, ()| f(Member::Str(name))),
_ => {
for m in self.iter() {
f(m);
}
Cursor::END
}
}
}
#[must_use]
pub fn memory_bytes(&self) -> usize {
match &self.body {
Body::Ints(s) => s.memory_bytes(),
Body::Packed(lp) => lp.byte_len(),
Body::Table(t) => t.memory_bytes(),
}
}
pub fn add(&mut self, member: &[u8], limits: &Limits) -> bool {
match &mut self.body {
Body::Table(t) => return t.insert(member, ()).is_ok_and(|old| old.is_none()),
Body::Packed(lp) => {
if lp.find(member, 1).is_some() {
return false;
}
if lp.len() < limits.max_listpack_entries
&& member.len() <= limits.max_listpack_value
{
lp.push(member);
return true;
}
}
Body::Ints(s) => {
if let Some(v) = parse_i64(member) {
if !s.add(v) {
return false;
}
if s.len() > limits.max_intset_entries {
self.become_table(0);
}
return true;
}
if self.ints_fit_a_listpack(member, limits) {
self.become_listpack();
self.push_new(member);
return true;
}
}
}
self.become_table(1);
self.push_new(member);
true
}
fn push_new(&mut self, member: &[u8]) {
match &mut self.body {
Body::Packed(lp) => lp.push(member),
Body::Table(t) => {
t.insert(member, ())
.expect("the table was sized for this one");
}
Body::Ints(_) => unreachable!("no promotion ever lands on an intset"),
}
}
pub fn remove(&mut self, member: &[u8]) -> bool {
match &mut self.body {
Body::Ints(s) => parse_i64(member).is_some_and(|v| s.remove(v)),
Body::Packed(lp) => match lp.find(member, 1) {
Some(at) => lp.delete(at, 1),
None => false,
},
Body::Table(t) => t.remove(member).is_some(),
}
}
pub fn remove_at(&mut self, index: usize) -> Option<Vec<u8>> {
match &mut self.body {
Body::Ints(s) => {
let v = s.get(index)?;
s.remove(v);
let mut out = Vec::with_capacity(i64_len(v));
Member::Int(v).write_to(&mut out);
Some(out)
}
Body::Packed(lp) => {
let out = lp.get(index)?.to_vec();
lp.delete(index, 1);
Some(out)
}
Body::Table(t) => t.take_at(index).map(|(name, ())| name),
}
}
pub fn drop_at(&mut self, index: usize) -> bool {
match &mut self.body {
Body::Ints(s) => match s.get(index) {
Some(v) => {
s.remove(v);
true
}
None => false,
},
Body::Packed(lp) => {
if index >= lp.len() {
return false;
}
lp.delete(index, 1);
true
}
Body::Table(t) => t.remove_at(index).is_some(),
}
}
fn ints_fit_a_listpack(&self, member: &[u8], limits: &Limits) -> bool {
let Body::Ints(s) = &self.body else {
return false;
};
let widest = s
.min()
.map(i64_len)
.unwrap_or(0)
.max(s.max().map(i64_len).unwrap_or(0));
s.len() < limits.max_listpack_entries
&& member.len() <= limits.max_listpack_value
&& widest <= limits.max_listpack_value
}
fn become_listpack(&mut self) {
let Body::Ints(s) = &self.body else {
return;
};
let mut lp = Listpack::new();
let mut buf = Vec::with_capacity(20);
for v in s.iter() {
buf.clear();
Member::Int(v).write_to(&mut buf);
lp.push(&buf);
}
self.body = Body::Packed(lp);
}
fn become_table(&mut self, extra: usize) {
let mut t = Elements::with_capacity(self.len() + extra);
let mut buf = Vec::with_capacity(20);
for m in self.iter() {
match m {
Member::Str(b) => {
t.insert(b, ()).expect("room, and every member was unique");
}
Member::Int(v) => {
buf.clear();
Member::Int(v).write_to(&mut buf);
t.insert(&buf, ())
.expect("room, and every member was unique");
}
}
}
self.body = Body::Table(t);
}
}
impl Default for Set {
fn default() -> Set {
Set::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn of(members: &[&str]) -> Set {
let mut s = Set::new();
for m in members {
assert!(s.add(m.as_bytes(), &Limits::DEFAULT), "{m} was new");
}
s
}
fn members(s: &Set) -> Vec<String> {
let mut v: Vec<String> = s
.iter()
.map(|m| String::from_utf8(m.to_vec()).expect("utf8 in these tests"))
.collect();
v.sort();
v
}
#[test]
fn a_new_set_is_an_empty_intset() {
let s = Set::new();
assert_eq!(s.encoding(), Encoding::Intset);
assert_eq!(s.len(), 0);
assert!(s.is_empty());
assert!(!s.contains(b"1"));
assert_eq!(s.at(0), None);
}
#[test]
fn integers_stay_an_intset_and_come_back_as_members() {
let s = of(&["1", "2", "3"]);
assert_eq!(s.encoding(), Encoding::Intset);
assert_eq!(members(&s), ["1", "2", "3"]);
assert!(s.contains(b"2"));
assert!(!s.contains(b"4"));
assert!(!s.contains(b"two"));
}
#[test]
fn adding_the_same_member_twice_says_so_in_all_three() {
let mut ints = of(&["1", "2"]);
assert!(!ints.add(b"2", &Limits::DEFAULT));
assert_eq!(ints.len(), 2);
let mut packed = of(&["a", "b"]);
assert_eq!(packed.encoding(), Encoding::Listpack);
assert!(!packed.add(b"b", &Limits::DEFAULT));
assert_eq!(packed.len(), 2);
let mut table = of(&["a", "b"]);
table.become_table(0);
assert!(!table.add(b"b", &Limits::DEFAULT));
assert_eq!(table.len(), 2);
}
#[test]
fn a_string_turns_a_small_intset_into_a_listpack() {
let mut s = of(&["1", "2", "3"]);
assert!(s.add(b"hello", &Limits::DEFAULT));
assert_eq!(s.encoding(), Encoding::Listpack);
assert_eq!(members(&s), ["1", "2", "3", "hello"]);
assert!(s.contains(b"1"), "the integers survived the rewrite");
assert!(s.contains(b"hello"));
}
#[test]
fn a_string_turns_a_big_intset_straight_into_a_table() {
let mut s = Set::new();
for i in 0..200 {
s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
}
assert_eq!(s.encoding(), Encoding::Intset);
assert_eq!(s.len(), 200);
assert!(s.add(b"hello", &Limits::DEFAULT));
assert_eq!(s.encoding(), Encoding::Hashtable);
assert_eq!(s.len(), 201);
assert!(s.contains(b"199"));
assert!(s.contains(b"hello"));
}
#[test]
fn an_intset_holds_five_hundred_and_twelve_and_converts_at_the_next_one() {
let mut s = Set::new();
for i in 0..512 {
s.add(i.to_string().as_bytes(), &Limits::DEFAULT);
}
assert_eq!(s.encoding(), Encoding::Intset, "512 is still an intset");
assert_eq!(s.len(), 512);
s.add(b"512", &Limits::DEFAULT);
assert_eq!(s.encoding(), Encoding::Hashtable, "513 is not");
assert_eq!(s.len(), 513);
for i in 0..513 {
assert!(s.contains(i.to_string().as_bytes()), "{i} survived");
}
}
#[test]
fn a_listpack_converts_at_a_hundred_and_twenty_eight_members() {
let mut s = of(&["x"]);
assert_eq!(s.encoding(), Encoding::Listpack);
for i in 0..127 {
s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
}
assert_eq!(s.len(), 128);
assert_eq!(s.encoding(), Encoding::Listpack, "128 is still a listpack");
s.add(b"one more", &Limits::DEFAULT);
assert_eq!(s.len(), 129);
assert_eq!(s.encoding(), Encoding::Hashtable);
assert!(s.contains(b"x"));
assert!(s.contains(b"m126"));
assert!(s.contains(b"one more"));
}
#[test]
fn a_long_member_converts_a_listpack_whatever_the_count() {
let mut s = of(&["a"]);
let long = vec![b'z'; 65];
assert!(s.add(&long, &Limits::DEFAULT));
assert_eq!(s.encoding(), Encoding::Hashtable, "65 is past 64");
assert_eq!(s.len(), 2);
assert!(s.contains(&long));
let mut ok = of(&["a"]);
ok.add(&[b'z'; 64], &Limits::DEFAULT);
assert_eq!(ok.encoding(), Encoding::Listpack, "64 fits");
}
#[test]
fn a_long_member_sends_an_intset_to_a_table_and_not_a_listpack() {
let mut s = of(&["1", "2"]);
assert!(s.add(&[b'z'; 65], &Limits::DEFAULT));
assert_eq!(s.encoding(), Encoding::Hashtable);
assert_eq!(s.len(), 3);
}
#[test]
fn the_limits_are_configuration_and_moving_them_moves_the_encodings() {
let tight = Limits {
max_intset_entries: 2,
max_listpack_entries: 2,
max_listpack_value: 3,
};
let mut s = Set::new();
s.add(b"1", &tight);
s.add(b"2", &tight);
assert_eq!(s.encoding(), Encoding::Intset);
s.add(b"3", &tight);
assert_eq!(s.encoding(), Encoding::Hashtable, "three is past two");
let mut t = Set::new();
t.add(b"abc", &tight);
assert_eq!(t.encoding(), Encoding::Listpack, "three characters fit");
t.add(b"defg", &tight);
assert_eq!(t.encoding(), Encoding::Hashtable, "four do not");
assert!(t.contains(b"abc"));
assert!(t.contains(b"defg"));
let mut u = Set::new();
u.add(b"abcd", &tight);
assert_eq!(u.encoding(), Encoding::Hashtable);
assert!(u.contains(b"abcd"));
}
#[test]
fn with_hint_picks_the_representation_up_front() {
let d = &Limits::DEFAULT;
assert_eq!(
Set::with_hint(b"1", 10, d).encoding(),
Encoding::Intset,
"an integer and few enough of them"
);
assert_eq!(
Set::with_hint(b"1", 1000, d).encoding(),
Encoding::Hashtable,
"an integer and too many"
);
assert_eq!(
Set::with_hint(b"x", 10, d).encoding(),
Encoding::Listpack,
"not an integer and few enough"
);
assert_eq!(
Set::with_hint(b"x", 1000, d).encoding(),
Encoding::Hashtable,
"not an integer and too many"
);
}
#[test]
fn removing_works_in_all_three_and_never_demotes() {
let mut ints = of(&["1", "2", "3"]);
assert!(ints.remove(b"2"));
assert!(!ints.remove(b"2"));
assert!(!ints.remove(b"nope"), "not an integer, so not a member");
assert_eq!(members(&ints), ["1", "3"]);
assert_eq!(ints.encoding(), Encoding::Intset);
let mut packed = of(&["a", "b", "c"]);
assert!(packed.remove(b"b"));
assert!(!packed.remove(b"b"));
assert_eq!(members(&packed), ["a", "c"]);
assert_eq!(packed.encoding(), Encoding::Listpack);
let mut table = of(&["a", "b", "c"]);
table.become_table(0);
assert!(table.remove(b"b"));
assert!(!table.remove(b"b"));
assert_eq!(members(&table), ["a", "c"]);
assert_eq!(
table.encoding(),
Encoding::Hashtable,
"down to two members and still a table"
);
}
#[test]
fn a_set_can_be_emptied_a_member_at_a_time() {
for mut s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
let all: Vec<Vec<u8>> = s.iter().map(|m| m.to_vec()).collect();
for m in &all {
assert!(s.remove(m));
}
assert!(s.is_empty());
assert_eq!(s.at(0), None);
}
}
#[test]
fn removing_by_position_hands_the_member_back() {
let mut table = of(&["a", "b", "c", "d"]);
table.become_table(0);
for mut s in [
of(&["10", "20", "30", "40"]),
of(&["a", "b", "c", "d"]),
table,
] {
let mut got = Vec::new();
while !s.is_empty() {
got.push(String::from_utf8(s.remove_at(0).expect("not empty")).expect("utf8"));
}
got.sort();
assert_eq!(got.len(), 4, "four members and no repeats");
assert_eq!(s.len(), 0);
assert_eq!(s.remove_at(0), None);
}
}
#[test]
fn an_integer_member_is_the_same_member_however_it_is_written() {
let mut s = of(&["42"]);
assert!(!s.add(b"42", &Limits::DEFAULT));
assert_eq!(s.len(), 1);
assert!(s.add(b"042", &Limits::DEFAULT));
assert_eq!(s.encoding(), Encoding::Listpack);
assert_eq!(members(&s), ["042", "42"]);
assert!(s.contains(b"42"));
assert!(s.contains(b"042"));
}
#[test]
fn the_small_bands_answer_a_scan_in_one_go() {
for s in [of(&["1", "2", "3"]), of(&["a", "b", "c"])] {
let mut seen = Vec::new();
let next = s.scan(Cursor::at(1, 0, 99), 1, |m| seen.push(m.to_vec()));
assert!(next.is_end(), "{:?} split a scan up", s.encoding());
assert_eq!(seen.len(), 3);
}
}
#[test]
fn the_table_band_walks_a_scan_in_windows_and_misses_nothing() {
let mut s = Set::new();
for i in 0..300 {
s.add(format!("m{i}").as_bytes(), &Limits::DEFAULT);
}
assert_eq!(s.encoding(), Encoding::Hashtable);
let mut seen = Vec::new();
let mut c = Cursor::START;
let mut turns = 0;
loop {
c = s.scan(c, 7, |m| seen.push(m.to_vec()));
turns += 1;
assert!(turns < 100, "the scan did not finish");
if c.is_end() {
break;
}
}
assert!(
turns > 1,
"a window of seven over three hundred took one turn"
);
seen.sort();
seen.dedup();
assert_eq!(seen.len(), 300, "every member came back at least once");
}
#[test]
fn a_conversion_loses_no_member_at_any_of_the_three_boundaries() {
let mut wide = Set::new();
for v in [i64::MIN, -1, 0, 1, i64::MAX] {
wide.add(v.to_string().as_bytes(), &Limits::DEFAULT);
}
wide.add(b"str", &Limits::DEFAULT);
assert_eq!(wide.encoding(), Encoding::Listpack);
for v in [i64::MIN, -1, 0, 1, i64::MAX] {
assert!(wide.contains(v.to_string().as_bytes()), "{v} survived");
}
wide.add(&[b'q'; 100], &Limits::DEFAULT);
assert_eq!(wide.encoding(), Encoding::Hashtable);
for v in [i64::MIN, -1, 0, 1, i64::MAX] {
assert!(
wide.contains(v.to_string().as_bytes()),
"{v} survived twice"
);
}
assert!(wide.contains(b"str"));
assert_eq!(wide.len(), 7);
}
}