1use yo_common::num::DIGITS_MAX;
59use yo_common::{Code, Error, Result};
60
61use crate::db::{Db, Holds};
62use crate::elem::Elements;
63use crate::keyspace::Keyspace;
64use crate::scan::Cursor;
65use crate::strings;
66use crate::value::{self, Kind};
67use crate::zset::{Added, Bound, Lex, Member, Zset};
68use crate::zsetops::{self, Aggregate, Op, Operand};
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
72pub enum Gate {
73 #[default]
75 Always,
76 IfMissing,
78 IfPresent,
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
84pub enum Move {
85 #[default]
87 Any,
88 Up,
90 Down,
92}
93
94#[derive(Debug, Clone, Copy, Default)]
101pub struct ZAdd {
102 pub gate: Gate,
104 pub only: Move,
106 pub changed: bool,
108}
109
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
112pub enum From {
113 Min,
115 Max,
117}
118
119#[derive(Debug, Clone, Copy)]
121pub enum By<'a> {
122 Rank {
124 start: i64,
126 stop: i64,
128 },
129 Score {
131 min: Bound,
133 max: Bound,
135 },
136 Lex {
139 min: Lex<'a>,
141 max: Lex<'a>,
143 },
144}
145
146#[derive(Debug, Clone, Copy)]
153pub struct Query<'a> {
154 pub by: By<'a>,
156 pub rev: bool,
158 pub offset: usize,
160 pub count: Option<usize>,
163}
164
165impl<'a> Query<'a> {
166 #[must_use]
168 pub const fn rank(start: i64, stop: i64) -> Query<'a> {
169 Query {
170 by: By::Rank { start, stop },
171 rev: false,
172 offset: 0,
173 count: None,
174 }
175 }
176
177 #[must_use]
179 pub const fn score(min: Bound, max: Bound) -> Query<'a> {
180 Query {
181 by: By::Score { min, max },
182 rev: false,
183 offset: 0,
184 count: None,
185 }
186 }
187
188 #[must_use]
190 pub const fn lex(min: Lex<'a>, max: Lex<'a>) -> Query<'a> {
191 Query {
192 by: By::Lex { min, max },
193 rev: false,
194 offset: 0,
195 count: None,
196 }
197 }
198
199 #[must_use]
201 pub const fn rev(mut self, rev: bool) -> Query<'a> {
202 self.rev = rev;
203 self
204 }
205
206 #[must_use]
208 pub const fn limit(mut self, offset: usize, count: Option<usize>) -> Query<'a> {
209 self.offset = offset;
210 self.count = count;
211 self
212 }
213}
214
215#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
222pub struct Window {
223 pub from: usize,
225 pub count: usize,
227 pub rev: bool,
229}
230
231fn nan() -> Error {
233 Error::new(Code::Invalid, "resulting score is not a number (NaN)")
234}
235
236impl Keyspace {
237 pub fn zadd<'m, I>(&mut self, key: &[u8], pairs: I, opts: ZAdd) -> Result<usize>
247 where
248 I: Iterator<Item = (f64, &'m [u8])> + Clone,
249 {
250 let (added, changed) = self.zadd_counts(key, pairs, opts)?;
251 Ok(if opts.changed { added + changed } else { added })
252 }
253
254 pub fn zadd_counts<'m, I>(&mut self, key: &[u8], pairs: I, opts: ZAdd) -> Result<(usize, usize)>
268 where
269 I: Iterator<Item = (f64, &'m [u8])> + Clone,
270 {
271 for (score, m) in pairs.clone() {
272 strings::check_len(key, m.len())?;
273 if score.is_nan() {
274 return Err(nan());
275 }
276 }
277 let at = match self.zset_slot(key)? {
278 Some(at) => at,
279 None => {
280 if opts.gate == Gate::IfPresent || pairs.clone().next().is_none() {
286 return Ok((0, 0));
287 }
288 self.new_zset(key)
289 }
290 };
291 let limits = self.zset_limits;
292 let z = self
293 .zsets
294 .get_mut(at)
295 .expect("the record points at its body");
296 let mut added = 0usize;
297 let mut changed = 0usize;
298 for (score, m) in pairs {
299 let Some(want) = gated(z, m, score, opts) else {
300 continue;
301 };
302 if z.add(m, score, &limits) == Added::Full {
305 continue;
306 }
307 match want {
308 Added::New => added += 1,
309 Added::Changed => changed += 1,
310 _ => {}
311 }
312 }
313 if z.is_empty() {
317 self.drop_key(key);
318 }
319 Ok((added, changed))
320 }
321
322 pub fn zincrby(
331 &mut self,
332 key: &[u8],
333 member: &[u8],
334 by: f64,
335 opts: ZAdd,
336 ) -> Result<Option<f64>> {
337 strings::check_len(key, member.len())?;
338 if by.is_nan() {
339 return Err(nan());
340 }
341 let at = match self.zset_slot(key)? {
342 Some(at) => at,
343 None => {
344 if opts.gate == Gate::IfPresent {
345 return Ok(None);
346 }
347 self.new_zset(key)
348 }
349 };
350 let limits = self.zset_limits;
351 let z = self
352 .zsets
353 .get_mut(at)
354 .expect("the record points at its body");
355 let now = z.score(member);
356 let want = now.unwrap_or(0.0) + by;
357 if want.is_nan() {
360 if z.is_empty() {
361 self.drop_key(key);
362 }
363 return Err(nan());
364 }
365 let allowed = match (now, opts.gate, opts.only) {
366 (Some(_), Gate::IfMissing, _) | (None, Gate::IfPresent, _) => false,
367 (Some(was), _, Move::Up) => want > was,
368 (Some(was), _, Move::Down) => want < was,
369 _ => true,
370 };
371 if !allowed || z.add(member, want, &limits) == Added::Full {
372 if z.is_empty() {
373 self.drop_key(key);
374 }
375 return Ok(None);
376 }
377 Ok(Some(want))
378 }
379
380 pub fn zcard(&mut self, key: &[u8]) -> Result<usize> {
382 Ok(match self.zset_slot(key)? {
383 Some(at) => self.zset_at(at).len(),
384 None => 0,
385 })
386 }
387
388 pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> Result<Option<f64>> {
390 Ok(match self.zset_slot(key)? {
391 Some(at) => self.zset_at(at).score(member),
392 None => None,
393 })
394 }
395
396 pub fn zmscore<'m>(
401 &mut self,
402 key: &[u8],
403 members: impl Iterator<Item = &'m [u8]>,
404 out: &mut Vec<Option<f64>>,
405 ) -> Result<()> {
406 out.clear();
407 let Some(at) = self.zset_slot(key)? else {
408 out.extend(members.map(|_| None));
409 return Ok(());
410 };
411 let z = self.zset_at(at);
412 out.extend(members.map(|m| z.score(m)));
413 Ok(())
414 }
415
416 pub fn zrem<'m>(
421 &mut self,
422 key: &[u8],
423 members: impl Iterator<Item = &'m [u8]>,
424 ) -> Result<usize> {
425 let Some(at) = self.zset_slot(key)? else {
426 return Ok(0);
427 };
428 let z = self
429 .zsets
430 .get_mut(at)
431 .expect("the record points at its body");
432 let mut gone = 0;
433 for m in members {
434 if z.remove(m) {
435 gone += 1;
436 }
437 }
438 if z.is_empty() {
439 self.drop_key(key);
440 }
441 Ok(gone)
442 }
443
444 pub fn zrank(&mut self, key: &[u8], member: &[u8], rev: bool) -> Result<Option<(usize, f64)>> {
449 let Some(at) = self.zset_slot(key)? else {
450 return Ok(None);
451 };
452 let z = self.zset_at(at);
453 let Some(rank) = z.rank(member) else {
454 return Ok(None);
455 };
456 let score = z.score(member).unwrap_or(0.0);
457 Ok(Some((if rev { z.len() - 1 - rank } else { rank }, score)))
458 }
459
460 pub fn zwindow(&mut self, key: &[u8], q: &Query<'_>) -> Result<Window> {
467 let Some(at) = self.zset_slot(key)? else {
468 return Ok(Window::default());
469 };
470 Ok(window(self.zset_at(at), q))
471 }
472
473 pub fn zwalk<F>(&mut self, key: &[u8], w: Window, f: F) -> Result<()>
479 where
480 F: FnMut(Member<'_>, f64),
481 {
482 let Some(at) = self.zset_slot(key)? else {
483 return Ok(());
484 };
485 self.zset_at(at).walk(w.from, w.count, w.rev, f);
486 Ok(())
487 }
488
489 pub fn zcount(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
491 Ok(self.zwindow(key, q)?.count)
492 }
493
494 pub fn zremrange(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
501 let Some(at) = self.zset_slot(key)? else {
502 return Ok(0);
503 };
504 let w = window(self.zset_at(at), q);
505 let z = self
506 .zsets
507 .get_mut(at)
508 .expect("the record points at its body");
509 let low = if w.rev { w.from + 1 - w.count } else { w.from };
512 for i in (0..w.count).rev() {
513 z.remove_at(low + i);
514 }
515 if z.is_empty() {
516 self.drop_key(key);
517 }
518 Ok(w.count)
519 }
520
521 pub fn zpop<F>(&mut self, key: &[u8], end: From, count: usize, mut f: F) -> Result<usize>
527 where
528 F: FnMut(Member<'_>, f64),
529 {
530 let Some(at) = self.zset_slot(key)? else {
531 return Ok(0);
532 };
533 let z = self
534 .zsets
535 .get_mut(at)
536 .expect("the record points at its body");
537 let count = count.min(z.len());
538 for _ in 0..count {
539 let rank = if end == From::Min { 0 } else { z.len() - 1 };
543 let Some((m, s)) = z.at(rank) else { break };
544 f(m, s);
545 z.remove_at(rank);
546 }
547 if z.is_empty() {
548 self.drop_key(key);
549 }
550 Ok(count)
551 }
552
553 pub fn zpop_one(&mut self, key: &[u8], end: From) -> Result<Option<(Vec<u8>, f64)>> {
560 let mut got = None;
561 let mut name = [0u8; yo_common::num::DIGITS_MAX];
562 self.zpop(key, end, 1, |m, s| {
563 got = Some((member_bytes(m, &mut name).to_vec(), s));
564 })?;
565 Ok(got)
566 }
567
568 pub fn zrandmember<F>(&mut self, key: &[u8], count: i64, mut f: F) -> Result<usize>
579 where
580 F: FnMut(Member<'_>, f64),
581 {
582 let Some(at) = self.zset_slot(key)? else {
583 return Ok(0);
584 };
585 let len = self.zset_at(at).len();
586 if len == 0 || count == 0 {
587 return Ok(0);
588 }
589 if count < 0 {
590 let want = count.unsigned_abs() as usize;
591 for _ in 0..want {
592 let pick = self.rng.below(len);
593 let Some((m, s)) = self.zset_at(at).pick(pick) else {
594 break;
595 };
596 f(m, s);
597 }
598 return Ok(want);
599 }
600 let want = (count as usize).min(len);
601 if want == len {
605 for i in 0..len {
606 let Some((m, s)) = self.zset_at(at).pick(i) else {
607 break;
608 };
609 f(m, s);
610 }
611 return Ok(len);
612 }
613 let mut rows = std::mem::take(&mut self.rows);
623 rows.clear();
624 yo_alloc::high_water(|| rows.extend(0..len));
625 for i in 0..want {
626 let pick = i + self.rng.below(len - i);
627 rows.swap(i, pick);
628 let Some((m, s)) = self.zset_at(at).pick(rows[i]) else {
629 break;
630 };
631 f(m, s);
632 }
633 self.rows = rows;
634 Ok(want)
635 }
636
637 pub fn zscan<F>(&mut self, key: &[u8], cursor: Cursor, count: usize, f: F) -> Result<Cursor>
643 where
644 F: FnMut(Member<'_>, f64),
645 {
646 let Some(at) = self.zset_slot(key)? else {
647 return Ok(Cursor::END);
648 };
649 Ok(self.zset_at(at).scan(cursor, count, f))
650 }
651
652 pub fn zsetop<'k, F>(
659 &mut self,
660 op: Op,
661 keys: impl Iterator<Item = &'k [u8]>,
662 weights: &[f64],
663 agg: Aggregate,
664 f: F,
665 ) -> Result<usize>
666 where
667 F: FnMut(Member<'_>, f64),
668 {
669 let slots = self.operand_slots(keys)?;
670 let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
671 let limits = self.zset_limits;
672 let Some(z) = Zset::from_elements(got, &limits) else {
673 return Ok(0);
674 };
675 let len = z.len();
676 z.walk(0, len, false, f);
677 Ok(len)
678 }
679
680 pub fn zsetop_store<'k>(
687 &mut self,
688 op: Op,
689 destination: &[u8],
690 keys: impl Iterator<Item = &'k [u8]>,
691 weights: &[f64],
692 agg: Aggregate,
693 ) -> Result<usize> {
694 let slots = self.operand_slots(keys)?;
695 let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
696 let limits = self.zset_limits;
697 Ok(self.put_zset(destination, Zset::from_elements(got, &limits)))
698 }
699
700 pub fn zintercard<'k>(
706 &mut self,
707 keys: impl Iterator<Item = &'k [u8]>,
708 limit: usize,
709 ) -> Result<usize> {
710 let slots = self.operand_slots(keys)?;
711 Ok(zsetops::intercard(&self.operands_of(&slots), limit))
712 }
713
714 pub fn zrangestore(
721 &mut self,
722 destination: &[u8],
723 source: &[u8],
724 q: &Query<'_>,
725 ) -> Result<usize> {
726 let built = match self.zset_slot(source)? {
727 None => None,
728 Some(at) => {
729 let z = self.zset_at(at);
730 let w = window(z, q);
731 let mut got = Elements::with_capacity(w.count.max(16));
732 let mut digits = [0u8; DIGITS_MAX];
733 let from = if w.rev { w.from + 1 - w.count } else { w.from };
738 z.walk(from, w.count, false, |m, s| {
739 let _ = got.insert(member_bytes(m, &mut digits), s);
740 });
741 let limits = self.zset_limits;
742 Zset::from_elements(got, &limits)
743 }
744 };
745 Ok(self.put_zset(destination, built))
746 }
747
748 fn operand_slots<'k>(
755 &mut self,
756 keys: impl Iterator<Item = &'k [u8]>,
757 ) -> Result<Vec<Option<(Kind, u32)>>> {
758 let mut out = Vec::with_capacity(keys.size_hint().0);
759 for key in keys {
760 out.push(self.live_slot_either(key, Kind::Zset, Kind::Set)?);
761 }
762 Ok(out)
763 }
764
765 fn operands_of(&self, slots: &[Option<(Kind, u32)>]) -> Vec<Operand<'_>> {
768 slots
769 .iter()
770 .map(|got| match got {
771 Some((Kind::Zset, at)) => Operand::Zset(self.zset_at(*at)),
772 Some((Kind::Set, at)) => {
773 Operand::Set(self.sets.get(*at).expect("the record points at its body"))
774 }
775 _ => Operand::Missing,
776 })
777 .collect()
778 }
779
780 pub(crate) fn put_zset(&mut self, key: &[u8], z: Option<Zset>) -> usize {
790 let Some(z) = z else {
791 self.drop_key(key);
792 return 0;
793 };
794 self.replacing(key, Kind::Zset);
795 self.free_body(key);
796 let len = z.len();
797 let at = self.zsets.insert(z);
798 let record = value::slot_record_len(false);
799 self.write_rec(key, record, |out| {
800 value::write_slot_record(out, Kind::Zset, at, None);
801 });
802 self.bodies += 1;
803 len
804 }
805
806 #[inline]
808 pub(crate) fn zset_slot(&mut self, key: &[u8]) -> Result<Option<u32>> {
809 self.live_slot(key, Kind::Zset)
810 }
811
812 #[inline]
817 pub(crate) fn zset_at(&self, at: u32) -> &Zset {
818 self.zsets.get(at).expect("the record points at its body")
819 }
820
821 fn new_zset(&mut self, key: &[u8]) -> u32 {
827 let at = yo_alloc::first_touch(|| self.zsets.insert(Zset::new()));
831 let len = value::slot_record_len(false);
832 self.write_rec(key, len, |out| {
833 value::write_slot_record(out, Kind::Zset, at, None);
834 });
835 self.bodies += 1;
836 at
837 }
838}
839
840type Home = (usize, Kind, u32);
843
844impl Db {
845 pub fn zsetop<'k, F>(
853 &self,
854 op: Op,
855 keys: impl Iterator<Item = &'k [u8]> + Clone,
856 weights: &[f64],
857 agg: Aggregate,
858 f: F,
859 ) -> Result<usize>
860 where
861 F: FnMut(Member<'_>, f64),
862 {
863 if let Some(home) = self.one_stripe(keys.clone()) {
864 return self.hold_stripe(home).zsetop(op, keys, weights, agg, f);
865 }
866 let mut held = self.hold_operands(keys.clone(), Some(0));
869 let slots = self.operand_slots(&mut held, keys)?;
870 let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
871 let limits = held.stripe(0).zset_limits;
872 let Some(z) = Zset::from_elements(got, &limits) else {
873 return Ok(0);
874 };
875 let len = z.len();
876 z.walk(0, len, false, f);
877 Ok(len)
878 }
879
880 pub fn zsetop_store<'k>(
887 &self,
888 op: Op,
889 destination: &'k [u8],
890 keys: impl Iterator<Item = &'k [u8]> + Clone,
891 weights: &[f64],
892 agg: Aggregate,
893 ) -> Result<usize> {
894 if let Some(home) = self.one_stripe(std::iter::once(destination).chain(keys.clone())) {
895 return self
896 .hold_stripe(home)
897 .zsetop_store(op, destination, keys, weights, agg);
898 }
899 let onto = self.stripe_of(destination);
900 let mut held = self.hold_operands(keys.clone(), Some(onto));
901 let slots = self.operand_slots(&mut held, keys)?;
902 let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
903 let limits = held.stripe(onto).zset_limits;
906 let built = Zset::from_elements(got, &limits);
907 Ok(held.stripe_mut(onto).put_zset(destination, built))
908 }
909
910 pub fn zintercard<'k>(
912 &self,
913 keys: impl Iterator<Item = &'k [u8]> + Clone,
914 limit: usize,
915 ) -> Result<usize> {
916 if let Some(home) = self.one_stripe(keys.clone()) {
917 return self.hold_stripe(home).zintercard(keys, limit);
918 }
919 let mut held = self.hold_operands(keys.clone(), None);
920 let slots = self.operand_slots(&mut held, keys)?;
921 Ok(zsetops::intercard(&operands_of(&held, &slots), limit))
922 }
923
924 pub fn zrangestore(&self, destination: &[u8], source: &[u8], q: &Query<'_>) -> Result<usize> {
932 let (onto, home) = (self.stripe_of(destination), self.stripe_of(source));
933 if onto == home {
934 return self.hold_stripe(onto).zrangestore(destination, source, q);
935 }
936 let mut held = self.hold_many([home, onto].into_iter());
942 let slot = held.stripe_mut(home).zset_slot(source)?;
946 let built = match slot {
947 None => None,
948 Some(at) => {
949 let z = held.stripe(home).zset_at(at);
950 let w = window(z, q);
951 let mut got = Elements::with_capacity(w.count.max(16));
952 let mut digits = [0u8; DIGITS_MAX];
953 let from = if w.rev { w.from + 1 - w.count } else { w.from };
954 z.walk(from, w.count, false, |m, s| {
955 let _ = got.insert(member_bytes(m, &mut digits), s);
956 });
957 let limits = held.stripe(onto).zset_limits;
958 Zset::from_elements(got, &limits)
959 }
960 };
961 Ok(held.stripe_mut(onto).put_zset(destination, built))
962 }
963
964 fn hold_operands<'k>(
972 &self,
973 keys: impl Iterator<Item = &'k [u8]>,
974 onto: Option<usize>,
975 ) -> Holds<'_> {
976 let named = keys.map(|key| self.stripe_of(key));
977 self.hold_many(named.chain(onto))
978 }
979
980 fn operand_slots<'k>(
989 &self,
990 held: &mut Holds<'_>,
991 keys: impl Iterator<Item = &'k [u8]>,
992 ) -> Result<Vec<Option<Home>>> {
993 let mut out = Vec::with_capacity(keys.size_hint().0);
994 for key in keys {
995 let stripe = self.stripe_of(key);
996 let got = held
997 .stripe_mut(stripe)
998 .live_slot_either(key, Kind::Zset, Kind::Set)?;
999 out.push(got.map(|(kind, at)| (stripe, kind, at)));
1000 }
1001 Ok(out)
1002 }
1003}
1004
1005fn operands_of<'h>(held: &'h Holds<'_>, slots: &[Option<Home>]) -> Vec<Operand<'h>> {
1007 slots
1008 .iter()
1009 .map(|got| match got {
1010 Some((stripe, Kind::Zset, at)) => Operand::Zset(held.stripe(*stripe).zset_at(*at)),
1011 Some((stripe, Kind::Set, at)) => Operand::Set(
1012 held.stripe(*stripe)
1013 .sets
1014 .get(*at)
1015 .expect("the record points at its body"),
1016 ),
1017 _ => Operand::Missing,
1018 })
1019 .collect()
1020}
1021
1022fn gated(z: &Zset, member: &[u8], score: f64, opts: ZAdd) -> Option<Added> {
1028 match z.score(member) {
1029 None => match opts.gate {
1030 Gate::IfPresent => None,
1031 _ => Some(Added::New),
1034 },
1035 Some(was) => {
1036 if opts.gate == Gate::IfMissing {
1037 return None;
1038 }
1039 let ok = match opts.only {
1040 Move::Any => true,
1041 Move::Up => score > was,
1042 Move::Down => score < was,
1043 };
1044 if !ok {
1045 return None;
1046 }
1047 Some(if score == was {
1050 Added::Same
1051 } else {
1052 Added::Changed
1053 })
1054 }
1055 }
1056}
1057
1058fn window(z: &Zset, q: &Query<'_>) -> Window {
1060 let len = z.len();
1061 let range = match q.by {
1062 By::Rank { start, stop } => {
1063 let (from, count) = rank_span(start, stop, len);
1067 if q.rev {
1068 return apply_limit(
1069 Window {
1070 from: len - from - 1,
1071 count,
1072 rev: true,
1073 },
1074 q,
1075 true,
1076 );
1077 }
1078 return apply_limit(
1079 Window {
1080 from,
1081 count,
1082 rev: false,
1083 },
1084 q,
1085 true,
1086 );
1087 }
1088 By::Score { min, max } => z.window_by_score(min, max),
1089 By::Lex { min, max } => z.window_by_lex(min, max),
1090 };
1091 let count = range.end - range.start;
1092 let from = if q.rev {
1093 range.end.saturating_sub(1)
1094 } else {
1095 range.start
1096 };
1097 apply_limit(
1098 Window {
1099 from,
1100 count,
1101 rev: q.rev,
1102 },
1103 q,
1104 false,
1105 )
1106}
1107
1108fn apply_limit(w: Window, q: &Query<'_>, skip: bool) -> Window {
1117 if skip {
1118 return w;
1119 }
1120 let offset = q.offset.min(w.count);
1121 let count = q.count.unwrap_or(usize::MAX).min(w.count - offset);
1122 let from = if w.rev {
1123 w.from.saturating_sub(offset)
1124 } else {
1125 w.from + offset
1126 };
1127 Window {
1128 from,
1129 count,
1130 rev: w.rev,
1131 }
1132}
1133
1134fn rank_span(start: i64, stop: i64, len: usize) -> (usize, usize) {
1141 if len == 0 {
1142 return (0, 0);
1143 }
1144 let len = len as i64;
1145 let from = if start < 0 {
1146 (len + start).max(0)
1147 } else {
1148 start.min(len)
1149 };
1150 let to = if stop < 0 {
1151 len + stop
1152 } else {
1153 stop.min(len - 1)
1154 };
1155 if to < from {
1156 return (from as usize, 0);
1157 }
1158 (from as usize, (to - from + 1) as usize)
1159}
1160
1161pub(crate) fn member_bytes<'a>(
1164 m: Member<'a>,
1165 digits: &'a mut [u8; yo_common::num::DIGITS_MAX],
1166) -> &'a [u8] {
1167 match m {
1168 Member::Str(s) => s,
1169 Member::Int(n) => yo_common::num::i64_digits(digits, n),
1170 }
1171}
1172
1173#[cfg(test)]
1174mod tests {
1175 use super::*;
1176 use crate::many;
1177 use yo_common::num::DIGITS_MAX;
1178
1179 fn ks() -> Keyspace {
1180 Keyspace::new()
1181 }
1182
1183 fn add(k: &mut Keyspace, key: &[u8], pairs: &[(f64, &[u8])]) -> usize {
1185 k.zadd(key, pairs.iter().copied(), ZAdd::default()).unwrap()
1186 }
1187
1188 fn names(k: &mut Keyspace, key: &[u8], q: &Query<'_>) -> Vec<String> {
1190 let w = k.zwindow(key, q).unwrap();
1191 let mut out = Vec::new();
1192 let mut digits = [0u8; DIGITS_MAX];
1193 k.zwalk(key, w, |m, _| {
1194 out.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1195 })
1196 .unwrap();
1197 assert_eq!(
1198 out.len(),
1199 w.count,
1200 "the window said {} and the walk gave {}",
1201 w.count,
1202 out.len()
1203 );
1204 out
1205 }
1206
1207 #[test]
1212 fn zrandmember_stops_allocating_once_its_buffer_is_grown() {
1213 let mut k = ks();
1214 add(
1215 &mut k,
1216 b"z",
1217 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1218 );
1219 k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1220 let (_, allocs) = crate::tally::counted(|| {
1221 for _ in 0..100 {
1222 k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1223 }
1224 });
1225 assert_eq!(
1226 allocs, 0,
1227 "zrandmember allocated {allocs} times in a hundred"
1228 );
1229 }
1230
1231 #[test]
1235 fn zrandmember_hands_its_buffer_back() {
1236 let mut k = ks();
1237 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1238 for _ in 0..3 {
1239 let mut seen = Vec::new();
1240 let n = k
1241 .zrandmember(b"z", 2, |m, _| {
1242 let mut digits = [0u8; DIGITS_MAX];
1243 seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1244 })
1245 .expect("a sorted set");
1246 assert_eq!(n, 2);
1247 assert_eq!(seen.len(), 2);
1248 assert_ne!(seen[0], seen[1], "a positive count draws without replacing");
1249 }
1250 }
1251
1252 #[test]
1253 fn a_missing_key_is_an_empty_sorted_set() {
1254 let mut k = ks();
1255 assert_eq!(k.zcard(b"nope").unwrap(), 0);
1256 assert_eq!(k.zscore(b"nope", b"m").unwrap(), None);
1257 assert_eq!(k.zrank(b"nope", b"m", false).unwrap(), None);
1258 assert_eq!(k.zrem(b"nope", [b"m".as_slice()].into_iter()).unwrap(), 0);
1259 assert_eq!(
1260 names(&mut k, b"nope", &Query::rank(0, -1)),
1261 Vec::<String>::new()
1262 );
1263 assert!(!k.exists(b"nope"));
1264 }
1265
1266 #[test]
1267 fn a_key_holding_something_else_is_a_wrongtype() {
1268 let mut k = ks();
1269 k.set(b"s", b"hello", strings::SetOptions::default())
1270 .unwrap();
1271 assert_eq!(k.zcard(b"s").unwrap_err().code(), Code::WrongType);
1272 assert_eq!(
1273 k.zadd(b"s", [(1.0, b"m".as_slice())].into_iter(), ZAdd::default())
1274 .unwrap_err()
1275 .code(),
1276 Code::WrongType
1277 );
1278 assert_eq!(k.zscore(b"s", b"m").unwrap_err().code(), Code::WrongType);
1279 }
1280
1281 #[test]
1282 fn adding_answers_how_many_were_new() {
1283 let mut k = ks();
1284 assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]), 2);
1285 assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (3.0, b"c")]), 1);
1286 assert_eq!(k.zcard(b"z").unwrap(), 3);
1287 assert_eq!(k.zscore(b"z", b"c").unwrap(), Some(3.0));
1288 assert_eq!(k.kind_of(b"z").map(Kind::name), Some("zset"));
1289 assert_eq!(k.encoding_name(b"z"), Some("listpack"));
1290 }
1291
1292 #[test]
1293 fn the_ch_flag_counts_moved_scores_too() {
1294 let mut k = ks();
1295 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1296 let ch = ZAdd {
1297 changed: true,
1298 ..ZAdd::default()
1299 };
1300 let pairs = [
1302 (9.0, b"a".as_slice()),
1303 (2.0, b"b".as_slice()),
1304 (3.0, b"c".as_slice()),
1305 ];
1306 assert_eq!(k.zadd(b"z", pairs.into_iter(), ch).unwrap(), 2);
1307 assert_eq!(k.zadd(b"z", pairs.into_iter(), ZAdd::default()).unwrap(), 0);
1308 }
1309
1310 #[test]
1311 fn the_gates_let_the_right_members_through() {
1312 let mut k = ks();
1313 add(&mut k, b"z", &[(1.0, b"a")]);
1314 let nx = ZAdd {
1315 gate: Gate::IfMissing,
1316 ..ZAdd::default()
1317 };
1318 let xx = ZAdd {
1319 gate: Gate::IfPresent,
1320 ..ZAdd::default()
1321 };
1322 assert_eq!(
1323 k.zadd(b"z", [(5.0, b"a".as_slice())].into_iter(), nx)
1324 .unwrap(),
1325 0
1326 );
1327 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(1.0));
1328 assert_eq!(
1329 k.zadd(b"z", [(5.0, b"b".as_slice())].into_iter(), nx)
1330 .unwrap(),
1331 1
1332 );
1333 assert_eq!(
1334 k.zadd(b"z", [(7.0, b"c".as_slice())].into_iter(), xx)
1335 .unwrap(),
1336 0
1337 );
1338 assert_eq!(k.zscore(b"z", b"c").unwrap(), None);
1339 assert_eq!(
1340 k.zadd(b"z", [(7.0, b"a".as_slice())].into_iter(), xx)
1341 .unwrap(),
1342 0
1343 );
1344 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(7.0));
1345 assert_eq!(
1347 k.zadd(b"gone", [(1.0, b"a".as_slice())].into_iter(), xx)
1348 .unwrap(),
1349 0
1350 );
1351 assert!(!k.exists(b"gone"));
1352 }
1353
1354 #[test]
1355 fn gt_and_lt_only_move_a_score_one_way() {
1356 let mut k = ks();
1357 add(&mut k, b"z", &[(5.0, b"a")]);
1358 let gt = ZAdd {
1359 only: Move::Up,
1360 changed: true,
1361 ..ZAdd::default()
1362 };
1363 let lt = ZAdd {
1364 only: Move::Down,
1365 changed: true,
1366 ..ZAdd::default()
1367 };
1368 assert_eq!(
1369 k.zadd(b"z", [(3.0, b"a".as_slice())].into_iter(), gt)
1370 .unwrap(),
1371 0
1372 );
1373 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(5.0));
1374 assert_eq!(
1375 k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), gt)
1376 .unwrap(),
1377 1
1378 );
1379 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(9.0));
1380 assert_eq!(
1381 k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), lt)
1382 .unwrap(),
1383 0
1384 );
1385 assert_eq!(
1386 k.zadd(b"z", [(2.0, b"a".as_slice())].into_iter(), lt)
1387 .unwrap(),
1388 1
1389 );
1390 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.0));
1391 assert_eq!(
1393 k.zadd(b"z", [(1.0, b"new".as_slice())].into_iter(), gt)
1394 .unwrap(),
1395 1
1396 );
1397 }
1398
1399 #[test]
1400 fn incrementing_adds_to_what_is_there_or_to_nothing() {
1401 let mut k = ks();
1402 let plain = ZAdd::default();
1403 assert_eq!(k.zincrby(b"z", b"a", 5.0, plain).unwrap(), Some(5.0));
1404 assert_eq!(k.zincrby(b"z", b"a", -2.5, plain).unwrap(), Some(2.5));
1405 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1406 let nx = ZAdd {
1407 gate: Gate::IfMissing,
1408 ..ZAdd::default()
1409 };
1410 assert_eq!(k.zincrby(b"z", b"a", 1.0, nx).unwrap(), None);
1411 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1412 let xx = ZAdd {
1413 gate: Gate::IfPresent,
1414 ..ZAdd::default()
1415 };
1416 assert_eq!(k.zincrby(b"z", b"never", 1.0, xx).unwrap(), None);
1417 assert!(k.zscore(b"z", b"never").unwrap().is_none());
1418 let gt = ZAdd {
1419 only: Move::Up,
1420 ..ZAdd::default()
1421 };
1422 assert_eq!(k.zincrby(b"z", b"a", -1.0, gt).unwrap(), None);
1423 assert_eq!(k.zincrby(b"z", b"a", 1.0, gt).unwrap(), Some(3.5));
1424 }
1425
1426 #[test]
1427 fn a_score_that_is_not_a_number_is_refused() {
1428 let mut k = ks();
1429 let plain = ZAdd::default();
1430 assert_eq!(
1431 k.zadd(b"z", [(f64::NAN, b"a".as_slice())].into_iter(), plain)
1432 .unwrap_err()
1433 .code(),
1434 Code::Invalid
1435 );
1436 assert!(!k.exists(b"z"));
1437 k.zincrby(b"z", b"a", f64::INFINITY, plain).unwrap();
1438 let err = k.zincrby(b"z", b"a", f64::NEG_INFINITY, plain).unwrap_err();
1439 assert_eq!(err.code(), Code::Invalid);
1440 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(f64::INFINITY));
1442 }
1443
1444 #[test]
1445 fn a_sorted_set_that_loses_its_last_member_loses_its_key() {
1446 let mut k = ks();
1447 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1448 assert_eq!(
1449 k.zrem(b"z", [b"a".as_slice(), b"b".as_slice()].into_iter())
1450 .unwrap(),
1451 2
1452 );
1453 assert!(!k.exists(b"z"));
1454 assert_eq!(k.zcard(b"z").unwrap(), 0);
1455 }
1456
1457 #[test]
1458 fn ranks_count_from_both_ends() {
1459 let mut k = ks();
1460 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1461 assert_eq!(k.zrank(b"z", b"a", false).unwrap(), Some((0, 1.0)));
1462 assert_eq!(k.zrank(b"z", b"c", false).unwrap(), Some((2, 3.0)));
1463 assert_eq!(k.zrank(b"z", b"a", true).unwrap(), Some((2, 1.0)));
1464 assert_eq!(k.zrank(b"z", b"c", true).unwrap(), Some((0, 3.0)));
1465 assert_eq!(k.zrank(b"z", b"nope", false).unwrap(), None);
1466 }
1467
1468 #[test]
1469 fn a_rank_range_clamps_every_way_it_can_be_wrong() {
1470 let mut k = ks();
1471 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1472 assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), ["a", "b", "c"]);
1473 assert_eq!(names(&mut k, b"z", &Query::rank(1, 1)), ["b"]);
1474 assert_eq!(names(&mut k, b"z", &Query::rank(-2, -1)), ["b", "c"]);
1475 assert_eq!(names(&mut k, b"z", &Query::rank(-99, 99)), ["a", "b", "c"]);
1476 assert_eq!(
1477 names(&mut k, b"z", &Query::rank(2, 1)),
1478 Vec::<String>::new()
1479 );
1480 assert_eq!(
1481 names(&mut k, b"z", &Query::rank(5, 9)),
1482 Vec::<String>::new()
1483 );
1484 assert_eq!(
1485 names(&mut k, b"z", &Query::rank(0, -1).rev(true)),
1486 ["c", "b", "a"]
1487 );
1488 assert_eq!(
1489 names(&mut k, b"z", &Query::rank(0, 1).rev(true)),
1490 ["c", "b"]
1491 );
1492 }
1493
1494 #[test]
1495 fn a_score_range_walks_either_way_and_takes_a_limit() {
1496 let mut k = ks();
1497 add(
1498 &mut k,
1499 b"z",
1500 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1501 );
1502 let all = Query::score(
1503 Bound::closed(f64::NEG_INFINITY),
1504 Bound::closed(f64::INFINITY),
1505 );
1506 assert_eq!(names(&mut k, b"z", &all), ["a", "b", "c", "d"]);
1507 assert_eq!(names(&mut k, b"z", &all.rev(true)), ["d", "c", "b", "a"]);
1508 assert_eq!(
1509 names(
1510 &mut k,
1511 b"z",
1512 &Query::score(Bound::closed(2.0), Bound::closed(3.0))
1513 ),
1514 ["b", "c"]
1515 );
1516 assert_eq!(
1517 names(
1518 &mut k,
1519 b"z",
1520 &Query::score(Bound::open(2.0), Bound::open(4.0))
1521 ),
1522 ["c"]
1523 );
1524 assert_eq!(names(&mut k, b"z", &all.limit(1, Some(2))), ["b", "c"]);
1527 assert_eq!(
1528 names(&mut k, b"z", &all.rev(true).limit(1, Some(2))),
1529 ["c", "b"]
1530 );
1531 assert_eq!(
1532 names(&mut k, b"z", &all.limit(9, Some(2))),
1533 Vec::<String>::new()
1534 );
1535 assert_eq!(names(&mut k, b"z", &all.limit(2, None)), ["c", "d"]);
1536 assert_eq!(
1537 k.zcount(b"z", &Query::score(Bound::closed(2.0), Bound::closed(3.0)))
1538 .unwrap(),
1539 2
1540 );
1541 }
1542
1543 #[test]
1544 fn a_member_range_orders_by_member_when_every_score_is_the_same() {
1545 let mut k = ks();
1546 add(
1547 &mut k,
1548 b"z",
1549 &[(0.0, b"a"), (0.0, b"b"), (0.0, b"c"), (0.0, b"d")],
1550 );
1551 assert_eq!(
1552 names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max)),
1553 ["a", "b", "c", "d"]
1554 );
1555 assert_eq!(
1556 names(&mut k, b"z", &Query::lex(Lex::Incl(b"b"), Lex::Incl(b"c"))),
1557 ["b", "c"]
1558 );
1559 assert_eq!(
1560 names(&mut k, b"z", &Query::lex(Lex::Excl(b"a"), Lex::Excl(b"d"))),
1561 ["b", "c"]
1562 );
1563 assert_eq!(
1564 names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max).rev(true)),
1565 ["d", "c", "b", "a"]
1566 );
1567 assert_eq!(
1568 k.zcount(b"z", &Query::lex(Lex::Incl(b"b"), Lex::Max))
1569 .unwrap(),
1570 3
1571 );
1572 }
1573
1574 #[test]
1575 fn removing_a_range_takes_out_exactly_what_the_walk_would_have_given() {
1576 let mut k = ks();
1577 for (q, left) in [
1578 (Query::rank(0, 1), vec!["c", "d", "e"]),
1579 (Query::rank(-2, -1), vec!["a", "b", "c"]),
1580 (
1581 Query::score(Bound::closed(2.0), Bound::closed(4.0)),
1582 vec!["a", "e"],
1583 ),
1584 (
1585 Query::lex(Lex::Incl(b"b"), Lex::Excl(b"d")),
1586 vec!["a", "d", "e"],
1587 ),
1588 (Query::rank(0, -1).rev(true), vec![]),
1589 ] {
1590 k.del(b"z");
1591 add(
1592 &mut k,
1593 b"z",
1594 &[
1595 (1.0, b"a"),
1596 (2.0, b"b"),
1597 (3.0, b"c"),
1598 (4.0, b"d"),
1599 (5.0, b"e"),
1600 ],
1601 );
1602 let want = names(&mut k, b"z", &q).len();
1603 assert_eq!(k.zremrange(b"z", &q).unwrap(), want, "{q:?}");
1604 assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), left, "{q:?}");
1605 }
1606 assert!(!k.exists(b"z"));
1608 }
1609
1610 #[test]
1611 fn popping_takes_from_the_end_it_was_told_to() {
1612 let mut k = ks();
1613 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1614 let mut got = Vec::new();
1615 let mut digits = [0u8; DIGITS_MAX];
1616 k.zpop(b"z", From::Min, 2, |m, s| {
1617 got.push((
1618 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1619 s,
1620 ));
1621 })
1622 .unwrap();
1623 assert_eq!(got, [("a".to_string(), 1.0), ("b".to_string(), 2.0)]);
1624 assert_eq!(
1625 k.zpop_one(b"z", From::Max).unwrap(),
1626 Some((b"c".to_vec(), 3.0))
1627 );
1628 assert!(!k.exists(b"z"));
1629 add(&mut k, b"z", &[(1.0, b"a")]);
1631 assert_eq!(k.zpop(b"z", From::Max, 99, |_, _| {}).unwrap(), 1);
1632 assert!(!k.exists(b"z"));
1633 assert_eq!(k.zpop_one(b"z", From::Min).unwrap(), None);
1634 }
1635
1636 #[test]
1637 fn a_random_draw_is_with_or_without_replacement_by_the_sign_of_the_count() {
1638 let mut k = ks();
1639 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1640 let mut digits = [0u8; DIGITS_MAX];
1641 let mut seen = Vec::new();
1642 k.zrandmember(b"z", 2, |m, _| {
1643 seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1644 })
1645 .unwrap();
1646 assert_eq!(seen.len(), 2);
1647 seen.sort();
1648 seen.dedup();
1649 assert_eq!(seen.len(), 2, "a positive count does not repeat a member");
1650 let mut all = Vec::new();
1652 k.zrandmember(b"z", 99, |m, _| {
1653 all.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1654 })
1655 .unwrap();
1656 all.sort();
1657 assert_eq!(all, ["a", "b", "c"]);
1658 let mut many = 0;
1660 k.zrandmember(b"z", -10, |_, _| many += 1).unwrap();
1661 assert_eq!(many, 10);
1662 assert_eq!(k.zrandmember(b"nope", 3, |_, _| {}).unwrap(), 0);
1663 }
1664
1665 #[test]
1666 fn a_scan_of_either_band_walks_every_member_once() {
1667 let mut k = ks();
1668 for entries in [8usize, 4096] {
1669 k.del(b"z");
1670 let pairs: Vec<(f64, Vec<u8>)> = (0..entries)
1671 .map(|i| (i as f64, format!("m{i:05}").into_bytes()))
1672 .collect();
1673 k.zadd(
1674 b"z",
1675 pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1676 ZAdd::default(),
1677 )
1678 .unwrap();
1679 let mut seen = Vec::new();
1680 let mut digits = [0u8; DIGITS_MAX];
1681 let mut cursor = Cursor::START;
1682 loop {
1683 cursor = k
1684 .zscan(b"z", cursor, 16, |m, _| {
1685 seen.push(
1686 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1687 );
1688 })
1689 .unwrap();
1690 if cursor.is_end() {
1691 break;
1692 }
1693 }
1694 seen.sort();
1695 seen.dedup();
1696 assert_eq!(seen.len(), entries, "{entries} members");
1697 }
1698 }
1699
1700 #[test]
1701 fn a_big_sorted_set_promotes_and_still_answers_every_rank() {
1702 let n = many(5_000);
1706 let (mid, low) = (n / 2, n / 5);
1707 let mut k = ks();
1708 let pairs: Vec<(f64, Vec<u8>)> = (0..n)
1709 .map(|i| (f64::from(i), format!("m{i:05}").into_bytes()))
1710 .collect();
1711 assert_eq!(
1712 k.zadd(
1713 b"z",
1714 pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1715 ZAdd::default()
1716 )
1717 .unwrap(),
1718 n as usize
1719 );
1720 assert_eq!(k.encoding_name(b"z"), Some("skiplist"));
1721 assert_eq!(k.zcard(b"z").unwrap(), n as usize);
1722 let at_mid = format!("m{mid:05}").into_bytes();
1723 assert_eq!(
1724 k.zrank(b"z", &at_mid, false).unwrap(),
1725 Some((mid as usize, f64::from(mid)))
1726 );
1727 let q = Query::score(
1728 Bound::closed(f64::from(low)),
1729 Bound::open(f64::from(low + 10)),
1730 );
1731 assert_eq!(k.zcount(b"z", &q).unwrap(), 10);
1732 assert_eq!(
1733 names(&mut k, b"z", &q).first().map(String::as_str),
1734 Some(format!("m{low:05}").as_str())
1735 );
1736 assert_eq!(
1738 k.zremrange(b"z", &Query::rank(0, i64::from(mid) - 1))
1739 .unwrap(),
1740 mid as usize
1741 );
1742 assert_eq!(k.zcard(b"z").unwrap(), mid as usize);
1743 assert_eq!(
1744 k.zrank(b"z", &at_mid, false).unwrap(),
1745 Some((0, f64::from(mid)))
1746 );
1747 assert_eq!(k.zrank(b"z", b"m00000", false).unwrap(), None);
1748 }
1749
1750 #[test]
1751 fn a_deadline_on_a_sorted_set_leaves_its_members_alone() {
1752 let mut k = ks();
1753 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1754 assert!(k.set_expiry(b"z", Some(u64::MAX / 2)));
1755 assert_eq!(k.zcard(b"z").unwrap(), 2);
1756 assert_eq!(k.zscore(b"z", b"b").unwrap(), Some(2.0));
1757 assert!(k.set_expiry(b"z", None));
1758 assert_eq!(k.zcard(b"z").unwrap(), 2);
1759 }
1760
1761 #[test]
1762 fn writing_a_string_over_a_sorted_set_gives_its_body_back() {
1763 let mut k = ks();
1764 add(&mut k, b"z", &[(1.0, b"a")]);
1765 let held = k.memory_bytes();
1766 k.set(b"z", b"now a string", strings::SetOptions::default())
1767 .unwrap();
1768 assert_eq!(k.kind_of(b"z").map(Kind::name), Some("string"));
1769 assert!(
1770 k.memory_bytes() < held,
1771 "the sorted set's body was not freed"
1772 );
1773 }
1774
1775 fn all(k: &mut Keyspace, key: &[u8]) -> Vec<(String, f64)> {
1777 let q = Query::rank(0, -1);
1778 let w = k.zwindow(key, &q).unwrap();
1779 let mut out = Vec::new();
1780 let mut digits = [0u8; DIGITS_MAX];
1781 k.zwalk(key, w, |m, s| {
1782 out.push((
1783 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1784 s,
1785 ));
1786 })
1787 .unwrap();
1788 out
1789 }
1790
1791 fn got(
1793 k: &mut Keyspace,
1794 op: Op,
1795 keys: &[&[u8]],
1796 weights: &[f64],
1797 agg: Aggregate,
1798 ) -> Vec<(String, f64)> {
1799 let mut out = Vec::new();
1800 let mut digits = [0u8; DIGITS_MAX];
1801 let n = k
1802 .zsetop(op, keys.iter().copied(), weights, agg, |m, s| {
1803 out.push((
1804 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1805 s,
1806 ));
1807 })
1808 .unwrap();
1809 assert_eq!(out.len(), n, "the count and the walk disagree");
1810 out
1811 }
1812
1813 #[test]
1814 fn a_union_adds_the_scores_of_a_member_that_is_in_both() {
1815 let mut k = ks();
1816 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1817 add(&mut k, b"b", &[(10.0, b"y"), (3.0, b"z")]);
1818 assert_eq!(
1819 got(&mut k, Op::Union, &[b"a", b"b"], &[], Aggregate::Sum),
1820 [("x".into(), 1.0), ("z".into(), 3.0), ("y".into(), 12.0)]
1821 );
1822 }
1823
1824 #[test]
1825 fn an_intersection_keeps_only_what_every_input_has() {
1826 let mut k = ks();
1827 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1828 add(&mut k, b"b", &[(5.0, b"y"), (5.0, b"z")]);
1829 add(&mut k, b"c", &[(7.0, b"z")]);
1830 assert_eq!(
1831 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Sum),
1832 [("z".into(), 15.0)]
1833 );
1834 assert_eq!(
1835 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Min),
1836 [("z".into(), 3.0)]
1837 );
1838 assert_eq!(
1839 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Max),
1840 [("z".into(), 7.0)]
1841 );
1842 }
1843
1844 #[test]
1845 fn a_difference_takes_the_first_input_and_removes_the_rest() {
1846 let mut k = ks();
1847 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1848 add(&mut k, b"b", &[(99.0, b"y")]);
1849 assert_eq!(
1850 got(&mut k, Op::Diff, &[b"a", b"b"], &[], Aggregate::Sum),
1851 [("x".into(), 1.0), ("z".into(), 3.0)]
1852 );
1853 }
1854
1855 #[test]
1856 fn a_missing_key_keeps_its_place_so_the_weights_stay_lined_up() {
1857 let mut k = ks();
1858 add(&mut k, b"a", &[(1.0, b"x")]);
1859 add(&mut k, b"c", &[(1.0, b"x")]);
1860 let out = got(
1863 &mut k,
1864 Op::Union,
1865 &[b"a", b"gone", b"c"],
1866 &[2.0, 10.0, 3.0],
1867 Aggregate::Sum,
1868 );
1869 assert_eq!(out, [("x".into(), 5.0)]);
1870 }
1871
1872 #[test]
1873 fn a_plain_set_counts_as_every_score_being_one() {
1874 let mut k = ks();
1875 add(&mut k, b"z", &[(5.0, b"x")]);
1876 k.sadd(b"s", [b"x".as_slice(), b"y".as_slice()].into_iter())
1877 .unwrap();
1878 assert_eq!(
1879 got(&mut k, Op::Union, &[b"z", b"s"], &[], Aggregate::Sum),
1880 [("y".into(), 1.0), ("x".into(), 6.0)]
1881 );
1882 assert_eq!(
1883 got(&mut k, Op::Inter, &[b"z", b"s"], &[], Aggregate::Min),
1884 [("x".into(), 1.0)]
1885 );
1886 }
1887
1888 #[test]
1889 fn a_store_writes_the_result_and_answers_its_size() {
1890 let mut k = ks();
1891 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1892 add(&mut k, b"b", &[(10.0, b"y")]);
1893 assert_eq!(
1894 k.zsetop_store(
1895 Op::Union,
1896 b"d",
1897 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1898 &[],
1899 Aggregate::Sum
1900 )
1901 .unwrap(),
1902 2
1903 );
1904 assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1905 assert_eq!(k.zscore(b"d", b"y").unwrap(), Some(12.0));
1906 assert_eq!(k.zrank(b"d", b"y", false).unwrap(), Some((1, 12.0)));
1907 }
1908
1909 #[test]
1910 fn a_store_onto_one_of_its_own_sources_still_reads_the_old_body() {
1911 let mut k = ks();
1912 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1913 add(&mut k, b"b", &[(10.0, b"y")]);
1914 assert_eq!(
1915 k.zsetop_store(
1916 Op::Union,
1917 b"a",
1918 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1919 &[],
1920 Aggregate::Sum
1921 )
1922 .unwrap(),
1923 2
1924 );
1925 assert_eq!(all(&mut k, b"a"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1926 }
1927
1928 #[test]
1929 fn a_store_with_nothing_in_it_deletes_the_destination() {
1930 let mut k = ks();
1931 add(&mut k, b"a", &[(1.0, b"x")]);
1932 add(&mut k, b"b", &[(1.0, b"y")]);
1933 add(&mut k, b"d", &[(1.0, b"old")]);
1934 assert_eq!(
1935 k.zsetop_store(
1936 Op::Inter,
1937 b"d",
1938 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1939 &[],
1940 Aggregate::Sum
1941 )
1942 .unwrap(),
1943 0
1944 );
1945 assert!(!k.exists(b"d"));
1946 }
1947
1948 #[test]
1949 fn a_store_clears_the_deadline_the_destination_was_carrying() {
1950 let mut k = ks();
1951 add(&mut k, b"a", &[(1.0, b"x")]);
1952 add(&mut k, b"d", &[(1.0, b"old")]);
1953 assert!(k.set_expiry(b"d", Some(u64::MAX / 2)));
1954 k.zsetop_store(
1955 Op::Union,
1956 b"d",
1957 [b"a".as_slice()].into_iter(),
1958 &[],
1959 Aggregate::Sum,
1960 )
1961 .unwrap();
1962 assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0)]);
1963 assert_eq!(k.expire_at(b"d"), None);
1964 }
1965
1966 #[test]
1967 fn the_algebra_refuses_a_key_holding_something_that_is_neither() {
1968 let mut k = ks();
1969 add(&mut k, b"a", &[(1.0, b"x")]);
1970 k.set(b"s", b"hello", strings::SetOptions::default())
1971 .unwrap();
1972 assert_eq!(
1973 k.zsetop(
1974 Op::Union,
1975 [b"a".as_slice(), b"s".as_slice()].into_iter(),
1976 &[],
1977 Aggregate::Sum,
1978 |_, _| {}
1979 )
1980 .unwrap_err()
1981 .code(),
1982 Code::WrongType
1983 );
1984 assert!(!k.exists(b"d"));
1987 assert_eq!(
1988 k.zsetop_store(
1989 Op::Union,
1990 b"d",
1991 [b"a".as_slice(), b"s".as_slice()].into_iter(),
1992 &[],
1993 Aggregate::Sum
1994 )
1995 .unwrap_err()
1996 .code(),
1997 Code::WrongType
1998 );
1999 assert!(!k.exists(b"d"));
2000 }
2001
2002 #[test]
2003 fn intercard_counts_without_building_anything_and_stops_at_the_limit() {
2004 let mut k = ks();
2005 add(
2006 &mut k,
2007 b"a",
2008 &[(1.0, b"w"), (2.0, b"x"), (3.0, b"y"), (4.0, b"z")],
2009 );
2010 add(&mut k, b"b", &[(1.0, b"x"), (1.0, b"y"), (1.0, b"z")]);
2011 assert_eq!(
2012 k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 0)
2013 .unwrap(),
2014 3
2015 );
2016 assert_eq!(
2017 k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 2)
2018 .unwrap(),
2019 2
2020 );
2021 assert_eq!(
2022 k.zintercard([b"a".as_slice(), b"gone".as_slice()].into_iter(), 0)
2023 .unwrap(),
2024 0
2025 );
2026 }
2027
2028 #[test]
2029 fn a_range_store_copies_a_window_and_leaves_the_source_alone() {
2030 let mut k = ks();
2031 add(
2032 &mut k,
2033 b"z",
2034 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2035 );
2036 assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(1, 2)).unwrap(), 2);
2037 assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2038 assert_eq!(k.zcard(b"z").unwrap(), 4);
2039 }
2040
2041 #[test]
2042 fn a_reverse_range_store_picks_the_same_members_and_orders_them_the_same() {
2043 let mut k = ks();
2044 add(
2045 &mut k,
2046 b"z",
2047 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2048 );
2049 assert_eq!(
2052 k.zrangestore(b"d", b"z", &Query::rank(0, 1).rev(true))
2053 .unwrap(),
2054 2
2055 );
2056 assert_eq!(all(&mut k, b"d"), [("c".into(), 3.0), ("d".into(), 4.0)]);
2057 }
2058
2059 #[test]
2060 fn a_range_store_by_score_takes_the_bounds_the_range_would_have() {
2061 let mut k = ks();
2062 add(
2063 &mut k,
2064 b"z",
2065 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2066 );
2067 let q = Query::score(Bound::closed(2.0), Bound::open(4.0));
2068 assert_eq!(k.zrangestore(b"d", b"z", &q).unwrap(), 2);
2069 assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2070 }
2071
2072 #[test]
2073 fn a_range_store_of_an_empty_window_deletes_the_destination() {
2074 let mut k = ks();
2075 add(&mut k, b"z", &[(1.0, b"a")]);
2076 add(&mut k, b"d", &[(1.0, b"old")]);
2077 assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(5, 9)).unwrap(), 0);
2078 assert!(!k.exists(b"d"));
2079 assert_eq!(
2080 k.zrangestore(b"d", b"gone", &Query::rank(0, -1)).unwrap(),
2081 0
2082 );
2083 assert!(!k.exists(b"d"));
2084 }
2085
2086 #[test]
2087 fn a_range_store_onto_its_own_source_keeps_the_window() {
2088 let mut k = ks();
2089 add(
2090 &mut k,
2091 b"z",
2092 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2093 );
2094 assert_eq!(k.zrangestore(b"z", b"z", &Query::rank(1, 2)).unwrap(), 2);
2095 assert_eq!(all(&mut k, b"z"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2096 }
2097
2098 #[test]
2099 fn a_big_result_comes_out_on_the_table_band_in_the_right_order() {
2100 let mut k = ks();
2101 let names: Vec<String> = (0..3000).map(|i| format!("member-{i:05}")).collect();
2102 for (i, name) in names.iter().enumerate() {
2103 add(&mut k, b"a", &[((i % 17) as f64, name.as_bytes())]);
2104 }
2105 for name in names.iter().step_by(3) {
2106 add(&mut k, b"b", &[(100.0, name.as_bytes())]);
2107 }
2108 let n = k
2109 .zsetop_store(
2110 Op::Union,
2111 b"d",
2112 [b"a".as_slice(), b"b".as_slice()].into_iter(),
2113 &[],
2114 Aggregate::Sum,
2115 )
2116 .unwrap();
2117 assert_eq!(n, 3000);
2118 assert_eq!(
2119 k.zset_encoding(b"d").map(crate::zset::Encoding::name),
2120 Some("skiplist")
2121 );
2122 let out = all(&mut k, b"d");
2123 assert_eq!(out.len(), 3000);
2124 let mut want = out.clone();
2125 want.sort_by(|x, y| x.1.partial_cmp(&y.1).unwrap().then_with(|| x.0.cmp(&y.0)));
2126 assert_eq!(out, want, "the result came out in the wrong order");
2127 for (i, name) in names.iter().enumerate() {
2128 let base = (i % 17) as f64;
2129 let want = if i % 3 == 0 { base + 100.0 } else { base };
2130 assert_eq!(k.zscore(b"d", name.as_bytes()).unwrap(), Some(want));
2131 }
2132 }
2133}