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