Skip to main content

yo_kv/
bitmaps.rs

1//! The bitmap commands, which are string commands wearing a different hat.
2//!
3//! A bitmap in Redis is a string, and that is not an implementation detail a
4//! caller can ignore: `SET k "A"` then `GETBIT k 1` answers 1, because `A` is
5//! `0x41` and the second bit from the top of that byte is set. So there is no
6//! bitmap type here either, and everything in this file works on the same
7//! string records [`strings`](crate::strings) writes. The kernels are in
8//! [`bits`]; this is where a key turns into bytes, where a write
9//! is allowed to grow a value and where Redis's edges live.
10//!
11//! Three of those edges are worth stating up front, because all three have been
12//! measured on a real server rather than reasoned about.
13//!
14//! A write always leaves the value `raw`. `SET n 12345` reports `int` and a
15//! `SETBIT n 0 0` that changes nothing at all still reports `raw` afterwards,
16//! because Redis unshares the object before it looks at a bit. A read does not:
17//! `GETBIT n 3` on the same key leaves it `int`. That is why the in place fast
18//! path below only takes a record that is already raw.
19//!
20//! A write creates the key and pads it with zero bytes, even when the bit being
21//! written is zero and the byte is past the end. `SETBIT nokey 0 0` on an empty
22//! database leaves a one byte string behind.
23//!
24//! A `BITFIELD` is checked all the way through before any of it runs, so a bad
25//! field type in the last subcommand leaves the key untouched and, if it was not
26//! there, uncreated. That ordering is the wire layer's to keep, and it is why
27//! [`Keyspace::bitfield`] takes a list of already parsed subcommands rather than
28//! words to parse.
29
30use crate::bits::{self, Field, Op, Overflow};
31use crate::db::Db;
32use crate::keyspace::Keyspace;
33use crate::lookups;
34use crate::strings::{STRING_MAX, check_len};
35use crate::value::{self, Kind, Str};
36use yo_common::num::{self, DIGITS_MAX};
37use yo_common::{Code, Error, Result};
38use yo_index::RawMap;
39
40/// What Redis says about an offset that is not a number or is off the end.
41const BAD_BIT_OFFSET: &str = "bit offset is not an integer or out of range";
42/// What Redis says when a write would make a string too long.
43const TOO_LONG: &str = "string exceeds maximum allowed size (proto-max-bulk-len)";
44
45/// The highest bit `SETBIT` and `GETBIT` take.
46///
47/// It is 4 Gi bits, which is 512 MiB, which is Redis's string ceiling. Ours is a
48/// segment and smaller than that, so a write between the two limits is refused
49/// by the length check with the "string exceeds maximum allowed size" sentence
50/// rather than by this one. Both are Redis's own sentences and the boundary
51/// between them is where we diverge.
52pub const BIT_OFFSET_MAX: u64 = 4 * 1024 * 1024 * 1024 - 1;
53
54/// Whether a range's two ends count bytes or bits.
55///
56/// `BITCOUNT` and `BITPOS` both take an optional `BYTE` or `BIT` word after
57/// their two indexes, and both default to `BYTE`. The word is only allowed once
58/// both indexes are there: `BITPOS k 0 5 BIT` is not a bit ranged search from
59/// bit five, it is an error, because `BIT` is read as the end index.
60#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
61pub enum Unit {
62    /// Indexes count bytes. The default.
63    #[default]
64    Byte,
65    /// Indexes count bits.
66    Bit,
67}
68
69/// One `BITFIELD` subcommand.
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub struct Sub {
72    /// Which of the three it is, and what it carries.
73    pub op: SubOp,
74    /// The width and signedness of the field.
75    pub field: Field,
76    /// Where the field starts, in bits.
77    ///
78    /// The `#n` form a client can send is `n` times the width, and multiplying
79    /// it out is the wire layer's job.
80    pub at: u64,
81    /// What to do if the value will not fit. Ignored by `GET`.
82    pub on: Overflow,
83}
84
85/// The three things a `BITFIELD` subcommand does.
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
87pub enum SubOp {
88    /// `GET`, which never writes and never creates the key.
89    Get,
90    /// `SET`, answering the value that was there before.
91    Set(i64),
92    /// `INCRBY`, answering the value afterwards.
93    Incr(i64),
94}
95
96impl SubOp {
97    /// Whether this one writes, which is what decides how far the value grows.
98    const fn writes(self) -> bool {
99        !matches!(self, SubOp::Get)
100    }
101}
102
103impl Keyspace {
104    /// `GETBIT key offset`.
105    ///
106    /// A missing key, and any offset past the end of a key that is there, read
107    /// as zero. Nothing is created and nothing is re-encoded.
108    pub fn getbit(&mut self, key: &[u8], offset: u64) -> Result<bool> {
109        if offset > BIT_OFFSET_MAX {
110            return Err(Error::new(Code::Invalid, BAD_BIT_OFFSET));
111        }
112        self.reap(key);
113        self.string_only(key)?;
114        // A bitmap is a string, so it can have been demoted like any other, and
115        // the bit being asked about is somewhere in it. Warmed rather than
116        // thawed: reading a bit out of a cold bitmap is a read like any other
117        // and the doorkeeper decides whether it earns its way back.
118        self.warm(key)?;
119        let mut digits = [0u8; DIGITS_MAX];
120        let bytes = self.bitmap(key, &mut digits);
121        let byte = (offset / 8) as usize;
122        Ok(bytes.get(byte).is_some_and(|b| b & mask(offset) != 0))
123    }
124
125    /// `SETBIT key offset value`, answering the bit that was there before and
126    /// whether the value had to get longer to hold the offset.
127    ///
128    /// The value grows to hold the offset, padded with zero bytes, and keeps
129    /// whatever deadline it had. A key that was not there is created, even when
130    /// the bit being written is zero.
131    ///
132    /// The reply is the first half and the second half is for the notification.
133    /// A real server only says `setbit` when the write did something, and doing
134    /// something means either the bit came out different or the value got
135    /// longer, so `SETBIT k 1 0` on a bit that was already zero says nothing
136    /// while `SETBIT k 1000 0` on a short value says it. The caller cannot work
137    /// the second half out from the reply, since a value that was created or
138    /// padded reads back as a zero bit either way.
139    pub fn setbit(&mut self, key: &[u8], offset: u64, bit: bool) -> Result<(bool, bool)> {
140        if offset > BIT_OFFSET_MAX {
141            return Err(Error::new(Code::Invalid, BAD_BIT_OFFSET));
142        }
143        let byte = (offset / 8) as usize;
144        check_len(key, byte + 1)?;
145        self.thaw(key)?;
146        let now = self.clock.now_ms();
147        let hash = RawMap::hash_of(key);
148
149        // The fast path: the key is there, it is raw already, and the byte is
150        // inside it, so the write is one probe and one byte. This is the shape a
151        // bitmap is used in, a fixed size map of ids that was sized once and is
152        // written to for the rest of its life, and it is the only path that does
153        // not touch the arena. The kind check sits inside the probe for the
154        // reason `INCR`'s does: the byte holding it is already loaded here.
155        let mut dead = false;
156        if let Some(rec) = self.map.value_mut_hashed(hash, key) {
157            if value::kind(rec) != Kind::String {
158                return Err(crate::keyspace::wrong_type());
159            }
160            if value::is_expired(rec, now) {
161                dead = true;
162            } else if let Some(b) = value::raw_in_place(rec).and_then(|it| it.get_mut(byte)) {
163                let had = *b & mask(offset) != 0;
164                if bit {
165                    *b |= mask(offset);
166                } else {
167                    *b &= !mask(offset);
168                }
169                // Nothing grew on this path by definition, since it is the one
170                // taken when the byte is already inside the value.
171                return Ok((had, false));
172            }
173        }
174        if dead {
175            self.reaped(key);
176        }
177
178        // The slow path, which is every first write to a key and every write
179        // that makes it longer. Through the one scratch buffer, the way `APPEND`
180        // and `SETRANGE` go, since the old bytes are needed in hand while
181        // `store_raw` wants the database.
182        let mut bytes = std::mem::take(&mut self.scratch);
183        bytes.clear();
184        let deadline = match self.map.get(key) {
185            Some(rec) => {
186                value::read(rec).write_to(&mut bytes);
187                value::expire_at(rec)
188            }
189            None => None,
190        };
191        // Whether the value got longer, which is not the same question as
192        // whether this path was taken. A key holding an int encoded value comes
193        // through here to be written out as digits even when the byte being
194        // written is already inside those digits, and that is not growth.
195        let grew = bytes.len() <= byte;
196        if grew {
197            bytes.resize(byte + 1, 0);
198        }
199        let had = bytes[byte] & mask(offset) != 0;
200        if bit {
201            bytes[byte] |= mask(offset);
202        } else {
203            bytes[byte] &= !mask(offset);
204        }
205        self.store_raw(key, &bytes, deadline);
206        self.scratch = bytes;
207        Ok((had, grew))
208    }
209
210    /// `BITCOUNT key [start end [BYTE | BIT]]`.
211    ///
212    /// A missing key, an empty string and a range that ends before it starts all
213    /// answer zero. The two indexes may be negative, counting from the end, and
214    /// both are clamped rather than refused.
215    pub fn bitcount(&mut self, key: &[u8], range: Option<(i64, i64, Unit)>) -> Result<u64> {
216        self.reap(key);
217        self.string_only(key)?;
218        self.warm(key)?;
219        let mut digits = [0u8; DIGITS_MAX];
220        let bytes = self.bitmap(key, &mut digits);
221        let Some((start, end, unit)) = range else {
222            return Ok(bits::count(bytes));
223        };
224        match window(bytes.len(), start, end, unit) {
225            Some((from, to)) => Ok(bits::count_range(bytes, from, to)),
226            None => Ok(0),
227        }
228    }
229
230    /// `BITPOS key bit [start [end [BYTE | BIT]]]`.
231    ///
232    /// Answers minus one when there is no such bit, with the one exception Redis
233    /// carved out: looking for a zero with no end index given, over a range that
234    /// is all ones, answers the first bit past the end of the string. The idea is
235    /// that a string is followed by an infinity of zeros unless the caller said
236    /// where to stop. Giving an explicit end turns that back into minus one, and
237    /// so does asking about a range that is empty once it has been clamped.
238    pub fn bitpos(
239        &mut self,
240        key: &[u8],
241        bit: bool,
242        start: Option<i64>,
243        end: Option<i64>,
244        unit: Unit,
245    ) -> Result<i64> {
246        self.reap(key);
247        self.string_only(key)?;
248        self.warm(key)?;
249        let here = self.map.get(key).is_some();
250        let mut digits = [0u8; DIGITS_MAX];
251        let bytes = self.bitmap(key, &mut digits);
252        if bytes.is_empty() {
253            // A missing key is all zeros, so a zero is at bit nought and a one is
254            // nowhere. An empty string that is really there answers minus one
255            // either way, since there is no bit nought to point at.
256            return Ok(if !bit && !here { 0 } else { -1 });
257        }
258        let all = bytes.len() as u64 * 8;
259        let (from, to) = match (start, end) {
260            (None, _) => (0, all),
261            (Some(s), None) => match window(bytes.len(), s, -1, unit) {
262                Some(r) => r,
263                None => return Ok(-1),
264            },
265            (Some(s), Some(e)) => match window(bytes.len(), s, e, unit) {
266                Some(r) => r,
267                None => return Ok(-1),
268            },
269        };
270        match bits::find(bytes, bit, from, to) {
271            Some(at) => Ok(at as i64),
272            None if !bit && end.is_none() => Ok(all as i64),
273            None => Ok(-1),
274        }
275    }
276
277    /// `BITOP op dest src [src ...]`, answering the length of the result.
278    ///
279    /// A result with no bytes in it deletes the destination, and any other
280    /// result creates it whatever it holds, so a `BITOP AND` over sources that
281    /// share nothing leaves a destination full of zero bytes rather than no
282    /// destination at all. Whatever the destination held before goes, even if
283    /// it was not a string: it is written and never read, so its type is not
284    /// part of the operation, and only the sources have to be strings. Sources that are shorter than the longest read as
285    /// zeros past their end, and a source that is not there reads as empty.
286    ///
287    /// # Panics
288    ///
289    /// If `srcs` is empty, or holds more than one key for [`Op::Not`]. Both are
290    /// refused with a message on the wire before this is called.
291    pub fn bitop<'k, I>(&mut self, op: Op, dest: &[u8], srcs: I) -> Result<usize>
292    where
293        I: Iterator<Item = &'k [u8]> + Clone,
294    {
295        for src in srcs.clone() {
296            self.reap(src);
297            self.string_only(src)?;
298            // Every source at once, so every one of them has to be in memory
299            // rather than in the one buffer a fault serves out of. `BITOP` over
300            // demoted sources brings them back, which is also what a client
301            // running it in a loop wants.
302            self.thaw(src)?;
303        }
304        // The sources have to be copied out before the destination can be
305        // written, since they are borrowed from the map and the write wants the
306        // database back. They go end to end into the scratch buffer with their
307        // boundaries in `rows`, and the result goes on the end of the same
308        // buffer, so a `BITOP` over any number of sources is one buffer and no
309        // allocation past whatever growing that buffer costs.
310        let mut flat = std::mem::take(&mut self.scratch);
311        let mut ends = std::mem::take(&mut self.rows);
312        flat.clear();
313        ends.clear();
314        let mut digits = [0u8; DIGITS_MAX];
315        for src in srcs.clone() {
316            let bytes = self.bitmap(src, &mut digits);
317            flat.extend_from_slice(bytes);
318            ends.push(flat.len());
319        }
320        // As long as the longest source, `NOT` included: complementing a source
321        // cannot make it longer, and there is only ever the one of them.
322        let len = bits::width(parts(&flat, &ends));
323        if len > STRING_MAX {
324            self.scratch = flat;
325            self.rows = ends;
326            return Err(Error::new(Code::Invalid, TOO_LONG));
327        }
328
329        let split = flat.len();
330        flat.resize(split + len, 0);
331        // The sources and the destination are in the same buffer, so they have
332        // to be split apart before one can be read while the other is written.
333        let (read, write) = flat.split_at_mut(split);
334        bits::combine(op, parts(read, &ends), write);
335
336        // The sources above are what a real server counts here. The destination
337        // is written and never read, so the lookups it takes are the ones
338        // Redis's `LOOKUP_WRITE` leaves out. See [`crate::lookups::quiet`].
339        let _quiet = lookups::quiet();
340        let outcome = if len == 0 {
341            self.del(dest);
342            Ok(0)
343        } else {
344            // Whatever the destination held and not only a string. It is
345            // never read, so its type is not part of the operation, and a real
346            // server writes over a list or a set here rather than refusing.
347            // The sources are the ones that have to be strings.
348            self.reap(dest);
349            self.replacing(dest, Kind::String);
350            self.store_raw(dest, &flat[split..], None);
351            Ok(len)
352        };
353        self.scratch = flat;
354        self.rows = ends;
355        outcome
356    }
357
358    /// `BITFIELD key [subcommand ...]`, answering one reply per subcommand.
359    ///
360    /// A `None` in the answers is the nil an `OVERFLOW FAIL` subcommand gives
361    /// when its value would not fit; that one does not write and the ones around
362    /// it still do. The subcommands are expected to have been checked already,
363    /// which is what makes it safe for this to be the point of no return.
364    ///
365    /// The value grows once, before anything runs, to hold the last bit any
366    /// writing subcommand touches. That happens even if every one of those
367    /// writes then fails its overflow check, which is Redis's behaviour and
368    /// falls out of it growing the string before it looks at the values.
369    pub fn bitfield(&mut self, key: &[u8], ops: &[Sub]) -> Result<Vec<Option<i64>>> {
370        let grow = ops.iter().filter(|s| s.op.writes()).map(reach).max();
371        let (out, _) = self.bitfield_with(key, grow, |bytes| {
372            ops.iter().map(|&sub| apply(bytes, sub).0).collect()
373        })?;
374        Ok(out)
375    }
376
377    /// `BITFIELD`, with the subcommands run against the value in place.
378    ///
379    /// This is the form the wire uses. It hands over the bytes and lets the
380    /// caller walk its own arguments a second time, calling [`apply`] on each,
381    /// which is what lets a `BITFIELD` with two hundred subcommands write two
382    /// hundred replies without a list of them existing anywhere.
383    ///
384    /// `grow` is how many bytes the value has to reach, which is the last byte
385    /// any writing subcommand touches, and `None` for a call that only reads.
386    /// The growing happens once and before anything runs, even if every one of
387    /// those writes then fails its overflow check, because that is what Redis
388    /// does: it makes the string long enough while it is looking up the key and
389    /// only then starts on the values. A call that only reads stores nothing,
390    /// which is what keeps `BITFIELD k GET u8 0` from turning an `embstr` into a
391    /// `raw`.
392    ///
393    /// The second half of the answer is whether the value did get longer, which
394    /// the notification wants for the reason [`Keyspace::setbit`] gives: a write
395    /// that grew the value counts as having done something even when every bit
396    /// it wrote came out the same as the one it replaced.
397    pub fn bitfield_with<T>(
398        &mut self,
399        key: &[u8],
400        grow: Option<usize>,
401        run: impl FnOnce(&mut [u8]) -> T,
402    ) -> Result<(T, bool)> {
403        self.reap(key);
404        self.string_only(key)?;
405        // Every path here materialises the value and most of them write it
406        // back, so this thaws rather than asking the doorkeeper about a value
407        // that is going to be resident when the command ends anyway.
408        self.thaw(key)?;
409        let need = grow.unwrap_or(0);
410        check_len(key, need)?;
411
412        // Every path materialises the value, including the read only one, so
413        // that an int encoded key reads as the digits it prints as.
414        let mut bytes = std::mem::take(&mut self.scratch);
415        bytes.clear();
416        let deadline = match self.map.get(key) {
417            Some(rec) => {
418                value::read(rec).write_to(&mut bytes);
419                value::expire_at(rec)
420            }
421            None => None,
422        };
423        let grew = bytes.len() < need;
424        if grew {
425            bytes.resize(need, 0);
426        }
427        let out = run(&mut bytes);
428        if grow.is_some() {
429            self.store_raw(key, &bytes, deadline);
430        }
431        self.scratch = bytes;
432        Ok((out, grew))
433    }
434
435    /// The bytes of a string key, as the bit commands want to see them.
436    ///
437    /// A missing key is empty, which is what every one of these commands treats
438    /// it as. An int encoded key is the digits it would print as, because that
439    /// is the string it is: `SET n 65` then `GETBIT n 1` is asking about the
440    /// character `6`. The digits are written into the caller's buffer so that the
441    /// ordinary case, a raw string, is still a borrow and not a copy.
442    fn bitmap<'a>(&'a self, key: &[u8], digits: &'a mut [u8; DIGITS_MAX]) -> &'a [u8] {
443        match self.peek(key) {
444            None => &[],
445            Some(Str::Bytes(b)) => b,
446            Some(Str::Int(n)) => num::i64_digits(digits, n),
447        }
448    }
449}
450
451impl Db {
452    /// `BITOP op dest src [src ...]` over a database of any width.
453    ///
454    /// Every key on one stripe is that one stripe's `BITOP`, which is every
455    /// `BITOP` on a database of one stripe and every `BITOP` whose keys were
456    /// hash tagged into the same place. That path is the old one, byte for byte.
457    ///
458    /// The rest is the same work with the reads spread out. Every stripe the
459    /// command names is held for the whole of it, the sources are copied into a
460    /// buffer this database owns rather than one a stripe owns, and they are
461    /// combined there and written to whichever stripe the destination is on.
462    /// Held together rather than one after the other, because an operand that
463    /// was written to after it had been read would leave a result that no
464    /// arrangement of these keys ever had.
465    ///
466    /// # Panics
467    ///
468    /// As [`Keyspace::bitop`].
469    pub fn bitop<'k, I>(&self, op: Op, dest: &'k [u8], srcs: I) -> Result<usize>
470    where
471        I: Iterator<Item = &'k [u8]> + Clone,
472    {
473        if let Some(home) = self.one_stripe(std::iter::once(dest).chain(srcs.clone())) {
474            return self.hold_stripe(home).bitop(op, dest, srcs);
475        }
476        // The buffers before the stripes, which is the order every command that
477        // wants both takes them in.
478        let mut spare = self.spare();
479        let spare = &mut *spare;
480        let (flat, ends) = (&mut spare.bytes, &mut spare.rows);
481        flat.clear();
482        ends.clear();
483        let onto = self.stripe_of(dest);
484        let mut held = self
485            .hold_many(std::iter::once(onto).chain(srcs.clone().map(|src| self.stripe_of(src))));
486        for src in srcs.clone() {
487            let stripe = held.stripe_mut(self.stripe_of(src));
488            stripe.reap(src);
489            stripe.string_only(src)?;
490            stripe.thaw(src)?;
491        }
492        let mut digits = [0u8; DIGITS_MAX];
493        for src in srcs.clone() {
494            let bytes = held.stripe(self.stripe_of(src)).bitmap(src, &mut digits);
495            flat.extend_from_slice(bytes);
496            ends.push(flat.len());
497        }
498        let len = bits::width(parts(flat, ends));
499        if len > STRING_MAX {
500            return Err(Error::new(Code::Invalid, TOO_LONG));
501        }
502
503        let split = flat.len();
504        flat.resize(split + len, 0);
505        let (read, write) = flat.split_at_mut(split);
506        bits::combine(op, parts(read, ends), write);
507
508        if len == 0 {
509            held.stripe_mut(onto).del(dest);
510            return Ok(0);
511        }
512        // Whatever the destination held, for the reason [`Keyspace::bitop`]
513        // gives: it is written and never read.
514        let stripe = held.stripe_mut(onto);
515        stripe.reap(dest);
516        stripe.replacing(dest, Kind::String);
517        stripe.store_raw(dest, &flat[split..], None);
518        Ok(len)
519    }
520}
521
522/// The sources of a `BITOP`, out of the buffer they were copied into.
523///
524/// The boundaries are the end of each source, so the first one starts at nought
525/// and each of the others starts where the one before it ended. Written as a
526/// zip over two views of the same list rather than as a running offset, because
527/// the iterator has to be cloneable and a clone of a running offset would carry
528/// whatever the original had reached.
529fn parts<'a>(flat: &'a [u8], ends: &'a [usize]) -> impl Iterator<Item = &'a [u8]> + Clone {
530    std::iter::once(0)
531        .chain(ends.iter().copied())
532        .zip(ends.iter().copied())
533        .map(|(from, to)| &flat[from..to])
534}
535
536/// Run one subcommand against a value, answering what the client is owed and
537/// whether it left the value different from how it found it.
538///
539/// `None` is the nil an `OVERFLOW FAIL` subcommand gives when its value would
540/// not fit; that one writes nothing and the ones around it still do. A `SET`
541/// answers what was there before and an `INCRBY` answers what is there now,
542/// which is not symmetry anybody would have chosen but is what Redis does.
543///
544/// The second half is what the notification wants, and no call site can work it
545/// out from the first: a `SET` answering the old value has not said what the new
546/// one is, and an `INCRBY` answering the new one has not said what the old one
547/// was. A `GET` never changes anything and a subcommand that failed its overflow
548/// check wrote nothing, so both of those are false.
549///
550/// The bytes have to be long enough already, which is [`reach`]'s job.
551#[must_use]
552pub fn apply(bytes: &mut [u8], sub: Sub) -> (Option<i64>, bool) {
553    let had = bits::get(bytes, sub.at, sub.field);
554    match sub.op {
555        SubOp::Get => (Some(had), false),
556        SubOp::Set(val) => match bits::setting(sub.field, val, sub.on) {
557            Some(next) => {
558                bits::set(bytes, sub.at, sub.field, next);
559                (Some(had), next != had)
560            }
561            None => (None, false),
562        },
563        SubOp::Incr(by) => match bits::adding(sub.field, had, by, sub.on) {
564            Some(next) => {
565                bits::set(bytes, sub.at, sub.field, next);
566                (Some(next), next != had)
567            }
568            None => (None, false),
569        },
570    }
571}
572
573/// How many bytes a value needs before `sub` can be written into it.
574#[must_use]
575pub const fn reach(sub: &Sub) -> usize {
576    (sub.field.last_bit(sub.at) / 8 + 1) as usize
577}
578
579/// The bit `offset` names inside its byte.
580///
581/// Bit zero is the top bit, which is the convention all of these commands use.
582#[inline]
583const fn mask(offset: u64) -> u8 {
584    0x80 >> (offset % 8)
585}
586
587/// A start and end index turned into a half open range of bits.
588///
589/// `None` for a range that holds nothing, which is what an empty value, an
590/// out of range start or a backwards range all come to. Negative indexes count
591/// from the end and both ends are clamped, so `BITCOUNT k -100 100` over a three
592/// byte string is the whole string rather than an error.
593fn window(len: usize, start: i64, end: i64, unit: Unit) -> Option<(u64, u64)> {
594    let items = match unit {
595        Unit::Byte => len as i64,
596        Unit::Bit => (len as i64).checked_mul(8)?,
597    };
598    if items == 0 {
599        return None;
600    }
601    // The two ends are not clamped the same way, and the difference is what
602    // makes `BITCOUNT k 10 20` over a three byte string answer zero rather than
603    // counting its last byte. A negative index counts back from the end and
604    // stops at the front, the end index is pulled back to the last item, and a
605    // start past the last item is left where it is so that the range comes out
606    // backwards and is thrown away below.
607    let back = |i: i64| if i < 0 { (items + i).max(0) } else { i };
608    let (from, to) = (back(start), back(end).min(items - 1));
609    if from > to {
610        return None;
611    }
612    let scale = match unit {
613        Unit::Byte => 8,
614        Unit::Bit => 1,
615    };
616    Some(((from * scale) as u64, ((to + 1) * scale) as u64))
617}
618
619/// The largest value a bit range can name, for a caller checking its own limit.
620///
621/// Nothing here uses it; it is the ceiling [`STRING_MAX`] imposes expressed in
622/// bits, which is what a client asking "how big can this bitmap be" wants.
623#[must_use]
624pub const fn max_bits() -> u64 {
625    STRING_MAX as u64 * 8
626}
627
628#[cfg(test)]
629mod tests {
630    use super::*;
631    use crate::keyspace::Keyspace;
632
633    fn db() -> Keyspace {
634        Keyspace::new()
635    }
636
637    /// The source list `bitop` takes, out of the keys a test wants to name.
638    fn keys<'k>(names: &'k [&'k [u8]]) -> impl Iterator<Item = &'k [u8]> + Clone {
639        names.iter().copied()
640    }
641
642    #[test]
643    fn a_bit_is_set_and_read_back() {
644        let mut db = db();
645        assert!(!db.setbit(b"k", 7, true).expect("a bit").0);
646        assert!(db.getbit(b"k", 7).expect("a bit"));
647        assert!(!db.getbit(b"k", 6).expect("a bit"));
648        assert_eq!(db.strlen(b"k").expect("a length"), 1);
649        assert_eq!(
650            db.get(b"k").expect("a value").expect("bytes").to_vec(),
651            b"\x01"
652        );
653        // The answer is what was there, not what is there now.
654        assert!(db.setbit(b"k", 7, false).expect("a bit").0);
655        assert!(!db.setbit(b"k", 7, false).expect("a bit").0);
656    }
657
658    #[test]
659    fn a_write_creates_and_pads_even_when_the_bit_is_zero() {
660        let mut db = db();
661        assert!(!db.setbit(b"k", 0, false).expect("a bit").0);
662        assert!(db.exists(b"k"));
663        assert_eq!(db.strlen(b"k").expect("a length"), 1);
664        db.setbit(b"k", 40, true).expect("a bit");
665        assert_eq!(db.strlen(b"k").expect("a length"), 6);
666    }
667
668    #[test]
669    fn a_write_leaves_the_value_raw_and_a_read_does_not() {
670        let mut db = db();
671        db.set_plain(b"n", b"12345").expect("a set");
672        assert_eq!(db.encoding(b"n"), Some(value::Encoding::Int));
673        // Reading a bit out of an int is reading a bit out of its digits.
674        assert!(db.getbit(b"n", 3).expect("a bit"));
675        assert_eq!(db.encoding(b"n"), Some(value::Encoding::Int));
676        // Writing one, even a write that changes nothing, does not leave an int.
677        assert!(!db.setbit(b"n", 0, false).expect("a bit").0);
678        assert_eq!(db.encoding(b"n"), Some(value::Encoding::Raw));
679        assert_eq!(
680            db.get(b"n").expect("a value").expect("bytes").to_vec(),
681            b"12345"
682        );
683    }
684
685    #[test]
686    fn a_write_keeps_the_deadline() {
687        let mut db = db();
688        db.setex(b"k", 100, b"abc").expect("a set");
689        db.setbit(b"k", 40, true).expect("a bit");
690        assert_eq!(db.strlen(b"k").expect("a length"), 6);
691        assert!(db.expire_at(b"k").is_some());
692        // And so does the fast path, which does not go near the deadline.
693        db.setbit(b"k", 1, true).expect("a bit");
694        assert!(db.expire_at(b"k").is_some());
695    }
696
697    #[test]
698    fn counting_takes_the_ranges_a_real_server_takes() {
699        let mut db = db();
700        db.set_plain(b"k", b"foobar").expect("a set");
701        let count = |db: &mut Keyspace, r| db.bitcount(b"k", r).expect("a count");
702        assert_eq!(count(&mut db, None), 26);
703        assert_eq!(count(&mut db, Some((0, 0, Unit::Byte))), 4);
704        assert_eq!(count(&mut db, Some((1, 1, Unit::Byte))), 6);
705        assert_eq!(count(&mut db, Some((0, -5, Unit::Byte))), 10);
706        assert_eq!(count(&mut db, Some((5, 30, Unit::Bit))), 17);
707        // Redis's own documentation says 22 for this one. A real 8.10.1 says 25,
708        // and 25 is what counting the first 44 bits of `foobar` by hand gives,
709        // so the documentation is wrong and this is not a divergence.
710        assert_eq!(count(&mut db, Some((0, -5, Unit::Bit))), 25);
711        // Clamped at both ends, empty when it is backwards.
712        assert_eq!(count(&mut db, Some((-100, 100, Unit::Byte))), 26);
713        assert_eq!(count(&mut db, Some((2, 1, Unit::Byte))), 0);
714        assert_eq!(count(&mut db, Some((5, 3, Unit::Bit))), 0);
715        // A start past the end is nothing, not the whole string.
716        assert_eq!(count(&mut db, Some((10, 20, Unit::Byte))), 0);
717        assert_eq!(db.bitcount(b"gone", None).expect("a count"), 0);
718    }
719
720    #[test]
721    fn searching_takes_the_ranges_a_real_server_takes() {
722        let mut db = db();
723        db.set_plain(b"ones", b"\xff\xff\xff").expect("a set");
724        db.set_plain(b"mix", b"\x00\xff\x00").expect("a set");
725        let pos = |db: &mut Keyspace, k: &[u8], bit, s, e| {
726            db.bitpos(k, bit, s, e, Unit::Byte).expect("a position")
727        };
728        assert_eq!(pos(&mut db, b"mix", true, None, None), 8);
729        assert_eq!(pos(&mut db, b"mix", false, None, None), 0);
730        assert_eq!(pos(&mut db, b"mix", true, Some(2), None), -1);
731        assert_eq!(pos(&mut db, b"mix", true, Some(-1), Some(-1)), -1);
732        assert_eq!(pos(&mut db, b"mix", false, Some(-100), None), 0);
733        // The one exception: no end given, all ones, so the answer is the first
734        // bit past the end of the string.
735        assert_eq!(pos(&mut db, b"ones", false, None, None), 24);
736        assert_eq!(pos(&mut db, b"ones", false, Some(-1), None), 24);
737        // An explicit end takes that away again.
738        assert_eq!(pos(&mut db, b"ones", false, Some(0), Some(-1)), -1);
739        assert_eq!(pos(&mut db, b"ones", false, Some(0), Some(100)), -1);
740        // And so does a range that is empty once it has been clamped.
741        assert_eq!(pos(&mut db, b"ones", false, Some(10), None), -1);
742        assert_eq!(pos(&mut db, b"ones", false, Some(3), None), -1);
743        assert_eq!(pos(&mut db, b"ones", true, Some(10), None), -1);
744        assert_eq!(pos(&mut db, b"ones", false, Some(2), Some(1)), -1);
745        assert_eq!(
746            db.bitpos(b"ones", false, Some(5), Some(20), Unit::Bit)
747                .expect("a position"),
748            -1
749        );
750    }
751
752    #[test]
753    fn searching_an_absent_or_empty_key() {
754        let mut db = db();
755        let pos = |db: &mut Keyspace, k: &[u8], bit| {
756            db.bitpos(k, bit, None, None, Unit::Byte)
757                .expect("a position")
758        };
759        // A key that is not there is all zeros, so a zero is at the front.
760        assert_eq!(pos(&mut db, b"gone", false), 0);
761        assert_eq!(pos(&mut db, b"gone", true), -1);
762        // A key that is there and empty has no bits at all.
763        db.set_plain(b"empty", b"").expect("a set");
764        assert_eq!(pos(&mut db, b"empty", false), -1);
765        assert_eq!(pos(&mut db, b"empty", true), -1);
766        assert_eq!(
767            db.bitcount(b"empty", Some((0, -1, Unit::Byte)))
768                .expect("a count"),
769            0
770        );
771    }
772
773    #[test]
774    fn combining_writes_a_destination_and_deletes_an_empty_one() {
775        let mut db = db();
776        db.set_plain(b"a", b"\xf0\x0f\xff").expect("a set");
777        db.set_plain(b"b", b"\xff\x00").expect("a set");
778        let n = db
779            .bitop(Op::And, b"d", keys(&[b"a", b"b"]))
780            .expect("a length");
781        assert_eq!(n, 3);
782        assert_eq!(
783            db.get(b"d").expect("a value").expect("bytes").to_vec(),
784            b"\xf0\x00\x00"
785        );
786        // A destination full of nothing is still a destination.
787        db.set_plain(b"z", b"\x00\x00").expect("a set");
788        let n = db
789            .bitop(Op::And, b"d", keys(&[b"a", b"z"]))
790            .expect("a length");
791        assert_eq!(n, 3);
792        assert!(db.exists(b"d"));
793        // Sources that are all missing take the destination with them.
794        let n = db
795            .bitop(Op::Or, b"d", keys(&[b"no1", b"no2"]))
796            .expect("a length");
797        assert_eq!(n, 0);
798        assert!(!db.exists(b"d"));
799    }
800
801    /// The destination is written and never read, so what it held before does
802    /// not have to be a string and does not have to be anything.
803    #[test]
804    fn combining_writes_over_a_destination_of_any_type() {
805        let mut db = db();
806        db.set_plain(b"a", b"abc").expect("a set");
807        db.sadd(b"d", [b"m".as_slice()].into_iter())
808            .expect("a member");
809        let n = db.bitop(Op::And, b"d", keys(&[b"a"])).expect("a length");
810        assert_eq!(n, 3);
811        assert_eq!(
812            db.get(b"d").expect("a value").expect("bytes").to_vec(),
813            b"abc"
814        );
815        // And a source that is not a string is still refused, which is what
816        // makes the line above the destination rather than the check going.
817        db.sadd(b"s", [b"m".as_slice()].into_iter())
818            .expect("a member");
819        assert!(db.bitop(Op::Or, b"d", keys(&[b"s"])).is_err());
820    }
821
822    #[test]
823    fn combining_reads_an_int_key_as_its_digits() {
824        let mut db = db();
825        db.set_plain(b"n", b"12345").expect("a set");
826        db.bitop(Op::Or, b"d", keys(&[b"n"])).expect("a length");
827        assert_eq!(
828            db.get(b"d").expect("a value").expect("bytes").to_vec(),
829            b"12345"
830        );
831    }
832
833    #[test]
834    fn a_field_is_read_written_and_incremented() {
835        let mut db = db();
836        let u8f = Field::new(false, 8).expect("a width");
837        let sub = |op, at| Sub {
838            op,
839            field: u8f,
840            at,
841            on: Overflow::Wrap,
842        };
843        let out = db
844            .bitfield(b"k", &[sub(SubOp::Set(255), 0), sub(SubOp::Get, 0)])
845            .expect("replies");
846        assert_eq!(out, vec![Some(0), Some(255)]);
847        assert_eq!(db.strlen(b"k").expect("a length"), 1);
848
849        let out = db
850            .bitfield(b"k", &[sub(SubOp::Incr(10), 0)])
851            .expect("replies");
852        assert_eq!(out, vec![Some(9)], "wrapped round");
853
854        // A failing write answers nothing and leaves the field alone, and the
855        // subcommands around it still run.
856        let fail = Sub {
857            on: Overflow::Fail,
858            ..sub(SubOp::Incr(250), 0)
859        };
860        let out = db
861            .bitfield(b"k", &[fail, sub(SubOp::Get, 0)])
862            .expect("replies");
863        assert_eq!(out, vec![None, Some(9)]);
864    }
865
866    #[test]
867    fn a_read_only_bitfield_creates_nothing_and_re_encodes_nothing() {
868        let mut db = db();
869        let f = Field::new(true, 16).expect("a width");
870        let get = Sub {
871            op: SubOp::Get,
872            field: f,
873            at: 0,
874            on: Overflow::Wrap,
875        };
876        assert_eq!(
877            db.bitfield(b"gone", &[get]).expect("replies"),
878            vec![Some(0)]
879        );
880        assert!(!db.exists(b"gone"));
881
882        db.set_plain(b"s", b"hello").expect("a set");
883        assert_eq!(db.encoding(b"s"), Some(value::Encoding::Embstr));
884        db.bitfield(b"s", &[get]).expect("replies");
885        assert_eq!(
886            db.encoding(b"s"),
887            Some(value::Encoding::Embstr),
888            "still short"
889        );
890    }
891
892    #[test]
893    fn a_write_grows_the_value_even_when_every_write_fails() {
894        let mut db = db();
895        let f = Field::new(false, 8).expect("a width");
896        let sub = Sub {
897            op: SubOp::Set(300),
898            field: f,
899            at: 64,
900            on: Overflow::Fail,
901        };
902        assert_eq!(db.bitfield(b"k", &[sub]).expect("replies"), vec![None]);
903        assert_eq!(db.strlen(b"k").expect("a length"), 9);
904    }
905
906    #[test]
907    fn a_bit_command_on_the_wrong_type_says_so() {
908        let mut db = db();
909        let member: &[u8] = b"x";
910        db.sadd(b"s", std::iter::once(member)).expect("a member");
911        assert!(db.getbit(b"s", 0).is_err());
912        assert!(db.setbit(b"s", 0, true).is_err());
913        assert!(db.bitcount(b"s", None).is_err());
914        assert!(db.bitpos(b"s", true, None, None, Unit::Byte).is_err());
915        assert!(db.bitop(Op::Or, b"d", keys(&[b"s"])).is_err());
916        let f = Field::new(false, 8).expect("a width");
917        let sub = Sub {
918            op: SubOp::Get,
919            field: f,
920            at: 0,
921            on: Overflow::Wrap,
922        };
923        assert!(db.bitfield(b"s", &[sub]).is_err());
924    }
925
926    #[test]
927    fn an_offset_past_the_end_of_the_world_is_refused() {
928        let mut db = db();
929        assert!(db.setbit(b"k", BIT_OFFSET_MAX + 1, true).is_err());
930        assert!(db.getbit(b"k", BIT_OFFSET_MAX + 1).is_err());
931        // And one inside Redis's limit but outside ours is refused too, with the
932        // other sentence. This is the divergence [`STRING_MAX`] is about.
933        assert!(db.setbit(b"k", BIT_OFFSET_MAX, true).is_err());
934        assert!(max_bits() < BIT_OFFSET_MAX);
935    }
936}