feat: bitmap commands (GETBIT, SETBIT, BITCOUNT, BITPOS, BITOP)

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")]

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 }]
         }

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,
         );

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<BitRange>,
+    },
+    /// BITPOS key bit [range]. Finds first set or clear bit position.
+    BitPos {
+        key: String,
+        bit: u8,
+        range: Option<BitRange>,
+    },
+    /// BITOP op destkey key [key ...]. Bitwise operation across strings.
+    BitOp {
+        op: BitOpKind,
+        dest: String,
+        keys: Vec<String>,
+    },
     /// 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)

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,
             );

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<String>,
+    },
     /// 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")?;

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<Bytes> = 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);