1#[cfg(not(feature = "std"))]
6use crate::nostd_prelude::*;
7use crate::{ScoreCompare, SetCondition, Store, StoreError};
8
9#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
24pub struct ZaddFlags {
25 condition: SetCondition,
26 compare: ScoreCompare,
27 ch: bool,
28}
29
30impl ZaddFlags {
31 pub fn new(condition: SetCondition, compare: ScoreCompare) -> Option<Self> {
34 let legal = condition != SetCondition::IfAbsent || compare == ScoreCompare::Any;
35 legal.then_some(Self { condition, compare, ch: false })
36 }
37
38 #[must_use]
44 pub fn with_ch(mut self, ch: bool) -> Self {
45 self.ch = ch;
46 self
47 }
48
49 pub fn condition(self) -> SetCondition {
51 self.condition
52 }
53
54 pub fn compare(self) -> ScoreCompare {
56 self.compare
57 }
58
59 pub fn ch(self) -> bool {
61 self.ch
62 }
63
64 fn vetoes_update(self, old: f64, new: f64) -> bool {
67 self.condition == SetCondition::IfAbsent
68 || match self.compare {
69 ScoreCompare::Greater => new <= old,
70 ScoreCompare::Less => new >= old,
71 ScoreCompare::Any => false,
72 }
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Default)]
88#[non_exhaustive]
89pub struct ZaddReport {
90 pub added: usize,
101 pub changed: usize,
113 pub applied: Vec<(f64, Vec<u8>)>,
129}
130
131impl Store {
132 pub fn zadd_flags(
135 &mut self,
136 key: &[u8],
137 pairs: &[(f64, &[u8])],
138 flags: ZaddFlags,
139 ) -> Result<ZaddReport, StoreError> {
140 let mut rep = ZaddReport::default();
141 for (score, m) in pairs {
142 match self.zscore(key, m)? {
143 Some(old) => {
144 if flags.vetoes_update(old, *score) {
145 continue;
146 }
147 if *score != old {
148 self.zadd(key, &[(*score, m)])?;
149 rep.changed += 1;
150 rep.applied.push((*score, m.to_vec()));
151 }
152 }
153 None => {
154 if flags.condition == SetCondition::IfPresent {
155 continue;
156 }
157 self.zadd(key, &[(*score, m)])?;
158 rep.added += 1;
159 rep.changed += 1;
160 rep.applied.push((*score, m.to_vec()));
161 }
162 }
163 }
164 Ok(rep)
165 }
166
167 pub fn zadd_incr(
170 &mut self,
171 key: &[u8],
172 delta: f64,
173 member: &[u8],
174 flags: ZaddFlags,
175 ) -> Result<Option<f64>, StoreError> {
176 match self.zscore(key, member)? {
177 Some(old) => {
178 if flags.condition == SetCondition::IfAbsent {
179 return Ok(None);
180 }
181 let next = old + delta;
182 if !next.is_finite() {
183 return Err(StoreError::NotFloat);
184 }
185 if flags.vetoes_update(old, next) {
186 return Ok(None);
187 }
188 self.zadd(key, &[(next, member)])?;
189 Ok(Some(next))
190 }
191 None => {
192 if flags.condition == SetCondition::IfPresent {
193 return Ok(None);
194 }
195 if !delta.is_finite() {
196 return Err(StoreError::NotFloat);
197 }
198 self.zadd(key, &[(delta, member)])?;
199 Ok(Some(delta))
200 }
201 }
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208
209 fn zf() -> ZaddFlags {
210 ZaddFlags::default()
211 }
212
213 fn flags(c: SetCondition, s: ScoreCompare) -> ZaddFlags {
214 ZaddFlags::new(c, s).unwrap()
215 }
216
217 #[test]
218 fn only_nx_with_a_comparison_is_refused() {
219 use ScoreCompare::{Any, Greater, Less};
220 use SetCondition::{Always, IfAbsent, IfPresent};
221 for c in [Always, IfAbsent, IfPresent] {
222 for s in [Any, Greater, Less] {
223 let refused = c == IfAbsent && s != Any;
224 assert_eq!(ZaddFlags::new(c, s).is_none(), refused, "{c:?} {s:?}");
225 }
226 }
227 }
228
229 #[test]
230 fn nx_only_adds() {
231 let mut s = Store::new();
232 s.zadd(b"z", &[(1.0, b"m".as_slice())]).unwrap();
233 let r = s
234 .zadd_flags(
235 b"z",
236 &[(9.0, b"m"), (2.0, b"n")],
237 flags(SetCondition::IfAbsent, ScoreCompare::Any),
238 )
239 .unwrap();
240 assert_eq!((r.added, r.changed), (1, 1));
241 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(1.0)); assert_eq!(s.zscore(b"z", b"n").unwrap(), Some(2.0));
243 }
244
245 #[test]
246 fn xx_only_updates() {
247 let mut s = Store::new();
248 s.zadd(b"z", &[(1.0, b"m".as_slice())]).unwrap();
249 let r = s
250 .zadd_flags(
251 b"z",
252 &[(9.0, b"m"), (2.0, b"n")],
253 flags(SetCondition::IfPresent, ScoreCompare::Any),
254 )
255 .unwrap();
256 assert_eq!((r.added, r.changed), (0, 1));
257 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(9.0));
258 assert_eq!(s.zscore(b"z", b"n").unwrap(), None); }
260
261 #[test]
262 fn gt_is_monotonic_heal() {
263 let mut s = Store::new();
264 s.zadd(b"z", &[(5.0, b"m".as_slice())]).unwrap();
265 let gt = flags(SetCondition::Always, ScoreCompare::Greater);
266 let r = s.zadd_flags(b"z", &[(3.0, b"m")], gt).unwrap();
268 assert_eq!(r.changed, 0);
269 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(5.0));
270 let r = s.zadd_flags(b"z", &[(7.0, b"m")], gt).unwrap();
272 assert_eq!(r.changed, 1);
273 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(7.0));
274 let r = s.zadd_flags(b"z", &[(1.0, b"new")], gt).unwrap();
276 assert_eq!(r.added, 1);
277 }
278
279 #[test]
280 fn lt_mirror() {
281 let mut s = Store::new();
282 s.zadd(b"z", &[(5.0, b"m".as_slice())]).unwrap();
283 let lt = flags(SetCondition::Always, ScoreCompare::Less);
284 assert_eq!(s.zadd_flags(b"z", &[(7.0, b"m")], lt).unwrap().changed, 0);
285 assert_eq!(s.zadd_flags(b"z", &[(3.0, b"m")], lt).unwrap().changed, 1);
286 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(3.0));
287 }
288
289 #[test]
290 fn applied_reflects_effect_only() {
291 let mut s = Store::new();
292 s.zadd(b"z", &[(5.0, b"a".as_slice()), (5.0, b"b".as_slice())]).unwrap();
293 let r = s
294 .zadd_flags(
295 b"z",
296 &[(9.0, b"a"), (1.0, b"b"), (5.0, b"c")],
297 flags(SetCondition::Always, ScoreCompare::Greater),
298 )
299 .unwrap();
300 assert_eq!(r.applied, vec![(9.0, b"a".to_vec()), (5.0, b"c".to_vec())]);
302 }
303
304 #[test]
305 fn incr_form_vetoes_to_none() {
306 let mut s = Store::new();
307 s.zadd(b"z", &[(5.0, b"m".as_slice())]).unwrap();
308 let gt = flags(SetCondition::Always, ScoreCompare::Greater);
309 assert_eq!(s.zadd_incr(b"z", -2.0, b"m", gt).unwrap(), None);
311 assert_eq!(s.zscore(b"z", b"m").unwrap(), Some(5.0));
312 assert_eq!(s.zadd_incr(b"z", 2.0, b"m", gt).unwrap(), Some(7.0));
313 let xx = flags(SetCondition::IfPresent, ScoreCompare::Any);
315 assert_eq!(s.zadd_incr(b"z", 1.0, b"nope", xx).unwrap(), None);
316 let nx = flags(SetCondition::IfAbsent, ScoreCompare::Any);
318 assert_eq!(s.zadd_incr(b"z", 1.0, b"m", nx).unwrap(), None);
319 }
320
321 #[test]
322 fn wrongtype_propagates() {
323 let mut s = Store::new();
324 s.set(b"str", b"v".to_vec(), None, SetCondition::Always);
325 assert!(s.zadd_flags(b"str", &[(1.0, b"m")], zf()).is_err());
326 assert!(s.zadd_incr(b"str", 1.0, b"m", zf()).is_err());
327 }
328}