Skip to main content

kevy_store/
zset_flags.rs

1//! `ZADD` condition flags (Redis 6.2): `NX` / `XX` / `GT` / `LT` /
2//! `CH` / `INCR`. Split from `zset.rs` (500-LOC rule). The no-flags
3//! hot path stays `zadd` / `zadd` — nothing here taxes it.
4
5#[cfg(not(feature = "std"))]
6use crate::nostd_prelude::*;
7use crate::{ScoreCompare, SetCondition, Store, StoreError};
8
9/// Parsed `ZADD` condition flags: `NX` / `XX` as a [`SetCondition`],
10/// `GT` / `LT` as a [`ScoreCompare`], and `CH`. Only Redis's legal
11/// combinations can be built — [`ZaddFlags::new`] refuses `NX` with
12/// `GT` / `LT` — so the store never has a combination to reject. `CH`
13/// only changes the *reply* (changed count instead of added count);
14/// [`ZaddReport`] carries both counts either way.
15///
16/// ```
17/// use kevy_store::{ScoreCompare, SetCondition, ZaddFlags};
18/// let f = ZaddFlags::new(SetCondition::IfPresent, ScoreCompare::Greater).unwrap().with_ch(true);
19/// assert_eq!((f.condition(), f.compare(), f.ch()), (SetCondition::IfPresent, ScoreCompare::Greater, true));
20/// assert!(ZaddFlags::new(SetCondition::IfAbsent, ScoreCompare::Less).is_none());
21/// assert_eq!(ZaddFlags::default(), ZaddFlags::new(SetCondition::Always, ScoreCompare::Any).unwrap());
22/// ```
23#[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    /// The flags for `condition` and `compare`, or `None` for the one
32    /// combination Redis refuses: `NX` together with `GT` or `LT`.
33    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    /// Set `CH`: the reply counts changed members instead of added ones.
39    ///
40    /// ```
41    /// assert!(kevy_store::ZaddFlags::default().with_ch(true).ch());
42    /// ```
43    #[must_use]
44    pub fn with_ch(mut self, ch: bool) -> Self {
45        self.ch = ch;
46        self
47    }
48
49    /// `NX` / `XX`, or neither.
50    pub fn condition(self) -> SetCondition {
51        self.condition
52    }
53
54    /// `GT` / `LT`, or neither.
55    pub fn compare(self) -> ScoreCompare {
56        self.compare
57    }
58
59    /// Whether `CH` was given.
60    pub fn ch(self) -> bool {
61        self.ch
62    }
63
64    /// Whether these flags veto replacing an existing member's `old`
65    /// score with `new`.
66    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/// Outcome of a flags-aware `ZADD`.
77///
78/// ```
79/// use kevy_store::{ScoreCompare, SetCondition, Store, ZaddFlags};
80/// let mut s = Store::new();
81/// s.zadd(b"z", &[(5.0, b"a".as_slice())])?;
82/// let gt = ZaddFlags::new(SetCondition::Always, ScoreCompare::Greater).unwrap();
83/// let r = s.zadd_flags(b"z", &[(1.0, b"a".as_slice()), (2.0, b"b")], gt)?;
84/// assert_eq!((r.added, r.changed), (1, 1)); // `a` vetoed, `b` added
85/// # Ok::<(), kevy_store::StoreError>(())
86/// ```
87#[derive(Debug, Clone, PartialEq, Default)]
88#[non_exhaustive]
89pub struct ZaddReport {
90    /// Members newly added.
91    ///
92    /// ```
93    /// use kevy_store::{Store, ZaddFlags};
94    /// let mut s = Store::new();
95    /// s.zadd(b"z", &[(1.0, b"a".as_slice())])?;
96    /// let r = s.zadd_flags(b"z", &[(2.0, b"a".as_slice()), (1.0, b"b")], ZaddFlags::default())?;
97    /// assert_eq!(r.added, 1); // only `b` is new
98    /// # Ok::<(), kevy_store::StoreError>(())
99    /// ```
100    pub added: usize,
101    /// Members added or whose score actually changed (`CH` reply).
102    ///
103    /// ```
104    /// use kevy_store::{Store, ZaddFlags};
105    /// let mut s = Store::new();
106    /// s.zadd(b"z", &[(1.0, b"a".as_slice()), (1.0, b"c")])?;
107    /// let pairs = [(2.0, b"a".as_slice()), (1.0, b"b"), (1.0, b"c")];
108    /// let r = s.zadd_flags(b"z", &pairs, ZaddFlags::default().with_ch(true))?;
109    /// assert_eq!(r.changed, 2); // `a` moved, `b` added, `c` unchanged
110    /// # Ok::<(), kevy_store::StoreError>(())
111    /// ```
112    pub changed: usize,
113    /// The `(score, member)` pairs actually applied, in input order —
114    /// vetoed pairs are absent. Lets an AOF writer log the *effect*
115    /// as a plain unconditional `ZADD` (deterministic on replay; a
116    /// conditional replayed against divergent state could veto
117    /// differently).
118    ///
119    /// ```
120    /// use kevy_store::{ScoreCompare, SetCondition, Store, ZaddFlags};
121    /// let mut s = Store::new();
122    /// s.zadd(b"z", &[(5.0, b"a".as_slice())])?;
123    /// let lt = ZaddFlags::new(SetCondition::Always, ScoreCompare::Less).unwrap();
124    /// let r = s.zadd_flags(b"z", &[(9.0, b"a".as_slice()), (3.0, b"b")], lt)?;
125    /// assert_eq!(r.applied, [(3.0, b"b".to_vec())]);
126    /// # Ok::<(), kevy_store::StoreError>(())
127    /// ```
128    pub applied: Vec<(f64, Vec<u8>)>,
129}
130
131impl Store {
132    /// Flags-aware `ZADD`: each pair is added or updated as `flags`
133    /// allow, and the report says what actually happened.
134    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    /// `ZADD … INCR` — a conditional `ZINCRBY`: returns the new score,
168    /// or `None` when the flags veto the operation (Redis replies nil).
169    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)); // untouched
242        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); // not added
259    }
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        // Stale (lower) score: vetoed.
267        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        // Newer (higher) score: applied.
271        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        // GT still ADDS missing members (only XX suppresses adds).
275        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        // a updated, b vetoed, c added.
301        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        // Negative delta under GT: next < old → nil, score untouched.
310        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        // XX on a missing member → nil.
314        let xx = flags(SetCondition::IfPresent, ScoreCompare::Any);
315        assert_eq!(s.zadd_incr(b"z", 1.0, b"nope", xx).unwrap(), None);
316        // NX on an existing member → nil.
317        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}