Bitmaps on strings

High bit first

// trafikey — tra::bitmap, the offset and bit parsers and the MSB-first read
// the canonical offset domain ends at 2^32 - 1
const MAX_OFFSET: u64 = u32::MAX as u64;

/// Offsets: ASCII decimal digits only (leading zeroes allowed), within 0..=2^32-1.
fn parse_offset(raw: &[u8]) -> Option<u64> {
    if raw.is_empty() || raw.len() > 32 {
        return None;
    }
    let mut value: u64 = 0;
    for &c in raw {
        if !c.is_ascii_digit() {
            return None;
        }
        value = value.checked_mul(10)?.checked_add(u64::from(c - b'0'))?;
    }
    (value <= MAX_OFFSET).then_some(value)
}
/// Bit values are exactly `0` or `1`.
fn parse_bit(raw: &[u8]) -> Option<u8> {
    match raw {
        b"0" => Some(0),
        b"1" => Some(1),
        _ => None,
    }
}
// …
/// MSB-first bit read; anything at or past the logical length is zero.
fn get_bit(bytes: &[u8], offset: u64) -> u8 {
    match bytes.get((offset >> 3) as usize) {
        Some(byte) => (byte >> (7 - (offset & 7))) & 1,
        None => 0,
    }
}

Offset N is bit 7 − (N mod 8) of byte N ÷ 8, the layout the server’s own bit commands use, over an offset grammar the module checks itself.

The one write

// trafikey — tra::bitmap::setbit, the one TRA command that writes a key
pub(crate) fn setbit(ctx: &Context, args: Vec<ValkeyString>) -> ValkeyResult {
    // The host enforces arity 4 before the handler runs; the else arm is its arity reply.
    let [_, name, offset, bit] = args.as_slice() else {
        return Err(ValkeyError::WrongArity);
    };
    let offset = parse_offset(offset.as_slice()).ok_or(INVALID_OFFSET)?;
    let bit = parse_bit(bit.as_slice()).ok_or(INVALID_ARGUMENT)?;
    let index = (offset >> 3) as usize;
    let (previous, grows) = with_string(ctx, name, |bytes| match bytes {
        None => (0, true),
        Some(b) => (get_bit(b, offset), index >= b.len()),
    })?;
    // Equal bit without growth is not a logical mutation: no write, no replication, no WATCH touch.
    if !grows && previous == bit {
        return Ok(ValkeyValue::Integer(i64::from(previous)));
    }
    let key = writable(raw::open_key(
        ctx.ctx,
        name.inner,
        raw::KeyMode::READ | raw::KeyMode::WRITE,
    ))?;
    let mut length = 0;
    // StringTruncate only grows here (zero-filled); in-place mutation keeps the TTL. The preload
    // gate refuses the load without it; an absent API is the host refusing to grow.
    raw::string_dma(key, &mut length, raw::KeyMode::READ);
    if length <= index {
        // SAFETY: `RedisModule_StringTruncate` is a C global that `RedisModule_InitAPI` wrote
        // during OnLoad; TRA.SETBIT runs on the main thread after it and nothing writes it again.
        let grown = unsafe { raw::RedisModule_StringTruncate }.is_some_and(|truncate| {
            // SAFETY: `VM_StringTruncate` dereferences `key`: it is the non-null
            // write-mode handle `writable` accepted, not yet closed.
            (unsafe { truncate(key, index + 1) }) == raw::Status::Ok as i32
        });
        if !grown {
            raw::close_key(key);
            return Err(LIMIT);
        }
    }
    let ptr = raw::string_dma(key, &mut length, raw::KeyMode::WRITE).cast::<u8>();
    if ptr.is_null() || length <= index {
        raw::close_key(key);
        return Err(LIMIT);
    }
    let mask = 1u8 << (7 - (offset & 7));
    // SAFETY: `ptr` starts the `length`-byte write-mode StringDMA buffer of the open key and
    // `index < length` (checked above), so the offset stays inside that one allocation.
    let target = unsafe { ptr.add(index) };
    // SAFETY: `target` is an in-bounds byte of the write-mode DMA buffer, valid until close_key;
    // no key write runs before close_key, so this is the only access.
    let byte = unsafe { &mut *target };
    if bit == 1 {
        *byte |= mask;
    } else {
        *byte &= !mask;
    }
    raw::close_key(key);
    ctx.notify_keyspace_event(NotifyEvent::STRING, "setbit", name);
    // … four lines: a clippy lint expectation for the next call, with its reason
    ctx.replicate_verbatim();
    Ok(ValkeyValue::Integer(i64::from(previous)))
}

// trafikey — emq/host.rs, the EMQ adapter's replication rule, for contrast
//! So this file calls `RedisModule_Call` itself, with `"vE"` for reads and `"v!E"` for writes,
//! and classifies the raw reply type: `!` replicates each write as its effect, and every `!` call
//! of one command reaches the AOF and replicas inside one MULTI/EXEC. Never `ReplicateVerbatim`,
//! never the `3`, `0`, `X` or `K` flags.

A changed bit is one write and one replicated command; an unchanged bit inside the string is a reply and nothing else.

Counting a range

// trafikey — tra::bitmap::bitcount, the argument shapes, then the count
pub(crate) fn bitcount(ctx: &Context, args: Vec<ValkeyString>) -> ValkeyResult {
    let (name, range) = match args.as_slice() {
        [_, name] => (name, None),
        [_, name, start, rest @ ..] if rest.len() <= 2 => {
            let start = parse_i64(start.as_slice()).ok_or(INVALID_ARGUMENT)?;
            let (end, bit_unit) = match rest {
                [] => (None, false),
                [end, unit @ ..] => {
                    let end = parse_i64(end.as_slice()).ok_or(INVALID_ARGUMENT)?;
                    let bit_unit = match unit.first().map(|u| u.as_slice()) {
                        None => false,
                        Some(u) if u.eq_ignore_ascii_case(b"BIT") => true,
                        Some(u) if u.eq_ignore_ascii_case(b"BYTE") => false,
                        Some(_) => return Err(INVALID_ARGUMENT),
                    };
                    (Some(end), bit_unit)
                }
            };
            (name, Some((start, end, bit_unit)))
        }
        _ => return Err(INVALID_ARGUMENT),
    };
    with_string(ctx, name, |b| {
        let n = count(b.unwrap_or(&[]), range);
        ValkeyValue::Integer(n.min(i64::MAX as u64) as i64)
    })
}
/// TRA.BITCOUNT's answer for a parsed range: no range counts the whole value, and a start
/// without an end ends at the last byte.
fn count(bytes: &[u8], range: Option<(i64, Option<i64>, bool)>) -> u64 {
    let last_byte = bytes.len() as i64 - 1;
    range.map_or_else(
        || popcount(bytes),
        |(start, end, bit_unit)| count_range(bytes, start, end.unwrap_or(last_byte), bit_unit),
    )
}

// trafikey — tra::bitmap::count_range, the inclusive window over the string
// two negative endpoints with start > end are empty; negatives count back from the end,
// both are raised to 0 and the end lowered to the last unit, and start > end is empty
fn count_range(bytes: &[u8], start: i64, end: i64, bit_unit: bool) -> u64 {
    if start < 0 && end < 0 && start > end {
        return 0;
    }
    let domain: i128 = bytes.len() as i128 * if bit_unit { 8 } else { 1 };
    let mut s = i128::from(start);
    let mut e = i128::from(end);
    if s < 0 {
        s += domain;
    }
    if e < 0 {
        e += domain;
    }
    if s < 0 {
        s = 0;
    }
    if e < 0 {
        e = 0;
    }
    if e >= domain {
        e = domain - 1;
    }
    if s > e {
        return 0;
    }
    // From here 0 <= s <= e < domain, so both windows below lie inside `bytes`.
    if !bit_unit {
        return bytes.get(s as usize..=e as usize).map_or(0, popcount);
    }
    let (first, last) = ((s >> 3) as usize, (e >> 3) as usize);
    let head: u8 = 0xFF >> (s & 7);
    let tail: u8 = 0xFF << (7 - (e & 7));
    match bytes.get(first..=last).unwrap_or_default() {
        [] => 0,
        [only] => u64::from((only & head & tail).count_ones()),
        [lo, middle @ .., hi] => {
            u64::from((lo & head).count_ones())
                + popcount(middle)
                + u64::from((hi & tail).count_ones())
        }
    }
}

Arguments first, then the key: a malformed range is refused before the string is read, and a well-formed range on a string always counts, an empty window counting 0.

Strings only

// trafikey — tra::bitmap::with_string, the one reader all four commands share
/// Reads argv[1] as an ordinary string: Ok(None) when missing, WRONGTYPE for any other type,
/// including types this crate does not know.
fn with_string<T>(
    ctx: &Context,
    name: &ValkeyString,
    f: impl FnOnce(Option<&[u8]>) -> T,
) -> Result<T, ValkeyError> {
    // The preload gate refuses the load without KeyType; an absent API is the host refusing.
    let Some(key_type) = (unsafe { raw::RedisModule_KeyType }) else {
        return Err(LIMIT);
    };
    // Read mode, as the crate's `open_key`. A missing key is NULL, and KeyType(NULL) is EMPTY.
    let key = raw::open_key(ctx.ctx, name.inner, raw::KeyMode::READ);
    // SAFETY: `key` is the handle open_key just returned, not yet closed; NULL reads as EMPTY.
    let out = match classify(unsafe { key_type(key) }) {
        Kind::Missing => Ok(f(None)),
        Kind::String => {
            let mut length = 0;
            let ptr = raw::string_dma(key, &mut length, raw::KeyMode::READ).cast::<u8>();
            if ptr.is_null() {
                Err(ValkeyError::WrongType)
            } else {
                // SAFETY: StringDMA returned a non-null pointer to `length` bytes of this open
                // string key; they stay valid until close_key below, and `f` runs before it.
                Ok(f(Some(unsafe { std::slice::from_raw_parts(ptr, length) })))
            }
        }
        Kind::Other => Err(ValkeyError::WrongType),
    };
    raw::close_key(key);
    out
}

// trafikey — tra::bitmap, the raw key-type classification and the write-handle guard
// what a key holds, read from the raw RedisModule_KeyType value:
/// Never through the crate's `KeyType: From<c_int>`: it unwraps, and a type newer than the crate
/// (Valkey 9.2 `PATH_HASH` = 8) panics, which `panic = "abort"` turns into a server abort.
#[derive(Debug, PartialEq, Eq)]
enum Kind {
    Missing,
    String,
    Other,
}
fn classify(raw_type: c_int) -> Kind {
    match raw_type {
        t if t == raw::REDISMODULE_KEYTYPE_EMPTY as c_int => Kind::Missing,
        t if t == raw::REDISMODULE_KEYTYPE_STRING as c_int => Kind::String,
        _ => Kind::Other,
    }
}

// the second Valkey 9.2 hazard: a write-mode open can return NULL while a background
// iteration holds the entry, so a NULL handle never reaches StringDMA or StringTruncate
fn writable(key: *mut raw::RedisModuleKey) -> Result<*mut raw::RedisModuleKey, ValkeyError> {
    if key.is_null() {
        Err(LIMIT)
    } else {
        Ok(key)
    }
}

// trafikey — the test that pins the mapping, Valkey 9.2's type 8 included
    // every raw key type maps without a panic; only EMPTY and STRING are readable
    #[test]
    fn raw_key_types_classify_without_panic() {
        assert_eq!(classify(0), Kind::Missing);
        assert_eq!(classify(1), Kind::String);
        // LIST..STREAM (2..=7), Valkey 9.2 PATH_HASH (8), and values no server defines yet.
        for t in [2, 3, 4, 5, 6, 7, 8, 9, 99, -1, c_int::MIN, c_int::MAX] {
            assert_eq!(classify(t), Kind::Other, "{t}");
        }
        assert_eq!(raw::REDISMODULE_KEYTYPE_EMPTY, 0);
        assert_eq!(raw::REDISMODULE_KEYTYPE_STRING, 1);
    }

# trafikey — clippy.toml, the Trafikist profile's key-type pack
# rstyler: Profile trafikist
disallowed-methods = [
    # pack: valkey-module key types, keyed to valkey-module 0.1.14
    { path = "valkey_module::key::ValkeyKey::key_type", reason = "aborts on a Valkey 9.2 Path Hash key; read the raw key type" },
    { path = "valkey_module::key::ValkeyKey::get_value", reason = "verify_type aborts on a Valkey 9.2 Path Hash key" },
    { path = "valkey_module::key::ValkeyKeyWritable::key_type", reason = "aborts on a Valkey 9.2 Path Hash key; read the raw key type" },
    { path = "valkey_module::key::ValkeyKeyWritable::get_value", reason = "verify_type aborts on a Valkey 9.2 Path Hash key" },
    { path = "valkey_module::key::ValkeyKeyWritable::set_value", reason = "verify_type aborts on a Valkey 9.2 Path Hash key" },
    { path = "valkey_module::key::verify_type", reason = "aborts on a Valkey 9.2 Path Hash key" },
    # … four more: Context::open_key, open_key_writable and their _with_flags forms
    # … three more packs: NUL bytes, call replies and replication
]

A string, a missing key or WRONGTYPE: the classification reads the raw type value, so a key type the binding does not define is a reply, not an abort.

Pattern & implementation

References