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    /// Membership. A missing key is `false`, not an error; a
11    /// wrong-typed key is an error.
12    pub fn sismember(&mut self, key: &[u8], member: &[u8]) -> Result<bool, StoreError> {
13        match self.live_entry(key) {
14            None => Ok(false),
15            Some(e) => match &e.value {
16                Value::Set(s) => Ok(s.contains(member)),
17                Value::SegSet(s) => Ok(s.contains_key(member)),
18                Value::SmallSetInline(s) => Ok(s.contains(member)),
19                _ => Err(StoreError::WrongType),
20            },
21        }
22    }
23
24    /// Member count. A missing key is 0, matching SCARD.
25    pub fn scard(&mut self, key: &[u8]) -> Result<usize, StoreError> {
26        match self.live_entry(key) {
27            None => Ok(0),
28            Some(e) => match &e.value {
29                Value::Set(s) => Ok(s.len()),
30                Value::SegSet(s) => Ok(s.len()),
31                Value::SmallSetInline(s) => Ok(s.len()),
32                _ => Err(StoreError::WrongType),
33            },
34        }
35    }
36
37    /// Every member, copied out. Unordered — a set has no order to
38    /// preserve, and callers that need one must sort.
39    pub fn smembers(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
40        match self.live_entry(key) {
41            None => Ok(Vec::new()),
42            Some(e) => match &e.value {
43                Value::Set(s) => Ok(s.iter().map(kevy_bytes::SmallBytes::to_vec).collect()),
44                Value::SegSet(s) => Ok(s.keys().map(kevy_bytes::SmallBytes::to_vec).collect()),
45                Value::SmallSetInline(s) => Ok(s.iter_slices().map(<[u8]>::to_vec).collect()),
46                _ => Err(StoreError::WrongType),
47            },
48        }
49    }
50
51    /// `SRANDMEMBER key count` — up to `count` DISTINCT arbitrary
52    /// members, not removed.
53    ///
54    /// Two regimes, as Redis has: when `count` is a small fraction of
55    /// the set, probe random slots and reject duplicates — O(count)
56    /// expected. When it is most of the set, rejection would thrash, so
57    /// copy the members out and shuffle a prefix instead.
58    pub fn srandmember(&mut self, key: &[u8], count: usize) -> Result<Vec<Vec<u8>>, StoreError> {
59        let mut draws: Vec<u64> =
60            (0..count.saturating_mul(3).max(8)).map(|_| self.rng.next_u64()).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.iter().map(|d| all[(*d as usize) % all.len()].clone()).collect())
122                }
123                Value::Set(s) => {
124                    if s.is_empty() {
125                        return Ok(Vec::new());
126                    }
127                    Ok(draws
128                        .iter()
129                        .filter_map(|d| {
130                            s.iter_from_slot(*d as usize).next().map(kevy_bytes::SmallBytes::to_vec)
131                        })
132                        .collect())
133                }
134                Value::SegSet(s) => {
135                    if s.is_empty() {
136                        return Ok(Vec::new());
137                    }
138                    Ok(draws
139                        .iter()
140                        .filter_map(|d| s.rand_entry(*d).map(|(m, ())| m.to_vec()))
141                        .collect())
142                }
143                _ => Err(StoreError::WrongType),
144            },
145        }
146    }
147
148    /// Snapshot of a set's members for cross-shard algebra (SINTER/etc.).
149    pub fn set_snapshot(&mut self, key: &[u8]) -> Result<Vec<Vec<u8>>, StoreError> {
150        self.smembers(key)
151    }
152}
153
154/// SRANDMEMBER over a sharded set: rejection-probe via the weighted
155/// random walk; degenerate huge counts fall back to the copy regime
156/// like the flat path.
157fn seg_srandmember(
158    s: &crate::seg_map::SegMap<()>,
159    count: usize,
160    draws: &mut Vec<u64>,
161) -> Vec<Vec<u8>> {
162    if count * 4 >= s.len() {
163        let mut all: Vec<Vec<u8>> = s.keys().map(kevy_bytes::SmallBytes::to_vec).collect();
164        let k = crate::set::shuffle_prefix(&mut all, count, draws);
165        all.truncate(k);
166        return all;
167    }
168    let mut out: Vec<Vec<u8>> = Vec::with_capacity(count);
169    for d in draws.iter() {
170        if out.len() == count {
171            break;
172        }
173        if let Some((m, ())) = s.rand_entry(*d)
174            && !out.contains(&m.to_vec())
175        {
176            out.push(m.to_vec());
177        }
178    }
179    out
180}
kevy_store/set_read.rs

kevy_store/
set_read.rs
