kevy_store/
zset.rs

1//! `Store` sorted-set commands.
2
3use crate::util::*;
4use crate::value::*;
5use crate::{Entry, Store, StoreError};
6
7impl Store {
8    // ---- sorted sets ---------------------------------------------------
9
10    fn zset_mut(&mut self, key: &[u8], create: bool) -> Result<Option<&mut ZSetData>, StoreError> {
11        if self.live_entry_mut(key).is_none() {
12            if !create {
13                return Ok(None);
14            }
15            self.insert_entry(
16                SmallBytes::from_slice(key),
17                Entry::new(Value::ZSet(Box::default()), None),
18            );
19        }
20        match &mut self.map.get_mut(key).expect("present").value {
21            Value::ZSet(z) => Ok(Some(z)),
22            _ => Err(StoreError::WrongType),
23        }
24    }
25
26    fn zset_ref(&mut self, key: &[u8]) -> Result<Option<&ZSetData>, StoreError> {
27        match self.live_entry(key) {
28            None => Ok(None),
29            Some(e) => match &e.value {
30                Value::ZSet(z) => Ok(Some(z)),
31                _ => Err(StoreError::WrongType),
32            },
33        }
34    }
35
36    fn drop_if_empty_zset(&mut self, key: &[u8]) {
37        let empty = matches!(self.map.get(key).map(|e| &e.value), Some(Value::ZSet(z)) if z.len() == 0);
38        if empty {
39            self.remove_entry(key);
40        }
41    }
42
43    /// `ZADD` — returns the count of newly-added members (updates don't count).
44    pub fn zadd(&mut self, key: &[u8], pairs: &[(f64, Vec<u8>)]) -> Result<usize, StoreError> {
45        let (added, delta) = {
46            let z = self.zset_mut(key, true)?.expect("created");
47            let mut a = 0usize;
48            let mut d: i64 = 0;
49            for (score, m) in pairs {
50                let smb = SmallBytes::from_slice(m);
51                let w = zset_member_weight(&smb) as i64;
52                if z.insert(m, *score) {
53                    a += 1;
54                    d += w;
55                }
56                // Updating an existing score reuses the same member entry —
57                // no weight delta (f64 score is a fixed 8 B already counted).
58            }
59            (a, d)
60        };
61        self.account_delta(key, delta);
62        Ok(added)
63    }
64
65    pub fn zscore(&mut self, key: &[u8], member: &[u8]) -> Result<Option<f64>, StoreError> {
66        Ok(self
67            .zset_ref(key)?
68            .and_then(|z| z.by_member.get(member).copied()))
69    }
70
71    pub fn zcard(&mut self, key: &[u8]) -> Result<usize, StoreError> {
72        Ok(self.zset_ref(key)?.map_or(0, |z| z.len()))
73    }
74
75    pub fn zrem(&mut self, key: &[u8], members: &[Vec<u8>]) -> Result<usize, StoreError> {
76        let (removed, delta) = {
77            let mut r = 0usize;
78            let mut d: i64 = 0;
79            if let Some(z) = self.zset_mut(key, false)? {
80                for m in members {
81                    if z.remove(m.as_slice()) {
82                        r += 1;
83                        d -= zset_member_weight(&SmallBytes::from_slice(m)) as i64;
84                    }
85                }
86            }
87            (r, d)
88        };
89        self.account_delta(key, delta);
90        self.drop_if_empty_zset(key);
91        Ok(removed)
92    }
93
94    /// `ZRANK` — 0-based position in ascending order (O(n) for now).
95    pub fn zrank(&mut self, key: &[u8], member: &[u8]) -> Result<Option<usize>, StoreError> {
96        Ok(self
97            .zset_ref(key)?
98            .and_then(|z| z.ordered().position(|(m, _)| m == member)))
99    }
100
101    /// `ZINCRBY` — add `incr` to a member's score (default 0), returns the new score.
102    pub fn zincrby(&mut self, key: &[u8], incr: f64, member: &[u8]) -> Result<f64, StoreError> {
103        let (next, delta) = {
104            let z = self.zset_mut(key, true)?.expect("created");
105            let cur = z.by_member.get(member).copied().unwrap_or(0.0);
106            let next = cur + incr;
107            let smb = SmallBytes::from_slice(member);
108            let is_new = !z.by_member.contains_key(member);
109            z.insert(member, next);
110            let d = if is_new { zset_member_weight(&smb) as i64 } else { 0 };
111            (next, d)
112        };
113        self.account_delta(key, delta);
114        Ok(next)
115    }
116
117    /// `ZRANGE key start stop` by rank.
118    pub fn zrange(
119        &mut self,
120        key: &[u8],
121        start: i64,
122        stop: i64,
123    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
124        match self.zset_ref(key)? {
125            None => Ok(Vec::new()),
126            Some(z) => Ok(match range_bounds(start, stop, z.len()) {
127                None => Vec::new(),
128                Some((s, e)) => z
129                    .ordered()
130                    .skip(s)
131                    .take(e - s + 1)
132                    .map(|(m, sc)| (m.to_vec(), sc))
133                    .collect(),
134            }),
135        }
136    }
137
138    /// `ZRANGEBYSCORE` — members with score in the (possibly exclusive) bounds.
139    pub fn zrange_by_score(
140        &mut self,
141        key: &[u8],
142        min: ScoreBound,
143        max: ScoreBound,
144    ) -> Result<Vec<(Vec<u8>, f64)>, StoreError> {
145        Ok(self.zset_ref(key)?.map_or(Vec::new(), |z| {
146            z.ordered()
147                .filter(|(_, sc)| min.ge_ok(*sc) && max.le_ok(*sc))
148                .map(|(m, sc)| (m.to_vec(), sc))
149                .collect()
150        }))
151    }
152
153    /// `ZCOUNT` — number of members with score in the bounds.
154    pub fn zcount(
155        &mut self,
156        key: &[u8],
157        min: ScoreBound,
158        max: ScoreBound,
159    ) -> Result<usize, StoreError> {
160        Ok(self.zset_ref(key)?.map_or(0, |z| {
161            z.ordered()
162                .filter(|(_, sc)| min.ge_ok(*sc) && max.le_ok(*sc))
163                .count()
164        }))
165    }
166}
kevy_store/zset.rs

kevy_store/
zset.rs
