glide/commands/
bitmap.rs

1// Copyright Valkey GLIDE Project Contributors - SPDX Identifier: Apache-2.0
2//! Bitmap commands. Mirrors Python's bitmap command surface.
3
4use crate::ValkeyResult;
5use crate::cmd::Cmd;
6use crate::executor::CommandExecutor;
7use crate::value::FromValkeyValue;
8use crate::value::ValkeyValue;
9use crate::write::ToValkeyArgs;
10use crate::write::ValkeyWrite;
11use async_trait::async_trait;
12
13/// Index unit for `BITCOUNT`/`BITPOS` range queries.
14///
15/// Mirrors Python `BitmapIndexType`.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum BitmapIndexType {
18    /// Interpret the range as byte offsets (`BYTE`).
19    Byte,
20    /// Interpret the range as bit offsets (`BIT`).
21    Bit,
22}
23
24impl BitmapIndexType {
25    fn as_arg(&self) -> &'static str {
26        match self {
27            BitmapIndexType::Byte => "BYTE",
28            BitmapIndexType::Bit => "BIT",
29        }
30    }
31}
32
33/// A signed or unsigned bit encoding for `BITFIELD`/`BITFIELD_RO`.
34///
35/// Mirrors Python `SignedEncoding`/`UnsignedEncoding`.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum BitEncoding {
38    /// Signed encoding of the given bit width (`i<width>`, `< 65` bits).
39    Signed(u32),
40    /// Unsigned encoding of the given bit width (`u<width>`, `< 64` bits).
41    Unsigned(u32),
42}
43
44impl ToValkeyArgs for BitEncoding {
45    fn write_valkey_args<W: ?Sized + ValkeyWrite>(&self, out: &mut W) {
46        match self {
47            BitEncoding::Signed(n) => out.write_arg_fmt(format_args!("i{n}")),
48            BitEncoding::Unsigned(n) => out.write_arg_fmt(format_args!("u{n}")),
49        }
50    }
51}
52
53/// A bit offset for `BITFIELD`/`BITFIELD_RO`.
54///
55/// Mirrors Python `BitOffset`/`BitOffsetMultiplier`.
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum BitFieldOffset {
58    /// A raw bit index offset.
59    Bit(u64),
60    /// An offset multiplied by the encoding width (`#<offset>`).
61    Multiplier(u64),
62}
63
64impl ToValkeyArgs for BitFieldOffset {
65    fn write_valkey_args<W: ?Sized + ValkeyWrite>(&self, out: &mut W) {
66        match self {
67            BitFieldOffset::Bit(n) => out.write_arg_fmt(n),
68            BitFieldOffset::Multiplier(n) => out.write_arg_fmt(format_args!("#{n}")),
69        }
70    }
71}
72
73/// Overflow behavior for `BITFIELD` `SET`/`INCRBY` subcommands.
74///
75/// Mirrors Python `BitOverflowControl`.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub enum BitOverflow {
78    /// Wrap around on overflow (`WRAP`).
79    Wrap,
80    /// Saturate at the min/max value on overflow (`SAT`).
81    Sat,
82    /// Return `None` on overflow (`FAIL`).
83    Fail,
84}
85
86impl ToValkeyArgs for BitOverflow {
87    fn write_valkey_args<W: ?Sized + ValkeyWrite>(&self, out: &mut W) {
88        out.write_arg(match self {
89            BitOverflow::Wrap => b"WRAP".as_slice(),
90            BitOverflow::Sat => b"SAT".as_slice(),
91            BitOverflow::Fail => b"FAIL".as_slice(),
92        });
93    }
94}
95
96/// A subcommand for `BITFIELD`/`BITFIELD_RO`.
97///
98/// Mirrors Python `BitFieldGet`/`BitFieldSet`/`BitFieldIncrBy`/`BitFieldOverflow`.
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub enum BitFieldSubcommand {
101    /// Read a value (`GET encoding offset`).
102    Get {
103        /// The bit encoding.
104        encoding: BitEncoding,
105        /// The bit offset.
106        offset: BitFieldOffset,
107    },
108    /// Write a value (`SET encoding offset value`).
109    Set {
110        /// The bit encoding.
111        encoding: BitEncoding,
112        /// The bit offset.
113        offset: BitFieldOffset,
114        /// The value to set.
115        value: i64,
116    },
117    /// Increment a value (`INCRBY encoding offset increment`).
118    IncrBy {
119        /// The bit encoding.
120        encoding: BitEncoding,
121        /// The bit offset.
122        offset: BitFieldOffset,
123        /// The amount to increment by.
124        increment: i64,
125    },
126    /// Set the overflow behavior for subsequent subcommands (`OVERFLOW`).
127    Overflow(BitOverflow),
128}
129
130impl ToValkeyArgs for BitFieldSubcommand {
131    fn write_valkey_args<W: ?Sized + ValkeyWrite>(&self, out: &mut W) {
132        match self {
133            BitFieldSubcommand::Get { encoding, offset } => {
134                out.write_arg(b"GET");
135                encoding.write_valkey_args(out);
136                offset.write_valkey_args(out);
137            }
138            BitFieldSubcommand::Set {
139                encoding,
140                offset,
141                value,
142            } => {
143                out.write_arg(b"SET");
144                encoding.write_valkey_args(out);
145                offset.write_valkey_args(out);
146                value.write_valkey_args(out);
147            }
148            BitFieldSubcommand::IncrBy {
149                encoding,
150                offset,
151                increment,
152            } => {
153                out.write_arg(b"INCRBY");
154                encoding.write_valkey_args(out);
155                offset.write_valkey_args(out);
156                increment.write_valkey_args(out);
157            }
158            BitFieldSubcommand::Overflow(overflow) => {
159                out.write_arg(b"OVERFLOW");
160                overflow.write_valkey_args(out);
161            }
162        }
163    }
164}
165
166/// Bitmap commands (`SETBIT`, `GETBIT`, `BITCOUNT`, `BITPOS`, `BITOP`).
167#[async_trait]
168pub trait BitmapCommands: CommandExecutor {
169    // TODO #7082: add a `bitcount` variant that takes the `BYTE`/`BIT` index unit
170    // (the `BitmapIndexType` type already exists), to match the other GLIDE clients.
171
172    /// Find the position of the first bit set to `bit` (`BITPOS`).
173    async fn bitpos<K: ToValkeyArgs + Send>(&self, key: K, bit: u8) -> ValkeyResult<i64> {
174        let mut cmd = Cmd::new();
175        cmd.arg("BITPOS").arg(key).arg(bit);
176        i64::from_owned_valkey_value(self.execute_command(cmd, None).await?)
177    }
178
179    /// Find the position of the first bit set to `bit` within a range
180    /// (`BITPOS key bit start end [BYTE|BIT]`).
181    async fn bitpos_range<K: ToValkeyArgs + Send>(
182        &self,
183        key: K,
184        bit: u8,
185        start: i64,
186        end: i64,
187        index_type: Option<BitmapIndexType>,
188    ) -> ValkeyResult<i64> {
189        let mut cmd = Cmd::new();
190        cmd.arg("BITPOS").arg(key).arg(bit).arg(start).arg(end);
191        if let Some(it) = index_type {
192            cmd.arg(it.as_arg());
193        }
194        i64::from_owned_valkey_value(self.execute_command(cmd, None).await?)
195    }
196
197    /// Perform arbitrary bit-field operations (`BITFIELD`). Returns one result
198    /// per non-`OVERFLOW` subcommand; a `None` indicates a `FAIL` overflow.
199    async fn bitfield<K: ToValkeyArgs + Send>(
200        &self,
201        key: K,
202        subcommands: &[BitFieldSubcommand],
203    ) -> ValkeyResult<Vec<Option<i64>>> {
204        let mut cmd = Cmd::new();
205        cmd.arg("BITFIELD").arg(key);
206        cmd.arg(subcommands);
207        parse_bitfield(self.execute_command(cmd, None).await?)
208    }
209
210    /// Read-only bit-field operations (`BITFIELD_RO`). Only `GET` subcommands are
211    /// permitted by the server.
212    async fn bitfield_readonly<K: ToValkeyArgs + Send>(
213        &self,
214        key: K,
215        subcommands: &[BitFieldSubcommand],
216    ) -> ValkeyResult<Vec<Option<i64>>> {
217        let mut cmd = Cmd::new();
218        cmd.arg("BITFIELD_RO").arg(key);
219        cmd.arg(subcommands);
220        parse_bitfield(self.execute_command(cmd, None).await?)
221    }
222}
223
224impl<T: CommandExecutor + ?Sized> BitmapCommands for T {}
225
226/// Parse a `BITFIELD` reply (array of ints, with `Nil` for `FAIL` overflow).
227fn parse_bitfield(v: ValkeyValue) -> ValkeyResult<Vec<Option<i64>>> {
228    match v {
229        ValkeyValue::Nil => Ok(Vec::new()),
230        ValkeyValue::Array(items) => items
231            .into_iter()
232            .map(|it| match it {
233                ValkeyValue::Nil => Ok(None),
234                other => Ok(Some(i64::from_owned_valkey_value(other)?)),
235            })
236            .collect(),
237        other => Ok(vec![Some(i64::from_owned_valkey_value(other)?)]),
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use crate::test_utils::assert_args;
245
246    #[test]
247    fn encoding_and_offset_args() {
248        assert_args(BitEncoding::Signed(8), &["i8"]);
249        assert_args(BitEncoding::Unsigned(16), &["u16"]);
250        assert_args(BitFieldOffset::Bit(5), &["5"]);
251        assert_args(BitFieldOffset::Multiplier(3), &["#3"]);
252    }
253
254    #[test]
255    fn overflow_args() {
256        assert_args(BitOverflow::Wrap, &["WRAP"]);
257        assert_args(BitOverflow::Sat, &["SAT"]);
258        assert_args(BitOverflow::Fail, &["FAIL"]);
259    }
260
261    #[test]
262    fn bitfield_subcommand_args() {
263        assert_args(
264            BitFieldSubcommand::Get {
265                encoding: BitEncoding::Unsigned(8),
266                offset: BitFieldOffset::Bit(0),
267            },
268            &["GET", "u8", "0"],
269        );
270        assert_args(
271            BitFieldSubcommand::Set {
272                encoding: BitEncoding::Signed(5),
273                offset: BitFieldOffset::Multiplier(1),
274                value: 12,
275            },
276            &["SET", "i5", "#1", "12"],
277        );
278        assert_args(
279            BitFieldSubcommand::IncrBy {
280                encoding: BitEncoding::Unsigned(4),
281                offset: BitFieldOffset::Bit(2),
282                increment: -3,
283            },
284            &["INCRBY", "u4", "2", "-3"],
285        );
286        assert_args(
287            BitFieldSubcommand::Overflow(BitOverflow::Sat),
288            &["OVERFLOW", "SAT"],
289        );
290    }
291
292    #[test]
293    fn parse_bitfield_handles_nil() {
294        let v = ValkeyValue::Array(vec![ValkeyValue::Int(1), ValkeyValue::Nil]);
295        assert_eq!(parse_bitfield(v).unwrap(), vec![Some(1), None]);
296    }
297}
glide/commands/bitmap.rs

glide/commands/
bitmap.rs
