kevy_store/
set.rs

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

kevy_store/
set.rs
