1use crate::{Store, StoreError};
6
7#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
11pub struct ZaddFlags {
12 pub nx: bool,
14 pub xx: bool,
16 pub gt: bool,
18 pub lt: bool,
20 pub ch: bool,
22}
23
24impl ZaddFlags {
25 pub fn valid(self) -> bool {
28 !(self.nx && (self.xx || self.gt || self.lt)) && !(self.gt && self.lt)
29 }
30}
31
32pub struct ZaddReport {
34 pub added: usize,
36 pub changed: usize,
38 pub applied: Vec<(f64, Vec<u8>)>,
44}
45
46impl Store {
47 pub fn zadd_flags_borrowed(
51 &mut self,
52 key: &[u8],
53 pairs: &[(f64, &[u8])],
54 flags: ZaddFlags,
55 ) -> Result<ZaddReport, StoreError> {
56 debug_assert!(flags.valid(), "caller must reject invalid flag combos");
57 let mut rep = ZaddReport { added: 0, changed: 0, applied: Vec::new() };
58 for (score, m) in pairs {
59 match self.zscore(key, m)? {
60 Some(old) => {
61 if flags.nx
62 || (flags.gt && *score <= old)
63 || (flags.lt && *score >= old)
64 {
65 continue;
66 }
67 if *score != old {
68 self.zadd_borrowed(key, &[(*score, m)])?;
69 rep.changed += 1;
70 rep.applied.push((*score, m.to_vec()));
71 }
72 }
73 None => {
74 if flags.xx {
75 continue;
76 }
77 self.zadd_borrowed(key, &[(*score, m)])?;
78 rep.added += 1;
79 rep.changed += 1;
80 rep.applied.push((*score, m.to_vec()));
81 }
82 }
83 }
84 Ok(rep)
85 }
86
87 pub fn zadd_incr(
90 &mut self,
91 key: &[u8],
92 delta: f64,
93 member: &[u8],
94 flags: ZaddFlags,
95 ) -> Result<Option<f64>, StoreError> {
96 debug_assert!(flags.valid(), "caller must reject invalid flag combos");
97 match self.zscore(key, member)? {
98 Some(old) => {
99 if flags.nx {
100 return Ok(None);
101 }
102 let next = old + delta;
103 if !next.is_finite() {
104 return Err(StoreError::NotFloat);
105 }
106 if (flags.gt && next <= old) || (flags.lt && next >= old) {
107 return Ok(None);
108 }
109 self.zadd_borrowed(key, &[(next, member)])?;
110 Ok(Some(next))
111 }
112 None => {
113 if flags.xx {
114 return Ok(None);
115 }
116 if !delta.is_finite() {
117 return Err(StoreError::NotFloat);
118 }
119 self.zadd_borrowed(key, &[(delta, member)])?;
120 Ok(Some(delta))
121 }
122 }
123 }
124}
125
126#[cfg(test)]
127mod tests {
128 use super::*;
129
130 fn zf() -> ZaddFlags {
131 ZaddFlags::default()
132 }
133
134 #[test]
135 fn validity_matrix() {
136 assert!(zf().valid());
137 assert!(ZaddFlags { gt: true, ch: true, ..zf() }.valid());
138 assert!(ZaddFlags { xx: true, gt: true, ..zf() }.valid());
139 assert!(!ZaddFlags { nx: true, xx: true, ..zf() }.valid());
140 assert!(!ZaddFlags { nx: true, gt: true, ..zf() }.valid());
141 assert!(!ZaddFlags { gt: true, lt: true, ..zf() }.valid());
142 }
143
144 #[test]
145 fn nx_only_adds() {
146 let mut s = Store::new();
147 s.zadd(b"z", &[(1.0, b"m".to_vec())]).unwrap();
148 let r = s
149 .zadd_flags_borrowed(b"z", &[(9.0, b"m"), (2.0, b"n")], ZaddFlags { nx: true, ..zf() })
150 .unwrap();
151 assert_eq!((r.added, r.changed), (1, 1));
152 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(1.0)); assert_eq!(s.zscore(b"z", b"n").unwrap(), Some(2.0));
154 }
155
156 #[test]
157 fn xx_only_updates() {
158 let mut s = Store::new();
159 s.zadd(b"z", &[(1.0, b"m".to_vec())]).unwrap();
160 let r = s
161 .zadd_flags_borrowed(b"z", &[(9.0, b"m"), (2.0, b"n")], ZaddFlags { xx: true, ..zf() })
162 .unwrap();
163 assert_eq!((r.added, r.changed), (0, 1));
164 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(9.0));
165 assert_eq!(s.zscore(b"z", b"n").unwrap(), None); }
167
168 #[test]
169 fn gt_is_monotonic_heal() {
170 let mut s = Store::new();
171 s.zadd(b"z", &[(5.0, b"m".to_vec())]).unwrap();
172 let gt = ZaddFlags { gt: true, ..zf() };
173 let r = s.zadd_flags_borrowed(b"z", &[(3.0, b"m")], gt).unwrap();
175 assert_eq!(r.changed, 0);
176 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(5.0));
177 let r = s.zadd_flags_borrowed(b"z", &[(7.0, b"m")], gt).unwrap();
179 assert_eq!(r.changed, 1);
180 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(7.0));
181 let r = s.zadd_flags_borrowed(b"z", &[(1.0, b"new")], gt).unwrap();
183 assert_eq!(r.added, 1);
184 }
185
186 #[test]
187 fn lt_mirror() {
188 let mut s = Store::new();
189 s.zadd(b"z", &[(5.0, b"m".to_vec())]).unwrap();
190 let lt = ZaddFlags { lt: true, ..zf() };
191 assert_eq!(s.zadd_flags_borrowed(b"z", &[(7.0, b"m")], lt).unwrap().changed, 0);
192 assert_eq!(s.zadd_flags_borrowed(b"z", &[(3.0, b"m")], lt).unwrap().changed, 1);
193 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(3.0));
194 }
195
196 #[test]
197 fn applied_reflects_effect_only() {
198 let mut s = Store::new();
199 s.zadd(b"z", &[(5.0, b"a".to_vec()), (5.0, b"b".to_vec())]).unwrap();
200 let r = s
201 .zadd_flags_borrowed(
202 b"z",
203 &[(9.0, b"a"), (1.0, b"b"), (5.0, b"c")],
204 ZaddFlags { gt: true, ..zf() },
205 )
206 .unwrap();
207 assert_eq!(
209 r.applied,
210 vec![(9.0, b"a".to_vec()), (5.0, b"c".to_vec())]
211 );
212 }
213
214 #[test]
215 fn incr_form_vetoes_to_none() {
216 let mut s = Store::new();
217 s.zadd(b"z", &[(5.0, b"m".to_vec())]).unwrap();
218 let gt = ZaddFlags { gt: true, ..zf() };
219 assert_eq!(s.zadd_incr(b"z", -2.0, b"m", gt).unwrap(), None);
221 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(5.0));
222 assert_eq!(s.zadd_incr(b"z", 2.0, b"m", gt).unwrap(), Some(7.0));
223 let xx = ZaddFlags { xx: true, ..zf() };
225 assert_eq!(s.zadd_incr(b"z", 1.0, b"nope", xx).unwrap(), None);
226 let nx = ZaddFlags { nx: true, ..zf() };
228 assert_eq!(s.zadd_incr(b"z", 1.0, b"m", nx).unwrap(), None);
229 }
230
231 #[test]
232 fn wrongtype_propagates() {
233 let mut s = Store::new();
234 s.set(b"str", b"v".to_vec(), None, false, false);
235 assert!(s.zadd_flags_borrowed(b"str", &[(1.0, b"m")], zf()).is_err());
236 assert!(s.zadd_incr(b"str", 1.0, b"m", zf()).is_err());
237 }
238}