kevy_store/
set.rs

1//! `Store` set commands.
2
3use crate::value::{SetData, SmallBytes, Value, set_member_weight};
4use crate::{Entry, Store, StoreError};
5use std::sync::Arc;
6
7impl Store {
8    // ---- sets ----------------------------------------------------------
9
10    fn set_mut(&mut self, key: &[u8], create: bool) -> Result<Option<&mut SetData>, 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::Set(Arc::default()), None),
18            );
19        }
20        match &mut self.map.get_mut(key).expect("present").value {
21            Value::Set(s) => Ok(Some(Arc::make_mut(s))),
22            _ => Err(StoreError::WrongType),
23        }
24    }
25
26    fn set_ref(&mut self, key: &[u8]) -> Result<Option<&SetData>, StoreError> {
27        match self.live_entry(key) {
28            None => Ok(None),
29            Some(e) => match &e.value {
30                Value::Set(s) => Ok(Some(s.as_ref())),
31                _ => Err(StoreError::WrongType),
32            },
33        }
34    }
35
36    fn drop_if_empty_set(&mut self, key: &[u8]) {
37        let empty = matches!(self.map.get(key).map(|e| &e.value), Some(Value::Set(s)) if s.is_empty());
38        if empty {
39            self.remove_entry(key);
40        }
41    }
42
43    /// `SADD` — returns the count of newly-added members.
44    pub fn sadd(&mut self, key: &[u8], members: &[Vec<u8>]) -> Result<usize, StoreError> {
45        let (added, delta) = {
46            let s = self.set_mut(key, true)?.expect("created");
47            let mut a = 0usize;
48            let mut d: i64 = 0;
49            for m in members {
50                let smb = SmallBytes::from_slice(m);
51                let w = set_member_weight(&smb) as i64;
52                if s.insert(smb) {
53                    a += 1;
54                    d += w;
55                }
56            }
57            (a, d)
58        };
59        self.account_delta(key, delta);
60        Ok(added)
61    }
62
63    /// `SREM` — returns the count removed (deleting an emptied key).
64    pub fn srem(&mut self, key: &[u8], members: &[Vec<u8>]) -> Result<usize, StoreError> {
65        let (removed, delta) = {
66            let mut r = 0usize;
67            let mut d: i64 = 0;
68            if let Some(s) = self.set_mut(key, false)? {
69                for m in members {
70                    if s.remove(m.as_slice()) {
71                        r += 1;
72                        d -= set_member_weight(&SmallBytes::from_slice(m)) as i64;
73                    }
74                }
75            }
76            (r, d)
77        };
78        self.account_delta(key, delta);
79        self.drop_if_empty_set(key);
80        Ok(removed)
81    }
82
83    pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> Result<bool, StoreError> {
84        Ok(self.set_ref(key)?.is_some_and(|s| s.contains(member)))
85    }
86
87    pub fn scard(&mut self, key: &[u8]) -> Result<usize, StoreError> {
88        Ok(self.set_ref(key)?.map_or(0, kevy_map::KevySet::len))
89    }
90
91    pub fn smembers(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
92        Ok(self
93            .set_ref(key)?
94            .map_or(Vec::new(), |s| s.iter().map(kevy_bytes::SmallBytes::to_vec).collect()))
95    }
96
97    /// `SPOP key count` — remove and return up to `count` arbitrary members.
98    pub fn spop(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
99        let (out, delta) = {
100            let mut o: Vec<Vec<u8>> = Vec::new();
101            let mut d: i64 = 0;
102            if let Some(s) = self.set_mut(key, false)? {
103                let take: Vec<Vec<u8>> = s.iter().take(count).map(kevy_bytes::SmallBytes::to_vec).collect();
104                for m in &take {
105                    if s.remove(m.as_slice()) {
106                        d -= set_member_weight(&SmallBytes::from_slice(m)) as i64;
107                    }
108                }
109                o = take;
110            }
111            (o, d)
112        };
113        self.account_delta(key, delta);
114        self.drop_if_empty_set(key);
115        Ok(out)
116    }
117
118    /// `SRANDMEMBER key count` — up to `count` arbitrary members, not removed.
119    pub fn srandmember(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
120        Ok(self
121            .set_ref(key)?
122            .map_or(Vec::new(), |s| {
123                s.iter().take(count).map(kevy_bytes::SmallBytes::to_vec).collect()
124            }))
125    }
126
127    /// Snapshot of a set's members for cross-shard algebra (SINTER/etc.).
128    pub fn set_snapshot(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
129        self.smembers(key)
130    }
131}
kevy_store/set.rs

kevy_store/
set.rs
