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_borrowed` — nothing here taxes it.
4
5use crate::{Store, StoreError};
6
7/// Parsed `ZADD` condition flags. `CH` only changes the *reply*
8/// (changed count instead of added count) — callers read
9/// [`ZaddReport::changed`] when set; the engine behavior is identical.
10#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
11pub struct ZaddFlags {
12    /// Only add new members; never update existing ones.
13    pub nx: bool,
14    /// Only update existing members; never add new ones.
15    pub xx: bool,
16    /// Only update when the new score is greater than the current.
17    pub gt: bool,
18    /// Only update when the new score is less than the current.
19    pub lt: bool,
20    /// Reply with changed (added + updated) instead of added.
21    pub ch: bool,
22}
23
24impl ZaddFlags {
25    /// Redis 6.2 rule: `GT`, `LT` and `NX` are mutually exclusive
26    /// (and `GT`+`LT` together are, too). `XX`+`NX` likewise.
27    pub fn valid(self) -> bool {
28        !(self.nx && (self.xx || self.gt || self.lt)) && !(self.gt && self.lt)
29    }
30}
31
32/// Outcome of a flags-aware `ZADD`.
33pub struct ZaddReport {
34    /// Members newly added.
35    pub added: usize,
36    /// Members added or whose score actually changed (`CH` reply).
37    pub changed: usize,
38    /// The `(score, member)` pairs actually applied, in input order —
39    /// vetoed pairs are absent. Lets an AOF writer log the *effect*
40    /// as a plain unconditional `ZADD` (deterministic on replay; a
41    /// conditional replayed against divergent state could veto
42    /// differently).
43    pub applied: Vec<(f64, Vec<u8>)>,
44}
45
46impl Store {
47    /// Flags-aware `ZADD`. Caller validates [`ZaddFlags::valid`] at
48    /// its input boundary (RESP parse / typed API) — invalid combos
49    /// here are a caller bug.
50    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    /// `ZADD … INCR` — a conditional `ZINCRBY`: returns the new score,
88    /// or `None` when the flags veto the operation (Redis replies nil).
89    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)); // untouched
153        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); // not added
166    }
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        // Stale (lower) score: vetoed.
174        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        // Newer (higher) score: applied.
178        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        // GT still ADDS missing members (only XX suppresses adds).
182        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        // a updated, b vetoed, c added.
208        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        // Negative delta under GT: next < old → nil, score untouched.
220        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        // XX on a missing member → nil.
224        let xx = ZaddFlags { xx: true, ..zf() };
225        assert_eq!(s.zadd_incr(b"z", 1.0, b"nope", xx).unwrap(), None);
226        // NX on an existing member → nil.
227        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}