use core::cmp::Ordering;
use core::ops::Range;
use yo_common::num::{DIGITS_MAX, DOUBLE_MAX, i64_digits, parse_f64, write_double};
use crate::elem::Elements;
use crate::listpack::{self, Listpack};
use crate::rank::Rank;
use crate::scan::Cursor;
pub type Member<'a> = listpack::Entry<'a>;
#[derive(Debug, Clone, Copy)]
pub struct Limits {
pub max_listpack_entries: usize,
pub max_listpack_value: usize,
}
impl Limits {
pub const DEFAULT: Limits = Limits {
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 {
Listpack,
Skiplist,
}
impl Encoding {
#[must_use]
pub const fn name(self) -> &'static str {
match self {
Encoding::Listpack => "listpack",
Encoding::Skiplist => "skiplist",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Added {
New,
Changed,
Same,
Full,
}
#[derive(Debug, Clone, Copy)]
pub struct Bound {
pub at: f64,
pub open: bool,
}
impl Bound {
#[must_use]
pub const fn closed(at: f64) -> Bound {
Bound { at, open: false }
}
#[must_use]
pub const fn open(at: f64) -> Bound {
Bound { at, open: true }
}
}
#[derive(Debug, Clone, Copy)]
pub enum Lex<'a> {
Min,
Max,
Incl(&'a [u8]),
Excl(&'a [u8]),
}
#[derive(Debug, Clone)]
struct Table {
members: Elements<f64>,
order: Rank,
}
#[derive(Debug, Clone)]
enum Body {
Packed(Listpack),
Table(Table),
}
#[derive(Debug, Clone)]
pub struct Zset {
body: Body,
}
impl Default for Zset {
fn default() -> Self {
Self::new()
}
}
#[inline]
fn cmp_score(a: f64, b: f64) -> Ordering {
a.partial_cmp(&b).unwrap_or(Ordering::Equal)
}
#[inline]
fn cmp_key(score: f64, member: &[u8], other_score: f64, other_member: &[u8]) -> Ordering {
match cmp_score(score, other_score) {
Ordering::Equal => member.cmp(other_member),
other => other,
}
}
fn score_of(entry: Member<'_>) -> f64 {
match entry {
Member::Int(n) => n as f64,
Member::Str(s) => parse_f64(s).unwrap_or(0.0),
}
}
fn bytes_of<'a>(entry: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> &'a [u8] {
match entry {
Member::Str(s) => s,
Member::Int(n) => i64_digits(digits, n),
}
}
impl Zset {
#[must_use]
pub fn new() -> Zset {
Zset {
body: Body::Packed(Listpack::new()),
}
}
#[must_use]
pub fn with_hint(hint: usize, limits: &Limits) -> Zset {
if hint <= limits.max_listpack_entries {
return Zset::new();
}
Zset {
body: Body::Table(Table {
members: Elements::with_capacity(hint),
order: Rank::new(),
}),
}
}
pub(crate) fn from_packed(lp: Listpack, limits: &Limits) -> Result<Zset, Listpack> {
let n = lp.len();
if n == 0 || !n.is_multiple_of(2) || n / 2 > limits.max_listpack_entries {
return Err(lp);
}
let ok = {
let mut walk = lp.iter();
let mut prev: Option<(f64, Member<'_>)> = None;
let mut before_buf = [0u8; DIGITS_MAX];
let mut member_buf = [0u8; DIGITS_MAX];
loop {
let Some(member) = walk.next() else {
break true;
};
let Some(entry) = walk.next() else {
break false;
};
let score = match entry {
Member::Int(v) => v as f64,
Member::Str(s) => match parse_f64(s) {
Some(v) => v,
None => break false,
},
};
let bytes = bytes_of(member, &mut member_buf);
if bytes.len() > limits.max_listpack_value {
break false;
}
if let Some((before, was)) = prev {
let was = bytes_of(was, &mut before_buf);
if cmp_key(before, was, score, bytes) != Ordering::Less {
break false;
}
}
prev = Some((score, member));
}
};
if ok {
Ok(Zset {
body: Body::Packed(lp),
})
} else {
Err(lp)
}
}
#[must_use]
pub fn len(&self) -> usize {
match &self.body {
Body::Packed(lp) => lp.len() / 2,
Body::Table(t) => t.members.len(),
}
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[must_use]
pub const fn encoding(&self) -> Encoding {
match &self.body {
Body::Packed(_) => Encoding::Listpack,
Body::Table(_) => Encoding::Skiplist,
}
}
#[inline]
pub(crate) fn packed_bytes(&self) -> Option<&[u8]> {
match &self.body {
Body::Packed(lp) => Some(lp.as_bytes()),
Body::Table(_) => None,
}
}
#[must_use]
pub fn memory_bytes(&self) -> usize {
match &self.body {
Body::Packed(lp) => lp.byte_len(),
Body::Table(t) => t.members.memory_bytes() + t.order.bytes(),
}
}
#[must_use]
pub fn score(&self, member: &[u8]) -> Option<f64> {
match &self.body {
Body::Packed(lp) => {
let at = lp.find(member, 2)?;
lp.get(at + 1).map(score_of)
}
Body::Table(t) => t.members.get(member).copied(),
}
}
pub fn add(&mut self, member: &[u8], score: f64, limits: &Limits) -> Added {
if let Body::Packed(lp) = &mut self.body {
if let Some(at) = lp.find(member, 2) {
let old = lp.get(at + 1).map_or(0.0, score_of);
if cmp_score(old, score) == Ordering::Equal {
return Added::Same;
}
lp.delete(at, 2);
packed_insert(lp, member, score);
return Added::Changed;
}
if lp.len() / 2 < limits.max_listpack_entries
&& member.len() <= limits.max_listpack_value
{
packed_insert(lp, member, score);
return Added::New;
}
self.promote();
}
let Body::Table(t) = &mut self.body else {
unreachable!("promoted above")
};
t.add(member, score)
}
pub fn remove(&mut self, member: &[u8]) -> bool {
match &mut self.body {
Body::Packed(lp) => match lp.find(member, 2) {
Some(at) => lp.delete(at, 2),
None => false,
},
Body::Table(t) => {
let Some(row) = t.members.index_of(member) else {
return false;
};
let score = *t.members.get(member).expect("just found");
let rank = t.rank_of(row as u32, score, member);
t.take(rank, row);
true
}
}
}
#[must_use]
pub fn rank(&self, member: &[u8]) -> Option<usize> {
match &self.body {
Body::Packed(lp) => lp.find(member, 2).map(|at| at / 2),
Body::Table(t) => {
let row = t.members.index_of(member)?;
let score = *t.members.get(member)?;
Some(t.rank_of(row as u32, score, member))
}
}
}
#[must_use]
pub fn at(&self, rank: usize) -> Option<(Member<'_>, f64)> {
match &self.body {
Body::Packed(lp) => {
let member = lp.get(rank * 2)?;
let score = lp.get(rank * 2 + 1).map(score_of)?;
Some((member, score))
}
Body::Table(t) => {
let row = t.order.row_at(rank)?;
let (name, score) = t.members.at(row as usize)?;
Some((Member::Str(name), *score))
}
}
}
pub fn remove_at(&mut self, rank: usize) -> bool {
match &mut self.body {
Body::Packed(lp) => lp.delete(rank * 2, 2),
Body::Table(t) => {
let Some(row) = t.order.row_at(rank) else {
return false;
};
t.take(rank, row as usize);
true
}
}
}
#[must_use]
pub fn pick(&self, at: usize) -> Option<(Member<'_>, f64)> {
match &self.body {
Body::Packed(_) => self.at(at),
Body::Table(t) => {
let (name, score) = t.members.at(at)?;
Some((Member::Str(name), *score))
}
}
}
pub fn walk<F: FnMut(Member<'_>, f64)>(&self, from: usize, count: usize, rev: bool, mut f: F) {
let len = self.len();
if from >= len || count == 0 {
return;
}
let count = count.min(if rev { from + 1 } else { len - from });
match &self.body {
Body::Packed(lp) => {
for i in 0..count {
let at = if rev { from - i } else { from + i };
let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
return;
};
f(m, score_of(s));
}
}
Body::Table(t) => {
if rev {
for row in t.order.iter_back_from(from).take(count) {
let Some((name, score)) = t.members.at(row as usize) else {
return;
};
f(Member::Str(name), *score);
}
} else {
for row in t.order.iter_from(from).take(count) {
let Some((name, score)) = t.members.at(row as usize) else {
return;
};
f(Member::Str(name), *score);
}
}
}
}
}
#[must_use]
pub fn window_by_score(&self, min: Bound, max: Bound) -> Range<usize> {
let start = self.seek(|score, _| {
let before = match cmp_score(score, min.at) {
Ordering::Less => true,
Ordering::Equal => min.open,
Ordering::Greater => false,
};
if before {
Ordering::Greater
} else {
Ordering::Less
}
});
let end = self.seek(|score, _| {
let inside = match cmp_score(score, max.at) {
Ordering::Less => true,
Ordering::Equal => !max.open,
Ordering::Greater => false,
};
if inside {
Ordering::Greater
} else {
Ordering::Less
}
});
start..end.max(start)
}
#[must_use]
pub fn window_by_lex(&self, min: Lex<'_>, max: Lex<'_>) -> Range<usize> {
let start = self.seek(|_, member| match min {
Lex::Min => Ordering::Less,
Lex::Max => Ordering::Greater,
Lex::Incl(at) => {
if member < at {
Ordering::Greater
} else {
Ordering::Less
}
}
Lex::Excl(at) => {
if member <= at {
Ordering::Greater
} else {
Ordering::Less
}
}
});
let end = self.seek(|_, member| match max {
Lex::Min => Ordering::Less,
Lex::Max => Ordering::Greater,
Lex::Incl(at) => {
if member <= at {
Ordering::Greater
} else {
Ordering::Less
}
}
Lex::Excl(at) => {
if member < at {
Ordering::Greater
} else {
Ordering::Less
}
}
});
start..end.max(start)
}
fn seek<F: FnMut(f64, &[u8]) -> Ordering>(&self, mut probe: F) -> usize {
match &self.body {
Body::Packed(lp) => {
let mut digits = [0u8; DIGITS_MAX];
let mut at = 0;
while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
let bytes = bytes_of(m, &mut digits);
if probe(score_of(s), bytes) != Ordering::Greater {
break;
}
at += 1;
}
at
}
Body::Table(t) => t.order.seek(|row| {
let (name, score) = t.members.at(row as usize).expect("a row the tree holds");
probe(*score, name)
}),
}
}
pub fn scan<F: FnMut(Member<'_>, f64)>(
&self,
cursor: Cursor,
count: usize,
mut f: F,
) -> Cursor {
match &self.body {
Body::Packed(lp) => {
for at in 0..lp.len() / 2 {
let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
break;
};
f(m, score_of(s));
}
Cursor::END
}
Body::Table(t) => t.members.scan(cursor, count, |name, score| {
f(Member::Str(name), *score);
}),
}
}
#[must_use]
pub fn from_elements(members: Elements<f64>, limits: &Limits) -> Option<Zset> {
let n = members.len();
if n == 0 {
return None;
}
let mut rows: Vec<u32> = (0..n as u32).collect();
rows.sort_unstable_by(|&a, &b| {
let (a_name, a_score) = members.at(a as usize).expect("in range");
let (b_name, b_score) = members.at(b as usize).expect("in range");
cmp_key(*a_score, a_name, *b_score, b_name)
});
let packable = n <= limits.max_listpack_entries
&& rows.iter().all(|&r| {
members.at(r as usize).expect("in range").0.len() <= limits.max_listpack_value
});
if packable {
let mut lp = Listpack::new();
let mut score_buf = [0u8; DOUBLE_MAX];
for &row in &rows {
let (name, score) = members.at(row as usize).expect("in range");
lp.push(name);
lp.push(write_double(&mut score_buf, *score));
}
return Some(Zset {
body: Body::Packed(lp),
});
}
let mut order = Rank::new();
for (at, &row) in rows.iter().enumerate() {
order.insert_at(at, row);
}
Some(Zset {
body: Body::Table(Table { members, order }),
})
}
fn promote(&mut self) {
let Body::Packed(lp) = &self.body else { return };
let n = lp.len() / 2;
let mut table = Table {
members: Elements::with_capacity(n.next_power_of_two().max(16)),
order: Rank::new(),
};
let mut digits = [0u8; DIGITS_MAX];
let mut steps = lp.iter();
while let (Some(m), Some(s)) = (steps.next(), steps.next()) {
let bytes = bytes_of(m, &mut digits);
let row = table.members.len() as u32;
if table.members.insert(bytes, score_of(s)).is_err() {
break;
}
table.order.insert_at(row as usize, row);
}
self.body = Body::Table(table);
}
}
fn packed_insert(lp: &mut Listpack, member: &[u8], score: f64) {
let mut score_buf = [0u8; DOUBLE_MAX];
let text = write_double(&mut score_buf, score);
let mut digits = [0u8; DIGITS_MAX];
let mut at = 0;
while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
let bytes = bytes_of(m, &mut digits);
if cmp_key(score, member, score_of(s), bytes) == Ordering::Less {
break;
}
at += 1;
}
if at * 2 == lp.len() {
lp.push(member);
lp.push(text);
} else {
lp.insert(at * 2, member);
lp.insert(at * 2 + 1, text);
}
}
impl Table {
fn rank_of(&self, row: u32, score: f64, member: &[u8]) -> usize {
let members = &self.members;
self.order.seek(|other| {
if other == row {
return Ordering::Equal;
}
let (name, at) = members.at(other as usize).expect("a row the tree holds");
cmp_key(score, member, *at, name)
})
}
fn add(&mut self, member: &[u8], score: f64) -> Added {
if let Some(row) = self.members.index_of(member) {
let old = *self.members.at(row).map_or(&0.0, |(_, s)| s);
if cmp_score(old, score) == Ordering::Equal {
return Added::Same;
}
let was = self.rank_of(row as u32, old, member);
self.order.remove_at(was);
if let Some(at) = self.members.at_mut(row) {
*at = score;
}
let now = self.rank_of(row as u32, score, member);
self.order.insert_at(now, row as u32);
return Added::Changed;
}
let row = self.members.len() as u32;
if self.members.insert(member, score).is_err() {
return Added::Full;
}
let at = self.rank_of(row, score, member);
self.order.insert_at(at, row);
Added::New
}
fn take(&mut self, rank: usize, row: usize) {
let last = self.members.len() - 1;
let moved = if last == row {
None
} else {
let (name, score) = self.members.at(last).expect("the last row");
let at = {
let members = &self.members;
let score = *score;
self.order.seek(|other| {
if other as usize == last {
return Ordering::Equal;
}
let (other_name, other_score) =
members.at(other as usize).expect("a row the tree holds");
cmp_key(score, name, *other_score, other_name)
})
};
Some(at)
};
self.order.remove_at(rank);
self.members.remove_at(row);
if let Some(at) = moved {
let at = if at > rank { at - 1 } else { at };
self.order.set_at(at, row as u32);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn packed(pairs: &[(&[u8], &str)]) -> Listpack {
let mut lp = Listpack::new();
for (member, score) in pairs {
lp.push(member);
lp.push(score.as_bytes());
}
lp
}
#[test]
fn a_payload_in_this_layout_is_taken_whole() {
let lp = packed(&[(b"a", "1"), (b"b", "2"), (b"c", "2.5")]);
let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("in order and inside the limits");
assert_eq!(z.encoding(), Encoding::Listpack);
assert_eq!(z.len(), 3);
assert_eq!(z.score(b"a"), Some(1.0));
assert_eq!(z.score(b"c"), Some(2.5));
assert_eq!(z.rank(b"b"), Some(1));
assert_eq!(z.score(b"missing"), None);
let lp = packed(&[(b"10", "1"), (b"9", "2")]);
let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("sorted by score");
assert_eq!(z.score(b"10"), Some(1.0));
assert_eq!(z.score(b"9"), Some(2.0));
assert_eq!(z.rank(b"9"), Some(1));
}
#[test]
fn a_blob_this_band_cannot_hold_is_handed_back() {
let small = Limits {
max_listpack_entries: 4,
max_listpack_value: 8,
};
for (why, lp) in [
("out of order by score", packed(&[(b"a", "2"), (b"b", "1")])),
(
"out of order by member",
packed(&[(b"b", "1"), (b"a", "1")]),
),
(
"the same member twice",
packed(&[(b"a", "1"), (b"a", "1"), (b"b", "2")]),
),
("a score that is not a number", packed(&[(b"a", "no")])),
(
"a member past the value limit",
packed(&[(b"aaaaaaaaaa", "1")]),
),
(
"more members than the band takes",
packed(&[
(b"a", "1"),
(b"b", "2"),
(b"c", "3"),
(b"d", "4"),
(b"e", "5"),
]),
),
("nothing in it at all", packed(&[])),
] {
assert!(
Zset::from_packed(lp, &small).is_err(),
"{why} should be handed back"
);
}
let mut odd = Listpack::new();
odd.push(b"a");
assert!(Zset::from_packed(odd, &Limits::DEFAULT).is_err());
}
#[test]
fn a_hint_past_the_band_starts_on_the_table() {
let big = Zset::with_hint(Limits::DEFAULT.max_listpack_entries + 1, &Limits::DEFAULT);
assert_eq!(big.encoding(), Encoding::Skiplist);
assert!(big.is_empty());
let small = Zset::with_hint(Limits::DEFAULT.max_listpack_entries, &Limits::DEFAULT);
assert_eq!(small.encoding(), Encoding::Listpack);
let mut z = Zset::with_hint(1_000_000, &Limits::DEFAULT);
z.add(b"b", 2.0, &Limits::DEFAULT);
z.add(b"a", 1.0, &Limits::DEFAULT);
assert_eq!(z.len(), 2);
assert_eq!(z.rank(b"a"), Some(0));
assert_eq!(z.rank(b"b"), Some(1));
}
#[test]
#[ignore = "a measurement, run it by name"]
fn measure_bytes_per_entry() {
let mut lp = Zset::new();
let mut lp_payload = 0usize;
for i in 0..128 {
let m = format!("member:{i:09}");
lp_payload += m.len();
lp.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
}
println!(
"packed n=128 total={} payload={lp_payload} overhead_per_entry={:.2}",
lp.memory_bytes(),
(lp.memory_bytes() as f64 - lp_payload as f64) / 128.0
);
for n in [10_000usize, 100_000, 600_000, 1_000_000] {
let (z, payload) = filled(n);
let total = z.memory_bytes();
let scores = n * 8;
let (slots, rows, names, tree) = match &z.body {
Body::Table(t) => (
t.members.slot_bytes(),
t.members.row_bytes(),
t.members.name_bytes(),
t.order.bytes(),
),
Body::Packed(_) => (0, 0, 0, 0),
};
let per = |b: usize| b as f64 / n as f64;
println!(
"table n={n} total={total} slots={:.2}/e rows={:.2}/e names={:.2}/e tree={:.2}/e overhead_per_entry={:.2}",
per(slots),
per(rows),
per(names),
per(tree),
(total as f64 - payload as f64 - scores as f64) / n as f64
);
}
}
fn filled(n: usize) -> (Zset, usize) {
let mut z = Zset::new();
let mut payload = 0usize;
for i in 0..n {
let m = format!("member:{i:09}");
payload += m.len();
z.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
}
(z, payload)
}
#[test]
fn a_large_sorted_set_does_not_hold_much_more_than_it_stores() {
let n = 40_000usize;
let (z, payload) = filled(n);
let Body::Table(t) = &z.body else {
panic!("forty thousand members is not a listpack");
};
assert!(
t.members.row_bytes() < n * 25,
"the row and score arrays hold {} for {n} members",
t.members.row_bytes()
);
assert!(
t.members.name_bytes() < payload + payload / 4,
"the name blob holds {} for {payload} bytes of names",
t.members.name_bytes()
);
assert!(
t.order.bytes() < n * 4,
"the tree holds {} for {n} members",
t.order.bytes()
);
}
fn built(pairs: &[(&str, f64)], limits: &Limits) -> Zset {
let mut z = Zset::new();
for (m, s) in pairs {
z.add(m.as_bytes(), *s, limits);
}
z
}
fn listed(z: &Zset) -> Vec<(String, f64)> {
let mut out = Vec::new();
let mut digits = [0u8; DIGITS_MAX];
z.walk(0, z.len(), false, |m, s| {
let bytes = bytes_of(m, &mut digits).to_vec();
out.push((String::from_utf8(bytes).unwrap(), s));
});
out
}
fn model(pairs: &[(&str, f64)]) -> Vec<(String, f64)> {
let mut last: Vec<(String, f64)> = Vec::new();
for (m, s) in pairs {
match last.iter_mut().find(|(name, _)| name == m) {
Some(row) => row.1 = *s,
None => last.push(((*m).to_string(), *s)),
}
}
last.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
last
}
const PACKED: Limits = Limits::DEFAULT;
const TABLE: Limits = Limits {
max_listpack_entries: 0,
max_listpack_value: 64,
};
#[test]
fn an_empty_set_answers_nothing() {
let z = Zset::new();
assert_eq!(z.len(), 0);
assert!(z.is_empty());
assert_eq!(z.encoding(), Encoding::Listpack);
assert_eq!(z.score(b"nobody"), None);
assert_eq!(z.rank(b"nobody"), None);
assert!(z.at(0).is_none());
assert_eq!(
z.window_by_score(Bound::closed(f64::NEG_INFINITY), Bound::closed(0.0)),
0..0
);
}
#[test]
fn both_bands_put_members_in_the_same_order() {
let pairs = [("c", 3.0), ("a", 1.0), ("b", 2.0), ("d", 2.0), ("e", -1.5)];
for limits in [&PACKED, &TABLE] {
let z = built(&pairs, limits);
assert_eq!(listed(&z), model(&pairs), "{:?}", z.encoding());
assert_eq!(z.len(), 5);
}
}
#[test]
fn a_tie_on_score_is_broken_by_the_member() {
for limits in [&PACKED, &TABLE] {
let pairs = [
("beta", 1.0),
("alpha", 1.0),
("gamma", 1.0),
("Alpha", 1.0),
];
let z = built(&pairs, limits);
let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
assert_eq!(names, ["Alpha", "alpha", "beta", "gamma"]);
}
}
#[test]
fn adding_a_member_again_moves_it_rather_than_adding_it() {
for limits in [&PACKED, &TABLE] {
let mut z = built(&[("a", 1.0), ("b", 2.0), ("c", 3.0)], limits);
assert_eq!(z.add(b"a", 1.0, limits), Added::Same);
assert_eq!(z.add(b"a", 9.0, limits), Added::Changed);
assert_eq!(z.len(), 3);
assert_eq!(z.score(b"a"), Some(9.0));
assert_eq!(z.rank(b"a"), Some(2));
assert_eq!(listed(&z).last().unwrap().0, "a");
}
}
#[test]
fn a_removal_leaves_every_other_rank_right() {
for limits in [&PACKED, &TABLE] {
let pairs: Vec<(&str, f64)> =
vec![("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0), ("e", 5.0)];
let mut z = built(&pairs, limits);
assert!(z.remove(b"c"));
assert!(!z.remove(b"c"));
assert_eq!(z.len(), 4);
assert_eq!(z.rank(b"a"), Some(0));
assert_eq!(z.rank(b"d"), Some(2));
assert_eq!(z.rank(b"e"), Some(3));
assert_eq!(z.score(b"c"), None);
}
}
#[test]
fn removing_from_the_middle_renumbers_the_last_element() {
let limits = &TABLE;
let mut z = Zset::new();
for i in 0..64u32 {
z.add(format!("m{i:03}").as_bytes(), f64::from(i), limits);
}
for i in 0..63u32 {
assert!(z.remove(format!("m{i:03}").as_bytes()));
assert_eq!(z.len() as u32, 63 - i);
for j in i + 1..64 {
let name = format!("m{j:03}");
assert_eq!(
z.score(name.as_bytes()),
Some(f64::from(j)),
"score of {name}"
);
assert_eq!(
z.rank(name.as_bytes()),
Some((j - i - 1) as usize),
"rank of {name} after {i}"
);
}
}
}
#[test]
fn a_set_promotes_when_it_outgrows_the_packed_band() {
let limits = Limits {
max_listpack_entries: 4,
max_listpack_value: 64,
};
let mut z = Zset::new();
for i in 0..4u32 {
z.add(format!("m{i}").as_bytes(), f64::from(i), &limits);
}
assert_eq!(z.encoding(), Encoding::Listpack);
z.add(b"m4", 4.0, &limits);
assert_eq!(z.encoding(), Encoding::Skiplist);
assert_eq!(z.len(), 5);
let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
assert_eq!(names, ["m0", "m1", "m2", "m3", "m4"]);
z.remove(b"m4");
z.remove(b"m3");
assert_eq!(z.encoding(), Encoding::Skiplist);
}
#[test]
fn a_member_too_long_for_the_packed_band_promotes_on_its_own() {
let limits = Limits {
max_listpack_entries: 128,
max_listpack_value: 8,
};
let mut z = Zset::new();
z.add(b"short", 1.0, &limits);
assert_eq!(z.encoding(), Encoding::Listpack);
z.add(b"a member well past eight bytes", 2.0, &limits);
assert_eq!(z.encoding(), Encoding::Skiplist);
assert_eq!(z.len(), 2);
assert_eq!(z.score(b"a member well past eight bytes"), Some(2.0));
}
#[test]
fn a_score_range_finds_both_of_its_ends() {
for limits in [&PACKED, &TABLE] {
let pairs = [("a", 1.0), ("b", 2.0), ("c", 2.0), ("d", 3.0), ("e", 4.0)];
let z = built(&pairs, limits);
assert_eq!(
z.window_by_score(Bound::closed(2.0), Bound::closed(3.0)),
1..4
);
assert_eq!(
z.window_by_score(Bound::open(2.0), Bound::closed(3.0)),
3..4
);
assert_eq!(
z.window_by_score(Bound::closed(2.0), Bound::open(3.0)),
1..3
);
assert_eq!(z.window_by_score(Bound::open(1.0), Bound::open(4.0)), 1..4);
assert_eq!(
z.window_by_score(
Bound::closed(f64::NEG_INFINITY),
Bound::closed(f64::INFINITY)
),
0..5
);
assert_eq!(
z.window_by_score(Bound::closed(3.0), Bound::closed(2.0)),
3..3
);
assert_eq!(
z.window_by_score(Bound::closed(9.0), Bound::closed(10.0)),
5..5
);
}
}
#[test]
fn a_member_range_orders_by_member_when_the_scores_are_equal() {
for limits in [&PACKED, &TABLE] {
let pairs = [("a", 0.0), ("b", 0.0), ("c", 0.0), ("d", 0.0), ("e", 0.0)];
let z = built(&pairs, limits);
assert_eq!(z.window_by_lex(Lex::Min, Lex::Max), 0..5);
assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Incl(b"d")), 1..4);
assert_eq!(z.window_by_lex(Lex::Excl(b"b"), Lex::Excl(b"d")), 2..3);
assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Excl(b"c")), 1..2);
assert_eq!(z.window_by_lex(Lex::Excl(b"e"), Lex::Max), 5..5);
assert_eq!(z.window_by_lex(Lex::Min, Lex::Excl(b"a")), 0..0);
}
}
#[test]
fn a_walk_can_go_backwards_and_stops_where_it_is_told() {
for limits in [&PACKED, &TABLE] {
let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0)];
let z = built(&pairs, limits);
let mut seen = Vec::new();
let mut digits = [0u8; DIGITS_MAX];
z.walk(3, 2, true, |m, _| {
seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
});
assert_eq!(seen, ["d", "c"]);
let mut seen = Vec::new();
z.walk(1, 99, false, |m, _| {
seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
});
assert_eq!(seen, ["b", "c", "d"]);
let mut count = 0;
z.walk(4, 1, false, |_, _| count += 1);
assert_eq!(count, 0);
}
}
#[test]
fn taking_from_a_rank_takes_the_right_one() {
for limits in [&PACKED, &TABLE] {
let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0)];
let mut z = built(&pairs, limits);
assert!(z.remove_at(0));
assert_eq!(z.len(), 2);
assert_eq!(z.rank(b"b"), Some(0));
assert!(z.remove_at(1));
assert_eq!(z.score(b"c"), None);
assert!(!z.remove_at(5));
}
}
#[test]
fn a_run_of_everything_agrees_with_a_model() {
for limits in [
&PACKED,
&Limits {
max_listpack_entries: 8,
max_listpack_value: 64,
},
] {
let mut z = Zset::new();
let mut model: Vec<(String, f64)> = Vec::new();
let mut seed = 0x8765_4321_9ABC_DEF0u64;
let mut roll = || {
seed ^= seed << 13;
seed ^= seed >> 7;
seed ^= seed << 17;
seed
};
for round in 0..3_000 {
let name = format!("m{:02}", roll() % 40);
let score = (roll() % 7) as f64 - 3.0;
if round % 5 == 4 {
let hit = z.remove(name.as_bytes());
let was = model.iter().position(|(m, _)| *m == name);
assert_eq!(hit, was.is_some());
if let Some(at) = was {
model.remove(at);
}
} else {
z.add(name.as_bytes(), score, limits);
match model.iter_mut().find(|(m, _)| *m == name) {
Some(row) => row.1 = score,
None => model.push((name, score)),
}
}
if round % 97 == 0 {
let mut want = model.clone();
want.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
assert_eq!(listed(&z), want, "round {round}");
for (at, (m, s)) in want.iter().enumerate() {
assert_eq!(z.rank(m.as_bytes()), Some(at), "rank of {m}");
assert_eq!(z.score(m.as_bytes()), Some(*s), "score of {m}");
}
}
}
}
}
#[test]
fn a_big_set_costs_what_the_tree_said_it_would() {
let limits = &TABLE;
let mut z = Zset::new();
let n = 100_000u32;
for i in 0..n {
z.add(format!("member:{i:08}").as_bytes(), f64::from(i), limits);
}
assert_eq!(z.len(), n as usize);
assert_eq!(z.rank(b"member:00050000"), Some(50_000));
assert_eq!(z.at(0).map(|(_, s)| s), Some(0.0));
assert_eq!(z.at(n as usize - 1).map(|(_, s)| s), Some(f64::from(n - 1)));
let Body::Table(t) = &z.body else {
panic!("a hundred thousand members is not a listpack")
};
let per_order = t.order.bytes() as f64 / f64::from(n);
assert!(per_order < 3.4, "{per_order} bytes an element in the tree");
let per = z.memory_bytes() as f64 / f64::from(n);
assert!(per < 70.0, "{per} bytes a member all in");
}
}