kevy_store/
zset_range.rs

1//! `Store` sorted-set range / pop / range-removal commands
2//! (`ZRANGE` / `ZRANGEBYSCORE` / `ZCOUNT` / `ZPOPMIN` / `ZREMRANGEBY*`).
3//! Split out of `zset.rs` to keep it under the 500-LOC house cap; the
4//! write-path core (`ZADD` / `ZREM` / `ZINCRBY`) stays there.
5
6#[cfg(not(feature = "std"))]
7use crate::nostd_prelude::*;
8use crate::util::range_bounds;
9use crate::value::{ScoreBound, Value};
10use crate::{Store, StoreError};
11
12impl Store {
13    /// `ZRANGE key start stop` by rank.
14    pub fn zrange(
15        &mut self,
16        key: &[u8],
17        start: i64,
18        stop: i64,
19    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
20        match self.live_entry(key) {
21            None => Ok(Vec::new()),
22            Some(e) => match &e.value {
23                Value::ZSet(z) => Ok(match range_bounds(start, stop, z.len()) {
24                    None => Vec::new(),
25                    // O(log N) seek to the start rank, then walk M items —
26                    // no skip-walk from the front.
27                    Some((s, end)) => z
28                        .ordered_from(s)
29                        .take(end - s + 1)
30                        .map(|(m, sc)| (m.to_vec(), sc))
31                        .collect(),
32                }),
33                Value::SegZSet(z) => Ok(match range_bounds(start, stop, z.len()) {
34                    None => Vec::new(),
35                    Some((s, end)) => z
36                        .ordered_from(s)
37                        .take(end - s + 1)
38                        .map(|(m, sc)| (m.to_vec(), sc))
39                        .collect(),
40                }),
41                Value::SmallZSetInline(z) => {
42                    let mut entries: Vec<(Vec<u8>, f64)> =
43                        z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect();
44                    entries.sort_by(|a, b| {
45                        a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))
46                    });
47                    Ok(match range_bounds(start, stop, entries.len()) {
48                        None => Vec::new(),
49                        Some((s, end)) => entries.into_iter().skip(s).take(end - s + 1).collect(),
50                    })
51                }
52                _ => Err(StoreError::WrongType),
53            },
54        }
55    }
56
57    /// `ZRANGEBYSCORE`.
58    pub fn zrange_by_score(
59        &mut self,
60        key: &[u8],
61        min: ScoreBound,
62        max: ScoreBound,
63    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
64        match self.live_entry(key) {
65            None => Ok(Vec::new()),
66            Some(e) => match &e.value {
67                Value::ZSet(z) => {
68                    // Two O(log N) rank descents bracket the score range,
69                    // then only the M matches are walked — no scan+filter.
70                    let lo = z.score_start_rank(&min);
71                    let hi = z.score_end_rank(&max);
72                    Ok(z.ordered_from(lo)
73                        .take(hi.saturating_sub(lo))
74                        .map(|(m, sc)| (m.to_vec(), sc))
75                        .collect())
76                }
77                Value::SegZSet(z) => {
78                    let lo = z.score_start_rank(&min);
79                    let hi = z.score_end_rank(&max);
80                    Ok(z.ordered_from(lo)
81                        .take(hi.saturating_sub(lo))
82                        .map(|(m, sc)| (m.to_vec(), sc))
83                        .collect())
84                }
85                Value::SmallZSetInline(z) => {
86                    let mut entries: Vec<(Vec<u8>, f64)> = z
87                        .iter()
88                        .filter(|(_, sc)| min.ge_ok(*sc) && max.le_ok(*sc))
89                        .map(|(m, sc)| (m.to_vec(), sc))
90                        .collect();
91                    entries.sort_by(|a, b| {
92                        a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0))
93                    });
94                    Ok(entries)
95                }
96                _ => Err(StoreError::WrongType),
97            },
98        }
99    }
100
101    /// `ZCOUNT`.
102    pub fn zcount(
103        &mut self,
104        key: &[u8],
105        min: ScoreBound,
106        max: ScoreBound,
107    ) -> Result<usize, StoreError> {
108        match self.live_entry(key) {
109            None => Ok(0),
110            Some(e) => match &e.value {
111                // Two rank descents — O(log N), no iteration at all.
112                Value::ZSet(z) => Ok(z
113                    .score_end_rank(&max)
114                    .saturating_sub(z.score_start_rank(&min))),
115                Value::SegZSet(z) => Ok(z
116                    .score_end_rank(&max)
117                    .saturating_sub(z.score_start_rank(&min))),
118                Value::SmallZSetInline(z) => Ok(z
119                    .iter()
120                    .filter(|(_, sc)| min.ge_ok(*sc) && max.le_ok(*sc))
121                    .count()),
122                _ => Err(StoreError::WrongType),
123            },
124        }
125    }
126
127    /// `ZPOPMIN key [count]` — pop and return the `count` lowest-scored
128    /// members (ascending by `(score, member)`). Returns `(member,
129    /// score)` pairs in pop order; empty when the key is absent / empty.
130    pub fn zpopmin(
131        &mut self,
132        key: &[u8],
133        count: usize,
134    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
135        if count == 0 {
136            // Validate type up-front so ZPOPMIN k 0 against a wrong-type
137            // key still reports WRONGTYPE (Redis behaviour).
138            if let Some(e) = self.live_entry(key) {
139                match &e.value {
140                    Value::ZSet(_) | Value::SegZSet(_) | Value::SmallZSetInline(_) => {}
141                    _ => return Err(StoreError::WrongType),
142                }
143            }
144            return Ok(Vec::new());
145        }
146        // Snapshot the lowest `count` members first (immutable borrow),
147        // then remove them via the shared zrem path (which handles the
148        // encoding, weight accounting, and empty-key cleanup uniformly).
149        let to_pop: Vec<(Vec<u8>, f64)> = match self.live_entry(key) {
150            None => return Ok(Vec::new()),
151            Some(e) => match &e.value {
152                Value::ZSet(z) => z
153                    .ordered()
154                    .take(count)
155                    .map(|(m, sc)| (m.to_vec(), sc))
156                    .collect(),
157                Value::SegZSet(z) => z
158                    .ordered()
159                    .take(count)
160                    .map(|(m, sc)| (m.to_vec(), sc))
161                    .collect(),
162                Value::SmallZSetInline(z) => {
163                    let mut entries: Vec<(Vec<u8>, f64)> =
164                        z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect();
165                    entries.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
166                    entries.into_iter().take(count).collect()
167                }
168                _ => return Err(StoreError::WrongType),
169            },
170        };
171        if to_pop.is_empty() {
172            return Ok(to_pop);
173        }
174        let borrowed: Vec<&[u8]> = to_pop.iter().map(|(m, _)| m.as_slice()).collect();
175        self.zrem(key, &borrowed)?;
176        Ok(to_pop)
177    }
178
179    /// `zpopmin_below` — pop up to `count` lowest-scored members
180    /// whose score is `< below` (strictly). The delayed-job primitive:
181    /// "pop everything that is due" in one atomic call (score = due
182    /// time, `below` = now). Absent key = empty; wrong type errors.
183    pub fn zpopmin_below(
184        &mut self,
185        key: &[u8],
186        below: f64,
187        count: usize,
188    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
189        if count == 0 {
190            if let Some(e) = self.live_entry(key) {
191                match &e.value {
192                    Value::ZSet(_) | Value::SegZSet(_) | Value::SmallZSetInline(_) => {}
193                    _ => return Err(StoreError::WrongType),
194                }
195            }
196            return Ok(Vec::new());
197        }
198        let to_pop: Vec<(Vec<u8>, f64)> = match self.live_entry(key) {
199            None => return Ok(Vec::new()),
200            Some(e) => match &e.value {
201                Value::ZSet(z) => z
202                    .ordered()
203                    .take_while(|(_, sc)| *sc < below)
204                    .take(count)
205                    .map(|(m, sc)| (m.to_vec(), sc))
206                    .collect(),
207                Value::SegZSet(z) => z
208                    .ordered()
209                    .take_while(|(_, sc)| *sc < below)
210                    .take(count)
211                    .map(|(m, sc)| (m.to_vec(), sc))
212                    .collect(),
213                Value::SmallZSetInline(z) => {
214                    let mut entries: Vec<(Vec<u8>, f64)> =
215                        z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect();
216                    entries.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
217                    entries
218                        .into_iter()
219                        .take_while(|(_, sc)| *sc < below)
220                        .take(count)
221                        .collect()
222                }
223                _ => return Err(StoreError::WrongType),
224            },
225        };
226        if to_pop.is_empty() {
227            return Ok(to_pop);
228        }
229        let borrowed: Vec<&[u8]> = to_pop.iter().map(|(m, _)| m.as_slice()).collect();
230        self.zrem(key, &borrowed)?;
231        Ok(to_pop)
232    }
233
234    /// `ZREMRANGEBYRANK key start stop` — remove members in the rank
235    /// range `[start, stop]` (inclusive, negative indices count from
236    /// the tail). Returns the number of members removed.
237    pub fn zrem_range_by_rank(
238        &mut self,
239        key: &[u8],
240        start: i64,
241        stop: i64,
242    ) -> Result<usize, StoreError> {
243        let to_remove: Vec<Vec<u8>> = match self.live_entry(key) {
244            None => return Ok(0),
245            Some(e) => match &e.value {
246                Value::ZSet(z) => match crate::util::range_bounds(start, stop, z.len()) {
247                    None => return Ok(0),
248                    // Seek to the start rank (O(log N)), collect the M hits.
249                    Some((s, end)) => z
250                        .ordered_from(s)
251                        .take(end - s + 1)
252                        .map(|(m, _)| m.to_vec())
253                        .collect(),
254                },
255                Value::SegZSet(z) => match crate::util::range_bounds(start, stop, z.len()) {
256                    None => return Ok(0),
257                    Some((s, end)) => z
258                        .ordered_from(s)
259                        .take(end - s + 1)
260                        .map(|(m, _)| m.to_vec())
261                        .collect(),
262                },
263                Value::SmallZSetInline(z) => {
264                    let mut entries: Vec<(Vec<u8>, f64)> =
265                        z.iter().map(|(m, sc)| (m.to_vec(), sc)).collect();
266                    entries.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
267                    match crate::util::range_bounds(start, stop, entries.len()) {
268                        None => return Ok(0),
269                        Some((s, end)) => entries
270                            .into_iter()
271                            .skip(s)
272                            .take(end - s + 1)
273                            .map(|(m, _)| m)
274                            .collect(),
275                    }
276                }
277                _ => return Err(StoreError::WrongType),
278            },
279        };
280        if to_remove.is_empty() {
281            return Ok(0);
282        }
283        let borrowed: Vec<&[u8]> = to_remove.iter().map(Vec::as_slice).collect();
284        self.zrem(key, &borrowed)
285    }
286
287    /// `ZREMRANGEBYSCORE key min max` — remove every member whose score
288    /// satisfies `min ≤ score ≤ max` (with `(` for exclusive bounds via
289    /// `ScoreBound`). Returns the number removed.
290    pub fn zrem_range_by_score(
291        &mut self,
292        key: &[u8],
293        min: ScoreBound,
294        max: ScoreBound,
295    ) -> Result<usize, StoreError> {
296        // Reuse zrange_by_score's bound logic to materialise the hit set
297        // — keeps inline / heap parity in one place.
298        let hits = self.zrange_by_score(key, min, max)?;
299        if hits.is_empty() {
300            // Still need to honour wrong-type errors that zrange_by_score
301            // already surfaced; here Ok([]) means empty match, not type
302            // mismatch, so it's safe to early-return.
303            return Ok(0);
304        }
305        let borrowed: Vec<&[u8]> = hits.iter().map(|(m, _)| m.as_slice()).collect();
306        self.zrem(key, &borrowed)
307    }
308
309    /// `ZREVRANGEBYSCORE` — `zrange_by_score` reversed. Bounds are
310    /// passed in the `(min, max)` order already (the caller is
311    /// responsible for swapping the user-facing `max first, min second`
312    /// at the dispatch layer).
313    pub fn zrev_range_by_score(
314        &mut self,
315        key: &[u8],
316        min: ScoreBound,
317        max: ScoreBound,
318    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
319        let mut v = self.zrange_by_score(key, min, max)?;
320        v.reverse();
321        Ok(v)
322    }
323}
kevy_store/zset_range.rs

kevy_store/
zset_range.rs
