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 for (score, m) in pairs.clone() {
251 strings::check_len(key, m.len())?;
252 if score.is_nan() {
253 return Err(nan());
254 }
255 }
256 let at = match self.zset_slot(key)? {
257 Some(at) => at,
258 None => {
259 if opts.gate == Gate::IfPresent || pairs.clone().next().is_none() {
265 return Ok(0);
266 }
267 self.new_zset(key)
268 }
269 };
270 let limits = self.zset_limits;
271 let z = self
272 .zsets
273 .get_mut(at)
274 .expect("the record points at its body");
275 let mut added = 0usize;
276 let mut changed = 0usize;
277 for (score, m) in pairs {
278 let Some(want) = gated(z, m, score, opts) else {
279 continue;
280 };
281 if z.add(m, score, &limits) == Added::Full {
284 continue;
285 }
286 match want {
287 Added::New => added += 1,
288 Added::Changed => changed += 1,
289 _ => {}
290 }
291 }
292 if z.is_empty() {
296 self.drop_key(key);
297 }
298 Ok(if opts.changed { added + changed } else { added })
299 }
300
301 pub fn zincrby(
310 &mut self,
311 key: &[u8],
312 member: &[u8],
313 by: f64,
314 opts: ZAdd,
315 ) -> Result<Option<f64>> {
316 strings::check_len(key, member.len())?;
317 if by.is_nan() {
318 return Err(nan());
319 }
320 let at = match self.zset_slot(key)? {
321 Some(at) => at,
322 None => {
323 if opts.gate == Gate::IfPresent {
324 return Ok(None);
325 }
326 self.new_zset(key)
327 }
328 };
329 let limits = self.zset_limits;
330 let z = self
331 .zsets
332 .get_mut(at)
333 .expect("the record points at its body");
334 let now = z.score(member);
335 let want = now.unwrap_or(0.0) + by;
336 if want.is_nan() {
339 if z.is_empty() {
340 self.drop_key(key);
341 }
342 return Err(nan());
343 }
344 let allowed = match (now, opts.gate, opts.only) {
345 (Some(_), Gate::IfMissing, _) | (None, Gate::IfPresent, _) => false,
346 (Some(was), _, Move::Up) => want > was,
347 (Some(was), _, Move::Down) => want < was,
348 _ => true,
349 };
350 if !allowed || z.add(member, want, &limits) == Added::Full {
351 if z.is_empty() {
352 self.drop_key(key);
353 }
354 return Ok(None);
355 }
356 Ok(Some(want))
357 }
358
359 pub fn zcard(&mut self, key: &[u8]) -> Result<usize> {
361 Ok(match self.zset_slot(key)? {
362 Some(at) => self.zset_at(at).len(),
363 None => 0,
364 })
365 }
366
367 pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> Result<Option<f64>> {
369 Ok(match self.zset_slot(key)? {
370 Some(at) => self.zset_at(at).score(member),
371 None => None,
372 })
373 }
374
375 pub fn zmscore<'m>(
380 &mut self,
381 key: &[u8],
382 members: impl Iterator<Item = &'m [u8]>,
383 out: &mut Vec<Option<f64>>,
384 ) -> Result<()> {
385 out.clear();
386 let Some(at) = self.zset_slot(key)? else {
387 out.extend(members.map(|_| None));
388 return Ok(());
389 };
390 let z = self.zset_at(at);
391 out.extend(members.map(|m| z.score(m)));
392 Ok(())
393 }
394
395 pub fn zrem<'m>(
400 &mut self,
401 key: &[u8],
402 members: impl Iterator<Item = &'m [u8]>,
403 ) -> Result<usize> {
404 let Some(at) = self.zset_slot(key)? else {
405 return Ok(0);
406 };
407 let z = self
408 .zsets
409 .get_mut(at)
410 .expect("the record points at its body");
411 let mut gone = 0;
412 for m in members {
413 if z.remove(m) {
414 gone += 1;
415 }
416 }
417 if z.is_empty() {
418 self.drop_key(key);
419 }
420 Ok(gone)
421 }
422
423 pub fn zrank(&mut self, key: &[u8], member: &[u8], rev: bool) -> Result<Option<(usize, f64)>> {
428 let Some(at) = self.zset_slot(key)? else {
429 return Ok(None);
430 };
431 let z = self.zset_at(at);
432 let Some(rank) = z.rank(member) else {
433 return Ok(None);
434 };
435 let score = z.score(member).unwrap_or(0.0);
436 Ok(Some((if rev { z.len() - 1 - rank } else { rank }, score)))
437 }
438
439 pub fn zwindow(&mut self, key: &[u8], q: &Query<'_>) -> Result<Window> {
446 let Some(at) = self.zset_slot(key)? else {
447 return Ok(Window::default());
448 };
449 Ok(window(self.zset_at(at), q))
450 }
451
452 pub fn zwalk<F>(&mut self, key: &[u8], w: Window, f: F) -> Result<()>
458 where
459 F: FnMut(Member<'_>, f64),
460 {
461 let Some(at) = self.zset_slot(key)? else {
462 return Ok(());
463 };
464 self.zset_at(at).walk(w.from, w.count, w.rev, f);
465 Ok(())
466 }
467
468 pub fn zcount(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
470 Ok(self.zwindow(key, q)?.count)
471 }
472
473 pub fn zremrange(&mut self, key: &[u8], q: &Query<'_>) -> Result<usize> {
480 let Some(at) = self.zset_slot(key)? else {
481 return Ok(0);
482 };
483 let w = window(self.zset_at(at), q);
484 let z = self
485 .zsets
486 .get_mut(at)
487 .expect("the record points at its body");
488 let low = if w.rev { w.from + 1 - w.count } else { w.from };
491 for i in (0..w.count).rev() {
492 z.remove_at(low + i);
493 }
494 if z.is_empty() {
495 self.drop_key(key);
496 }
497 Ok(w.count)
498 }
499
500 pub fn zpop<F>(&mut self, key: &[u8], end: From, count: usize, mut f: F) -> Result<usize>
506 where
507 F: FnMut(Member<'_>, f64),
508 {
509 let Some(at) = self.zset_slot(key)? else {
510 return Ok(0);
511 };
512 let z = self
513 .zsets
514 .get_mut(at)
515 .expect("the record points at its body");
516 let count = count.min(z.len());
517 for _ in 0..count {
518 let rank = if end == From::Min { 0 } else { z.len() - 1 };
522 let Some((m, s)) = z.at(rank) else { break };
523 f(m, s);
524 z.remove_at(rank);
525 }
526 if z.is_empty() {
527 self.drop_key(key);
528 }
529 Ok(count)
530 }
531
532 pub fn zpop_one(&mut self, key: &[u8], end: From) -> Result<Option<(Vec<u8>, f64)>> {
539 let mut got = None;
540 let mut name = [0u8; yo_common::num::DIGITS_MAX];
541 self.zpop(key, end, 1, |m, s| {
542 got = Some((member_bytes(m, &mut name).to_vec(), s));
543 })?;
544 Ok(got)
545 }
546
547 pub fn zrandmember<F>(&mut self, key: &[u8], count: i64, mut f: F) -> Result<usize>
558 where
559 F: FnMut(Member<'_>, f64),
560 {
561 let Some(at) = self.zset_slot(key)? else {
562 return Ok(0);
563 };
564 let len = self.zset_at(at).len();
565 if len == 0 || count == 0 {
566 return Ok(0);
567 }
568 if count < 0 {
569 let want = count.unsigned_abs() as usize;
570 for _ in 0..want {
571 let pick = self.rng.below(len);
572 let Some((m, s)) = self.zset_at(at).pick(pick) else {
573 break;
574 };
575 f(m, s);
576 }
577 return Ok(want);
578 }
579 let want = (count as usize).min(len);
580 if want == len {
584 for i in 0..len {
585 let Some((m, s)) = self.zset_at(at).pick(i) else {
586 break;
587 };
588 f(m, s);
589 }
590 return Ok(len);
591 }
592 let mut rows = std::mem::take(&mut self.rows);
602 rows.clear();
603 yo_alloc::high_water(|| rows.extend(0..len));
604 for i in 0..want {
605 let pick = i + self.rng.below(len - i);
606 rows.swap(i, pick);
607 let Some((m, s)) = self.zset_at(at).pick(rows[i]) else {
608 break;
609 };
610 f(m, s);
611 }
612 self.rows = rows;
613 Ok(want)
614 }
615
616 pub fn zscan<F>(&mut self, key: &[u8], cursor: Cursor, count: usize, f: F) -> Result<Cursor>
622 where
623 F: FnMut(Member<'_>, f64),
624 {
625 let Some(at) = self.zset_slot(key)? else {
626 return Ok(Cursor::END);
627 };
628 Ok(self.zset_at(at).scan(cursor, count, f))
629 }
630
631 pub fn zsetop<'k, F>(
638 &mut self,
639 op: Op,
640 keys: impl Iterator<Item = &'k [u8]>,
641 weights: &[f64],
642 agg: Aggregate,
643 f: F,
644 ) -> Result<usize>
645 where
646 F: FnMut(Member<'_>, f64),
647 {
648 let slots = self.operand_slots(keys)?;
649 let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
650 let limits = self.zset_limits;
651 let Some(z) = Zset::from_elements(got, &limits) else {
652 return Ok(0);
653 };
654 let len = z.len();
655 z.walk(0, len, false, f);
656 Ok(len)
657 }
658
659 pub fn zsetop_store<'k>(
666 &mut self,
667 op: Op,
668 destination: &[u8],
669 keys: impl Iterator<Item = &'k [u8]>,
670 weights: &[f64],
671 agg: Aggregate,
672 ) -> Result<usize> {
673 let slots = self.operand_slots(keys)?;
674 let got = zsetops::gather(op, &self.operands_of(&slots), weights, agg);
675 let limits = self.zset_limits;
676 Ok(self.put_zset(destination, Zset::from_elements(got, &limits)))
677 }
678
679 pub fn zintercard<'k>(
685 &mut self,
686 keys: impl Iterator<Item = &'k [u8]>,
687 limit: usize,
688 ) -> Result<usize> {
689 let slots = self.operand_slots(keys)?;
690 Ok(zsetops::intercard(&self.operands_of(&slots), limit))
691 }
692
693 pub fn zrangestore(
700 &mut self,
701 destination: &[u8],
702 source: &[u8],
703 q: &Query<'_>,
704 ) -> Result<usize> {
705 let built = match self.zset_slot(source)? {
706 None => None,
707 Some(at) => {
708 let z = self.zset_at(at);
709 let w = window(z, q);
710 let mut got = Elements::with_capacity(w.count.max(16));
711 let mut digits = [0u8; DIGITS_MAX];
712 let from = if w.rev { w.from + 1 - w.count } else { w.from };
717 z.walk(from, w.count, false, |m, s| {
718 let _ = got.insert(member_bytes(m, &mut digits), s);
719 });
720 let limits = self.zset_limits;
721 Zset::from_elements(got, &limits)
722 }
723 };
724 Ok(self.put_zset(destination, built))
725 }
726
727 fn operand_slots<'k>(
734 &mut self,
735 keys: impl Iterator<Item = &'k [u8]>,
736 ) -> Result<Vec<Option<(Kind, u32)>>> {
737 let mut out = Vec::with_capacity(keys.size_hint().0);
738 for key in keys {
739 out.push(self.live_slot_either(key, Kind::Zset, Kind::Set)?);
740 }
741 Ok(out)
742 }
743
744 fn operands_of(&self, slots: &[Option<(Kind, u32)>]) -> Vec<Operand<'_>> {
747 slots
748 .iter()
749 .map(|got| match got {
750 Some((Kind::Zset, at)) => Operand::Zset(self.zset_at(*at)),
751 Some((Kind::Set, at)) => {
752 Operand::Set(self.sets.get(*at).expect("the record points at its body"))
753 }
754 _ => Operand::Missing,
755 })
756 .collect()
757 }
758
759 pub(crate) fn put_zset(&mut self, key: &[u8], z: Option<Zset>) -> usize {
769 let Some(z) = z else {
770 self.drop_key(key);
771 return 0;
772 };
773 self.free_body(key);
774 let len = z.len();
775 let at = self.zsets.insert(z);
776 let record = value::slot_record_len(false);
777 self.write_rec(key, record, |out| {
778 value::write_slot_record(out, Kind::Zset, at, None);
779 });
780 self.bodies += 1;
781 len
782 }
783
784 #[inline]
786 pub(crate) fn zset_slot(&mut self, key: &[u8]) -> Result<Option<u32>> {
787 self.live_slot(key, Kind::Zset)
788 }
789
790 #[inline]
795 pub(crate) fn zset_at(&self, at: u32) -> &Zset {
796 self.zsets.get(at).expect("the record points at its body")
797 }
798
799 fn new_zset(&mut self, key: &[u8]) -> u32 {
805 let at = yo_alloc::first_touch(|| self.zsets.insert(Zset::new()));
809 let len = value::slot_record_len(false);
810 self.write_rec(key, len, |out| {
811 value::write_slot_record(out, Kind::Zset, at, None);
812 });
813 self.bodies += 1;
814 at
815 }
816}
817
818type Home = (usize, Kind, u32);
821
822impl Db {
823 pub fn zsetop<'k, F>(
831 &self,
832 op: Op,
833 keys: impl Iterator<Item = &'k [u8]> + Clone,
834 weights: &[f64],
835 agg: Aggregate,
836 f: F,
837 ) -> Result<usize>
838 where
839 F: FnMut(Member<'_>, f64),
840 {
841 if let Some(home) = self.one_stripe(keys.clone()) {
842 return self.hold_stripe(home).zsetop(op, keys, weights, agg, f);
843 }
844 let mut held = self.hold_operands(keys.clone(), Some(0));
847 let slots = self.operand_slots(&mut held, keys)?;
848 let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
849 let limits = held.stripe(0).zset_limits;
850 let Some(z) = Zset::from_elements(got, &limits) else {
851 return Ok(0);
852 };
853 let len = z.len();
854 z.walk(0, len, false, f);
855 Ok(len)
856 }
857
858 pub fn zsetop_store<'k>(
865 &self,
866 op: Op,
867 destination: &'k [u8],
868 keys: impl Iterator<Item = &'k [u8]> + Clone,
869 weights: &[f64],
870 agg: Aggregate,
871 ) -> Result<usize> {
872 if let Some(home) = self.one_stripe(std::iter::once(destination).chain(keys.clone())) {
873 return self
874 .hold_stripe(home)
875 .zsetop_store(op, destination, keys, weights, agg);
876 }
877 let onto = self.stripe_of(destination);
878 let mut held = self.hold_operands(keys.clone(), Some(onto));
879 let slots = self.operand_slots(&mut held, keys)?;
880 let got = zsetops::gather(op, &operands_of(&held, &slots), weights, agg);
881 let limits = held.stripe(onto).zset_limits;
884 let built = Zset::from_elements(got, &limits);
885 Ok(held.stripe_mut(onto).put_zset(destination, built))
886 }
887
888 pub fn zintercard<'k>(
890 &self,
891 keys: impl Iterator<Item = &'k [u8]> + Clone,
892 limit: usize,
893 ) -> Result<usize> {
894 if let Some(home) = self.one_stripe(keys.clone()) {
895 return self.hold_stripe(home).zintercard(keys, limit);
896 }
897 let mut held = self.hold_operands(keys.clone(), None);
898 let slots = self.operand_slots(&mut held, keys)?;
899 Ok(zsetops::intercard(&operands_of(&held, &slots), limit))
900 }
901
902 pub fn zrangestore(&self, destination: &[u8], source: &[u8], q: &Query<'_>) -> Result<usize> {
910 let (onto, home) = (self.stripe_of(destination), self.stripe_of(source));
911 if onto == home {
912 return self.hold_stripe(onto).zrangestore(destination, source, q);
913 }
914 let mut held = self.hold_many([home, onto].into_iter());
920 let slot = held.stripe_mut(home).zset_slot(source)?;
924 let built = match slot {
925 None => None,
926 Some(at) => {
927 let z = held.stripe(home).zset_at(at);
928 let w = window(z, q);
929 let mut got = Elements::with_capacity(w.count.max(16));
930 let mut digits = [0u8; DIGITS_MAX];
931 let from = if w.rev { w.from + 1 - w.count } else { w.from };
932 z.walk(from, w.count, false, |m, s| {
933 let _ = got.insert(member_bytes(m, &mut digits), s);
934 });
935 let limits = held.stripe(onto).zset_limits;
936 Zset::from_elements(got, &limits)
937 }
938 };
939 Ok(held.stripe_mut(onto).put_zset(destination, built))
940 }
941
942 fn hold_operands<'k>(
950 &self,
951 keys: impl Iterator<Item = &'k [u8]>,
952 onto: Option<usize>,
953 ) -> Holds<'_> {
954 let named = keys.map(|key| self.stripe_of(key));
955 self.hold_many(named.chain(onto))
956 }
957
958 fn operand_slots<'k>(
967 &self,
968 held: &mut Holds<'_>,
969 keys: impl Iterator<Item = &'k [u8]>,
970 ) -> Result<Vec<Option<Home>>> {
971 let mut out = Vec::with_capacity(keys.size_hint().0);
972 for key in keys {
973 let stripe = self.stripe_of(key);
974 let got = held
975 .stripe_mut(stripe)
976 .live_slot_either(key, Kind::Zset, Kind::Set)?;
977 out.push(got.map(|(kind, at)| (stripe, kind, at)));
978 }
979 Ok(out)
980 }
981}
982
983fn operands_of<'h>(held: &'h Holds<'_>, slots: &[Option<Home>]) -> Vec<Operand<'h>> {
985 slots
986 .iter()
987 .map(|got| match got {
988 Some((stripe, Kind::Zset, at)) => Operand::Zset(held.stripe(*stripe).zset_at(*at)),
989 Some((stripe, Kind::Set, at)) => Operand::Set(
990 held.stripe(*stripe)
991 .sets
992 .get(*at)
993 .expect("the record points at its body"),
994 ),
995 _ => Operand::Missing,
996 })
997 .collect()
998}
999
1000fn gated(z: &Zset, member: &[u8], score: f64, opts: ZAdd) -> Option<Added> {
1006 match z.score(member) {
1007 None => match opts.gate {
1008 Gate::IfPresent => None,
1009 _ => Some(Added::New),
1012 },
1013 Some(was) => {
1014 if opts.gate == Gate::IfMissing {
1015 return None;
1016 }
1017 let ok = match opts.only {
1018 Move::Any => true,
1019 Move::Up => score > was,
1020 Move::Down => score < was,
1021 };
1022 if !ok {
1023 return None;
1024 }
1025 Some(if score == was {
1028 Added::Same
1029 } else {
1030 Added::Changed
1031 })
1032 }
1033 }
1034}
1035
1036fn window(z: &Zset, q: &Query<'_>) -> Window {
1038 let len = z.len();
1039 let range = match q.by {
1040 By::Rank { start, stop } => {
1041 let (from, count) = rank_span(start, stop, len);
1045 if q.rev {
1046 return apply_limit(
1047 Window {
1048 from: len - from - 1,
1049 count,
1050 rev: true,
1051 },
1052 q,
1053 true,
1054 );
1055 }
1056 return apply_limit(
1057 Window {
1058 from,
1059 count,
1060 rev: false,
1061 },
1062 q,
1063 true,
1064 );
1065 }
1066 By::Score { min, max } => z.window_by_score(min, max),
1067 By::Lex { min, max } => z.window_by_lex(min, max),
1068 };
1069 let count = range.end - range.start;
1070 let from = if q.rev {
1071 range.end.saturating_sub(1)
1072 } else {
1073 range.start
1074 };
1075 apply_limit(
1076 Window {
1077 from,
1078 count,
1079 rev: q.rev,
1080 },
1081 q,
1082 false,
1083 )
1084}
1085
1086fn apply_limit(w: Window, q: &Query<'_>, skip: bool) -> Window {
1095 if skip {
1096 return w;
1097 }
1098 let offset = q.offset.min(w.count);
1099 let count = q.count.unwrap_or(usize::MAX).min(w.count - offset);
1100 let from = if w.rev {
1101 w.from.saturating_sub(offset)
1102 } else {
1103 w.from + offset
1104 };
1105 Window {
1106 from,
1107 count,
1108 rev: w.rev,
1109 }
1110}
1111
1112fn rank_span(start: i64, stop: i64, len: usize) -> (usize, usize) {
1119 if len == 0 {
1120 return (0, 0);
1121 }
1122 let len = len as i64;
1123 let from = if start < 0 {
1124 (len + start).max(0)
1125 } else {
1126 start.min(len)
1127 };
1128 let to = if stop < 0 {
1129 len + stop
1130 } else {
1131 stop.min(len - 1)
1132 };
1133 if to < from {
1134 return (from as usize, 0);
1135 }
1136 (from as usize, (to - from + 1) as usize)
1137}
1138
1139pub(crate) fn member_bytes<'a>(
1142 m: Member<'a>,
1143 digits: &'a mut [u8; yo_common::num::DIGITS_MAX],
1144) -> &'a [u8] {
1145 match m {
1146 Member::Str(s) => s,
1147 Member::Int(n) => yo_common::num::i64_digits(digits, n),
1148 }
1149}
1150
1151#[cfg(test)]
1152mod tests {
1153 use super::*;
1154 use yo_common::num::DIGITS_MAX;
1155
1156 fn ks() -> Keyspace {
1157 Keyspace::new()
1158 }
1159
1160 fn add(k: &mut Keyspace, key: &[u8], pairs: &[(f64, &[u8])]) -> usize {
1162 k.zadd(key, pairs.iter().copied(), ZAdd::default()).unwrap()
1163 }
1164
1165 fn names(k: &mut Keyspace, key: &[u8], q: &Query<'_>) -> Vec<String> {
1167 let w = k.zwindow(key, q).unwrap();
1168 let mut out = Vec::new();
1169 let mut digits = [0u8; DIGITS_MAX];
1170 k.zwalk(key, w, |m, _| {
1171 out.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1172 })
1173 .unwrap();
1174 assert_eq!(
1175 out.len(),
1176 w.count,
1177 "the window said {} and the walk gave {}",
1178 w.count,
1179 out.len()
1180 );
1181 out
1182 }
1183
1184 #[test]
1189 fn zrandmember_stops_allocating_once_its_buffer_is_grown() {
1190 let mut k = ks();
1191 add(
1192 &mut k,
1193 b"z",
1194 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1195 );
1196 k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1197 let (_, allocs) = crate::tally::counted(|| {
1198 for _ in 0..100 {
1199 k.zrandmember(b"z", 2, |_, _| {}).expect("a sorted set");
1200 }
1201 });
1202 assert_eq!(
1203 allocs, 0,
1204 "zrandmember allocated {allocs} times in a hundred"
1205 );
1206 }
1207
1208 #[test]
1212 fn zrandmember_hands_its_buffer_back() {
1213 let mut k = ks();
1214 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1215 for _ in 0..3 {
1216 let mut seen = Vec::new();
1217 let n = k
1218 .zrandmember(b"z", 2, |m, _| {
1219 let mut digits = [0u8; DIGITS_MAX];
1220 seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1221 })
1222 .expect("a sorted set");
1223 assert_eq!(n, 2);
1224 assert_eq!(seen.len(), 2);
1225 assert_ne!(seen[0], seen[1], "a positive count draws without replacing");
1226 }
1227 }
1228
1229 #[test]
1230 fn a_missing_key_is_an_empty_sorted_set() {
1231 let mut k = ks();
1232 assert_eq!(k.zcard(b"nope").unwrap(), 0);
1233 assert_eq!(k.zscore(b"nope", b"m").unwrap(), None);
1234 assert_eq!(k.zrank(b"nope", b"m", false).unwrap(), None);
1235 assert_eq!(k.zrem(b"nope", [b"m".as_slice()].into_iter()).unwrap(), 0);
1236 assert_eq!(
1237 names(&mut k, b"nope", &Query::rank(0, -1)),
1238 Vec::<String>::new()
1239 );
1240 assert!(!k.exists(b"nope"));
1241 }
1242
1243 #[test]
1244 fn a_key_holding_something_else_is_a_wrongtype() {
1245 let mut k = ks();
1246 k.set(b"s", b"hello", strings::SetOptions::default())
1247 .unwrap();
1248 assert_eq!(k.zcard(b"s").unwrap_err().code(), Code::WrongType);
1249 assert_eq!(
1250 k.zadd(b"s", [(1.0, b"m".as_slice())].into_iter(), ZAdd::default())
1251 .unwrap_err()
1252 .code(),
1253 Code::WrongType
1254 );
1255 assert_eq!(k.zscore(b"s", b"m").unwrap_err().code(), Code::WrongType);
1256 }
1257
1258 #[test]
1259 fn adding_answers_how_many_were_new() {
1260 let mut k = ks();
1261 assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]), 2);
1262 assert_eq!(add(&mut k, b"z", &[(1.0, b"a"), (3.0, b"c")]), 1);
1263 assert_eq!(k.zcard(b"z").unwrap(), 3);
1264 assert_eq!(k.zscore(b"z", b"c").unwrap(), Some(3.0));
1265 assert_eq!(k.kind_of(b"z").map(Kind::name), Some("zset"));
1266 assert_eq!(k.encoding_name(b"z"), Some("listpack"));
1267 }
1268
1269 #[test]
1270 fn the_ch_flag_counts_moved_scores_too() {
1271 let mut k = ks();
1272 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1273 let ch = ZAdd {
1274 changed: true,
1275 ..ZAdd::default()
1276 };
1277 let pairs = [
1279 (9.0, b"a".as_slice()),
1280 (2.0, b"b".as_slice()),
1281 (3.0, b"c".as_slice()),
1282 ];
1283 assert_eq!(k.zadd(b"z", pairs.into_iter(), ch).unwrap(), 2);
1284 assert_eq!(k.zadd(b"z", pairs.into_iter(), ZAdd::default()).unwrap(), 0);
1285 }
1286
1287 #[test]
1288 fn the_gates_let_the_right_members_through() {
1289 let mut k = ks();
1290 add(&mut k, b"z", &[(1.0, b"a")]);
1291 let nx = ZAdd {
1292 gate: Gate::IfMissing,
1293 ..ZAdd::default()
1294 };
1295 let xx = ZAdd {
1296 gate: Gate::IfPresent,
1297 ..ZAdd::default()
1298 };
1299 assert_eq!(
1300 k.zadd(b"z", [(5.0, b"a".as_slice())].into_iter(), nx)
1301 .unwrap(),
1302 0
1303 );
1304 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(1.0));
1305 assert_eq!(
1306 k.zadd(b"z", [(5.0, b"b".as_slice())].into_iter(), nx)
1307 .unwrap(),
1308 1
1309 );
1310 assert_eq!(
1311 k.zadd(b"z", [(7.0, b"c".as_slice())].into_iter(), xx)
1312 .unwrap(),
1313 0
1314 );
1315 assert_eq!(k.zscore(b"z", b"c").unwrap(), None);
1316 assert_eq!(
1317 k.zadd(b"z", [(7.0, b"a".as_slice())].into_iter(), xx)
1318 .unwrap(),
1319 0
1320 );
1321 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(7.0));
1322 assert_eq!(
1324 k.zadd(b"gone", [(1.0, b"a".as_slice())].into_iter(), xx)
1325 .unwrap(),
1326 0
1327 );
1328 assert!(!k.exists(b"gone"));
1329 }
1330
1331 #[test]
1332 fn gt_and_lt_only_move_a_score_one_way() {
1333 let mut k = ks();
1334 add(&mut k, b"z", &[(5.0, b"a")]);
1335 let gt = ZAdd {
1336 only: Move::Up,
1337 changed: true,
1338 ..ZAdd::default()
1339 };
1340 let lt = ZAdd {
1341 only: Move::Down,
1342 changed: true,
1343 ..ZAdd::default()
1344 };
1345 assert_eq!(
1346 k.zadd(b"z", [(3.0, b"a".as_slice())].into_iter(), gt)
1347 .unwrap(),
1348 0
1349 );
1350 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(5.0));
1351 assert_eq!(
1352 k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), gt)
1353 .unwrap(),
1354 1
1355 );
1356 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(9.0));
1357 assert_eq!(
1358 k.zadd(b"z", [(9.0, b"a".as_slice())].into_iter(), lt)
1359 .unwrap(),
1360 0
1361 );
1362 assert_eq!(
1363 k.zadd(b"z", [(2.0, b"a".as_slice())].into_iter(), lt)
1364 .unwrap(),
1365 1
1366 );
1367 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.0));
1368 assert_eq!(
1370 k.zadd(b"z", [(1.0, b"new".as_slice())].into_iter(), gt)
1371 .unwrap(),
1372 1
1373 );
1374 }
1375
1376 #[test]
1377 fn incrementing_adds_to_what_is_there_or_to_nothing() {
1378 let mut k = ks();
1379 let plain = ZAdd::default();
1380 assert_eq!(k.zincrby(b"z", b"a", 5.0, plain).unwrap(), Some(5.0));
1381 assert_eq!(k.zincrby(b"z", b"a", -2.5, plain).unwrap(), Some(2.5));
1382 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1383 let nx = ZAdd {
1384 gate: Gate::IfMissing,
1385 ..ZAdd::default()
1386 };
1387 assert_eq!(k.zincrby(b"z", b"a", 1.0, nx).unwrap(), None);
1388 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(2.5));
1389 let xx = ZAdd {
1390 gate: Gate::IfPresent,
1391 ..ZAdd::default()
1392 };
1393 assert_eq!(k.zincrby(b"z", b"never", 1.0, xx).unwrap(), None);
1394 assert!(k.zscore(b"z", b"never").unwrap().is_none());
1395 let gt = ZAdd {
1396 only: Move::Up,
1397 ..ZAdd::default()
1398 };
1399 assert_eq!(k.zincrby(b"z", b"a", -1.0, gt).unwrap(), None);
1400 assert_eq!(k.zincrby(b"z", b"a", 1.0, gt).unwrap(), Some(3.5));
1401 }
1402
1403 #[test]
1404 fn a_score_that_is_not_a_number_is_refused() {
1405 let mut k = ks();
1406 let plain = ZAdd::default();
1407 assert_eq!(
1408 k.zadd(b"z", [(f64::NAN, b"a".as_slice())].into_iter(), plain)
1409 .unwrap_err()
1410 .code(),
1411 Code::Invalid
1412 );
1413 assert!(!k.exists(b"z"));
1414 k.zincrby(b"z", b"a", f64::INFINITY, plain).unwrap();
1415 let err = k.zincrby(b"z", b"a", f64::NEG_INFINITY, plain).unwrap_err();
1416 assert_eq!(err.code(), Code::Invalid);
1417 assert_eq!(k.zscore(b"z", b"a").unwrap(), Some(f64::INFINITY));
1419 }
1420
1421 #[test]
1422 fn a_sorted_set_that_loses_its_last_member_loses_its_key() {
1423 let mut k = ks();
1424 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1425 assert_eq!(
1426 k.zrem(b"z", [b"a".as_slice(), b"b".as_slice()].into_iter())
1427 .unwrap(),
1428 2
1429 );
1430 assert!(!k.exists(b"z"));
1431 assert_eq!(k.zcard(b"z").unwrap(), 0);
1432 }
1433
1434 #[test]
1435 fn ranks_count_from_both_ends() {
1436 let mut k = ks();
1437 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1438 assert_eq!(k.zrank(b"z", b"a", false).unwrap(), Some((0, 1.0)));
1439 assert_eq!(k.zrank(b"z", b"c", false).unwrap(), Some((2, 3.0)));
1440 assert_eq!(k.zrank(b"z", b"a", true).unwrap(), Some((2, 1.0)));
1441 assert_eq!(k.zrank(b"z", b"c", true).unwrap(), Some((0, 3.0)));
1442 assert_eq!(k.zrank(b"z", b"nope", false).unwrap(), None);
1443 }
1444
1445 #[test]
1446 fn a_rank_range_clamps_every_way_it_can_be_wrong() {
1447 let mut k = ks();
1448 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1449 assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), ["a", "b", "c"]);
1450 assert_eq!(names(&mut k, b"z", &Query::rank(1, 1)), ["b"]);
1451 assert_eq!(names(&mut k, b"z", &Query::rank(-2, -1)), ["b", "c"]);
1452 assert_eq!(names(&mut k, b"z", &Query::rank(-99, 99)), ["a", "b", "c"]);
1453 assert_eq!(
1454 names(&mut k, b"z", &Query::rank(2, 1)),
1455 Vec::<String>::new()
1456 );
1457 assert_eq!(
1458 names(&mut k, b"z", &Query::rank(5, 9)),
1459 Vec::<String>::new()
1460 );
1461 assert_eq!(
1462 names(&mut k, b"z", &Query::rank(0, -1).rev(true)),
1463 ["c", "b", "a"]
1464 );
1465 assert_eq!(
1466 names(&mut k, b"z", &Query::rank(0, 1).rev(true)),
1467 ["c", "b"]
1468 );
1469 }
1470
1471 #[test]
1472 fn a_score_range_walks_either_way_and_takes_a_limit() {
1473 let mut k = ks();
1474 add(
1475 &mut k,
1476 b"z",
1477 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1478 );
1479 let all = Query::score(
1480 Bound::closed(f64::NEG_INFINITY),
1481 Bound::closed(f64::INFINITY),
1482 );
1483 assert_eq!(names(&mut k, b"z", &all), ["a", "b", "c", "d"]);
1484 assert_eq!(names(&mut k, b"z", &all.rev(true)), ["d", "c", "b", "a"]);
1485 assert_eq!(
1486 names(
1487 &mut k,
1488 b"z",
1489 &Query::score(Bound::closed(2.0), Bound::closed(3.0))
1490 ),
1491 ["b", "c"]
1492 );
1493 assert_eq!(
1494 names(
1495 &mut k,
1496 b"z",
1497 &Query::score(Bound::open(2.0), Bound::open(4.0))
1498 ),
1499 ["c"]
1500 );
1501 assert_eq!(names(&mut k, b"z", &all.limit(1, Some(2))), ["b", "c"]);
1504 assert_eq!(
1505 names(&mut k, b"z", &all.rev(true).limit(1, Some(2))),
1506 ["c", "b"]
1507 );
1508 assert_eq!(
1509 names(&mut k, b"z", &all.limit(9, Some(2))),
1510 Vec::<String>::new()
1511 );
1512 assert_eq!(names(&mut k, b"z", &all.limit(2, None)), ["c", "d"]);
1513 assert_eq!(
1514 k.zcount(b"z", &Query::score(Bound::closed(2.0), Bound::closed(3.0)))
1515 .unwrap(),
1516 2
1517 );
1518 }
1519
1520 #[test]
1521 fn a_member_range_orders_by_member_when_every_score_is_the_same() {
1522 let mut k = ks();
1523 add(
1524 &mut k,
1525 b"z",
1526 &[(0.0, b"a"), (0.0, b"b"), (0.0, b"c"), (0.0, b"d")],
1527 );
1528 assert_eq!(
1529 names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max)),
1530 ["a", "b", "c", "d"]
1531 );
1532 assert_eq!(
1533 names(&mut k, b"z", &Query::lex(Lex::Incl(b"b"), Lex::Incl(b"c"))),
1534 ["b", "c"]
1535 );
1536 assert_eq!(
1537 names(&mut k, b"z", &Query::lex(Lex::Excl(b"a"), Lex::Excl(b"d"))),
1538 ["b", "c"]
1539 );
1540 assert_eq!(
1541 names(&mut k, b"z", &Query::lex(Lex::Min, Lex::Max).rev(true)),
1542 ["d", "c", "b", "a"]
1543 );
1544 assert_eq!(
1545 k.zcount(b"z", &Query::lex(Lex::Incl(b"b"), Lex::Max))
1546 .unwrap(),
1547 3
1548 );
1549 }
1550
1551 #[test]
1552 fn removing_a_range_takes_out_exactly_what_the_walk_would_have_given() {
1553 let mut k = ks();
1554 for (q, left) in [
1555 (Query::rank(0, 1), vec!["c", "d", "e"]),
1556 (Query::rank(-2, -1), vec!["a", "b", "c"]),
1557 (
1558 Query::score(Bound::closed(2.0), Bound::closed(4.0)),
1559 vec!["a", "e"],
1560 ),
1561 (
1562 Query::lex(Lex::Incl(b"b"), Lex::Excl(b"d")),
1563 vec!["a", "d", "e"],
1564 ),
1565 (Query::rank(0, -1).rev(true), vec![]),
1566 ] {
1567 k.del(b"z");
1568 add(
1569 &mut k,
1570 b"z",
1571 &[
1572 (1.0, b"a"),
1573 (2.0, b"b"),
1574 (3.0, b"c"),
1575 (4.0, b"d"),
1576 (5.0, b"e"),
1577 ],
1578 );
1579 let want = names(&mut k, b"z", &q).len();
1580 assert_eq!(k.zremrange(b"z", &q).unwrap(), want, "{q:?}");
1581 assert_eq!(names(&mut k, b"z", &Query::rank(0, -1)), left, "{q:?}");
1582 }
1583 assert!(!k.exists(b"z"));
1585 }
1586
1587 #[test]
1588 fn popping_takes_from_the_end_it_was_told_to() {
1589 let mut k = ks();
1590 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1591 let mut got = Vec::new();
1592 let mut digits = [0u8; DIGITS_MAX];
1593 k.zpop(b"z", From::Min, 2, |m, s| {
1594 got.push((
1595 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1596 s,
1597 ));
1598 })
1599 .unwrap();
1600 assert_eq!(got, [("a".to_string(), 1.0), ("b".to_string(), 2.0)]);
1601 assert_eq!(
1602 k.zpop_one(b"z", From::Max).unwrap(),
1603 Some((b"c".to_vec(), 3.0))
1604 );
1605 assert!(!k.exists(b"z"));
1606 add(&mut k, b"z", &[(1.0, b"a")]);
1608 assert_eq!(k.zpop(b"z", From::Max, 99, |_, _| {}).unwrap(), 1);
1609 assert!(!k.exists(b"z"));
1610 assert_eq!(k.zpop_one(b"z", From::Min).unwrap(), None);
1611 }
1612
1613 #[test]
1614 fn a_random_draw_is_with_or_without_replacement_by_the_sign_of_the_count() {
1615 let mut k = ks();
1616 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c")]);
1617 let mut digits = [0u8; DIGITS_MAX];
1618 let mut seen = Vec::new();
1619 k.zrandmember(b"z", 2, |m, _| {
1620 seen.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1621 })
1622 .unwrap();
1623 assert_eq!(seen.len(), 2);
1624 seen.sort();
1625 seen.dedup();
1626 assert_eq!(seen.len(), 2, "a positive count does not repeat a member");
1627 let mut all = Vec::new();
1629 k.zrandmember(b"z", 99, |m, _| {
1630 all.push(String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap());
1631 })
1632 .unwrap();
1633 all.sort();
1634 assert_eq!(all, ["a", "b", "c"]);
1635 let mut many = 0;
1637 k.zrandmember(b"z", -10, |_, _| many += 1).unwrap();
1638 assert_eq!(many, 10);
1639 assert_eq!(k.zrandmember(b"nope", 3, |_, _| {}).unwrap(), 0);
1640 }
1641
1642 #[test]
1643 fn a_scan_of_either_band_walks_every_member_once() {
1644 let mut k = ks();
1645 for entries in [8usize, 4096] {
1646 k.del(b"z");
1647 let pairs: Vec<(f64, Vec<u8>)> = (0..entries)
1648 .map(|i| (i as f64, format!("m{i:05}").into_bytes()))
1649 .collect();
1650 k.zadd(
1651 b"z",
1652 pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1653 ZAdd::default(),
1654 )
1655 .unwrap();
1656 let mut seen = Vec::new();
1657 let mut digits = [0u8; DIGITS_MAX];
1658 let mut cursor = Cursor::START;
1659 loop {
1660 cursor = k
1661 .zscan(b"z", cursor, 16, |m, _| {
1662 seen.push(
1663 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1664 );
1665 })
1666 .unwrap();
1667 if cursor.is_end() {
1668 break;
1669 }
1670 }
1671 seen.sort();
1672 seen.dedup();
1673 assert_eq!(seen.len(), entries, "{entries} members");
1674 }
1675 }
1676
1677 #[test]
1678 fn a_big_sorted_set_promotes_and_still_answers_every_rank() {
1679 let mut k = ks();
1680 let pairs: Vec<(f64, Vec<u8>)> = (0..5_000)
1681 .map(|i| (f64::from(i), format!("m{i:05}").into_bytes()))
1682 .collect();
1683 assert_eq!(
1684 k.zadd(
1685 b"z",
1686 pairs.iter().map(|(s, m)| (*s, m.as_slice())),
1687 ZAdd::default()
1688 )
1689 .unwrap(),
1690 5_000
1691 );
1692 assert_eq!(k.encoding_name(b"z"), Some("skiplist"));
1693 assert_eq!(k.zcard(b"z").unwrap(), 5_000);
1694 assert_eq!(
1695 k.zrank(b"z", b"m02500", false).unwrap(),
1696 Some((2_500, 2500.0))
1697 );
1698 let q = Query::score(Bound::closed(1000.0), Bound::open(1010.0));
1699 assert_eq!(k.zcount(b"z", &q).unwrap(), 10);
1700 assert_eq!(
1701 names(&mut k, b"z", &q).first().map(String::as_str),
1702 Some("m01000")
1703 );
1704 assert_eq!(k.zremrange(b"z", &Query::rank(0, 2_499)).unwrap(), 2_500);
1706 assert_eq!(k.zcard(b"z").unwrap(), 2_500);
1707 assert_eq!(k.zrank(b"z", b"m02500", false).unwrap(), Some((0, 2500.0)));
1708 assert_eq!(k.zrank(b"z", b"m00000", false).unwrap(), None);
1709 }
1710
1711 #[test]
1712 fn a_deadline_on_a_sorted_set_leaves_its_members_alone() {
1713 let mut k = ks();
1714 add(&mut k, b"z", &[(1.0, b"a"), (2.0, b"b")]);
1715 assert!(k.set_expiry(b"z", Some(u64::MAX / 2)));
1716 assert_eq!(k.zcard(b"z").unwrap(), 2);
1717 assert_eq!(k.zscore(b"z", b"b").unwrap(), Some(2.0));
1718 assert!(k.set_expiry(b"z", None));
1719 assert_eq!(k.zcard(b"z").unwrap(), 2);
1720 }
1721
1722 #[test]
1723 fn writing_a_string_over_a_sorted_set_gives_its_body_back() {
1724 let mut k = ks();
1725 add(&mut k, b"z", &[(1.0, b"a")]);
1726 let held = k.memory_bytes();
1727 k.set(b"z", b"now a string", strings::SetOptions::default())
1728 .unwrap();
1729 assert_eq!(k.kind_of(b"z").map(Kind::name), Some("string"));
1730 assert!(
1731 k.memory_bytes() < held,
1732 "the sorted set's body was not freed"
1733 );
1734 }
1735
1736 fn all(k: &mut Keyspace, key: &[u8]) -> Vec<(String, f64)> {
1738 let q = Query::rank(0, -1);
1739 let w = k.zwindow(key, &q).unwrap();
1740 let mut out = Vec::new();
1741 let mut digits = [0u8; DIGITS_MAX];
1742 k.zwalk(key, w, |m, s| {
1743 out.push((
1744 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1745 s,
1746 ));
1747 })
1748 .unwrap();
1749 out
1750 }
1751
1752 fn got(
1754 k: &mut Keyspace,
1755 op: Op,
1756 keys: &[&[u8]],
1757 weights: &[f64],
1758 agg: Aggregate,
1759 ) -> Vec<(String, f64)> {
1760 let mut out = Vec::new();
1761 let mut digits = [0u8; DIGITS_MAX];
1762 let n = k
1763 .zsetop(op, keys.iter().copied(), weights, agg, |m, s| {
1764 out.push((
1765 String::from_utf8(member_bytes(m, &mut digits).to_vec()).unwrap(),
1766 s,
1767 ));
1768 })
1769 .unwrap();
1770 assert_eq!(out.len(), n, "the count and the walk disagree");
1771 out
1772 }
1773
1774 #[test]
1775 fn a_union_adds_the_scores_of_a_member_that_is_in_both() {
1776 let mut k = ks();
1777 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1778 add(&mut k, b"b", &[(10.0, b"y"), (3.0, b"z")]);
1779 assert_eq!(
1780 got(&mut k, Op::Union, &[b"a", b"b"], &[], Aggregate::Sum),
1781 [("x".into(), 1.0), ("z".into(), 3.0), ("y".into(), 12.0)]
1782 );
1783 }
1784
1785 #[test]
1786 fn an_intersection_keeps_only_what_every_input_has() {
1787 let mut k = ks();
1788 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1789 add(&mut k, b"b", &[(5.0, b"y"), (5.0, b"z")]);
1790 add(&mut k, b"c", &[(7.0, b"z")]);
1791 assert_eq!(
1792 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Sum),
1793 [("z".into(), 15.0)]
1794 );
1795 assert_eq!(
1796 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Min),
1797 [("z".into(), 3.0)]
1798 );
1799 assert_eq!(
1800 got(&mut k, Op::Inter, &[b"a", b"b", b"c"], &[], Aggregate::Max),
1801 [("z".into(), 7.0)]
1802 );
1803 }
1804
1805 #[test]
1806 fn a_difference_takes_the_first_input_and_removes_the_rest() {
1807 let mut k = ks();
1808 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y"), (3.0, b"z")]);
1809 add(&mut k, b"b", &[(99.0, b"y")]);
1810 assert_eq!(
1811 got(&mut k, Op::Diff, &[b"a", b"b"], &[], Aggregate::Sum),
1812 [("x".into(), 1.0), ("z".into(), 3.0)]
1813 );
1814 }
1815
1816 #[test]
1817 fn a_missing_key_keeps_its_place_so_the_weights_stay_lined_up() {
1818 let mut k = ks();
1819 add(&mut k, b"a", &[(1.0, b"x")]);
1820 add(&mut k, b"c", &[(1.0, b"x")]);
1821 let out = got(
1824 &mut k,
1825 Op::Union,
1826 &[b"a", b"gone", b"c"],
1827 &[2.0, 10.0, 3.0],
1828 Aggregate::Sum,
1829 );
1830 assert_eq!(out, [("x".into(), 5.0)]);
1831 }
1832
1833 #[test]
1834 fn a_plain_set_counts_as_every_score_being_one() {
1835 let mut k = ks();
1836 add(&mut k, b"z", &[(5.0, b"x")]);
1837 k.sadd(b"s", [b"x".as_slice(), b"y".as_slice()].into_iter())
1838 .unwrap();
1839 assert_eq!(
1840 got(&mut k, Op::Union, &[b"z", b"s"], &[], Aggregate::Sum),
1841 [("y".into(), 1.0), ("x".into(), 6.0)]
1842 );
1843 assert_eq!(
1844 got(&mut k, Op::Inter, &[b"z", b"s"], &[], Aggregate::Min),
1845 [("x".into(), 1.0)]
1846 );
1847 }
1848
1849 #[test]
1850 fn a_store_writes_the_result_and_answers_its_size() {
1851 let mut k = ks();
1852 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1853 add(&mut k, b"b", &[(10.0, b"y")]);
1854 assert_eq!(
1855 k.zsetop_store(
1856 Op::Union,
1857 b"d",
1858 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1859 &[],
1860 Aggregate::Sum
1861 )
1862 .unwrap(),
1863 2
1864 );
1865 assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1866 assert_eq!(k.zscore(b"d", b"y").unwrap(), Some(12.0));
1867 assert_eq!(k.zrank(b"d", b"y", false).unwrap(), Some((1, 12.0)));
1868 }
1869
1870 #[test]
1871 fn a_store_onto_one_of_its_own_sources_still_reads_the_old_body() {
1872 let mut k = ks();
1873 add(&mut k, b"a", &[(1.0, b"x"), (2.0, b"y")]);
1874 add(&mut k, b"b", &[(10.0, b"y")]);
1875 assert_eq!(
1876 k.zsetop_store(
1877 Op::Union,
1878 b"a",
1879 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1880 &[],
1881 Aggregate::Sum
1882 )
1883 .unwrap(),
1884 2
1885 );
1886 assert_eq!(all(&mut k, b"a"), [("x".into(), 1.0), ("y".into(), 12.0)]);
1887 }
1888
1889 #[test]
1890 fn a_store_with_nothing_in_it_deletes_the_destination() {
1891 let mut k = ks();
1892 add(&mut k, b"a", &[(1.0, b"x")]);
1893 add(&mut k, b"b", &[(1.0, b"y")]);
1894 add(&mut k, b"d", &[(1.0, b"old")]);
1895 assert_eq!(
1896 k.zsetop_store(
1897 Op::Inter,
1898 b"d",
1899 [b"a".as_slice(), b"b".as_slice()].into_iter(),
1900 &[],
1901 Aggregate::Sum
1902 )
1903 .unwrap(),
1904 0
1905 );
1906 assert!(!k.exists(b"d"));
1907 }
1908
1909 #[test]
1910 fn a_store_clears_the_deadline_the_destination_was_carrying() {
1911 let mut k = ks();
1912 add(&mut k, b"a", &[(1.0, b"x")]);
1913 add(&mut k, b"d", &[(1.0, b"old")]);
1914 assert!(k.set_expiry(b"d", Some(u64::MAX / 2)));
1915 k.zsetop_store(
1916 Op::Union,
1917 b"d",
1918 [b"a".as_slice()].into_iter(),
1919 &[],
1920 Aggregate::Sum,
1921 )
1922 .unwrap();
1923 assert_eq!(all(&mut k, b"d"), [("x".into(), 1.0)]);
1924 assert_eq!(k.expire_at(b"d"), None);
1925 }
1926
1927 #[test]
1928 fn the_algebra_refuses_a_key_holding_something_that_is_neither() {
1929 let mut k = ks();
1930 add(&mut k, b"a", &[(1.0, b"x")]);
1931 k.set(b"s", b"hello", strings::SetOptions::default())
1932 .unwrap();
1933 assert_eq!(
1934 k.zsetop(
1935 Op::Union,
1936 [b"a".as_slice(), b"s".as_slice()].into_iter(),
1937 &[],
1938 Aggregate::Sum,
1939 |_, _| {}
1940 )
1941 .unwrap_err()
1942 .code(),
1943 Code::WrongType
1944 );
1945 assert!(!k.exists(b"d"));
1948 assert_eq!(
1949 k.zsetop_store(
1950 Op::Union,
1951 b"d",
1952 [b"a".as_slice(), b"s".as_slice()].into_iter(),
1953 &[],
1954 Aggregate::Sum
1955 )
1956 .unwrap_err()
1957 .code(),
1958 Code::WrongType
1959 );
1960 assert!(!k.exists(b"d"));
1961 }
1962
1963 #[test]
1964 fn intercard_counts_without_building_anything_and_stops_at_the_limit() {
1965 let mut k = ks();
1966 add(
1967 &mut k,
1968 b"a",
1969 &[(1.0, b"w"), (2.0, b"x"), (3.0, b"y"), (4.0, b"z")],
1970 );
1971 add(&mut k, b"b", &[(1.0, b"x"), (1.0, b"y"), (1.0, b"z")]);
1972 assert_eq!(
1973 k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 0)
1974 .unwrap(),
1975 3
1976 );
1977 assert_eq!(
1978 k.zintercard([b"a".as_slice(), b"b".as_slice()].into_iter(), 2)
1979 .unwrap(),
1980 2
1981 );
1982 assert_eq!(
1983 k.zintercard([b"a".as_slice(), b"gone".as_slice()].into_iter(), 0)
1984 .unwrap(),
1985 0
1986 );
1987 }
1988
1989 #[test]
1990 fn a_range_store_copies_a_window_and_leaves_the_source_alone() {
1991 let mut k = ks();
1992 add(
1993 &mut k,
1994 b"z",
1995 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
1996 );
1997 assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(1, 2)).unwrap(), 2);
1998 assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
1999 assert_eq!(k.zcard(b"z").unwrap(), 4);
2000 }
2001
2002 #[test]
2003 fn a_reverse_range_store_picks_the_same_members_and_orders_them_the_same() {
2004 let mut k = ks();
2005 add(
2006 &mut k,
2007 b"z",
2008 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2009 );
2010 assert_eq!(
2013 k.zrangestore(b"d", b"z", &Query::rank(0, 1).rev(true))
2014 .unwrap(),
2015 2
2016 );
2017 assert_eq!(all(&mut k, b"d"), [("c".into(), 3.0), ("d".into(), 4.0)]);
2018 }
2019
2020 #[test]
2021 fn a_range_store_by_score_takes_the_bounds_the_range_would_have() {
2022 let mut k = ks();
2023 add(
2024 &mut k,
2025 b"z",
2026 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2027 );
2028 let q = Query::score(Bound::closed(2.0), Bound::open(4.0));
2029 assert_eq!(k.zrangestore(b"d", b"z", &q).unwrap(), 2);
2030 assert_eq!(all(&mut k, b"d"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2031 }
2032
2033 #[test]
2034 fn a_range_store_of_an_empty_window_deletes_the_destination() {
2035 let mut k = ks();
2036 add(&mut k, b"z", &[(1.0, b"a")]);
2037 add(&mut k, b"d", &[(1.0, b"old")]);
2038 assert_eq!(k.zrangestore(b"d", b"z", &Query::rank(5, 9)).unwrap(), 0);
2039 assert!(!k.exists(b"d"));
2040 assert_eq!(
2041 k.zrangestore(b"d", b"gone", &Query::rank(0, -1)).unwrap(),
2042 0
2043 );
2044 assert!(!k.exists(b"d"));
2045 }
2046
2047 #[test]
2048 fn a_range_store_onto_its_own_source_keeps_the_window() {
2049 let mut k = ks();
2050 add(
2051 &mut k,
2052 b"z",
2053 &[(1.0, b"a"), (2.0, b"b"), (3.0, b"c"), (4.0, b"d")],
2054 );
2055 assert_eq!(k.zrangestore(b"z", b"z", &Query::rank(1, 2)).unwrap(), 2);
2056 assert_eq!(all(&mut k, b"z"), [("b".into(), 2.0), ("c".into(), 3.0)]);
2057 }
2058
2059 #[test]
2060 fn a_big_result_comes_out_on_the_table_band_in_the_right_order() {
2061 let mut k = ks();
2062 let names: Vec<String> = (0..3000).map(|i| format!("member-{i:05}")).collect();
2063 for (i, name) in names.iter().enumerate() {
2064 add(&mut k, b"a", &[((i % 17) as f64, name.as_bytes())]);
2065 }
2066 for name in names.iter().step_by(3) {
2067 add(&mut k, b"b", &[(100.0, name.as_bytes())]);
2068 }
2069 let n = k
2070 .zsetop_store(
2071 Op::Union,
2072 b"d",
2073 [b"a".as_slice(), b"b".as_slice()].into_iter(),
2074 &[],
2075 Aggregate::Sum,
2076 )
2077 .unwrap();
2078 assert_eq!(n, 3000);
2079 assert_eq!(
2080 k.zset_encoding(b"d").map(crate::zset::Encoding::name),
2081 Some("skiplist")
2082 );
2083 let out = all(&mut k, b"d");
2084 assert_eq!(out.len(), 3000);
2085 let mut want = out.clone();
2086 want.sort_by(|x, y| x.1.partial_cmp(&y.1).unwrap().then_with(|| x.0.cmp(&y.0)));
2087 assert_eq!(out, want, "the result came out in the wrong order");
2088 for (i, name) in names.iter().enumerate() {
2089 let base = (i % 17) as f64;
2090 let want = if i % 3 == 0 { base + 100.0 } else { base };
2091 assert_eq!(k.zscore(b"d", name.as_bytes()).unwrap(), Some(want));
2092 }
2093 }
2094}