From 67a9f32206fb5b018081cb51fc60b1ef73ccb7e6 Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Thu, 26 Feb 2026 15:28:43 -0500 Subject: [PATCH] feat: bitmap commands (GETBIT, SETBIT, BITCOUNT, BITPOS, BITOP) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit implements all five redis-compatible bitmap commands backed by Value::String storage with big-endian bit ordering (bit 0 = MSB of byte 0). no new value type — bitmaps reuse the existing string type. - GETBIT / SETBIT with auto-extension and old-value return - BITCOUNT with byte-range and bit-range (BIT unit) support - BITPOS finds first set/clear bit; returns -1 when not found - BITOP (AND, OR, XOR, NOT) operates across arbitrary source keys, result length equals the longest source - AOF persistence for SETBIT and BITOP (tags 33 and 34) - recovery replay in ember-persistence without cross-crate dep (BitOpKind encoded as u8 in AofRecord) - fix pre-existing spawn_shard arity mismatch in shard unit tests (expired_tx param added in #313 but tests not updated) - 15 integration tests + 20 unit tests covering edge cases --- crates/ember-core/src/keyspace/bitmap.rs | 713 ++++++++++++++++++ crates/ember-core/src/keyspace/mod.rs | 1 + crates/ember-core/src/shard/aof.rs | 18 + crates/ember-core/src/shard/blocking.rs | 5 + crates/ember-core/src/shard/mod.rs | 53 ++ crates/ember-core/src/shard/persistence.rs | 5 + crates/ember-persistence/src/aof.rs | 58 ++ crates/ember-persistence/src/recovery.rs | 81 ++ .../ember-protocol/src/command/attributes.rs | 18 +- crates/ember-protocol/src/command/mod.rs | 61 ++ crates/ember-protocol/src/command/parse.rs | 141 +++- crates/ember-server/src/connection/execute.rs | 57 ++ crates/ember-server/src/replication.rs | 26 +- tests/integration/src/bitmap.rs | 297 ++++++++ tests/integration/src/main.rs | 1 + 15 files changed, 1531 insertions(+), 4 deletions(-) create mode 100644 crates/ember-core/src/keyspace/bitmap.rs create mode 100644 tests/integration/src/bitmap.rs diff --git a/crates/ember-core/src/keyspace/bitmap.rs b/crates/ember-core/src/keyspace/bitmap.rs new file mode 100644 index 00000000..a4e853f1 --- /dev/null +++ b/crates/ember-core/src/keyspace/bitmap.rs @@ -0,0 +1,713 @@ +//! Bitmap operations on string-typed keys. +//! +//! Bitmaps are not a separate data type — they use `Value::String(Bytes)` +//! with big-endian bit ordering (Redis compatible): bit 0 is the most +//! significant bit of byte 0. + +use ember_protocol::command::{BitOpKind, BitRange, BitRangeUnit}; + +use super::*; + +impl Keyspace { + /// Returns the bit at `offset` in the string stored at `key`. + /// + /// Bit ordering is big-endian: byte `offset / 8`, bit position + /// `7 - (offset % 8)`. Returns 0 for missing keys or offsets + /// beyond the string's length. + pub fn getbit(&mut self, key: &str, offset: u64) -> Result { + if self.remove_if_expired(key) { + return Ok(0); + } + match self.entries.get(key) { + None => Ok(0), + Some(e) => match &e.value { + Value::String(data) => { + let byte_idx = (offset / 8) as usize; + if byte_idx >= data.len() { + return Ok(0); + } + let bit_pos = 7 - (offset % 8) as u32; + Ok((data[byte_idx] >> bit_pos) & 1) + } + _ => Err(WrongType), + }, + } + } + + /// Sets the bit at `offset` to `value` (0 or 1). + /// + /// Extends the string with zero bytes when `offset` reaches beyond the + /// current string length. Returns the old bit value. + /// + /// Follows the same memory-tracking pattern as `setrange`. + pub fn setbit(&mut self, key: &str, offset: u64, value: u8) -> Result { + self.remove_if_expired(key); + + let byte_idx = (offset / 8) as usize; + let bit_pos = 7 - (offset % 8) as u32; + let mask = 1u8 << bit_pos; + + let (existing, expire) = match self.entries.get(key) { + Some(entry) => match &entry.value { + Value::String(data) => { + let expire = time::remaining_ms(entry.expires_at_ms).map(Duration::from_millis); + (data.clone(), expire) + } + _ => return Err(WriteError::WrongType), + }, + None => (Bytes::new(), None), + }; + + let old_bit = if byte_idx < existing.len() { + (existing[byte_idx] >> bit_pos) & 1 + } else { + 0 + }; + + // build the new buffer: existing data + zero-padding if needed + let new_len = existing.len().max(byte_idx + 1); + let mut buf = existing.to_vec(); + buf.resize(new_len, 0); + + if value == 1 { + buf[byte_idx] |= mask; + } else { + buf[byte_idx] &= !mask; + } + + match self.set(key.to_owned(), Bytes::from(buf), expire, false, false) { + SetResult::Ok | SetResult::Blocked => Ok(old_bit), + SetResult::OutOfMemory => Err(WriteError::OutOfMemory), + } + } + + /// Counts set bits in the string at `key`. + /// + /// When `range` is `None`, counts bits across the entire string. + /// When `range` is `Some(r)`, restricts to the given byte or bit range + /// (negative indices count from the end of the string). + /// + /// Returns 0 for missing keys. + pub fn bitcount(&mut self, key: &str, range: Option) -> Result { + if self.remove_if_expired(key) { + return Ok(0); + } + let data = match self.entries.get(key) { + None => return Ok(0), + Some(e) => match &e.value { + Value::String(b) => b.clone(), + _ => return Err(WrongType), + }, + }; + + match range { + None => Ok(data.iter().map(|b| b.count_ones() as u64).sum()), + Some(r) if r.unit == BitRangeUnit::Bit => { + // bit-granularity: count each individual bit in [start_bit, end_bit] + let len_bits = data.len() as i64 * 8; + let start = normalize_bit_index(r.start, len_bits).min(len_bits); + let end = normalize_bit_index(r.end, len_bits).min(len_bits - 1); + if start > end { + return Ok(0); + } + let mut count = 0u64; + for bit_idx in start..=end { + let byte_idx = (bit_idx / 8) as usize; + let bit_pos = 7 - (bit_idx % 8) as u32; + count += ((data[byte_idx] >> bit_pos) & 1) as u64; + } + Ok(count) + } + Some(r) => { + let slice = bit_range_slice(&data, r); + Ok(slice.iter().map(|b| b.count_ones() as u64).sum()) + } + } + } + + /// Returns the position of the first bit equal to `bit` (0 or 1). + /// + /// `range` works the same as for `bitcount`. When no `end` is given for + /// `BITPOS 1`, the search covers the whole string; for `BITPOS 0`, it + /// also covers the virtual zero bits beyond the string's end. + /// + /// Returns -1 if the bit is not found (except for `BITPOS 0` on a missing + /// key, which returns 0). + pub fn bitpos( + &mut self, + key: &str, + bit: u8, + range: Option, + ) -> Result { + if self.remove_if_expired(key) { + // missing key: BITPOS 0 → 0, BITPOS 1 → -1 + return Ok(if bit == 0 { 0 } else { -1 }); + } + let data = match self.entries.get(key) { + None => { + return Ok(if bit == 0 { 0 } else { -1 }); + } + Some(e) => match &e.value { + Value::String(b) => b.clone(), + _ => return Err(WrongType), + }, + }; + + // determine whether the caller constrained the end boundary + let has_explicit_end = range.map(|r| r.end != -1).unwrap_or(false); + + let (slice, bit_offset) = match range { + None => (&data[..], 0i64), + Some(r) if r.unit == BitRangeUnit::Bit => { + // bit-granularity range: resolve to an inclusive bit range + let len_bits = data.len() as i64 * 8; + let start = normalize_bit_index(r.start, len_bits).min(len_bits); + let end = normalize_bit_index(r.end, len_bits).min(len_bits - 1); + if start > end { + return Ok(-1); + } + // search bit-by-bit within the resolved range + for bit_idx in start..=end { + let byte_idx = (bit_idx / 8) as usize; + let bit_pos = 7 - (bit_idx % 8) as u32; + let found = (data[byte_idx] >> bit_pos) & 1; + if found == bit { + return Ok(bit_idx); + } + } + return Ok(-1); + } + Some(r) => { + // byte-granularity range + let (s, e) = resolve_byte_range(r.start, r.end, data.len()); + if s >= data.len() { + return Ok(-1); + } + let end = e.min(data.len() - 1); + (&data[s..=end], (s as i64) * 8) + } + }; + + // scan bytes for the first matching bit + for (i, &byte) in slice.iter().enumerate() { + let b = if bit == 1 { byte } else { !byte }; + if b != 0 { + let bit_in_byte = b.leading_zeros() as i64; + return Ok(bit_offset + (i as i64) * 8 + bit_in_byte); + } + } + + // not found — for BITPOS 0 without an explicit end, the answer is the + // first virtual bit past the end of the string. + if bit == 0 && !has_explicit_end { + Ok((data.len() as i64) * 8) + } else { + Ok(-1) + } + } + + /// Performs a bitwise operation across `keys` and stores the result in `dest`. + /// + /// Returns the length of the result string (equal to the longest source). + /// Missing keys are treated as zero-filled strings of the same length. + /// `NOT` requires exactly one source key (enforced at parse time). + pub fn bitop( + &mut self, + op: BitOpKind, + dest: String, + keys: &[String], + ) -> Result { + // collect source bytes — type-check each before mutating anything + let mut sources: Vec = Vec::with_capacity(keys.len()); + for key in keys { + self.remove_if_expired(key); + match self.entries.get(key.as_str()) { + None => sources.push(Bytes::new()), + Some(e) => match &e.value { + Value::String(b) => sources.push(b.clone()), + _ => return Err(WriteError::WrongType), + }, + } + } + + let result_len = sources.iter().map(|s| s.len()).max().unwrap_or(0); + let mut result = vec![0u8; result_len]; + + match op { + BitOpKind::Not => { + // NOT of the single source; bytes beyond the source length → 0xFF + let src = sources.first().map(|b| b.as_ref()).unwrap_or(&[]); + for (i, b) in result.iter_mut().enumerate() { + *b = if i < src.len() { !src[i] } else { 0xFF }; + } + } + BitOpKind::And => { + // initialize from first source; bytes beyond any source → AND with 0 + if let Some(first) = sources.first() { + for (i, b) in result.iter_mut().enumerate() { + *b = if i < first.len() { first[i] } else { 0 }; + } + } + for src in sources.iter().skip(1) { + for (i, b) in result.iter_mut().enumerate() { + let s = if i < src.len() { src[i] } else { 0 }; + *b &= s; + } + } + } + BitOpKind::Or => { + for src in &sources { + for (i, b) in result.iter_mut().enumerate() { + if i < src.len() { + *b |= src[i]; + } + } + } + } + BitOpKind::Xor => { + for src in &sources { + for (i, b) in result.iter_mut().enumerate() { + if i < src.len() { + *b ^= src[i]; + } + } + } + } + } + + // store the result — set() handles memory accounting and version bumping. + // it also implicitly removes any prior value at dest (including wrong-type keys). + match self.set(dest, Bytes::from(result), None, false, false) { + SetResult::Ok | SetResult::Blocked => Ok(result_len), + SetResult::OutOfMemory => Err(WriteError::OutOfMemory), + } + } +} + +/// Resolves a byte-granularity range `[start, end]` (Redis semantics) into +/// a concrete `(start_byte, end_byte)` pair against a string of `len` bytes. +/// +/// Negative indices count from the end (-1 = last byte). The returned range +/// is NOT yet clamped to `[0, len)` — callers must do that. +fn resolve_byte_range(start: i64, end: i64, len: usize) -> (usize, usize) { + let len = len as i64; + let s = if start < 0 { + (len + start).max(0) + } else { + start + } as usize; + let e = if end < 0 { (len + end).max(0) } else { end } as usize; + (s, e) +} + +/// Resolves a signed bit index against `len_bits` (negative = from end). +fn normalize_bit_index(idx: i64, len_bits: i64) -> i64 { + if idx < 0 { + (len_bits + idx).max(0) + } else { + idx + } +} + +/// Returns the sub-slice of `data` described by `range`. +/// +/// Handles both byte and bit-unit ranges. For bit-unit ranges, the slice is +/// rounded to the containing bytes (bit searches happen inside the caller). +fn bit_range_slice(data: &[u8], range: BitRange) -> &[u8] { + match range.unit { + BitRangeUnit::Byte => { + let (s, e) = resolve_byte_range(range.start, range.end, data.len()); + if s >= data.len() { + return &[]; + } + let end = e.min(data.len() - 1); + if s > end { + &[] + } else { + &data[s..=end] + } + } + BitRangeUnit::Bit => { + // for BITCOUNT with BIT range, convert to byte boundaries (inclusive) + let len_bits = data.len() as i64 * 8; + let start_bit = normalize_bit_index(range.start, len_bits).min(len_bits); + let end_bit = normalize_bit_index(range.end, len_bits).min(len_bits - 1); + if start_bit > end_bit || data.is_empty() { + return &[]; + } + // return the byte slice containing all bits in [start_bit, end_bit] + // (individual bit masking happens in the caller for bitcount) + let start_byte = (start_bit / 8) as usize; + let end_byte = (end_bit / 8) as usize; + &data[start_byte..=end_byte.min(data.len() - 1)] + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // --- getbit --- + + #[test] + fn getbit_missing_key_returns_zero() { + let mut ks = Keyspace::new(); + assert_eq!(ks.getbit("nope", 0).unwrap(), 0); + assert_eq!(ks.getbit("nope", 100).unwrap(), 0); + } + + #[test] + fn getbit_reads_msb_first() { + let mut ks = Keyspace::new(); + // byte 0xFF: all bits set + ks.set("k".into(), Bytes::from(vec![0xFF]), None, false, false); + for offset in 0..8 { + assert_eq!(ks.getbit("k", offset).unwrap(), 1, "offset {offset}"); + } + // byte 0x00: no bits set + ks.set("k".into(), Bytes::from(vec![0x00]), None, false, false); + for offset in 0..8 { + assert_eq!(ks.getbit("k", offset).unwrap(), 0, "offset {offset}"); + } + } + + #[test] + fn getbit_big_endian_ordering() { + let mut ks = Keyspace::new(); + // 0x80 = 0b10000000: only bit 0 (MSB) is set + ks.set("k".into(), Bytes::from(vec![0x80]), None, false, false); + assert_eq!(ks.getbit("k", 0).unwrap(), 1); + assert_eq!(ks.getbit("k", 1).unwrap(), 0); + // 0x01 = 0b00000001: only bit 7 (LSB) is set + ks.set("k".into(), Bytes::from(vec![0x01]), None, false, false); + assert_eq!(ks.getbit("k", 7).unwrap(), 1); + assert_eq!(ks.getbit("k", 0).unwrap(), 0); + } + + #[test] + fn getbit_beyond_string_returns_zero() { + let mut ks = Keyspace::new(); + ks.set("k".into(), Bytes::from(vec![0xFF]), None, false, false); + // bit 8 is byte 1, which doesn't exist + assert_eq!(ks.getbit("k", 8).unwrap(), 0); + } + + #[test] + fn getbit_wrong_type() { + let mut ks = Keyspace::new(); + ks.lpush("list", &[Bytes::from("a")]).unwrap(); + assert!(ks.getbit("list", 0).is_err()); + } + + // --- setbit --- + + #[test] + fn setbit_returns_old_bit() { + let mut ks = Keyspace::new(); + // new key: old bit is 0 + assert_eq!(ks.setbit("k", 7, 1).unwrap(), 0); + // now bit 7 is set, returns 1 + assert_eq!(ks.setbit("k", 7, 1).unwrap(), 1); + // clear it: returns 1 + assert_eq!(ks.setbit("k", 7, 0).unwrap(), 1); + // now it's 0 again + assert_eq!(ks.setbit("k", 7, 0).unwrap(), 0); + } + + #[test] + fn setbit_roundtrip_with_getbit() { + let mut ks = Keyspace::new(); + ks.setbit("k", 10, 1).unwrap(); + assert_eq!(ks.getbit("k", 10).unwrap(), 1); + assert_eq!(ks.getbit("k", 0).unwrap(), 0); + } + + #[test] + fn setbit_extends_string() { + let mut ks = Keyspace::new(); + // setting bit 15 requires 2 bytes + ks.setbit("k", 15, 1).unwrap(); + let val = match ks.get("k").unwrap() { + Some(Value::String(b)) => b, + other => panic!("expected String, got {other:?}"), + }; + assert_eq!(val.len(), 2); + assert_eq!(ks.getbit("k", 15).unwrap(), 1); + } + + #[test] + fn setbit_preserves_ttl() { + let mut ks = Keyspace::new(); + ks.set( + "k".into(), + Bytes::from(vec![0u8]), + Some(Duration::from_secs(60)), + false, + false, + ); + ks.setbit("k", 0, 1).unwrap(); + assert!(matches!(ks.ttl("k"), TtlResult::Seconds(_))); + } + + #[test] + fn setbit_wrong_type() { + let mut ks = Keyspace::new(); + ks.lpush("list", &[Bytes::from("a")]).unwrap(); + assert!(ks.setbit("list", 0, 1).is_err()); + } + + // --- bitcount --- + + #[test] + fn bitcount_missing_key_returns_zero() { + let mut ks = Keyspace::new(); + assert_eq!(ks.bitcount("nope", None).unwrap(), 0); + } + + #[test] + fn bitcount_full_string() { + let mut ks = Keyspace::new(); + // 0xFF has 8 set bits, 0x0F has 4 + ks.set( + "k".into(), + Bytes::from(vec![0xFF, 0x0F]), + None, + false, + false, + ); + assert_eq!(ks.bitcount("k", None).unwrap(), 12); + } + + #[test] + fn bitcount_byte_range() { + let mut ks = Keyspace::new(); + ks.set( + "k".into(), + Bytes::from(vec![0xFF, 0x00, 0xFF]), + None, + false, + false, + ); + // only byte 0 + assert_eq!( + ks.bitcount( + "k", + Some(BitRange { + start: 0, + end: 0, + unit: BitRangeUnit::Byte + }) + ) + .unwrap(), + 8 + ); + // bytes 0 and 1 + assert_eq!( + ks.bitcount( + "k", + Some(BitRange { + start: 0, + end: 1, + unit: BitRangeUnit::Byte + }) + ) + .unwrap(), + 8 + ); + } + + #[test] + fn bitcount_bit_range() { + let mut ks = Keyspace::new(); + // 0xFF: bits 0-7 all set + ks.set("k".into(), Bytes::from(vec![0xFF]), None, false, false); + assert_eq!( + ks.bitcount( + "k", + Some(BitRange { + start: 0, + end: 7, + unit: BitRangeUnit::Bit + }) + ) + .unwrap(), + 8 + ); + assert_eq!( + ks.bitcount( + "k", + Some(BitRange { + start: 0, + end: 3, + unit: BitRangeUnit::Bit + }) + ) + .unwrap(), + 4 + ); + } + + #[test] + fn bitcount_wrong_type() { + let mut ks = Keyspace::new(); + ks.lpush("list", &[Bytes::from("a")]).unwrap(); + assert!(ks.bitcount("list", None).is_err()); + } + + // --- bitpos --- + + #[test] + fn bitpos_missing_key_bit1_returns_minus_one() { + let mut ks = Keyspace::new(); + assert_eq!(ks.bitpos("nope", 1, None).unwrap(), -1); + } + + #[test] + fn bitpos_missing_key_bit0_returns_zero() { + let mut ks = Keyspace::new(); + assert_eq!(ks.bitpos("nope", 0, None).unwrap(), 0); + } + + #[test] + fn bitpos_find_first_set_bit() { + let mut ks = Keyspace::new(); + // 0x00 0x01: first set bit is at position 15 (LSB of byte 1) + ks.set( + "k".into(), + Bytes::from(vec![0x00, 0x01]), + None, + false, + false, + ); + assert_eq!(ks.bitpos("k", 1, None).unwrap(), 15); + } + + #[test] + fn bitpos_find_first_clear_bit_in_all_ones() { + let mut ks = Keyspace::new(); + // all bytes 0xFF: first clear bit is at position 16 (past end) + ks.set( + "k".into(), + Bytes::from(vec![0xFF, 0xFF]), + None, + false, + false, + ); + assert_eq!(ks.bitpos("k", 0, None).unwrap(), 16); + } + + #[test] + fn bitpos_wrong_type() { + let mut ks = Keyspace::new(); + ks.lpush("list", &[Bytes::from("a")]).unwrap(); + assert!(ks.bitpos("list", 1, None).is_err()); + } + + // --- bitop --- + + #[test] + fn bitop_and() { + let mut ks = Keyspace::new(); + ks.set( + "a".into(), + Bytes::from(vec![0xFF, 0x0F]), + None, + false, + false, + ); + ks.set( + "b".into(), + Bytes::from(vec![0x0F, 0xFF]), + None, + false, + false, + ); + let len = ks + .bitop(BitOpKind::And, "dest".into(), &["a".into(), "b".into()]) + .unwrap(); + assert_eq!(len, 2); + let val = match ks.get("dest").unwrap() { + Some(Value::String(b)) => b, + other => panic!("expected String, got {other:?}"), + }; + assert_eq!(val[0], 0x0F); + assert_eq!(val[1], 0x0F); + } + + #[test] + fn bitop_or() { + let mut ks = Keyspace::new(); + ks.set("a".into(), Bytes::from(vec![0xF0]), None, false, false); + ks.set("b".into(), Bytes::from(vec![0x0F]), None, false, false); + ks.bitop(BitOpKind::Or, "dest".into(), &["a".into(), "b".into()]) + .unwrap(); + let val = match ks.get("dest").unwrap() { + Some(Value::String(b)) => b, + other => panic!("expected String, got {other:?}"), + }; + assert_eq!(val[0], 0xFF); + } + + #[test] + fn bitop_xor() { + let mut ks = Keyspace::new(); + ks.set("a".into(), Bytes::from(vec![0xFF]), None, false, false); + ks.set("b".into(), Bytes::from(vec![0xFF]), None, false, false); + ks.bitop(BitOpKind::Xor, "dest".into(), &["a".into(), "b".into()]) + .unwrap(); + let val = match ks.get("dest").unwrap() { + Some(Value::String(b)) => b, + other => panic!("expected String, got {other:?}"), + }; + assert_eq!(val[0], 0x00); + } + + #[test] + fn bitop_not() { + let mut ks = Keyspace::new(); + ks.set( + "src".into(), + Bytes::from(vec![0xF0, 0x0F]), + None, + false, + false, + ); + let len = ks + .bitop(BitOpKind::Not, "dest".into(), &["src".into()]) + .unwrap(); + assert_eq!(len, 2); + let val = match ks.get("dest").unwrap() { + Some(Value::String(b)) => b, + other => panic!("expected String, got {other:?}"), + }; + assert_eq!(val[0], 0x0F); + assert_eq!(val[1], 0xF0); + } + + #[test] + fn bitop_wrong_type() { + let mut ks = Keyspace::new(); + ks.lpush("list", &[Bytes::from("a")]).unwrap(); + assert!(ks + .bitop(BitOpKind::And, "dest".into(), &["list".into()]) + .is_err()); + } + + #[test] + fn bitop_extends_to_longest_source() { + let mut ks = Keyspace::new(); + ks.set( + "a".into(), + Bytes::from(vec![0xFF, 0xFF, 0xFF]), + None, + false, + false, + ); + ks.set("b".into(), Bytes::from(vec![0xFF]), None, false, false); + let len = ks + .bitop(BitOpKind::Or, "dest".into(), &["a".into(), "b".into()]) + .unwrap(); + assert_eq!(len, 3); + } +} diff --git a/crates/ember-core/src/keyspace/mod.rs b/crates/ember-core/src/keyspace/mod.rs index 374efe1e..18122cf8 100644 --- a/crates/ember-core/src/keyspace/mod.rs +++ b/crates/ember-core/src/keyspace/mod.rs @@ -21,6 +21,7 @@ use crate::time; use crate::types::sorted_set::{ScoreBound, SortedSet, ZAddFlags}; use crate::types::{self, normalize_range, Value}; +mod bitmap; mod hash; mod list; #[cfg(feature = "protobuf")] diff --git a/crates/ember-core/src/shard/aof.rs b/crates/ember-core/src/shard/aof.rs index 29d1c31e..b529dfba 100644 --- a/crates/ember-core/src/shard/aof.rs +++ b/crates/ember-core/src/shard/aof.rs @@ -143,6 +143,24 @@ pub(super) fn to_aof_records( (ShardRequest::SetRange { key, offset, value }, ShardResponse::Len(_)) => { smallvec![AofRecord::SetRange { key, offset, value }] } + // SETBIT: record the offset + bit value for replay + (ShardRequest::SetBit { key, offset, value }, ShardResponse::Integer(_)) => { + smallvec![AofRecord::SetBit { key, offset, value }] + } + // BITOP: record the operation + dest + source keys for replay + (ShardRequest::BitOp { op, dest, keys }, ShardResponse::Integer(_)) => { + let op_byte: u8 = match op { + BitOpKind::And => 0, + BitOpKind::Or => 1, + BitOpKind::Xor => 2, + BitOpKind::Not => 3, + }; + smallvec![AofRecord::BitOp { + op: op_byte, + dest, + keys, + }] + } (ShardRequest::Rename { key, newkey }, ShardResponse::Ok) => { smallvec![AofRecord::Rename { key, newkey }] } diff --git a/crates/ember-core/src/shard/blocking.rs b/crates/ember-core/src/shard/blocking.rs index 9a91280f..0398039e 100644 --- a/crates/ember-core/src/shard/blocking.rs +++ b/crates/ember-core/src/shard/blocking.rs @@ -154,6 +154,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -195,6 +196,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -240,6 +242,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -276,6 +279,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -313,6 +317,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index a8e76184..764f7611 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -68,6 +68,7 @@ use crate::keyspace::{ }; use crate::types::sorted_set::{ScoreBound, ZAddFlags}; use crate::types::Value; +use ember_protocol::command::{BitOpKind, BitRange}; /// How often the shard runs active expiration. 100ms matches /// Redis's hz=10 default and keeps CPU overhead negligible. @@ -156,6 +157,34 @@ pub enum ShardRequest { offset: usize, value: Bytes, }, + /// GETBIT key offset. Returns the bit at `offset` (0 or 1). Big-endian ordering. + GetBit { + key: String, + offset: u64, + }, + /// SETBIT key offset value. Sets the bit at `offset` to 0 or 1. Returns old bit. + SetBit { + key: String, + offset: u64, + value: u8, + }, + /// BITCOUNT key [range]. Counts set bits, optionally restricted to a range. + BitCount { + key: String, + range: Option, + }, + /// BITPOS key bit [range]. Finds first set or clear bit position. + BitPos { + key: String, + bit: u8, + range: Option, + }, + /// BITOP op destkey key [key ...]. Bitwise operation across strings. + BitOp { + op: BitOpKind, + dest: String, + keys: Vec, + }, /// Returns all keys matching a glob pattern in this shard. Keys { pattern: String, @@ -687,6 +716,8 @@ impl ShardRequest { | ShardRequest::DecrBy { .. } | ShardRequest::IncrByFloat { .. } | ShardRequest::Append { .. } + | ShardRequest::SetBit { .. } + | ShardRequest::BitOp { .. } | ShardRequest::Del { .. } | ShardRequest::Unlink { .. } | ShardRequest::Rename { .. } @@ -1609,6 +1640,28 @@ fn dispatch( ShardRequest::SetRange { key, offset, value } => { write_result_len(ks.setrange(key, *offset, value)) } + ShardRequest::GetBit { key, offset } => match ks.getbit(key, *offset) { + Ok(bit) => ShardResponse::Integer(bit as i64), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::SetBit { key, offset, value } => match ks.setbit(key, *offset, *value) { + Ok(old_bit) => ShardResponse::Integer(old_bit as i64), + Err(WriteError::WrongType) => ShardResponse::WrongType, + Err(WriteError::OutOfMemory) => ShardResponse::OutOfMemory, + }, + ShardRequest::BitCount { key, range } => match ks.bitcount(key, *range) { + Ok(count) => ShardResponse::Integer(count as i64), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::BitPos { key, bit, range } => match ks.bitpos(key, *bit, *range) { + Ok(pos) => ShardResponse::Integer(pos), + Err(_) => ShardResponse::WrongType, + }, + ShardRequest::BitOp { op, dest, keys } => match ks.bitop(*op, dest.clone(), keys) { + Ok(len) => ShardResponse::Integer(len as i64), + Err(WriteError::WrongType) => ShardResponse::WrongType, + Err(WriteError::OutOfMemory) => ShardResponse::OutOfMemory, + }, ShardRequest::Keys { pattern } => { let keys = ks.keys(pattern); ShardResponse::StringArray(keys) diff --git a/crates/ember-core/src/shard/persistence.rs b/crates/ember-core/src/shard/persistence.rs index f4efdc0b..d52778d3 100644 --- a/crates/ember-core/src/shard/persistence.rs +++ b/crates/ember-core/src/shard/persistence.rs @@ -243,6 +243,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -281,6 +282,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -313,6 +315,7 @@ mod tests { None, None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -387,6 +390,7 @@ mod tests { Some(pcfg.clone()), None, None, + None, #[cfg(feature = "protobuf")] None, ); @@ -436,6 +440,7 @@ mod tests { Some(pcfg), None, None, + None, #[cfg(feature = "protobuf")] None, ); diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 1fc4b41d..6872d8d9 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -105,6 +105,10 @@ const TAG_LINSERT: u8 = 30; const TAG_LREM: u8 = 31; const TAG_SETRANGE: u8 = 32; +// bitmap +const TAG_SETBIT: u8 = 33; +const TAG_BITOP: u8 = 34; + // vector #[cfg(feature = "vector")] const TAG_VADD: u8 = 25; @@ -203,6 +207,16 @@ pub enum AofRecord { offset: usize, value: Bytes, }, + /// SETBIT key offset value. Replays the bit mutation verbatim. + SetBit { key: String, offset: u64, value: u8 }, + /// BITOP op destkey key [key ...]. Replays the bitwise operation. + /// + /// `op` is stored as a raw byte: 0=AND, 1=OR, 2=XOR, 3=NOT. + BitOp { + op: u8, + dest: String, + keys: Vec, + }, /// RENAME key newkey. Rename { key: String, newkey: String }, /// COPY source destination [REPLACE]. @@ -279,6 +293,8 @@ impl AofRecord { AofRecord::DecrBy { .. } => TAG_DECRBY, AofRecord::Append { .. } => TAG_APPEND, AofRecord::SetRange { .. } => TAG_SETRANGE, + AofRecord::SetBit { .. } => TAG_SETBIT, + AofRecord::BitOp { .. } => TAG_BITOP, AofRecord::Rename { .. } => TAG_RENAME, AofRecord::Copy { .. } => TAG_COPY, #[cfg(feature = "vector")] @@ -374,6 +390,13 @@ impl AofRecord { AofRecord::SetRange { key, value, .. } => { 1 + LEN_PREFIX + key.len() + 8 + LEN_PREFIX + value.len() } + // 1 tag + 4 key-len + key + 8 offset + 1 value + AofRecord::SetBit { key, .. } => 1 + LEN_PREFIX + key.len() + 8 + 1, + // 1 tag + 1 op + 4 dest-len + dest + 4 count + (4 key-len + key) * n + AofRecord::BitOp { dest, keys, .. } => { + let keys_size: usize = keys.iter().map(|k| LEN_PREFIX + k.len()).sum(); + 1 + 1 + LEN_PREFIX + dest.len() + 4 + keys_size + } AofRecord::Rename { key, newkey } => { 1 + LEN_PREFIX + key.len() + LEN_PREFIX + newkey.len() } @@ -560,6 +583,24 @@ impl AofRecord { format::write_bytes(&mut buf, value)?; } + // key + offset (as i64) + bit value (u8) + AofRecord::SetBit { key, offset, value } => { + format::write_bytes(&mut buf, key.as_bytes())?; + // offset is u64 but fits in i64 in practice (max bit offset < 2^32) + format::write_i64(&mut buf, *offset as i64)?; + format::write_u8(&mut buf, *value)?; + } + + // op byte + dest + key list + AofRecord::BitOp { op, dest, keys } => { + format::write_u8(&mut buf, *op)?; + format::write_bytes(&mut buf, dest.as_bytes())?; + format::write_len(&mut buf, keys.len())?; + for key in keys { + format::write_bytes(&mut buf, key.as_bytes())?; + } + } + // key + newkey AofRecord::Rename { key, newkey } => { format::write_bytes(&mut buf, key.as_bytes())?; @@ -813,6 +854,23 @@ impl AofRecord { replace, }) } + TAG_SETBIT => { + let key = read_string(&mut cursor, "key")?; + let offset = format::read_i64(&mut cursor)? as u64; + let value = format::read_u8(&mut cursor)?; + Ok(AofRecord::SetBit { key, offset, value }) + } + TAG_BITOP => { + let op = format::read_u8(&mut cursor)?; + if op > 3 { + return Err(FormatError::InvalidData(format!( + "BITOP: unknown op byte {op} in AOF record" + ))); + } + let dest = read_string(&mut cursor, "dest")?; + let keys = read_string_list(&mut cursor, "key")?; + Ok(AofRecord::BitOp { op, dest, keys }) + } #[cfg(feature = "vector")] TAG_VADD => { let key = read_string(&mut cursor, "key")?; diff --git a/crates/ember-persistence/src/recovery.rs b/crates/ember-persistence/src/recovery.rs index 90f7c23f..d3f48cc2 100644 --- a/crates/ember-persistence/src/recovery.rs +++ b/crates/ember-persistence/src/recovery.rs @@ -553,6 +553,87 @@ fn replay_aof( *data = Bytes::from(buf); } } + AofRecord::SetBit { key, offset, value } => { + let entry = map + .entry(key) + .or_insert_with(|| (RecoveredValue::String(Bytes::new()), -1)); + if let RecoveredValue::String(ref mut data) = entry.0 { + let byte_idx = (offset / 8) as usize; + let bit_pos = 7 - (offset % 8) as u32; + let mask = 1u8 << bit_pos; + let new_len = data.len().max(byte_idx + 1); + let mut buf = data.to_vec(); + buf.resize(new_len, 0); + if value == 1 { + buf[byte_idx] |= mask; + } else { + buf[byte_idx] &= !mask; + } + *data = Bytes::from(buf); + } + } + AofRecord::BitOp { op, dest, keys } => { + // op byte: 0=AND, 1=OR, 2=XOR, 3=NOT (matches aof.rs encoding) + let sources: Vec = keys + .iter() + .map(|k| { + map.get(k) + .and_then(|(v, _)| { + if let RecoveredValue::String(b) = v { + Some(b.clone()) + } else { + None + } + }) + .unwrap_or_default() + }) + .collect(); + let result_len = sources.iter().map(|s| s.len()).max().unwrap_or(0); + let mut result = vec![0u8; result_len]; + match op { + 3 => { + // NOT + let src = sources.first().map(|b| b.as_ref()).unwrap_or(&[]); + for (i, b) in result.iter_mut().enumerate() { + *b = if i < src.len() { !src[i] } else { 0xFF }; + } + } + 0 => { + // AND + if let Some(first) = sources.first() { + for (i, b) in result.iter_mut().enumerate() { + *b = if i < first.len() { first[i] } else { 0 }; + } + } + for src in sources.iter().skip(1) { + for (i, b) in result.iter_mut().enumerate() { + *b &= if i < src.len() { src[i] } else { 0 }; + } + } + } + 1 => { + // OR + for src in &sources { + for (i, b) in result.iter_mut().enumerate() { + if i < src.len() { + *b |= src[i]; + } + } + } + } + _ => { + // XOR (op == 2) + for src in &sources { + for (i, b) in result.iter_mut().enumerate() { + if i < src.len() { + *b ^= src[i]; + } + } + } + } + } + map.insert(dest, (RecoveredValue::String(Bytes::from(result)), -1)); + } AofRecord::Rename { key, newkey } => { if let Some(entry) = map.remove(&key) { map.insert(newkey, entry); diff --git a/crates/ember-protocol/src/command/attributes.rs b/crates/ember-protocol/src/command/attributes.rs index b8894ad5..af072621 100644 --- a/crates/ember-protocol/src/command/attributes.rs +++ b/crates/ember-protocol/src/command/attributes.rs @@ -42,6 +42,13 @@ impl Command { Command::GetRange { .. } => "getrange", Command::SetRange { .. } => "setrange", + // bitmaps + Command::GetBit { .. } => "getbit", + Command::SetBit { .. } => "setbit", + Command::BitCount { .. } => "bitcount", + Command::BitPos { .. } => "bitpos", + Command::BitOp { .. } => "bitop", + // key lifecycle Command::Del { .. } => "del", Command::Unlink { .. } => "unlink", @@ -244,6 +251,8 @@ impl Command { | Command::IncrBy { .. } | Command::DecrBy { .. } | Command::IncrByFloat { .. } + | Command::SetBit { .. } + | Command::BitOp { .. } // key lifecycle | Command::Del { .. } | Command::Unlink { .. } @@ -352,13 +361,18 @@ impl Command { Command::Get { .. } | Command::MGet { .. } | Command::Strlen { .. } - | Command::GetRange { .. } => READ | STRING | FAST, + | Command::GetRange { .. } + | Command::GetBit { .. } + | Command::BitCount { .. } + | Command::BitPos { .. } => READ | STRING | FAST, // string — writes Command::Set { .. } | Command::MSet { .. } | Command::Append { .. } - | Command::SetRange { .. } => WRITE | STRING | SLOW, + | Command::SetRange { .. } + | Command::SetBit { .. } + | Command::BitOp { .. } => WRITE | STRING | SLOW, Command::Incr { .. } | Command::Decr { .. } | Command::IncrBy { .. } diff --git a/crates/ember-protocol/src/command/mod.rs b/crates/ember-protocol/src/command/mod.rs index d2e00ca0..317679d8 100644 --- a/crates/ember-protocol/src/command/mod.rs +++ b/crates/ember-protocol/src/command/mod.rs @@ -73,6 +73,41 @@ pub enum Command { value: Bytes, }, + /// GETBIT `key` `offset`. Returns the bit at `offset` in the string stored at key. + /// + /// Bit ordering is big-endian (Redis compatible): byte 0 holds bits 0–7, MSB first. + /// Returns 0 if the key does not exist. + GetBit { key: String, offset: u64 }, + + /// SETBIT `key` `offset` `value`. Sets or clears the bit at `offset`. + /// + /// The string is automatically grown to accommodate the offset. + /// Returns the original bit value. + SetBit { key: String, offset: u64, value: u8 }, + + /// BITCOUNT `key` \[start end \[BYTE\|BIT\]\]. Counts set bits in the string, + /// optionally restricted to a byte or bit range. + BitCount { + key: String, + range: Option, + }, + + /// BITPOS `key` `bit` \[start \[end \[BYTE\|BIT\]\]\]. Returns the position of the + /// first set (`bit=1`) or clear (`bit=0`) bit in the string. + BitPos { + key: String, + bit: u8, + range: Option, + }, + + /// BITOP `operation` `destkey` `key` \[key ...\]. Performs a bitwise operation + /// across source strings and stores the result in `destkey`. + BitOp { + op: BitOpKind, + dest: String, + keys: Vec, + }, + /// KEYS `pattern`. Returns all keys matching a glob pattern. Keys { pattern: String }, @@ -790,6 +825,32 @@ pub enum Command { Unknown(String), } +/// Unit for BITCOUNT and BITPOS range arguments. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BitRangeUnit { + /// Byte-granularity range (default for Redis). + Byte, + /// Bit-granularity range (Redis 7.0+). + Bit, +} + +/// Range argument for BITCOUNT and BITPOS. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BitRange { + pub start: i64, + pub end: i64, + pub unit: BitRangeUnit, +} + +/// Operation kind for BITOP. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BitOpKind { + And, + Or, + Xor, + Not, +} + /// A score bound for sorted set range queries (ZRANGEBYSCORE, ZCOUNT, etc.). /// /// Redis supports `-inf`, `+inf`, inclusive (default), and exclusive diff --git a/crates/ember-protocol/src/command/parse.rs b/crates/ember-protocol/src/command/parse.rs index a0743b9f..4bc6eb08 100644 --- a/crates/ember-protocol/src/command/parse.rs +++ b/crates/ember-protocol/src/command/parse.rs @@ -8,7 +8,7 @@ use bytes::Bytes; use crate::error::ProtocolError; use crate::types::Frame; -use super::{Command, ScoreBound, SetExpire, ZAddFlags}; +use super::{BitOpKind, BitRange, BitRangeUnit, Command, ScoreBound, SetExpire, ZAddFlags}; /// Maximum number of dimensions in a vector. 65,536 is generous for any /// real-world embedding model (OpenAI: 1536, Cohere: 4096) while preventing @@ -89,6 +89,11 @@ impl Command { "PSETEX" => parse_psetex(&frames[1..]), "GETRANGE" | "SUBSTR" => parse_getrange(&frames[1..]), "SETRANGE" => parse_setrange(&frames[1..]), + "GETBIT" => parse_getbit(&frames[1..]), + "SETBIT" => parse_setbit(&frames[1..]), + "BITCOUNT" => parse_bitcount(&frames[1..]), + "BITPOS" => parse_bitpos(&frames[1..]), + "BITOP" => parse_bitop(&frames[1..]), "KEYS" => parse_keys(&frames[1..]), "RENAME" => parse_rename(&frames[1..]), "DEL" => parse_del(&frames[1..]), @@ -608,6 +613,140 @@ fn parse_setrange(args: &[Frame]) -> Result { Ok(Command::SetRange { key, offset, value }) } +/// GETBIT key offset. +fn parse_getbit(args: &[Frame]) -> Result { + if args.len() != 2 { + return Err(wrong_arity("GETBIT")); + } + let key = extract_string(&args[0])?; + let offset = parse_u64(&args[1], "GETBIT")?; + Ok(Command::GetBit { key, offset }) +} + +/// SETBIT key offset value. +fn parse_setbit(args: &[Frame]) -> Result { + if args.len() != 3 { + return Err(wrong_arity("SETBIT")); + } + let key = extract_string(&args[0])?; + let offset = parse_u64(&args[1], "SETBIT")?; + let raw = parse_u64(&args[2], "SETBIT")?; + if raw > 1 { + return Err(ProtocolError::InvalidCommandFrame( + "SETBIT: bit value must be 0 or 1".into(), + )); + } + Ok(Command::SetBit { + key, + offset, + value: raw as u8, + }) +} + +/// Parses an optional `[start end [BYTE|BIT]]` suffix for BITCOUNT / BITPOS. +/// +/// Accepts 0, 2, or 3 trailing arguments. Returns `None` when there are none. +fn parse_bit_range(args: &[Frame], cmd: &str) -> Result, ProtocolError> { + match args.len() { + 0 => Ok(None), + 2 | 3 => { + let start = parse_i64(&args[0], cmd)?; + let end = parse_i64(&args[1], cmd)?; + let unit = if args.len() == 3 { + let mut kw = [0u8; MAX_KEYWORD_LEN]; + match uppercase_arg(&args[2], &mut kw)? { + "BYTE" => BitRangeUnit::Byte, + "BIT" => BitRangeUnit::Bit, + other => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "{cmd}: invalid unit '{other}', expected BYTE or BIT" + ))); + } + } + } else { + BitRangeUnit::Byte + }; + Ok(Some(BitRange { start, end, unit })) + } + _ => Err(ProtocolError::InvalidCommandFrame(format!( + "{cmd}: wrong number of arguments" + ))), + } +} + +/// BITCOUNT key [start end [BYTE|BIT]]. +fn parse_bitcount(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(wrong_arity("BITCOUNT")); + } + let key = extract_string(&args[0])?; + let range = parse_bit_range(&args[1..], "BITCOUNT")?; + Ok(Command::BitCount { key, range }) +} + +/// BITPOS key bit [start [end [BYTE|BIT]]]. +/// +/// Redis allows 1, 2, or 3 trailing args (start, start+end, start+end+unit). +/// No trailing args means "search the whole string". +fn parse_bitpos(args: &[Frame]) -> Result { + if args.len() < 2 { + return Err(wrong_arity("BITPOS")); + } + let key = extract_string(&args[0])?; + let raw = parse_u64(&args[1], "BITPOS")?; + if raw > 1 { + return Err(ProtocolError::InvalidCommandFrame( + "BITPOS: bit value must be 0 or 1".into(), + )); + } + let bit = raw as u8; + let range = match args.len() - 2 { + 0 => None, + 1 => { + let start = parse_i64(&args[2], "BITPOS")?; + Some(BitRange { + start, + end: -1, + unit: BitRangeUnit::Byte, + }) + } + 2 | 3 => parse_bit_range(&args[2..], "BITPOS")?, + _ => { + return Err(ProtocolError::InvalidCommandFrame( + "BITPOS: wrong number of arguments".into(), + )) + } + }; + Ok(Command::BitPos { key, bit, range }) +} + +/// BITOP AND|OR|XOR|NOT destkey key [key ...]. +fn parse_bitop(args: &[Frame]) -> Result { + if args.len() < 3 { + return Err(wrong_arity("BITOP")); + } + let mut kw = [0u8; MAX_KEYWORD_LEN]; + let op = match uppercase_arg(&args[0], &mut kw)? { + "AND" => BitOpKind::And, + "OR" => BitOpKind::Or, + "XOR" => BitOpKind::Xor, + "NOT" => BitOpKind::Not, + other => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "BITOP: unknown operation '{other}'" + ))); + } + }; + let dest = extract_string(&args[1])?; + let keys = extract_strings(&args[2..])?; + if op == BitOpKind::Not && keys.len() != 1 { + return Err(ProtocolError::InvalidCommandFrame( + "BITOP NOT must be called with a single source key".into(), + )); + } + Ok(Command::BitOp { op, dest, keys }) +} + fn parse_keys(args: &[Frame]) -> Result { if args.len() != 1 { return Err(wrong_arity("KEYS")); diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index af280bf8..c498986b 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -317,6 +317,63 @@ pub(super) async fn execute( } } + Command::GetBit { key, offset } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::GetBit { key, offset }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Integer(bit)) => Frame::Integer(bit), + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::SetBit { key, offset, value } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::SetBit { key, offset, value }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Integer(old_bit)) => Frame::Integer(old_bit), + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::BitCount { key, range } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::BitCount { key, range }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Integer(n)) => Frame::Integer(n), + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::BitPos { key, bit, range } => { + let idx = engine.shard_for_key(&key); + let req = ShardRequest::BitPos { key, bit, range }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Integer(pos)) => Frame::Integer(pos), + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + + Command::BitOp { op, dest, keys } => { + let idx = engine.shard_for_key(&dest); + let req = ShardRequest::BitOp { op, dest, keys }; + match engine.send_to_shard(idx, req).await { + Ok(ShardResponse::Integer(len)) => Frame::Integer(len), + Ok(ShardResponse::WrongType) => wrongtype_error(), + Ok(ShardResponse::OutOfMemory) => oom_error(), + Ok(other) => Frame::Error(format!("ERR unexpected shard response: {other:?}")), + Err(e) => Frame::Error(format!("ERR {e}")), + } + } + Command::IncrByFloat { key, delta } => { let idx = engine.shard_for_key(&key); let req = ShardRequest::IncrByFloat { key, delta }; diff --git a/crates/ember-server/src/replication.rs b/crates/ember-server/src/replication.rs index 35823cd2..a201a9c9 100644 --- a/crates/ember-server/src/replication.rs +++ b/crates/ember-server/src/replication.rs @@ -822,6 +822,25 @@ pub fn aof_record_to_shard_request(record: &AofRecord) -> Option { offset: *offset, value: value.clone(), }), + AofRecord::SetBit { key, offset, value } => Some(ShardRequest::SetBit { + key: key.clone(), + offset: *offset, + value: *value, + }), + AofRecord::BitOp { op, dest, keys } => { + use ember_protocol::command::BitOpKind; + let op_kind = match op { + 0 => BitOpKind::And, + 1 => BitOpKind::Or, + 2 => BitOpKind::Xor, + _ => BitOpKind::Not, + }; + Some(ShardRequest::BitOp { + op: op_kind, + dest: dest.clone(), + keys: keys.clone(), + }) + } AofRecord::Rename { key, newkey } => Some(ShardRequest::Rename { key: key.clone(), newkey: newkey.clone(), @@ -907,7 +926,12 @@ fn primary_key_for_request(req: &ShardRequest) -> Option<&str> { | ShardRequest::HIncrBy { key, .. } | ShardRequest::SAdd { key, .. } | ShardRequest::SRem { key, .. } - | ShardRequest::Rename { key, .. } => Some(key), + | ShardRequest::Rename { key, .. } + | ShardRequest::SetBit { key, .. } + | ShardRequest::GetBit { key, .. } + | ShardRequest::BitCount { key, .. } + | ShardRequest::BitPos { key, .. } => Some(key), + ShardRequest::BitOp { dest, .. } => Some(dest), _ => None, } } diff --git a/tests/integration/src/bitmap.rs b/tests/integration/src/bitmap.rs new file mode 100644 index 00000000..3f77cd7f --- /dev/null +++ b/tests/integration/src/bitmap.rs @@ -0,0 +1,297 @@ +//! Integration tests for bitmap commands: GETBIT, SETBIT, BITCOUNT, BITPOS, BITOP. + +use ember_protocol::Frame; + +use crate::helpers::TestServer; + +// --- GETBIT / SETBIT --- + +#[tokio::test] +async fn setbit_and_getbit_round_trip() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // bit not set yet — should return 0 + assert_eq!(c.get_int(&["GETBIT", "bits", "7"]).await, 0); + + // set bit 7 → old value is 0 + assert_eq!(c.get_int(&["SETBIT", "bits", "7", "1"]).await, 0); + + // now bit 7 is 1 + assert_eq!(c.get_int(&["GETBIT", "bits", "7"]).await, 1); + + // clear it + assert_eq!(c.get_int(&["SETBIT", "bits", "7", "0"]).await, 1); + assert_eq!(c.get_int(&["GETBIT", "bits", "7"]).await, 0); +} + +#[tokio::test] +async fn setbit_auto_extends_string() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // set bit 100 — the string has to grow to at least 13 bytes + c.get_int(&["SETBIT", "bits", "100", "1"]).await; + assert_eq!(c.get_int(&["GETBIT", "bits", "100"]).await, 1); + + // surrounding bits should be 0 + assert_eq!(c.get_int(&["GETBIT", "bits", "99"]).await, 0); + assert_eq!(c.get_int(&["GETBIT", "bits", "101"]).await, 0); +} + +#[tokio::test] +async fn getbit_on_missing_key_returns_zero() { + let server = TestServer::start(); + let mut c = server.connect().await; + + assert_eq!(c.get_int(&["GETBIT", "nokey", "0"]).await, 0); + assert_eq!(c.get_int(&["GETBIT", "nokey", "999"]).await, 0); +} + +#[tokio::test] +async fn setbit_big_endian_ordering() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // set the MSB of byte 0 (bit 0) + c.get_int(&["SETBIT", "bits", "0", "1"]).await; + + // the underlying string should be "\x80" + match c.cmd(&["GET", "bits"]).await { + Frame::Bulk(data) => assert_eq!(&data[..], &[0x80u8]), + other => panic!("expected Bulk, got {other:?}"), + } +} + +// --- BITCOUNT --- + +#[tokio::test] +async fn bitcount_all_bits() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // set all 8 bits via SETBIT to get a fully-set byte + for i in 0..8 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + assert_eq!(c.get_int(&["BITCOUNT", "bits"]).await, 8); + + // a key with no bits set (auto-created, all zeros) + c.get_int(&["SETBIT", "zero", "0", "0"]).await; + assert_eq!(c.get_int(&["BITCOUNT", "zero"]).await, 0); +} + +#[tokio::test] +async fn bitcount_byte_range() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // build two bytes: byte 0 = 0xFF, byte 1 = 0x00 + for i in 0..8 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + // byte 1 stays at 0x00 (auto-extended to accommodate bit 15) + c.get_int(&["SETBIT", "bits", "15", "0"]).await; + + // byte 0 only: 8 set bits + assert_eq!(c.get_int(&["BITCOUNT", "bits", "0", "0"]).await, 8); + + // byte 1 only: 0 set bits + assert_eq!(c.get_int(&["BITCOUNT", "bits", "1", "1"]).await, 0); + + // both bytes: 8 set bits total + assert_eq!(c.get_int(&["BITCOUNT", "bits", "0", "1"]).await, 8); +} + +#[tokio::test] +async fn bitcount_bit_range_unit() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // set all 8 bits to get 0xFF + for i in 0..8 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + + // bits 0–3 → 4 set bits + assert_eq!(c.get_int(&["BITCOUNT", "bits", "0", "3", "BIT"]).await, 4); + // bits 0–7 → 8 set bits + assert_eq!(c.get_int(&["BITCOUNT", "bits", "0", "7", "BIT"]).await, 8); +} + +#[tokio::test] +async fn bitcount_missing_key_returns_zero() { + let server = TestServer::start(); + let mut c = server.connect().await; + + assert_eq!(c.get_int(&["BITCOUNT", "nokey"]).await, 0); +} + +// --- BITPOS --- + +#[tokio::test] +async fn bitpos_first_set_bit() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // byte 0 = 0x00, byte 1 = 0xFF: first set bit is at position 8 + // extend to 2 bytes with byte1 bit 8 set + c.get_int(&["SETBIT", "bits", "15", "0"]).await; // creates 2-byte string of zeros + for i in 8..16 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + assert_eq!(c.get_int(&["BITPOS", "bits", "1"]).await, 8); +} + +#[tokio::test] +async fn bitpos_first_clear_bit() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // byte 0 = 0xFF, byte 1 = 0x00: first clear bit is at position 8 + for i in 0..8 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + // extend to 2 bytes (byte 1 stays 0x00) + c.get_int(&["SETBIT", "bits", "15", "0"]).await; + assert_eq!(c.get_int(&["BITPOS", "bits", "0"]).await, 8); +} + +#[tokio::test] +async fn bitpos_missing_key_clear() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // missing key: BITPOS 0 → 0 + assert_eq!(c.get_int(&["BITPOS", "nokey", "0"]).await, 0); + // missing key: BITPOS 1 → -1 + let resp = c.cmd(&["BITPOS", "nokey", "1"]).await; + match resp { + Frame::Integer(n) => assert_eq!(n, -1), + other => panic!("expected Integer(-1), got {other:?}"), + } +} + +#[tokio::test] +async fn bitpos_with_byte_range() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // bytes 0,1 = 0x00, byte 2 = 0xFF + // extend to 3 bytes with first two bytes zero + c.get_int(&["SETBIT", "bits", "23", "0"]).await; // creates 3-byte string of zeros + for i in 16..24 { + c.get_int(&["SETBIT", "bits", &i.to_string(), "1"]).await; + } + // first set bit in byte range [2, 2] is position 16 + assert_eq!(c.get_int(&["BITPOS", "bits", "1", "2", "2"]).await, 16); +} + +// --- BITOP --- + +#[tokio::test] +async fn bitop_and() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // a = 0xFF (all bits set) + for i in 0..8 { + c.get_int(&["SETBIT", "a", &i.to_string(), "1"]).await; + } + // b = 0x0F (lower nibble set: bits 4-7) + for i in 4..8 { + c.get_int(&["SETBIT", "b", &i.to_string(), "1"]).await; + } + c.get_int(&["SETBIT", "b", "0", "0"]).await; // ensure b is 1 byte + + // 0xFF AND 0x0F = 0x0F → 4 set bits + let len = c.get_int(&["BITOP", "AND", "dest", "a", "b"]).await; + assert_eq!(len, 1); + assert_eq!(c.get_int(&["BITCOUNT", "dest"]).await, 4); +} + +#[tokio::test] +async fn bitop_or() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // a = 0xF0 (upper nibble: bits 0-3) + for i in 0..4 { + c.get_int(&["SETBIT", "a", &i.to_string(), "1"]).await; + } + c.get_int(&["SETBIT", "a", "7", "0"]).await; // ensure a is 1 byte + + // b = 0x0F (lower nibble: bits 4-7) + for i in 4..8 { + c.get_int(&["SETBIT", "b", &i.to_string(), "1"]).await; + } + c.get_int(&["SETBIT", "b", "0", "0"]).await; // ensure b is 1 byte + + // 0xF0 OR 0x0F = 0xFF → 8 set bits + let len = c.get_int(&["BITOP", "OR", "dest", "a", "b"]).await; + assert_eq!(len, 1); + assert_eq!(c.get_int(&["BITCOUNT", "dest"]).await, 8); +} + +#[tokio::test] +async fn bitop_xor() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // both a and b = 0xFF + for i in 0..8 { + c.get_int(&["SETBIT", "a", &i.to_string(), "1"]).await; + c.get_int(&["SETBIT", "b", &i.to_string(), "1"]).await; + } + + // 0xFF XOR 0xFF = 0x00 → 0 set bits + let len = c.get_int(&["BITOP", "XOR", "dest", "a", "b"]).await; + assert_eq!(len, 1); + assert_eq!(c.get_int(&["BITCOUNT", "dest"]).await, 0); +} + +#[tokio::test] +async fn bitop_not() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // src = 0xFF + for i in 0..8 { + c.get_int(&["SETBIT", "src", &i.to_string(), "1"]).await; + } + + // NOT 0xFF = 0x00 → 0 set bits + let len = c.get_int(&["BITOP", "NOT", "dest", "src"]).await; + assert_eq!(len, 1); + assert_eq!(c.get_int(&["BITCOUNT", "dest"]).await, 0); +} + +#[tokio::test] +async fn bitop_result_length_matches_longest_source() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // a is 3 bytes (bit 23 set), b is 1 byte + c.get_int(&["SETBIT", "a", "0", "1"]).await; + c.get_int(&["SETBIT", "a", "23", "1"]).await; // forces 3-byte string + c.get_int(&["SETBIT", "b", "4", "1"]).await; // 1-byte string + + let len = c.get_int(&["BITOP", "AND", "dest", "a", "b"]).await; + assert_eq!(len, 3); +} + +#[tokio::test] +async fn bitop_missing_source_treated_as_zeros() { + let server = TestServer::start(); + let mut c = server.connect().await; + + // a = 0xFF + for i in 0..8 { + c.get_int(&["SETBIT", "a", &i.to_string(), "1"]).await; + } + // "missing" key does not exist — treated as 0x00 + + // AND with a zero-filled ghost → result is all zeros + c.get_int(&["BITOP", "AND", "dest", "a", "missing"]).await; + assert_eq!(c.get_int(&["BITCOUNT", "dest"]).await, 0); +} diff --git a/tests/integration/src/main.rs b/tests/integration/src/main.rs index 65b036b3..702ca894 100644 --- a/tests/integration/src/main.rs +++ b/tests/integration/src/main.rs @@ -2,6 +2,7 @@ mod helpers; mod auth; mod basic_operations; +mod bitmap; mod cli; mod client_typed_api; mod cluster;