kevy_store/
zset.rs

1//! `Store` sorted-set write-path commands (`ZADD` / `ZREM` / `ZINCRBY` /
2//! `ZSCORE` / `ZCARD` / `ZRANK`). The range / pop / range-removal family
3//! lives in `zset_range.rs` (500-LOC house cap).
4
5#[cfg(not(feature = "std"))]
6use crate::nostd_prelude::*;
7use crate::small_zset::{self, AddResult as ZAddResult, SmallZSetData};
8use crate::zset_seg::{SegZSetData, Z_PROMOTE};
9use crate::value::{ZSetData, SmallBytes, Value, zset_member_weight};
10use crate::{Entry, Store, StoreError};
11use alloc::sync::Arc;
12
13impl Store {
14    // ---- sorted sets ---------------------------------------------------
15
16    /// Borrow the key's zset mutably; promotes inline → flat, and flat →
17    /// segmented at the threshold (so ZINCRBY-only workloads cross the
18    /// boundary too).
19    fn zset_mut(&mut self, key: &[u8], create: bool) -> Result<Option<ZRefMut<'_>>, StoreError> {
20        if self.live_entry_mut(key).is_none() {
21            if !create {
22                return Ok(None);
23            }
24            self.insert_entry(
25                SmallBytes::from_slice(key),
26                Entry::new(Value::ZSet(Arc::default()), None),
27            );
28        }
29        // A.8: see hash.rs::hash_mut — promote out-of-scope.
30        let needs = match self.map.get(key).map(|e| &e.value) {
31            Some(Value::SmallZSetInline(_)) => true,
32            Some(Value::ZSet(z)) => z.len() >= Z_PROMOTE,
33            _ => false,
34        };
35        if needs {
36            self.promote_zset_encoding(key);
37        }
38        match &mut self.map.get_mut(key).expect("present").value {
39            Value::ZSet(z) => Ok(Some(ZRefMut::Flat(Arc::make_mut(z)))),
40            Value::SegZSet(z) => Ok(Some(ZRefMut::Seg(Arc::make_mut(z)))),
41            _ => Err(StoreError::WrongType),
42        }
43    }
44
45    /// One promotion step: inline → flat, or flat-at-threshold →
46    /// segmented. Reweighs the entry.
47    fn promote_zset_encoding(&mut self, key: &[u8]) {
48        let Some(e) = self.map.get_mut(key) else { return };
49        match &mut e.value {
50            Value::SmallZSetInline(s) => {
51                e.value = Value::ZSet(Arc::new(small_zset::promote(s)));
52            }
53            Value::ZSet(z) => {
54                e.value = Value::SegZSet(Arc::new(SegZSetData::from_flat(z)));
55            }
56            _ => return,
57        }
58        self.reweigh_entry(key);
59    }
60
61    /// A.8: read the key's zset slot for ZADD. None when absent.
62    fn zset_value_for_set(&mut self, key: &[u8]) -> Result<Option<&mut Value>, StoreError> {
63        match self.live_entry_mut(key) {
64            None => Ok(None),
65            Some(e) => match &e.value {
66                Value::ZSet(_) | Value::SegZSet(_) | Value::SmallZSetInline(_) => {
67                    Ok(Some(&mut e.value))
68                }
69                _ => Err(StoreError::WrongType),
70            },
71        }
72    }
73
74    fn drop_if_empty_zset(&mut self, key: &[u8]) {
75        let empty = match self.map.get(key).map(|e| &e.value) {
76            Some(Value::ZSet(z)) => z.len() == 0,
77            Some(Value::SegZSet(z)) => z.is_empty(),
78            Some(Value::SmallZSetInline(z)) => z.is_empty(),
79            _ => false,
80        };
81        if empty {
82            self.remove_entry(key);
83        }
84    }
85
86    /// `ZADD` — returns the count of newly-added members. Borrowed
87    /// argv: no per-member allocation; routes through the
88    /// encoding-switch path.
89    pub fn zadd(
90        &mut self,
91        key: &[u8],
92        pairs: &[(f64, &[u8])],
93    ) -> Result<usize, StoreError> {
94        if pairs.is_empty() {
95            return Ok(0);
96        }
97        let mut added = 0usize;
98        let mut delta: i64 = 0;
99        for (score, m) in pairs {
100            match self.zadd_one(key, m, *score)? {
101                ZaddOutcome::AddedInline => added += 1,
102                ZaddOutcome::UpdatedInline => {}
103                ZaddOutcome::AddedHeap(w) => {
104                    added += 1;
105                    delta += w;
106                }
107                ZaddOutcome::UpdatedHeap => {}
108            }
109        }
110        self.account_delta(key, delta);
111        Ok(added)
112    }
113
114    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> Result<Option<f64>, StoreError> {
115        match self.live_entry(key) {
116            None => Ok(None),
117            Some(e) => match &e.value {
118                Value::ZSet(z) => Ok(z.by_member.get(member).copied()),
119                Value::SegZSet(z) => Ok(z.score_of(member)),
120                Value::SmallZSetInline(z) => Ok(z.score(member)),
121                _ => Err(StoreError::WrongType),
122            },
123        }
124    }
125
126    pub fn zcard(&mut self, key: &[u8]) -> Result<usize, StoreError> {
127        match self.live_entry(key) {
128            None => Ok(0),
129            Some(e) => match &e.value {
130                Value::ZSet(z) => Ok(z.len()),
131                Value::SegZSet(z) => Ok(z.len()),
132                Value::SmallZSetInline(z) => Ok(z.len()),
133                _ => Err(StoreError::WrongType),
134            },
135        }
136    }
137
138    /// `ZREM` — returns the count of members removed.
139    pub fn zrem(
140        &mut self,
141        key: &[u8],
142        members: &[&[u8]],
143    ) -> Result<usize, StoreError> {
144        let (removed, delta) = {
145            let mut r = 0usize;
146            let mut d: i64 = 0;
147            if let Some(e) = self.live_entry_mut(key) {
148                match &mut e.value {
149                    Value::ZSet(z) => {
150                        // G-A3: hoist Arc::make_mut OUT of loop.
151                        let z = Arc::make_mut(z);
152                        for m in members {
153                            if z.remove(m) {
154                                r += 1;
155                                d -= zset_member_weight(&SmallBytes::from_slice(m)) as i64;
156                            }
157                        }
158                    }
159                    Value::SegZSet(z) => {
160                        let z = Arc::make_mut(z);
161                        for m in members {
162                            if z.remove(m) {
163                                r += 1;
164                                d -= zset_member_weight(&SmallBytes::from_slice(m)) as i64;
165                            }
166                        }
167                    }
168                    Value::SmallZSetInline(z) => {
169                        for m in members {
170                            if z.try_remove(m) {
171                                r += 1;
172                            }
173                        }
174                    }
175                    _ => return Err(StoreError::WrongType),
176                }
177            }
178            (r, d)
179        };
180        self.account_delta(key, delta);
181        self.drop_if_empty_zset(key);
182        Ok(removed)
183    }
184
185    /// `ZRANK` — 0-based position in ascending order. O(log N): a hash
186    /// lookup for the score, then one order-statistic tree descent.
187    pub fn zrank(&mut self, key: &[u8], member: &[u8]) -> Result<Option<usize>, StoreError> {
188        match self.live_entry(key) {
189            None => Ok(None),
190            Some(e) => match &e.value {
191                Value::ZSet(z) => Ok(z
192                    .by_member
193                    .get(member)
194                    .copied()
195                    .and_then(|sc| z.rank_of(member, sc))),
196                Value::SegZSet(z) => {
197                    Ok(z.score_of(member).and_then(|sc| z.rank_of(member, sc)))
198                }
199                Value::SmallZSetInline(z) => {
200                    // Inline holds at most 2 entries; sort by score (then
201                    // bytes) so ZRANK matches ZRANGE order.
202                    let mut entries: Vec<(&[u8], f64)> = z.iter().collect();
203                    entries.sort_by(|a, b| {
204                        a.1.total_cmp(&b.1).then_with(|| a.0.cmp(b.0))
205                    });
206                    Ok(entries.iter().position(|(m, _)| *m == member))
207                }
208                _ => Err(StoreError::WrongType),
209            },
210        }
211    }
212
213    /// `ZINCRBY` — add `incr` to a member's score; returns the new score.
214    pub fn zincrby(&mut self, key: &[u8], incr: f64, member: &[u8]) -> Result<f64, StoreError> {
215        let mut z = self.zset_mut(key, true)?.expect("created");
216        let cur = z.score_of(member).unwrap_or(0.0);
217        let next = cur + incr;
218        let smb = SmallBytes::from_slice(member);
219        let is_new = !z.contains_member(member);
220        z.insert(member, next);
221        let d = if is_new { zset_member_weight(&smb) as i64 } else { 0 };
222        self.account_delta(key, d);
223        Ok(next)
224    }
225
226    /// A.8 core: set one `(member, score)` pair via encoding-switch.
227    fn zadd_one(&mut self, key: &[u8], m: &[u8], score: f64) -> Result<ZaddOutcome, StoreError> {
228        if self.zset_value_for_set(key)?.is_none() {
229            return Ok(self.zadd_create(key, m, score));
230        }
231        let v = self.zset_value_for_set(key)?.expect("present and a zset");
232        match v {
233            Value::SmallZSetInline(z) => match z.try_set(m, score) {
234                ZAddResult::Added => Ok(ZaddOutcome::AddedInline),
235                ZAddResult::Updated => Ok(ZaddOutcome::UpdatedInline),
236                ZAddResult::NoRoom => {
237                    let outcome = promote_inline_zset_and_add(v, m, score);
238                    self.reweigh_entry(key);
239                    Ok(outcome)
240                }
241            },
242            Value::ZSet(z) if z.len() >= Z_PROMOTE => {
243                let is_new = promote_flat_zset_and_add(v, m, score);
244                self.reweigh_entry(key);
245                // Reweighed from scratch — swallow the per-member delta.
246                if is_new {
247                    Ok(ZaddOutcome::AddedHeap(0))
248                } else {
249                    Ok(ZaddOutcome::UpdatedHeap)
250                }
251            }
252            Value::ZSet(z) => {
253                let z = Arc::make_mut(z);
254                let smb = SmallBytes::from_slice(m);
255                let w = zset_member_weight(&smb) as i64;
256                if z.insert(m, score) {
257                    Ok(ZaddOutcome::AddedHeap(w))
258                } else {
259                    Ok(ZaddOutcome::UpdatedHeap)
260                }
261            }
262            Value::SegZSet(z) => {
263                let z = Arc::make_mut(z);
264                let smb = SmallBytes::from_slice(m);
265                let w = zset_member_weight(&smb) as i64;
266                if z.insert(m, score) {
267                    Ok(ZaddOutcome::AddedHeap(w))
268                } else {
269                    Ok(ZaddOutcome::UpdatedHeap)
270                }
271            }
272            _ => Err(StoreError::WrongType),
273        }
274    }
275
276    /// Create a fresh entry holding one `(member, score)` pair.
277    fn zadd_create(&mut self, key: &[u8], m: &[u8], score: f64) -> ZaddOutcome {
278        if let Some(inline) = SmallZSetData::with_one(m, score) {
279            self.insert_entry(
280                SmallBytes::from_slice(key),
281                Entry::new(Value::SmallZSetInline(inline), None),
282            );
283            ZaddOutcome::AddedInline
284        } else {
285            let mut z = ZSetData::default();
286            z.insert(m, score);
287            self.insert_entry(
288                SmallBytes::from_slice(key),
289                Entry::new(Value::ZSet(Arc::new(z)), None),
290            );
291            ZaddOutcome::AddedInline
292        }
293    }
294}
295
296/// Inline zset out of room: promote to the flat heap encoding, then
297/// set the spilling pair. Caller reweighs the entry.
298fn promote_inline_zset_and_add(v: &mut Value, m: &[u8], score: f64) -> ZaddOutcome {
299    let Value::SmallZSetInline(z) = v else { unreachable!("matched inline") };
300    let mut promoted = small_zset::promote(z);
301    let smb = SmallBytes::from_slice(m);
302    let is_new = !promoted.by_member.contains_key(m);
303    let w = zset_member_weight(&smb) as i64;
304    promoted.insert(m, score);
305    *v = Value::ZSet(Arc::new(promoted));
306    if is_new {
307        ZaddOutcome::AddedHeap(w)
308    } else {
309        ZaddOutcome::UpdatedHeap
310    }
311}
312
313/// Flat zset at the threshold: segment, then set. One-time
314/// O(Z_PROMOTE) rebuild (or clone, if a view pins it now). Returns
315/// whether the member was new; caller reweighs.
316fn promote_flat_zset_and_add(v: &mut Value, m: &[u8], score: f64) -> bool {
317    let Value::ZSet(z) = v else { unreachable!("matched ZSet") };
318    let mut seg = SegZSetData::from_flat(z);
319    let is_new = seg.insert(m, score);
320    *v = Value::SegZSet(Arc::new(seg));
321    is_new
322}
323
324/// A mutable borrow of either heap zset encoding — the read-modify-
325/// write entry point (`zincrby`) stays encoding-blind.
326enum ZRefMut<'a> {
327    Flat(&'a mut ZSetData),
328    Seg(&'a mut SegZSetData),
329}
330
331impl ZRefMut<'_> {
332    fn score_of(&self, member: &[u8]) -> Option<f64> {
333        match self {
334            Self::Flat(z) => z.by_member.get(member).copied(),
335            Self::Seg(z) => z.score_of(member),
336        }
337    }
338    fn contains_member(&self, member: &[u8]) -> bool {
339        match self {
340            Self::Flat(z) => z.by_member.contains_key(member),
341            Self::Seg(z) => z.contains_member(member),
342        }
343    }
344    fn insert(&mut self, member: &[u8], score: f64) -> bool {
345        match self {
346            Self::Flat(z) => z.insert(member, score),
347            Self::Seg(z) => z.insert(member, score),
348        }
349    }
350}
351
352enum ZaddOutcome {
353    AddedInline,
354    UpdatedInline,
355    AddedHeap(i64),
356    UpdatedHeap,
357}
kevy_store/zset.rs

kevy_store/
zset.rs
