1use core::cmp::Ordering;
56use core::ops::Range;
57
58use yo_common::num::{DIGITS_MAX, DOUBLE_MAX, i64_digits, parse_f64, write_double};
59
60use crate::elem::Elements;
61use crate::frozen::{self, Broken};
62use crate::listpack::{self, Listpack};
63use crate::rank::Rank;
64use crate::scan::Cursor;
65
66pub type Member<'a> = listpack::Entry<'a>;
68
69const FORM_PACKED: u8 = 1;
71const FORM_MEMBERS: u8 = 2;
73
74#[derive(Debug, Clone, Copy)]
76pub struct Limits {
77 pub max_listpack_entries: usize,
79 pub max_listpack_value: usize,
81}
82
83impl Limits {
84 pub const DEFAULT: Limits = Limits {
86 max_listpack_entries: 128,
87 max_listpack_value: 64,
88 };
89}
90
91impl Default for Limits {
92 fn default() -> Limits {
93 Limits::DEFAULT
94 }
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
99pub enum Encoding {
100 Listpack,
102 Skiplist,
109}
110
111impl Encoding {
112 #[must_use]
114 pub const fn name(self) -> &'static str {
115 match self {
116 Encoding::Listpack => "listpack",
117 Encoding::Skiplist => "skiplist",
118 }
119 }
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
124pub enum Added {
125 New,
127 Changed,
129 Same,
131 Full,
133}
134
135#[derive(Debug, Clone, Copy)]
137pub struct Bound {
138 pub at: f64,
140 pub open: bool,
143}
144
145impl Bound {
146 #[must_use]
148 pub const fn closed(at: f64) -> Bound {
149 Bound { at, open: false }
150 }
151
152 #[must_use]
154 pub const fn open(at: f64) -> Bound {
155 Bound { at, open: true }
156 }
157}
158
159#[derive(Debug, Clone, Copy)]
161pub enum Lex<'a> {
162 Min,
164 Max,
166 Incl(&'a [u8]),
168 Excl(&'a [u8]),
170}
171
172#[derive(Debug, Clone)]
174struct Table {
175 members: Elements<f64>,
176 order: Rank,
177}
178
179#[derive(Debug, Clone)]
181enum Body {
182 Packed(Listpack),
184 Table(Table),
185}
186
187#[derive(Debug, Clone)]
189pub struct Zset {
190 body: Body,
191}
192
193impl Default for Zset {
194 fn default() -> Self {
195 Self::new()
196 }
197}
198
199#[inline]
206fn cmp_score(a: f64, b: f64) -> Ordering {
207 a.partial_cmp(&b).unwrap_or(Ordering::Equal)
208}
209
210#[inline]
212fn cmp_key(score: f64, member: &[u8], other_score: f64, other_member: &[u8]) -> Ordering {
213 match cmp_score(score, other_score) {
214 Ordering::Equal => member.cmp(other_member),
215 other => other,
216 }
217}
218
219fn score_of(entry: Member<'_>) -> f64 {
221 match entry {
222 Member::Int(n) => n as f64,
223 Member::Str(s) => parse_f64(s).unwrap_or(0.0),
226 }
227}
228
229fn bytes_of<'a>(entry: Member<'a>, digits: &'a mut [u8; DIGITS_MAX]) -> &'a [u8] {
231 match entry {
232 Member::Str(s) => s,
233 Member::Int(n) => i64_digits(digits, n),
234 }
235}
236
237impl Zset {
238 #[must_use]
240 pub fn new() -> Zset {
241 Zset {
242 body: Body::Packed(Listpack::new()),
243 }
244 }
245
246 #[must_use]
259 pub fn with_hint(hint: usize, limits: &Limits) -> Zset {
260 if hint <= limits.max_listpack_entries {
261 return Zset::new();
262 }
263 Zset {
264 body: Body::Table(Table {
265 members: Elements::with_capacity(hint),
266 order: Rank::new(),
267 }),
268 }
269 }
270
271 pub(crate) fn from_packed(lp: Listpack, limits: &Limits) -> Result<Zset, Listpack> {
294 let n = lp.len();
295 if n == 0 || !n.is_multiple_of(2) || n / 2 > limits.max_listpack_entries {
296 return Err(lp);
297 }
298 let ok = {
299 let mut walk = lp.iter();
300 let mut prev: Option<(f64, Member<'_>)> = None;
301 let mut before_buf = [0u8; DIGITS_MAX];
302 let mut member_buf = [0u8; DIGITS_MAX];
303 loop {
304 let Some(member) = walk.next() else {
305 break true;
306 };
307 let Some(entry) = walk.next() else {
310 break false;
311 };
312 let score = match entry {
313 Member::Int(v) => v as f64,
314 Member::Str(s) => match parse_f64(s) {
319 Some(v) => v,
320 None => break false,
321 },
322 };
323 let bytes = bytes_of(member, &mut member_buf);
324 if bytes.len() > limits.max_listpack_value {
325 break false;
326 }
327 if let Some((before, was)) = prev {
328 let was = bytes_of(was, &mut before_buf);
329 if cmp_key(before, was, score, bytes) != Ordering::Less {
330 break false;
331 }
332 }
333 prev = Some((score, member));
334 }
335 };
336 if ok {
337 Ok(Zset {
338 body: Body::Packed(lp),
339 })
340 } else {
341 Err(lp)
342 }
343 }
344
345 pub fn freeze(&self, out: &mut Vec<u8>) {
356 match &self.body {
357 Body::Packed(lp) => {
358 out.push(FORM_PACKED);
359 out.extend_from_slice(lp.as_bytes());
360 }
361 Body::Table(t) => {
362 out.push(FORM_MEMBERS);
363 frozen::put_uint(out, t.members.len() as u64);
364 for row in t.order.iter_from(0) {
368 let Some((name, score)) = t.members.at(row as usize) else {
369 continue;
370 };
371 frozen::put_bytes(out, name);
372 frozen::put_f64(out, *score);
373 }
374 }
375 }
376 }
377
378 pub fn thaw(bytes: &[u8]) -> Result<Zset, Broken> {
384 let mut cut = frozen::Cut::new(bytes);
385 match cut.byte()? {
386 FORM_PACKED => Ok(Zset {
387 body: Body::Packed(Listpack::from_bytes(cut.rest()).map_err(|_| Broken::Body)?),
388 }),
389 FORM_MEMBERS => {
390 let n = usize::try_from(cut.uint()?).map_err(|_| Broken::Short)?;
391 if n > cut.rest().len() {
395 return Err(Broken::Body);
396 }
397 let mut table = Table {
398 members: Elements::with_capacity(n),
399 order: Rank::new(),
400 };
401 for _ in 0..n {
402 let name = cut.bytes()?;
403 let score = cut.f64()?;
404 let row = table.members.len() as u32;
405 if !matches!(table.members.insert(name, score), Ok(None)) {
409 return Err(Broken::Body);
410 }
411 table.order.insert_at(row as usize, row);
414 }
415 Ok(Zset {
416 body: Body::Table(table),
417 })
418 }
419 _ => Err(Broken::Form),
420 }
421 }
422
423 #[must_use]
425 pub fn len(&self) -> usize {
426 match &self.body {
427 Body::Packed(lp) => lp.len() / 2,
428 Body::Table(t) => t.members.len(),
429 }
430 }
431
432 #[must_use]
434 pub fn is_empty(&self) -> bool {
435 self.len() == 0
436 }
437
438 #[must_use]
440 pub const fn encoding(&self) -> Encoding {
441 match &self.body {
442 Body::Packed(_) => Encoding::Listpack,
443 Body::Table(_) => Encoding::Skiplist,
444 }
445 }
446
447 #[inline]
454 pub(crate) fn packed_bytes(&self) -> Option<&[u8]> {
455 match &self.body {
456 Body::Packed(lp) => Some(lp.as_bytes()),
457 Body::Table(_) => None,
458 }
459 }
460
461 #[must_use]
463 pub fn memory_bytes(&self) -> usize {
464 match &self.body {
465 Body::Packed(lp) => lp.byte_len(),
466 Body::Table(t) => t.members.memory_bytes() + t.order.bytes(),
467 }
468 }
469
470 #[must_use]
474 pub fn score(&self, member: &[u8]) -> Option<f64> {
475 match &self.body {
476 Body::Packed(lp) => {
477 let at = lp.find(member, 2)?;
478 lp.get(at + 1).map(score_of)
479 }
480 Body::Table(t) => t.members.get(member).copied(),
481 }
482 }
483
484 pub fn add(&mut self, member: &[u8], score: f64, limits: &Limits) -> Added {
486 if let Body::Packed(lp) = &mut self.body {
487 if let Some(at) = lp.find(member, 2) {
488 let old = lp.get(at + 1).map_or(0.0, score_of);
489 if cmp_score(old, score) == Ordering::Equal {
490 return Added::Same;
491 }
492 lp.delete(at, 2);
493 packed_insert(lp, member, score);
494 return Added::Changed;
495 }
496 if lp.len() / 2 < limits.max_listpack_entries
497 && member.len() <= limits.max_listpack_value
498 {
499 packed_insert(lp, member, score);
500 return Added::New;
501 }
502 self.promote();
503 }
504 let Body::Table(t) = &mut self.body else {
505 unreachable!("promoted above")
506 };
507 t.add(member, score)
508 }
509
510 pub fn remove(&mut self, member: &[u8]) -> bool {
514 match &mut self.body {
515 Body::Packed(lp) => match lp.find(member, 2) {
516 Some(at) => lp.delete(at, 2),
517 None => false,
518 },
519 Body::Table(t) => {
520 let Some(row) = t.members.index_of(member) else {
521 return false;
522 };
523 let score = *t.members.get(member).expect("just found");
524 let rank = t.rank_of(row as u32, score, member);
525 t.take(rank, row);
526 true
527 }
528 }
529 }
530
531 #[must_use]
535 pub fn rank(&self, member: &[u8]) -> Option<usize> {
536 match &self.body {
537 Body::Packed(lp) => lp.find(member, 2).map(|at| at / 2),
538 Body::Table(t) => {
539 let row = t.members.index_of(member)?;
540 let score = *t.members.get(member)?;
541 Some(t.rank_of(row as u32, score, member))
542 }
543 }
544 }
545
546 #[must_use]
552 pub fn at(&self, rank: usize) -> Option<(Member<'_>, f64)> {
553 match &self.body {
554 Body::Packed(lp) => {
555 let member = lp.get(rank * 2)?;
556 let score = lp.get(rank * 2 + 1).map(score_of)?;
557 Some((member, score))
558 }
559 Body::Table(t) => {
560 let row = t.order.row_at(rank)?;
561 let (name, score) = t.members.at(row as usize)?;
562 Some((Member::Str(name), *score))
563 }
564 }
565 }
566
567 pub fn remove_at(&mut self, rank: usize) -> bool {
569 match &mut self.body {
570 Body::Packed(lp) => lp.delete(rank * 2, 2),
571 Body::Table(t) => {
572 let Some(row) = t.order.row_at(rank) else {
573 return false;
574 };
575 t.take(rank, row as usize);
576 true
577 }
578 }
579 }
580
581 #[must_use]
587 pub fn pick(&self, at: usize) -> Option<(Member<'_>, f64)> {
588 match &self.body {
589 Body::Packed(_) => self.at(at),
590 Body::Table(t) => {
591 let (name, score) = t.members.at(at)?;
592 Some((Member::Str(name), *score))
593 }
594 }
595 }
596
597 pub fn walk<F: FnMut(Member<'_>, f64)>(&self, from: usize, count: usize, rev: bool, mut f: F) {
603 let len = self.len();
604 if from >= len || count == 0 {
605 return;
606 }
607 let count = count.min(if rev { from + 1 } else { len - from });
608 match &self.body {
609 Body::Packed(lp) => {
610 for i in 0..count {
611 let at = if rev { from - i } else { from + i };
612 let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
613 return;
614 };
615 f(m, score_of(s));
616 }
617 }
618 Body::Table(t) => {
619 if rev {
620 for row in t.order.iter_back_from(from).take(count) {
621 let Some((name, score)) = t.members.at(row as usize) else {
622 return;
623 };
624 f(Member::Str(name), *score);
625 }
626 } else {
627 for row in t.order.iter_from(from).take(count) {
628 let Some((name, score)) = t.members.at(row as usize) else {
629 return;
630 };
631 f(Member::Str(name), *score);
632 }
633 }
634 }
635 }
636 }
637
638 #[must_use]
644 pub fn window_by_score(&self, min: Bound, max: Bound) -> Range<usize> {
645 let start = self.seek(|score, _| {
646 let before = match cmp_score(score, min.at) {
648 Ordering::Less => true,
649 Ordering::Equal => min.open,
650 Ordering::Greater => false,
651 };
652 if before {
653 Ordering::Greater
654 } else {
655 Ordering::Less
656 }
657 });
658 let end = self.seek(|score, _| {
659 let inside = match cmp_score(score, max.at) {
660 Ordering::Less => true,
661 Ordering::Equal => !max.open,
662 Ordering::Greater => false,
663 };
664 if inside {
665 Ordering::Greater
666 } else {
667 Ordering::Less
668 }
669 });
670 start..end.max(start)
671 }
672
673 #[must_use]
678 pub fn window_by_lex(&self, min: Lex<'_>, max: Lex<'_>) -> Range<usize> {
679 let start = self.seek(|_, member| match min {
680 Lex::Min => Ordering::Less,
681 Lex::Max => Ordering::Greater,
682 Lex::Incl(at) => {
683 if member < at {
684 Ordering::Greater
685 } else {
686 Ordering::Less
687 }
688 }
689 Lex::Excl(at) => {
690 if member <= at {
691 Ordering::Greater
692 } else {
693 Ordering::Less
694 }
695 }
696 });
697 let end = self.seek(|_, member| match max {
698 Lex::Min => Ordering::Less,
699 Lex::Max => Ordering::Greater,
700 Lex::Incl(at) => {
701 if member <= at {
702 Ordering::Greater
703 } else {
704 Ordering::Less
705 }
706 }
707 Lex::Excl(at) => {
708 if member < at {
709 Ordering::Greater
710 } else {
711 Ordering::Less
712 }
713 }
714 });
715 start..end.max(start)
716 }
717
718 fn seek<F: FnMut(f64, &[u8]) -> Ordering>(&self, mut probe: F) -> usize {
724 match &self.body {
725 Body::Packed(lp) => {
726 let mut digits = [0u8; DIGITS_MAX];
727 let mut at = 0;
728 while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
729 let bytes = bytes_of(m, &mut digits);
730 if probe(score_of(s), bytes) != Ordering::Greater {
731 break;
732 }
733 at += 1;
734 }
735 at
736 }
737 Body::Table(t) => t.order.seek(|row| {
738 let (name, score) = t.members.at(row as usize).expect("a row the tree holds");
739 probe(*score, name)
740 }),
741 }
742 }
743
744 pub fn scan<F: FnMut(Member<'_>, f64)>(
746 &self,
747 cursor: Cursor,
748 count: usize,
749 mut f: F,
750 ) -> Cursor {
751 match &self.body {
752 Body::Packed(lp) => {
753 for at in 0..lp.len() / 2 {
754 let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) else {
755 break;
756 };
757 f(m, score_of(s));
758 }
759 Cursor::END
760 }
761 Body::Table(t) => t.members.scan(cursor, count, |name, score| {
762 f(Member::Str(name), *score);
763 }),
764 }
765 }
766
767 #[must_use]
784 pub fn from_elements(members: Elements<f64>, limits: &Limits) -> Option<Zset> {
785 let n = members.len();
786 if n == 0 {
787 return None;
788 }
789 let mut rows: Vec<u32> = (0..n as u32).collect();
793 rows.sort_unstable_by(|&a, &b| {
794 let (a_name, a_score) = members.at(a as usize).expect("in range");
795 let (b_name, b_score) = members.at(b as usize).expect("in range");
796 cmp_key(*a_score, a_name, *b_score, b_name)
797 });
798 let packable = n <= limits.max_listpack_entries
801 && rows.iter().all(|&r| {
802 members.at(r as usize).expect("in range").0.len() <= limits.max_listpack_value
803 });
804 if packable {
805 let mut lp = Listpack::new();
806 let mut score_buf = [0u8; DOUBLE_MAX];
807 for &row in &rows {
808 let (name, score) = members.at(row as usize).expect("in range");
809 lp.push(name);
810 lp.push(write_double(&mut score_buf, *score));
811 }
812 return Some(Zset {
813 body: Body::Packed(lp),
814 });
815 }
816 let mut order = Rank::new();
817 for (at, &row) in rows.iter().enumerate() {
820 order.insert_at(at, row);
821 }
822 Some(Zset {
823 body: Body::Table(Table { members, order }),
824 })
825 }
826
827 fn promote(&mut self) {
829 let Body::Packed(lp) = &self.body else { return };
830 let n = lp.len() / 2;
831 let mut table = Table {
832 members: Elements::with_capacity(n.next_power_of_two().max(16)),
833 order: Rank::new(),
834 };
835 let mut digits = [0u8; DIGITS_MAX];
843 let mut steps = lp.iter();
844 while let (Some(m), Some(s)) = (steps.next(), steps.next()) {
845 let bytes = bytes_of(m, &mut digits);
846 let row = table.members.len() as u32;
847 if table.members.insert(bytes, score_of(s)).is_err() {
848 break;
849 }
850 table.order.insert_at(row as usize, row);
851 }
852 self.body = Body::Table(table);
853 }
854}
855
856fn packed_insert(lp: &mut Listpack, member: &[u8], score: f64) {
858 let mut score_buf = [0u8; DOUBLE_MAX];
862 let text = write_double(&mut score_buf, score);
863 let mut digits = [0u8; DIGITS_MAX];
864 let mut at = 0;
865 while let (Some(m), Some(s)) = (lp.get(at * 2), lp.get(at * 2 + 1)) {
866 let bytes = bytes_of(m, &mut digits);
867 if cmp_key(score, member, score_of(s), bytes) == Ordering::Less {
868 break;
869 }
870 at += 1;
871 }
872 if at * 2 == lp.len() {
873 lp.push(member);
874 lp.push(text);
875 } else {
876 lp.insert(at * 2, member);
877 lp.insert(at * 2 + 1, text);
878 }
879}
880
881impl Table {
882 fn rank_of(&self, row: u32, score: f64, member: &[u8]) -> usize {
887 let members = &self.members;
888 self.order.seek(|other| {
889 if other == row {
890 return Ordering::Equal;
891 }
892 let (name, at) = members.at(other as usize).expect("a row the tree holds");
893 cmp_key(score, member, *at, name)
894 })
895 }
896
897 fn add(&mut self, member: &[u8], score: f64) -> Added {
898 if let Some(row) = self.members.index_of(member) {
899 let old = *self.members.at(row).map_or(&0.0, |(_, s)| s);
900 if cmp_score(old, score) == Ordering::Equal {
901 return Added::Same;
902 }
903 let was = self.rank_of(row as u32, old, member);
907 self.order.remove_at(was);
908 if let Some(at) = self.members.at_mut(row) {
909 *at = score;
910 }
911 let now = self.rank_of(row as u32, score, member);
912 self.order.insert_at(now, row as u32);
913 return Added::Changed;
914 }
915 let row = self.members.len() as u32;
916 if self.members.insert(member, score).is_err() {
917 return Added::Full;
918 }
919 let at = self.rank_of(row, score, member);
920 self.order.insert_at(at, row);
921 Added::New
922 }
923
924 fn take(&mut self, rank: usize, row: usize) {
926 let last = self.members.len() - 1;
927 let moved = if last == row {
931 None
932 } else {
933 let (name, score) = self.members.at(last).expect("the last row");
934 let at = {
938 let members = &self.members;
939 let score = *score;
940 self.order.seek(|other| {
941 if other as usize == last {
942 return Ordering::Equal;
943 }
944 let (other_name, other_score) =
945 members.at(other as usize).expect("a row the tree holds");
946 cmp_key(score, name, *other_score, other_name)
947 })
948 };
949 Some(at)
950 };
951 self.order.remove_at(rank);
952 self.members.remove_at(row);
953 if let Some(at) = moved {
954 let at = if at > rank { at - 1 } else { at };
957 self.order.set_at(at, row as u32);
958 }
959 }
960}
961
962#[cfg(test)]
963mod tests {
964 use super::*;
965
966 fn packed(pairs: &[(&[u8], &str)]) -> Listpack {
968 let mut lp = Listpack::new();
969 for (member, score) in pairs {
970 lp.push(member);
971 lp.push(score.as_bytes());
972 }
973 lp
974 }
975
976 #[test]
983 fn a_payload_in_this_layout_is_taken_whole() {
984 let lp = packed(&[(b"a", "1"), (b"b", "2"), (b"c", "2.5")]);
985 let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("in order and inside the limits");
986 assert_eq!(z.encoding(), Encoding::Listpack);
987 assert_eq!(z.len(), 3);
988 assert_eq!(z.score(b"a"), Some(1.0));
989 assert_eq!(z.score(b"c"), Some(2.5));
990 assert_eq!(z.rank(b"b"), Some(1));
991 assert_eq!(z.score(b"missing"), None);
992
993 let lp = packed(&[(b"10", "1"), (b"9", "2")]);
998 let z = Zset::from_packed(lp, &Limits::DEFAULT).expect("sorted by score");
999 assert_eq!(z.score(b"10"), Some(1.0));
1000 assert_eq!(z.score(b"9"), Some(2.0));
1001 assert_eq!(z.rank(b"9"), Some(1));
1002 }
1003
1004 #[test]
1011 fn a_blob_this_band_cannot_hold_is_handed_back() {
1012 let small = Limits {
1013 max_listpack_entries: 4,
1014 max_listpack_value: 8,
1015 };
1016 for (why, lp) in [
1017 ("out of order by score", packed(&[(b"a", "2"), (b"b", "1")])),
1018 (
1019 "out of order by member",
1020 packed(&[(b"b", "1"), (b"a", "1")]),
1021 ),
1022 (
1023 "the same member twice",
1024 packed(&[(b"a", "1"), (b"a", "1"), (b"b", "2")]),
1025 ),
1026 ("a score that is not a number", packed(&[(b"a", "no")])),
1027 (
1028 "a member past the value limit",
1029 packed(&[(b"aaaaaaaaaa", "1")]),
1030 ),
1031 (
1032 "more members than the band takes",
1033 packed(&[
1034 (b"a", "1"),
1035 (b"b", "2"),
1036 (b"c", "3"),
1037 (b"d", "4"),
1038 (b"e", "5"),
1039 ]),
1040 ),
1041 ("nothing in it at all", packed(&[])),
1042 ] {
1043 assert!(
1044 Zset::from_packed(lp, &small).is_err(),
1045 "{why} should be handed back"
1046 );
1047 }
1048
1049 let mut odd = Listpack::new();
1052 odd.push(b"a");
1053 assert!(Zset::from_packed(odd, &Limits::DEFAULT).is_err());
1054 }
1055
1056 #[test]
1058 fn a_hint_past_the_band_starts_on_the_table() {
1059 let big = Zset::with_hint(Limits::DEFAULT.max_listpack_entries + 1, &Limits::DEFAULT);
1060 assert_eq!(big.encoding(), Encoding::Skiplist);
1061 assert!(big.is_empty());
1062
1063 let small = Zset::with_hint(Limits::DEFAULT.max_listpack_entries, &Limits::DEFAULT);
1065 assert_eq!(small.encoding(), Encoding::Listpack);
1066
1067 let mut z = Zset::with_hint(1_000_000, &Limits::DEFAULT);
1070 z.add(b"b", 2.0, &Limits::DEFAULT);
1071 z.add(b"a", 1.0, &Limits::DEFAULT);
1072 assert_eq!(z.len(), 2);
1073 assert_eq!(z.rank(b"a"), Some(0));
1074 assert_eq!(z.rank(b"b"), Some(1));
1075 }
1076
1077 #[test]
1088 #[ignore = "a measurement, run it by name"]
1089 fn measure_bytes_per_entry() {
1090 let mut lp = Zset::new();
1092 let mut lp_payload = 0usize;
1093 for i in 0..128 {
1094 let m = format!("member:{i:09}");
1095 lp_payload += m.len();
1096 lp.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
1097 }
1098 println!(
1099 "packed n=128 total={} payload={lp_payload} overhead_per_entry={:.2}",
1100 lp.memory_bytes(),
1101 (lp.memory_bytes() as f64 - lp_payload as f64) / 128.0
1102 );
1103 for n in [10_000usize, 100_000, 600_000, 1_000_000] {
1107 let (z, payload) = filled(n);
1108 let total = z.memory_bytes();
1109 let scores = n * 8;
1110 let (slots, rows, names, tree) = match &z.body {
1111 Body::Table(t) => (
1112 t.members.slot_bytes(),
1113 t.members.row_bytes(),
1114 t.members.name_bytes(),
1115 t.order.bytes(),
1116 ),
1117 Body::Packed(_) => (0, 0, 0, 0),
1118 };
1119 let per = |b: usize| b as f64 / n as f64;
1120 println!(
1121 "table n={n} total={total} slots={:.2}/e rows={:.2}/e names={:.2}/e tree={:.2}/e overhead_per_entry={:.2}",
1122 per(slots),
1123 per(rows),
1124 per(names),
1125 per(tree),
1126 (total as f64 - payload as f64 - scores as f64) / n as f64
1127 );
1128 }
1129 }
1130
1131 fn filled(n: usize) -> (Zset, usize) {
1134 let mut z = Zset::new();
1135 let mut payload = 0usize;
1136 for i in 0..n {
1137 let m = format!("member:{i:09}");
1138 payload += m.len();
1139 z.add(m.as_bytes(), i as f64, &Limits::DEFAULT);
1140 }
1141 (z, payload)
1142 }
1143
1144 #[test]
1153 fn a_large_sorted_set_does_not_hold_much_more_than_it_stores() {
1154 let n = 40_000usize;
1155 let (z, payload) = filled(n);
1156 let Body::Table(t) = &z.body else {
1157 panic!("forty thousand members is not a listpack");
1158 };
1159 assert!(
1163 t.members.row_bytes() < n * 25,
1164 "the row and score arrays hold {} for {n} members",
1165 t.members.row_bytes()
1166 );
1167 assert!(
1168 t.members.name_bytes() < payload + payload / 4,
1169 "the name blob holds {} for {payload} bytes of names",
1170 t.members.name_bytes()
1171 );
1172 assert!(
1178 t.order.bytes() < n * 4,
1179 "the tree holds {} for {n} members",
1180 t.order.bytes()
1181 );
1182 }
1183
1184 fn built(pairs: &[(&str, f64)], limits: &Limits) -> Zset {
1186 let mut z = Zset::new();
1187 for (m, s) in pairs {
1188 z.add(m.as_bytes(), *s, limits);
1189 }
1190 z
1191 }
1192
1193 fn listed(z: &Zset) -> Vec<(String, f64)> {
1195 let mut out = Vec::new();
1196 let mut digits = [0u8; DIGITS_MAX];
1197 z.walk(0, z.len(), false, |m, s| {
1198 let bytes = bytes_of(m, &mut digits).to_vec();
1199 out.push((String::from_utf8(bytes).unwrap(), s));
1200 });
1201 out
1202 }
1203
1204 fn model(pairs: &[(&str, f64)]) -> Vec<(String, f64)> {
1206 let mut last: Vec<(String, f64)> = Vec::new();
1207 for (m, s) in pairs {
1208 match last.iter_mut().find(|(name, _)| name == m) {
1209 Some(row) => row.1 = *s,
1210 None => last.push(((*m).to_string(), *s)),
1211 }
1212 }
1213 last.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
1214 last
1215 }
1216
1217 const PACKED: Limits = Limits::DEFAULT;
1218 const TABLE: Limits = Limits {
1219 max_listpack_entries: 0,
1220 max_listpack_value: 64,
1221 };
1222
1223 #[test]
1224 fn an_empty_set_answers_nothing() {
1225 let z = Zset::new();
1226 assert_eq!(z.len(), 0);
1227 assert!(z.is_empty());
1228 assert_eq!(z.encoding(), Encoding::Listpack);
1229 assert_eq!(z.score(b"nobody"), None);
1230 assert_eq!(z.rank(b"nobody"), None);
1231 assert!(z.at(0).is_none());
1232 assert_eq!(
1233 z.window_by_score(Bound::closed(f64::NEG_INFINITY), Bound::closed(0.0)),
1234 0..0
1235 );
1236 }
1237
1238 #[test]
1239 fn both_bands_put_members_in_the_same_order() {
1240 let pairs = [("c", 3.0), ("a", 1.0), ("b", 2.0), ("d", 2.0), ("e", -1.5)];
1241 for limits in [&PACKED, &TABLE] {
1242 let z = built(&pairs, limits);
1243 assert_eq!(listed(&z), model(&pairs), "{:?}", z.encoding());
1244 assert_eq!(z.len(), 5);
1245 }
1246 }
1247
1248 #[test]
1249 fn a_tie_on_score_is_broken_by_the_member() {
1250 for limits in [&PACKED, &TABLE] {
1251 let pairs = [
1252 ("beta", 1.0),
1253 ("alpha", 1.0),
1254 ("gamma", 1.0),
1255 ("Alpha", 1.0),
1256 ];
1257 let z = built(&pairs, limits);
1258 let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
1259 assert_eq!(names, ["Alpha", "alpha", "beta", "gamma"]);
1260 }
1261 }
1262
1263 #[test]
1264 fn adding_a_member_again_moves_it_rather_than_adding_it() {
1265 for limits in [&PACKED, &TABLE] {
1266 let mut z = built(&[("a", 1.0), ("b", 2.0), ("c", 3.0)], limits);
1267 assert_eq!(z.add(b"a", 1.0, limits), Added::Same);
1268 assert_eq!(z.add(b"a", 9.0, limits), Added::Changed);
1269 assert_eq!(z.len(), 3);
1270 assert_eq!(z.score(b"a"), Some(9.0));
1271 assert_eq!(z.rank(b"a"), Some(2));
1272 assert_eq!(listed(&z).last().unwrap().0, "a");
1273 }
1274 }
1275
1276 #[test]
1277 fn a_removal_leaves_every_other_rank_right() {
1278 for limits in [&PACKED, &TABLE] {
1279 let pairs: Vec<(&str, f64)> =
1280 vec![("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0), ("e", 5.0)];
1281 let mut z = built(&pairs, limits);
1282 assert!(z.remove(b"c"));
1283 assert!(!z.remove(b"c"));
1284 assert_eq!(z.len(), 4);
1285 assert_eq!(z.rank(b"a"), Some(0));
1286 assert_eq!(z.rank(b"d"), Some(2));
1287 assert_eq!(z.rank(b"e"), Some(3));
1288 assert_eq!(z.score(b"c"), None);
1289 }
1290 }
1291
1292 #[test]
1296 fn removing_from_the_middle_renumbers_the_last_element() {
1297 let limits = &TABLE;
1298 let mut z = Zset::new();
1299 for i in 0..64u32 {
1300 z.add(format!("m{i:03}").as_bytes(), f64::from(i), limits);
1301 }
1302 for i in 0..63u32 {
1305 assert!(z.remove(format!("m{i:03}").as_bytes()));
1306 assert_eq!(z.len() as u32, 63 - i);
1307 for j in i + 1..64 {
1310 let name = format!("m{j:03}");
1311 assert_eq!(
1312 z.score(name.as_bytes()),
1313 Some(f64::from(j)),
1314 "score of {name}"
1315 );
1316 assert_eq!(
1317 z.rank(name.as_bytes()),
1318 Some((j - i - 1) as usize),
1319 "rank of {name} after {i}"
1320 );
1321 }
1322 }
1323 }
1324
1325 #[test]
1326 fn a_set_promotes_when_it_outgrows_the_packed_band() {
1327 let limits = Limits {
1328 max_listpack_entries: 4,
1329 max_listpack_value: 64,
1330 };
1331 let mut z = Zset::new();
1332 for i in 0..4u32 {
1333 z.add(format!("m{i}").as_bytes(), f64::from(i), &limits);
1334 }
1335 assert_eq!(z.encoding(), Encoding::Listpack);
1336 z.add(b"m4", 4.0, &limits);
1337 assert_eq!(z.encoding(), Encoding::Skiplist);
1338 assert_eq!(z.len(), 5);
1339 let names: Vec<String> = listed(&z).into_iter().map(|(m, _)| m).collect();
1340 assert_eq!(names, ["m0", "m1", "m2", "m3", "m4"]);
1341 z.remove(b"m4");
1343 z.remove(b"m3");
1344 assert_eq!(z.encoding(), Encoding::Skiplist);
1345 }
1346
1347 #[test]
1348 fn a_member_too_long_for_the_packed_band_promotes_on_its_own() {
1349 let limits = Limits {
1350 max_listpack_entries: 128,
1351 max_listpack_value: 8,
1352 };
1353 let mut z = Zset::new();
1354 z.add(b"short", 1.0, &limits);
1355 assert_eq!(z.encoding(), Encoding::Listpack);
1356 z.add(b"a member well past eight bytes", 2.0, &limits);
1357 assert_eq!(z.encoding(), Encoding::Skiplist);
1358 assert_eq!(z.len(), 2);
1359 assert_eq!(z.score(b"a member well past eight bytes"), Some(2.0));
1360 }
1361
1362 #[test]
1363 fn a_score_range_finds_both_of_its_ends() {
1364 for limits in [&PACKED, &TABLE] {
1365 let pairs = [("a", 1.0), ("b", 2.0), ("c", 2.0), ("d", 3.0), ("e", 4.0)];
1366 let z = built(&pairs, limits);
1367 assert_eq!(
1368 z.window_by_score(Bound::closed(2.0), Bound::closed(3.0)),
1369 1..4
1370 );
1371 assert_eq!(
1372 z.window_by_score(Bound::open(2.0), Bound::closed(3.0)),
1373 3..4
1374 );
1375 assert_eq!(
1376 z.window_by_score(Bound::closed(2.0), Bound::open(3.0)),
1377 1..3
1378 );
1379 assert_eq!(z.window_by_score(Bound::open(1.0), Bound::open(4.0)), 1..4);
1380 assert_eq!(
1381 z.window_by_score(
1382 Bound::closed(f64::NEG_INFINITY),
1383 Bound::closed(f64::INFINITY)
1384 ),
1385 0..5
1386 );
1387 assert_eq!(
1389 z.window_by_score(Bound::closed(3.0), Bound::closed(2.0)),
1390 3..3
1391 );
1392 assert_eq!(
1393 z.window_by_score(Bound::closed(9.0), Bound::closed(10.0)),
1394 5..5
1395 );
1396 }
1397 }
1398
1399 #[test]
1400 fn a_member_range_orders_by_member_when_the_scores_are_equal() {
1401 for limits in [&PACKED, &TABLE] {
1402 let pairs = [("a", 0.0), ("b", 0.0), ("c", 0.0), ("d", 0.0), ("e", 0.0)];
1403 let z = built(&pairs, limits);
1404 assert_eq!(z.window_by_lex(Lex::Min, Lex::Max), 0..5);
1405 assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Incl(b"d")), 1..4);
1406 assert_eq!(z.window_by_lex(Lex::Excl(b"b"), Lex::Excl(b"d")), 2..3);
1407 assert_eq!(z.window_by_lex(Lex::Incl(b"b"), Lex::Excl(b"c")), 1..2);
1408 assert_eq!(z.window_by_lex(Lex::Excl(b"e"), Lex::Max), 5..5);
1409 assert_eq!(z.window_by_lex(Lex::Min, Lex::Excl(b"a")), 0..0);
1410 }
1411 }
1412
1413 #[test]
1414 fn a_walk_can_go_backwards_and_stops_where_it_is_told() {
1415 for limits in [&PACKED, &TABLE] {
1416 let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0), ("d", 4.0)];
1417 let z = built(&pairs, limits);
1418 let mut seen = Vec::new();
1419 let mut digits = [0u8; DIGITS_MAX];
1420 z.walk(3, 2, true, |m, _| {
1421 seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
1422 });
1423 assert_eq!(seen, ["d", "c"]);
1424 let mut seen = Vec::new();
1425 z.walk(1, 99, false, |m, _| {
1426 seen.push(String::from_utf8(bytes_of(m, &mut digits).to_vec()).unwrap());
1427 });
1428 assert_eq!(seen, ["b", "c", "d"]);
1429 let mut count = 0;
1431 z.walk(4, 1, false, |_, _| count += 1);
1432 assert_eq!(count, 0);
1433 }
1434 }
1435
1436 #[test]
1437 fn taking_from_a_rank_takes_the_right_one() {
1438 for limits in [&PACKED, &TABLE] {
1439 let pairs = [("a", 1.0), ("b", 2.0), ("c", 3.0)];
1440 let mut z = built(&pairs, limits);
1441 assert!(z.remove_at(0));
1442 assert_eq!(z.len(), 2);
1443 assert_eq!(z.rank(b"b"), Some(0));
1444 assert!(z.remove_at(1));
1445 assert_eq!(z.score(b"c"), None);
1446 assert!(!z.remove_at(5));
1447 }
1448 }
1449
1450 #[test]
1453 fn a_run_of_everything_agrees_with_a_model() {
1454 for limits in [
1455 &PACKED,
1456 &Limits {
1457 max_listpack_entries: 8,
1458 max_listpack_value: 64,
1459 },
1460 ] {
1461 let mut z = Zset::new();
1462 let mut model: Vec<(String, f64)> = Vec::new();
1463 let mut seed = 0x8765_4321_9ABC_DEF0u64;
1464 let mut roll = || {
1465 seed ^= seed << 13;
1466 seed ^= seed >> 7;
1467 seed ^= seed << 17;
1468 seed
1469 };
1470 for round in 0..3_000 {
1471 let name = format!("m{:02}", roll() % 40);
1472 let score = (roll() % 7) as f64 - 3.0;
1475 if round % 5 == 4 {
1476 let hit = z.remove(name.as_bytes());
1477 let was = model.iter().position(|(m, _)| *m == name);
1478 assert_eq!(hit, was.is_some());
1479 if let Some(at) = was {
1480 model.remove(at);
1481 }
1482 } else {
1483 z.add(name.as_bytes(), score, limits);
1484 match model.iter_mut().find(|(m, _)| *m == name) {
1485 Some(row) => row.1 = score,
1486 None => model.push((name, score)),
1487 }
1488 }
1489 if round % 97 == 0 {
1490 let mut want = model.clone();
1491 want.sort_by(|a, b| cmp_key(a.1, a.0.as_bytes(), b.1, b.0.as_bytes()));
1492 assert_eq!(listed(&z), want, "round {round}");
1493 for (at, (m, s)) in want.iter().enumerate() {
1494 assert_eq!(z.rank(m.as_bytes()), Some(at), "rank of {m}");
1495 assert_eq!(z.score(m.as_bytes()), Some(*s), "score of {m}");
1496 }
1497 }
1498 }
1499 }
1500 }
1501
1502 #[test]
1503 fn a_big_set_costs_what_the_tree_said_it_would() {
1504 let limits = &TABLE;
1505 let mut z = Zset::new();
1506 let n = 100_000u32;
1507 for i in 0..n {
1508 z.add(format!("member:{i:08}").as_bytes(), f64::from(i), limits);
1509 }
1510 assert_eq!(z.len(), n as usize);
1511 assert_eq!(z.rank(b"member:00050000"), Some(50_000));
1512 assert_eq!(z.at(0).map(|(_, s)| s), Some(0.0));
1513 assert_eq!(z.at(n as usize - 1).map(|(_, s)| s), Some(f64::from(n - 1)));
1514 let Body::Table(t) = &z.body else {
1522 panic!("a hundred thousand members is not a listpack")
1523 };
1524 let per_order = t.order.bytes() as f64 / f64::from(n);
1525 assert!(per_order < 3.4, "{per_order} bytes an element in the tree");
1526 let per = z.memory_bytes() as f64 / f64::from(n);
1527 assert!(per < 70.0, "{per} bytes a member all in");
1528 }
1529
1530 fn round_trip(z: &Zset) -> Zset {
1532 let mut buf = Vec::new();
1533 z.freeze(&mut buf);
1534 let back = Zset::thaw(&buf).expect("what freeze wrote");
1535 assert_eq!(back.len(), z.len(), "the member count");
1536 assert_eq!(back.encoding(), z.encoding(), "the band");
1537 assert_eq!(listed(&back), listed(z), "the members in rank order");
1538 back
1539 }
1540
1541 #[test]
1542 fn a_frozen_sorted_set_comes_back_in_the_band_it_left() {
1543 let pairs: Vec<(String, f64)> = (0..40)
1544 .map(|i| (format!("member:{i:04}"), f64::from(i) * 1.5 - 12.0))
1545 .collect();
1546 let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1547
1548 let packed = round_trip(&built(&refs, &PACKED));
1549 assert_eq!(packed.encoding(), Encoding::Listpack);
1550 let table = round_trip(&built(&refs, &TABLE));
1551 assert_eq!(table.encoding(), Encoding::Skiplist);
1552
1553 assert_eq!(listed(&packed), model(&refs));
1556 assert_eq!(listed(&table), model(&refs));
1557
1558 round_trip(&Zset::new());
1559 }
1560
1561 #[test]
1562 fn a_score_survives_the_trip_exactly() {
1563 let scores = [
1568 0.0,
1569 -0.0,
1570 1.0 / 3.0,
1571 -1.0 / 3.0,
1572 f64::MIN_POSITIVE,
1573 f64::MAX,
1574 f64::MIN,
1575 9_007_199_254_740_993.0,
1576 f64::INFINITY,
1577 f64::NEG_INFINITY,
1578 ];
1579 let mut z = Zset::new();
1580 for (i, s) in scores.iter().enumerate() {
1581 z.add(format!("member:{i:04}").as_bytes(), *s, &TABLE);
1582 }
1583 let back = round_trip(&z);
1584 for (i, s) in scores.iter().enumerate() {
1585 let member = format!("member:{i:04}");
1586 assert_eq!(
1587 back.score(member.as_bytes()).map(f64::to_bits),
1588 Some(s.to_bits()),
1589 "the score of {member}"
1590 );
1591 }
1592 }
1593
1594 #[test]
1595 fn a_sorted_set_that_came_back_still_takes_members_and_ranks_them() {
1596 let pairs: Vec<(String, f64)> = (0..200)
1597 .map(|i| (format!("member:{i:04}"), f64::from(i)))
1598 .collect();
1599 let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1600 let mut back = round_trip(&built(&refs, &TABLE));
1601
1602 assert_eq!(back.add(b"middle", 99.5, &TABLE), Added::New);
1605 assert_eq!(back.add(b"first", -1.0, &TABLE), Added::New);
1606 assert_eq!(back.add(b"last", 1000.0, &TABLE), Added::New);
1607 assert_eq!(back.add(b"member:0000", 500.0, &TABLE), Added::Changed);
1608
1609 assert_eq!(back.len(), 203);
1610 assert_eq!(back.rank(b"first"), Some(0));
1611 assert_eq!(back.rank(b"last"), Some(202));
1612 assert_eq!(back.rank(b"middle"), Some(100));
1615 assert_eq!(back.score(b"member:0000"), Some(500.0));
1616 assert!(back.remove(b"middle"));
1617 assert_eq!(back.rank(b"last"), Some(201));
1618 }
1619
1620 #[test]
1621 fn a_frozen_sorted_set_that_arrives_damaged_is_an_error_and_not_a_panic() {
1622 let pairs: Vec<(String, f64)> = (0..200)
1623 .map(|i| (format!("member:{i:04}"), f64::from(i)))
1624 .collect();
1625 let refs: Vec<(&str, f64)> = pairs.iter().map(|(m, s)| (m.as_str(), *s)).collect();
1626 let z = built(&refs, &TABLE);
1627 let mut buf = Vec::new();
1628 z.freeze(&mut buf);
1629
1630 assert!(Zset::thaw(&[]).is_err(), "nothing at all");
1631 assert!(Zset::thaw(&[99]).is_err(), "a form nobody wrote");
1632 for cut in 1..buf.len().min(64) {
1633 assert!(Zset::thaw(&buf[..cut]).is_err(), "cut at {cut}");
1634 }
1635 let mut lying = vec![FORM_MEMBERS];
1637 frozen::put_uint(&mut lying, u64::MAX);
1638 assert!(Zset::thaw(&lying).is_err(), "a count nobody could hold");
1639 let mut twice = vec![FORM_MEMBERS];
1642 frozen::put_uint(&mut twice, 2);
1643 for _ in 0..2 {
1644 frozen::put_bytes(&mut twice, b"member");
1645 frozen::put_f64(&mut twice, 1.0);
1646 }
1647 assert!(Zset::thaw(&twice).is_err(), "a member written twice");
1648 }
1649}