kevy_store/
set_read.rs

1//! `Store` set read commands — split from `set.rs` when the SegSet
2//! arms pushed it against the 500-LOC cap.
3
4#[cfg(not(feature = "std"))]
5use crate::nostd_prelude::*;
6use crate::value::Value;
7use crate::{Store, StoreError};
8
9impl Store {
10    pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> Result<bool, StoreError> {
11        match self.live_entry(key) {
12            None => Ok(false),
13            Some(e) => match &e.value {
14                Value::Set(s) => Ok(s.contains(member)),
15                Value::SegSet(s) => Ok(s.contains_key(member)),
16                Value::SmallSetInline(s) => Ok(s.contains(member)),
17                _ => Err(StoreError::WrongType),
18            },
19        }
20    }
21
22    pub fn scard(&mut self, key: &[u8]) -> Result<usize, StoreError> {
23        match self.live_entry(key) {
24            None => Ok(0),
25            Some(e) => match &e.value {
26                Value::Set(s) => Ok(s.len()),
27                Value::SegSet(s) => Ok(s.len()),
28                Value::SmallSetInline(s) => Ok(s.len()),
29                _ => Err(StoreError::WrongType),
30            },
31        }
32    }
33
34    pub fn smembers(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
35        match self.live_entry(key) {
36            None => Ok(Vec::new()),
37            Some(e) => match &e.value {
38                Value::Set(s) => {
39                    Ok(s.iter().map(kevy_bytes::SmallBytes::to_vec).collect())
40                }
41                Value::SegSet(s) => Ok(s.keys().map(kevy_bytes::SmallBytes::to_vec).collect()),
42                Value::SmallSetInline(s) => {
43                    Ok(s.iter_slices().map(<[u8]>::to_vec).collect())
44                }
45                _ => Err(StoreError::WrongType),
46            },
47        }
48    }
49
50    /// `SRANDMEMBER key count` — up to `count` DISTINCT arbitrary
51    /// members, not removed.
52    ///
53    /// Two regimes, as Redis has: when `count` is a small fraction of
54    /// the set, probe random slots and reject duplicates — O(count)
55    /// expected. When it is most of the set, rejection would thrash, so
56    /// copy the members out and shuffle a prefix instead.
57    pub fn srandmember(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
58        let mut draws: Vec<u64> = (0..count.saturating_mul(3).max(8))
59            .map(|_| self.rng.next_u64())
60            .collect();
61        match self.live_entry(key) {
62            None => Ok(Vec::new()),
63            Some(e) => match &e.value {
64                Value::SmallSetInline(s) => {
65                    let mut all: Vec<Vec<u8>> = s.iter_slices().map(<[u8]>::to_vec).collect();
66                    let k = crate::set::shuffle_prefix(&mut all, count, &mut draws);
67                    all.truncate(k);
68                    Ok(all)
69                }
70                Value::Set(s) => {
71                    let n = s.len();
72                    if count >= n {
73                        return Ok(s.iter().map(kevy_bytes::SmallBytes::to_vec).collect());
74                    }
75                    if count * 4 >= n {
76                        // Wanting most of the set: copying beats rejecting.
77                        let mut all: Vec<Vec<u8>> =
78                            s.iter().map(kevy_bytes::SmallBytes::to_vec).collect();
79                        let k = crate::set::shuffle_prefix(&mut all, count, &mut draws);
80                        all.truncate(k);
81                        return Ok(all);
82                    }
83                    let mut out: Vec<Vec<u8>> = Vec::with_capacity(count);
84                    for slot in &draws {
85                        if out.len() == count {
86                            break;
87                        }
88                        if let Some(m) = s
89                            .iter_from_slot(*slot as usize)
90                            .next()
91                            .map(kevy_bytes::SmallBytes::to_vec)
92                            && !out.contains(&m)
93                        {
94                            out.push(m);
95                        }
96                    }
97                    Ok(out)
98                }
99                Value::SegSet(s) => Ok(seg_srandmember(s, count, &mut draws)),
100                _ => Err(StoreError::WrongType),
101            },
102        }
103    }
104
105    /// `SRANDMEMBER key -count` — exactly `count` members, WITH
106    /// repetition.
107    pub fn srandmember_with_repeats(
108        &mut self,
109        key: &[u8],
110        count: usize,
111    ) -> Result<Vec<Vec<u8>>, StoreError> {
112        let draws: Vec<u64> = (0..count).map(|_| self.rng.next_u64()).collect();
113        match self.live_entry(key) {
114            None => Ok(Vec::new()),
115            Some(e) => match &e.value {
116                Value::SmallSetInline(s) => {
117                    let all: Vec<Vec<u8>> = s.iter_slices().map(<[u8]>::to_vec).collect();
118                    if all.is_empty() {
119                        return Ok(Vec::new());
120                    }
121                    Ok(draws
122                        .iter()
123                        .map(|d| all[(*d as usize) % all.len()].clone())
124                        .collect())
125                }
126                Value::Set(s) => {
127                    if s.is_empty() {
128                        return Ok(Vec::new());
129                    }
130                    Ok(draws
131                        .iter()
132                        .filter_map(|d| {
133                            s.iter_from_slot(*d as usize)
134                                .next()
135                                .map(kevy_bytes::SmallBytes::to_vec)
136                        })
137                        .collect())
138                }
139                Value::SegSet(s) => {
140                    if s.is_empty() {
141                        return Ok(Vec::new());
142                    }
143                    Ok(draws
144                        .iter()
145                        .filter_map(|d| s.rand_entry(*d).map(|(m, ())| m.to_vec()))
146                        .collect())
147                }
148                _ => Err(StoreError::WrongType),
149            },
150        }
151    }
152
153    /// Snapshot of a set's members for cross-shard algebra (SINTER/etc.).
154    pub fn set_snapshot(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
155        self.smembers(key)
156    }
157}
158
159/// SRANDMEMBER over a sharded set: rejection-probe via the weighted
160/// random walk; degenerate huge counts fall back to the copy regime
161/// like the flat path.
162fn seg_srandmember(
163    s: &crate::seg_map::SegMap<()>,
164    count: usize,
165    draws: &mut Vec<u64>,
166) -> Vec<Vec<u8>> {
167    if count * 4 >= s.len() {
168        let mut all: Vec<Vec<u8>> = s.keys().map(kevy_bytes::SmallBytes::to_vec).collect();
169        let k = crate::set::shuffle_prefix(&mut all, count, draws);
170        all.truncate(k);
171        return all;
172    }
173    let mut out: Vec<Vec<u8>> = Vec::with_capacity(count);
174    for d in draws.iter() {
175        if out.len() == count {
176            break;
177        }
178        if let Some((m, ())) = s.rand_entry(*d)
179            && !out.contains(&m.to_vec())
180        {
181            out.push(m.to_vec());
182        }
183    }
184    out
185}
kevy_store/set_read.rs

kevy_store/
set_read.rs
