kevy_store/
bitmap.rs

1//! Bitmap ops on string-typed values — `SETBIT` / `GETBIT` /
2//! `BITCOUNT`. Redis treats strings as byte arrays addressed at the
3//! bit level; this module exposes those reads / writes against the
4//! existing string value encodings (`Value::Str` / `Value::ArcBulk` /
5//! `Value::Int`).
6//!
7//! Split out from `string.rs` to keep that file under the 500-LOC
8//! house rule.
9
10use std::borrow::Cow;
11use std::num::NonZeroU64;
12use std::sync::Arc;
13
14use crate::value::{SmallBytes, Value};
15use crate::{Entry, Store, StoreError};
16
17impl Store {
18    /// `GETBIT key offset` — read the bit at `offset` (MSB-first
19    /// within each byte, matching Redis). Returns `0` for missing
20    /// key or offset past the end. Errors on wrong type.
21    pub fn getbit(&mut self, key: &[u8], offset: u64) -> Result<u8, StoreError> {
22        let bytes = match self.get(key)? {
23            Some(cow) => cow,
24            None => return Ok(0),
25        };
26        let byte_idx = (offset / 8) as usize;
27        let bit_idx = 7 - (offset % 8) as u8;
28        if byte_idx >= bytes.len() {
29            return Ok(0);
30        }
31        Ok((bytes[byte_idx] >> bit_idx) & 1)
32    }
33
34    /// `SETBIT key offset value` — set the bit at `offset` to `value`
35    /// (0 or 1). Extends the underlying string with zero-padding if
36    /// `offset / 8 >= current_len`. Returns the PREVIOUS bit value.
37    /// Errors on wrong type or `value > 1`.
38    pub fn setbit(
39        &mut self,
40        key: &[u8],
41        offset: u64,
42        value: u8,
43    ) -> Result<u8, StoreError> {
44        if value > 1 {
45            return Err(StoreError::OutOfRange);
46        }
47        let byte_idx = (offset / 8) as usize;
48        let bit_idx = 7 - (offset % 8) as u8;
49
50        // Read current bytes (Cow); compute previous bit; extend +
51        // write back. We collect into a fresh Vec each time — bitmaps
52        // tend to be hot-write so SmallBytes shrink-fit is moot.
53        let mut owned: Vec<u8> = match self.get(key)? {
54            Some(Cow::Borrowed(b)) => b.to_vec(),
55            Some(Cow::Owned(v)) => v,
56            None => Vec::new(),
57        };
58        if byte_idx >= owned.len() {
59            owned.resize(byte_idx + 1, 0);
60        }
61        let prev = (owned[byte_idx] >> bit_idx) & 1;
62        if value == 1 {
63            owned[byte_idx] |= 1 << bit_idx;
64        } else {
65            owned[byte_idx] &= !(1u8 << bit_idx);
66        }
67        // Store back. Always use the byte-array encoding (never int).
68        let new_val = if owned.is_empty() {
69            Value::Str(SmallBytes::from_slice(&[]))
70        } else {
71            Value::ArcBulk(Arc::new(owned.into_boxed_slice()))
72        };
73        // Take any existing TTL, re-attach to the new entry. Entry
74        // stores `expire_at_ns: Option<NonZeroU64>` (absolute ns).
75        let ttl_ns = self
76            .live_entry(key)
77            .and_then(|e| e.expire_at_ns.map(NonZeroU64::get));
78        self.insert_entry(
79            SmallBytes::from_slice(key),
80            Entry::new(new_val, ttl_ns),
81        );
82        Ok(prev)
83    }
84
85    /// `BITCOUNT key [start end [BYTE|BIT]]` — count set bits.
86    /// `start`/`end` are byte offsets (inclusive, negative-from-tail
87    /// like Redis). `None` for both = whole string.
88    pub fn bitcount(
89        &mut self,
90        key: &[u8],
91        range: Option<(i64, i64)>,
92    ) -> Result<u64, StoreError> {
93        let bytes = match self.get(key)? {
94            Some(cow) => cow,
95            None => return Ok(0),
96        };
97        if bytes.is_empty() {
98            return Ok(0);
99        }
100        let len = bytes.len() as i64;
101        let (s, e) = match range {
102            None => (0, (len - 1) as usize),
103            Some((start, end)) => {
104                let norm = |x: i64| -> i64 {
105                    if x < 0 { (len + x).max(0) } else { x.min(len - 1) }
106                };
107                let s = norm(start);
108                let e = norm(end);
109                if s > e {
110                    return Ok(0);
111                }
112                (s as usize, e as usize)
113            }
114        };
115        Ok(bytes[s..=e]
116            .iter()
117            .map(|b| u64::from(b.count_ones()))
118            .sum())
119    }
120}
kevy_store/bitmap.rs

kevy_store/
bitmap.rs
