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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
//! Bitmap ops on string-typed values — `SETBIT` / `GETBIT` /
//! `BITCOUNT`. Redis treats strings as byte arrays addressed at the
//! bit level; this module exposes those reads / writes against the
//! existing string value encodings (`Value::Str` / `Value::ArcBulk` /
//! `Value::Int`).
//!
//! Split out from `string.rs` to keep that file under the 500-LOC
//! house rule.
use std::borrow::Cow;
use std::num::NonZeroU64;
use std::sync::Arc;
use crate::value::{SmallBytes, Value};
use crate::{Entry, Store, StoreError};
impl Store {
/// `GETBIT key offset` — read the bit at `offset` (MSB-first
/// within each byte, matching Redis). Returns `0` for missing
/// key or offset past the end. Errors on wrong type.
pub fn getbit(&mut self, key: &[u8], offset: u64) -> Result<u8, StoreError> {
let bytes = match self.get(key)? {
Some(cow) => cow,
None => return Ok(0),
};
let byte_idx = (offset / 8) as usize;
let bit_idx = 7 - (offset % 8) as u8;
if byte_idx >= bytes.len() {
return Ok(0);
}
Ok((bytes[byte_idx] >> bit_idx) & 1)
}
/// `SETBIT key offset value` — set the bit at `offset` to `value`
/// (0 or 1). Extends the underlying string with zero-padding if
/// `offset / 8 >= current_len`. Returns the PREVIOUS bit value.
/// Errors on wrong type or `value > 1`.
pub fn setbit(
&mut self,
key: &[u8],
offset: u64,
value: u8,
) -> Result<u8, StoreError> {
if value > 1 {
return Err(StoreError::OutOfRange);
}
let byte_idx = (offset / 8) as usize;
let bit_idx = 7 - (offset % 8) as u8;
// Read current bytes (Cow); compute previous bit; extend +
// write back. We collect into a fresh Vec each time — bitmaps
// tend to be hot-write so SmallBytes shrink-fit is moot.
let mut owned: Vec<u8> = match self.get(key)? {
Some(Cow::Borrowed(b)) => b.to_vec(),
Some(Cow::Owned(v)) => v,
None => Vec::new(),
};
if byte_idx >= owned.len() {
owned.resize(byte_idx + 1, 0);
}
let prev = (owned[byte_idx] >> bit_idx) & 1;
if value == 1 {
owned[byte_idx] |= 1 << bit_idx;
} else {
owned[byte_idx] &= !(1u8 << bit_idx);
}
// Store back. Always use the byte-array encoding (never int).
let new_val = if owned.is_empty() {
Value::Str(SmallBytes::from_slice(&[]))
} else {
Value::ArcBulk(Arc::new(owned.into_boxed_slice()))
};
// Take any existing TTL, re-attach to the new entry. Entry
// stores `expire_at_ns: Option<NonZeroU64>` (absolute ns).
let ttl_ns = self
.live_entry(key)
.and_then(|e| e.expire_at_ns.map(NonZeroU64::get));
self.insert_entry(
SmallBytes::from_slice(key),
Entry::new(new_val, ttl_ns),
);
Ok(prev)
}
/// `BITCOUNT key [start end [BYTE|BIT]]` — count set bits.
/// `start`/`end` are byte offsets (inclusive, negative-from-tail
/// like Redis). `None` for both = whole string.
pub fn bitcount(
&mut self,
key: &[u8],
range: Option<(i64, i64)>,
) -> Result<u64, StoreError> {
let bytes = match self.get(key)? {
Some(cow) => cow,
None => return Ok(0),
};
if bytes.is_empty() {
return Ok(0);
}
let len = bytes.len() as i64;
let (s, e) = match range {
None => (0, (len - 1) as usize),
Some((start, end)) => {
let norm = |x: i64| -> i64 {
if x < 0 { (len + x).max(0) } else { x.min(len - 1) }
};
let s = norm(start);
let e = norm(end);
if s > e {
return Ok(0);
}
(s as usize, e as usize)
}
};
Ok(bytes[s..=e]
.iter()
.map(|b| u64::from(b.count_ones()))
.sum())
}
/// `BITPOS key bit [start [end]]` — return the position (bit
/// index, MSB-first) of the first bit equal to `bit` (0 or 1)
/// in the byte range `[start, end]` (inclusive, Redis-style
/// negative indexing). Returns `None` (Redis `-1`) when not
/// found. Errors with `OutOfRange` if `bit` > 1.
pub fn bitpos(
&mut self,
key: &[u8],
bit: u8,
range: Option<(i64, i64)>,
) -> Result<Option<u64>, StoreError> {
if bit > 1 {
return Err(StoreError::OutOfRange);
}
let bytes = match self.get(key)? {
Some(cow) => cow,
None => return Ok(if bit == 0 { Some(0) } else { None }),
};
if bytes.is_empty() {
return Ok(if bit == 0 { Some(0) } else { None });
}
let len = bytes.len() as i64;
let (s, e) = match range {
None => (0usize, (len - 1) as usize),
Some((start, end)) => {
let norm = |x: i64| -> i64 {
if x < 0 { (len + x).max(0) } else { x.min(len - 1) }
};
let s = norm(start);
let e = norm(end);
if s > e {
return Ok(None);
}
(s as usize, e as usize)
}
};
for (i, &b) in bytes[s..=e].iter().enumerate() {
let target_mask = if bit == 1 { b } else { !b };
if target_mask != 0 {
let bit_in_byte = target_mask.leading_zeros() as u64;
let byte_idx = (s + i) as u64;
return Ok(Some(byte_idx * 8 + bit_in_byte));
}
}
Ok(None)
}
/// `GETRANGE key start end` — substring with Redis-style
/// negative indexing; `[start, end]` inclusive. Returns empty
/// `Vec` when key absent or range out of bounds.
pub fn getrange(
&mut self,
key: &[u8],
start: i64,
end: i64,
) -> Result<Vec<u8>, StoreError> {
let bytes = match self.get(key)? {
Some(cow) => cow,
None => return Ok(Vec::new()),
};
if bytes.is_empty() {
return Ok(Vec::new());
}
let len = bytes.len() as i64;
let norm = |x: i64| -> i64 {
if x < 0 { (len + x).max(0) } else { x.min(len - 1) }
};
let s = norm(start) as usize;
let e = norm(end) as usize;
if s > e {
return Ok(Vec::new());
}
Ok(bytes[s..=e].to_vec())
}
/// `SETRANGE key offset value` — overwrite bytes at `offset`
/// with `value`. Extends the string with zero padding if
/// `offset > len`. Returns the new total length. Preserves
/// any existing TTL.
pub fn setrange(
&mut self,
key: &[u8],
offset: u64,
value: &[u8],
) -> Result<usize, StoreError> {
let offset = offset as usize;
let mut owned: Vec<u8> = match self.get(key)? {
Some(Cow::Borrowed(b)) => b.to_vec(),
Some(Cow::Owned(v)) => v,
None => Vec::new(),
};
let needed = offset + value.len();
if needed > owned.len() {
owned.resize(needed, 0);
}
owned[offset..offset + value.len()].copy_from_slice(value);
let new_len = owned.len();
let new_val = if owned.is_empty() {
Value::Str(SmallBytes::from_slice(&[]))
} else {
Value::ArcBulk(Arc::new(owned.into_boxed_slice()))
};
let ttl_ns = self
.live_entry(key)
.and_then(|e| e.expire_at_ns.map(NonZeroU64::get));
self.insert_entry(
SmallBytes::from_slice(key),
Entry::new(new_val, ttl_ns),
);
Ok(new_len)
}
}