pub mod bitops;
pub mod conf;
pub mod meta;
use std::collections::hash_map::Entry;
use std::str;
use rapidhash::RapidHashMap as HashMap;
pub use bitops::{
ArrayBitfieldBitmap, BITMAP_SEGMENT_BITS, BITMAP_SEGMENT_BYTES, bit_op_exec, bit_op_exec_into,
bitfield_op_calc, expand_bitmap_segment, find_bit_in_byte_lsb, find_bit_in_byte_msb,
get_bit_from_bytes, get_bit_lsb, normalize_bit_range_to_byte_mask, normalize_range,
normalize_to_byte_range_with_padding_mask, raw_bitpos, raw_bitpos_lsb, raw_popcount,
segment_byte_offset_for_bit, segment_index_for_bit, set_bit_in_bytes, set_bit_lsb,
signed_bitfield_plus, string_bitcount, string_bitpos, unsigned_bitfield_plus,
};
pub use conf::{
BitOp, BitfieldEncoding, BitfieldOpType, BitfieldOperation, BitfieldOverflow, BitfieldValue,
parse_bitfield_offset,
};
pub use meta::BitmapMeta;
use crate::db::WeDb;
use crate::error::{ERR_WRONG_TYPE, Error, Result};
use crate::key_composer::KeyComposer;
use crate::meta::current_now_ms;
use crate::string::{decode_string_value, encode_string_value, is_string_expired};
pub const MAX_BITMAP_TO_STRING_BYTES: usize = 512 * 1024 * 1024;
struct SegmentCacheStore<'a> {
db: &'a WeDb,
kc: KeyComposer<'a>,
k_str: &'a str,
cache: HashMap<u32, (bool, Vec<u8>)>,
}
impl<'a> SegmentCacheStore<'a> {
#[inline]
fn new(db: &'a WeDb, kc: KeyComposer<'a>, k_str: &'a str) -> Self {
Self {
db,
kc,
k_str,
cache: HashMap::default(),
}
}
#[inline]
fn get(&mut self, seg_idx: u32) -> Result<&[u8]> {
let entry = match self.cache.entry(seg_idx) {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => {
let seg_k = self.kc.bm_segment(self.k_str, seg_idx);
let bytes = self
.db
.data
.get(seg_k.as_bytes())?
.map(|v| v.to_vec())
.unwrap_or_default();
e.insert((false, bytes))
}
};
Ok(&entry.1)
}
#[inline]
fn get_mut(&mut self, seg_idx: u32, min_bytes: usize) -> Result<&mut Vec<u8>> {
let entry = match self.cache.entry(seg_idx) {
Entry::Occupied(e) => e.into_mut(),
Entry::Vacant(e) => {
let seg_k = self.kc.bm_segment(self.k_str, seg_idx);
let bytes = self
.db
.data
.get(seg_k.as_bytes())?
.map(|v| v.to_vec())
.unwrap_or_default();
e.insert((false, bytes))
}
};
expand_bitmap_segment(&mut entry.1, min_bytes);
entry.0 = true;
Ok(&mut entry.1)
}
#[inline]
fn dirty_segments(&self) -> impl Iterator<Item = (u32, &[u8])> {
self.cache
.iter()
.filter_map(|(&seg_idx, (dirty, seg_data))| {
if *dirty {
Some((seg_idx, seg_data.as_slice()))
} else {
None
}
})
}
}
impl WeDb {
pub fn setbit<K: AsRef<[u8]>>(&self, key: K, offset: u64, bit: u8) -> Result<u8> {
if bit > 1 {
return Err(Error::invalid_data(
"ERR bit is out of range, must be 0 or 1",
));
}
if offset > u32::MAX as u64 {
return Err(Error::invalid_data(
"ERR bit offset is not an integer or out of range",
));
}
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let bm_meta_k = kc.bm_meta(k_str);
let cur_meta_opt = self
.meta
.get(bm_meta_k.as_bytes())?
.and_then(|b| BitmapMeta::decode(&b));
if let Some(mut meta) = cur_meta_opt
&& !meta.is_expired(now_ms)
{
let seg_idx = segment_index_for_bit(offset);
let bit_offset_in_seg = (offset % (BITMAP_SEGMENT_BITS as u64)) as usize;
let byte_idx_in_seg = bit_offset_in_seg >> 3;
let seg_k = kc.bm_segment(k_str, seg_idx);
let seg_slice_opt = self.data.get(seg_k.as_bytes())?;
let old_bit = seg_slice_opt
.as_deref()
.map(|s| get_bit_lsb(s, bit_offset_in_seg))
.unwrap_or(0);
let used_size = segment_byte_offset_for_bit(offset) as u64 + byte_idx_in_seg as u64 + 1;
let bitmap_size = meta.base.size.max(used_size);
if let Some(ref seg_slice) = seg_slice_opt
&& old_bit == bit
&& meta.base.size == bitmap_size
&& byte_idx_in_seg < seg_slice.len()
{
return Ok(old_bit);
}
let mut seg = seg_slice_opt.map(|v| v.to_vec()).unwrap_or_default();
expand_bitmap_segment(&mut seg, byte_idx_in_seg + 1);
set_bit_lsb(&mut seg, bit_offset_in_seg, bit);
meta.base.size = bitmap_size;
let mut batch = self.db.batch();
batch.insert(&self.data, seg_k.as_bytes(), &seg);
batch.insert(&self.meta, bm_meta_k.as_bytes(), meta.encode());
batch.commit()?;
return Ok(old_bit);
}
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
let byte_idx = (offset >> 3) as usize;
let old_bit = get_bit_from_bytes(val, offset as usize);
if old_bit == bit && byte_idx < val.len() {
return Ok(old_bit);
}
let mut str_bytes = val.to_vec();
if byte_idx >= str_bytes.len() {
str_bytes.resize(byte_idx + 1, 0);
}
set_bit_in_bytes(&mut str_bytes, offset as usize, bit);
let enc_val = encode_string_value(&str_bytes, expire_at);
self.data.insert(&*raw_k, enc_val)?;
return Ok(old_bit);
}
}
let seg_idx = segment_index_for_bit(offset);
let bit_offset_in_seg = (offset % (BITMAP_SEGMENT_BITS as u64)) as usize;
let byte_idx_in_seg = bit_offset_in_seg >> 3;
let mut seg = Vec::new();
expand_bitmap_segment(&mut seg, byte_idx_in_seg + 1);
let old_bit = set_bit_lsb(&mut seg, bit_offset_in_seg, bit);
let used_size = segment_byte_offset_for_bit(offset) as u64 + byte_idx_in_seg as u64 + 1;
let meta = BitmapMeta::new_with_version(0, used_size);
let seg_k = kc.bm_segment(k_str, seg_idx);
let mut batch = self.db.batch();
batch.insert(&self.data, seg_k.as_bytes(), &seg);
batch.insert(&self.meta, bm_meta_k.as_bytes(), meta.encode());
batch.commit()?;
Ok(old_bit)
}
pub fn getbit<K: AsRef<[u8]>>(&self, key: K, offset: u64) -> Result<u8> {
if offset > u32::MAX as u64 {
return Ok(0);
}
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let bm_meta_k = kc.bm_meta(k_str);
if let Some(m_bytes) = self.meta.get(bm_meta_k.as_bytes())?
&& let Some(meta) = BitmapMeta::decode(&m_bytes)
{
if meta.is_expired(now_ms) {
return Ok(0);
}
let seg_idx = segment_index_for_bit(offset);
let bit_offset_in_seg = (offset % (BITMAP_SEGMENT_BITS as u64)) as usize;
let seg_k = kc.bm_segment(k_str, seg_idx);
if let Some(seg) = self.data.get(seg_k.as_bytes())? {
return Ok(get_bit_lsb(&seg, bit_offset_in_seg));
}
return Ok(0);
}
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
return Ok(get_bit_from_bytes(val, offset as usize));
}
}
Ok(0)
}
pub fn bitcount<K: AsRef<[u8]>>(
&self,
key: K,
start: Option<i64>,
end: Option<i64>,
) -> Result<u64> {
self.bitcount_opts(key, start, end, false)
}
pub fn bitcount_opts<K: AsRef<[u8]>>(
&self,
key: K,
start: Option<i64>,
end: Option<i64>,
is_bit_index: bool,
) -> Result<u64> {
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let bm_meta_k = kc.bm_meta(k_str);
if let Some(m_bytes) = self.meta.get(bm_meta_k.as_bytes())?
&& let Some(meta) = BitmapMeta::decode(&m_bytes)
{
if meta.is_expired(now_ms) || meta.is_empty() {
return Ok(0);
}
let length = if is_bit_index {
(meta.base.size * 8) as i64
} else {
meta.base.size as i64
};
let s = start.unwrap_or(0);
let e = end.unwrap_or(-1);
if s < 0 && e < 0 && s > e {
return Ok(0);
}
let (norm_s, norm_e) = normalize_range(s, e, length);
if norm_s > norm_e {
return Ok(0);
}
let (start_byte, stop_byte, first_mask, last_mask) =
normalize_to_byte_range_with_padding_mask(is_bit_index, norm_s, norm_e);
let first_seg = start_byte / BITMAP_SEGMENT_BYTES;
let last_seg = stop_byte / BITMAP_SEGMENT_BYTES;
let mut total_cnt = 0u64;
for seg_idx in first_seg..=last_seg {
let seg_k = kc.bm_segment(k_str, seg_idx as u32);
if let Some(seg) = self.data.get(seg_k.as_bytes())? {
let seg_offset = seg_idx * BITMAP_SEGMENT_BYTES;
let seg_start = start_byte.saturating_sub(seg_offset);
let seg_stop = if stop_byte < seg_offset + BITMAP_SEGMENT_BYTES {
stop_byte - seg_offset
} else {
BITMAP_SEGMENT_BYTES - 1
};
if seg_start < seg.len() {
let actual_stop = seg_stop.min(seg.len() - 1);
if seg_start <= actual_stop {
let bytes = &seg[seg_start..=actual_stop];
let cnt = raw_popcount(bytes);
let mut mask_cnt = 0u64;
if first_seg == last_seg
&& seg_idx == first_seg
&& seg_start == actual_stop
{
let combined_mask = (first_mask | last_mask).reverse_bits();
if combined_mask != 0 {
mask_cnt +=
(seg[seg_start] & combined_mask).count_ones() as u64;
}
} else {
if first_mask != 0 && seg_idx == first_seg && seg_start < seg.len()
{
let reversed_first_mask = first_mask.reverse_bits();
mask_cnt +=
(seg[seg_start] & reversed_first_mask).count_ones() as u64;
}
if last_mask != 0
&& seg_idx == last_seg
&& actual_stop == seg_stop
&& actual_stop < seg.len()
{
let reversed_last_mask = last_mask.reverse_bits();
mask_cnt +=
(seg[actual_stop] & reversed_last_mask).count_ones() as u64;
}
}
total_cnt += cnt.saturating_sub(mask_cnt);
}
}
}
}
return Ok(total_cnt);
}
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
return Ok(string_bitcount(val, start, end, is_bit_index));
}
}
Ok(0)
}
pub fn bitpos<K: AsRef<[u8]>>(
&self,
key: K,
bit: u8,
start: Option<i64>,
end: Option<i64>,
) -> Result<i64> {
let stop_given = end.is_some();
self.bitpos_opts(key, bit, start, end, stop_given, false)
}
pub fn bitpos_opts<K: AsRef<[u8]>>(
&self,
key: K,
bit: u8,
start: Option<i64>,
end: Option<i64>,
stop_given: bool,
is_bit_index: bool,
) -> Result<i64> {
if bit > 1 {
return Err(Error::invalid_data(
"ERR bit is out of range, must be 0 or 1",
));
}
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let bm_meta_k = kc.bm_meta(k_str);
if let Some(m_bytes) = self.meta.get(bm_meta_k.as_bytes())?
&& let Some(meta) = BitmapMeta::decode(&m_bytes)
{
if meta.is_expired(now_ms) || meta.is_empty() {
return Ok(if bit == 0 { 0 } else { -1 });
}
let length = if is_bit_index {
(meta.base.size * 8) as i64
} else {
meta.base.size as i64
};
let s = start.unwrap_or(0);
let e = end.unwrap_or(-1);
let (norm_s, norm_e) = normalize_range(s, e, length);
if norm_s > norm_e {
return Ok(-1);
}
let u_start = norm_s as usize;
let u_stop = norm_e as usize;
let byte_start = if is_bit_index { u_start / 8 } else { u_start };
let byte_stop = if is_bit_index { u_stop / 8 } else { u_stop };
let start_seg = byte_start / BITMAP_SEGMENT_BYTES;
let stop_seg = byte_stop / BITMAP_SEGMENT_BYTES;
for seg_idx in start_seg..=stop_seg {
let seg_k = kc.bm_segment(k_str, seg_idx as u32);
let seg_opt = self.data.get(seg_k.as_bytes())?;
let seg_offset_bytes = seg_idx * BITMAP_SEGMENT_BYTES;
let seg_start_byte = byte_start.saturating_sub(seg_offset_bytes);
let seg_stop_byte = if byte_stop < seg_offset_bytes + BITMAP_SEGMENT_BYTES {
byte_stop - seg_offset_bytes
} else {
BITMAP_SEGMENT_BYTES - 1
};
if let Some(seg) = seg_opt {
let seg_slice = &seg[..];
if seg_start_byte < seg_slice.len() {
let actual_stop = seg_stop_byte.min(seg_slice.len() - 1);
if is_bit_index {
for (b_idx, &b) in seg_slice[..=actual_stop]
.iter()
.enumerate()
.skip(seg_start_byte)
{
let start_bit = if seg_idx == start_seg && b_idx == seg_start_byte {
u_start % 8
} else {
0
};
let stop_bit = if seg_idx == stop_seg && b_idx == seg_stop_byte {
u_stop % 8
} else {
7
};
if let Some(bit_idx) =
find_bit_in_byte_lsb(b, bit, start_bit, stop_bit)
{
let abs_pos = ((seg_offset_bytes + b_idx) * 8 + bit_idx) as i64;
return Ok(abs_pos);
}
}
} else if let Some(rel_pos) =
raw_bitpos_lsb(&seg_slice[seg_start_byte..=actual_stop], bit)
{
let abs_pos =
((seg_offset_bytes + seg_start_byte) * 8 + rel_pos) as i64;
return Ok(abs_pos);
}
}
if bit == 0 && seg_slice.len() <= seg_stop_byte {
let start_byte_in_seg = seg_start_byte.max(seg_slice.len());
let first_zero_bit = (seg_offset_bytes + start_byte_in_seg) * 8;
let abs_pos = if seg_idx == start_seg {
u_start.max(first_zero_bit) as i64
} else {
first_zero_bit as i64
};
if is_bit_index && abs_pos > norm_e {
return Ok(-1);
}
return Ok(abs_pos);
}
} else if bit == 0 {
let pos_in_seg = if seg_idx == start_seg {
if is_bit_index {
u_start.saturating_sub(seg_offset_bytes * 8)
} else {
seg_start_byte * 8
}
} else {
0
};
let abs_pos = (seg_offset_bytes * 8 + pos_in_seg) as i64;
if is_bit_index && abs_pos > norm_e {
return Ok(-1);
}
return Ok(abs_pos);
}
}
return Ok(if stop_given && bit == 0 {
-1
} else if bit == 0 {
(meta.base.size * 8) as i64
} else {
-1
});
}
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
return Ok(string_bitpos(
val,
bit,
start,
end,
stop_given,
is_bit_index,
));
}
}
Ok(if bit == 0 { 0 } else { -1 })
}
pub fn bitop<K: AsRef<[u8]>, S: AsRef<[u8]>>(
&self,
op: BitOp,
dest_key: K,
src_keys: &[S],
) -> Result<usize> {
if src_keys.is_empty() || (op == BitOp::Not && src_keys.len() != 1) {
return Err(Error::invalid_data(
"ERR syntax error in BITOP or wrong number of arguments",
));
}
let kc = KeyComposer::new("default");
let dest_str = str::from_utf8(dest_key.as_ref()).unwrap_or("");
let dest_meta_k = kc.bm_meta(dest_str);
let now_ms = current_now_ms();
let mut src_metas: Vec<(&str, BitmapMeta)> = Vec::with_capacity(src_keys.len());
let mut max_bitmap_size = 0u64;
for sk in src_keys {
let sk_str = str::from_utf8(sk.as_ref()).unwrap_or("");
let raw_sk = kc.raw_key(sk_str);
if let Some(raw) = self.data.get(raw_sk.as_bytes())? {
let (expire_at, _) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
return Err(Error::wrong_type(ERR_WRONG_TYPE));
}
}
let sm_k = kc.bm_meta(sk_str);
if let Some(m_bytes) = self.meta.get(sm_k.as_bytes())?
&& let Some(m) = BitmapMeta::decode(&m_bytes)
&& !m.is_expired(now_ms)
{
max_bitmap_size = max_bitmap_size.max(m.base.size);
src_metas.push((sk_str, m));
}
}
let mut batch = self.db.batch();
let old_dest_meta = self
.meta
.get(dest_meta_k.as_bytes())?
.and_then(|b| BitmapMeta::decode(&b));
if max_bitmap_size == 0 {
batch.remove(&self.meta, dest_meta_k.as_bytes());
if let Some(old_m) = old_dest_meta {
let old_stop_seg =
(old_m.base.size.saturating_sub(1) as usize) / BITMAP_SEGMENT_BYTES;
for seg_idx in 0..=old_stop_seg {
let seg_k = kc.bm_segment(dest_str, seg_idx as u32);
batch.remove(&self.data, seg_k.as_bytes());
}
}
batch.commit()?;
return Ok(0);
}
let can_skip_op = op == BitOp::And && src_metas.len() != src_keys.len();
if can_skip_op {
if let Some(old_m) = old_dest_meta {
let old_stop_seg =
(old_m.base.size.saturating_sub(1) as usize) / BITMAP_SEGMENT_BYTES;
for seg_idx in 0..=old_stop_seg {
let seg_k = kc.bm_segment(dest_str, seg_idx as u32);
batch.remove(&self.data, seg_k.as_bytes());
}
}
let dest_meta = BitmapMeta::new_with_version(0, max_bitmap_size);
batch.insert(&self.meta, dest_meta_k.as_bytes(), dest_meta.encode());
batch.commit()?;
return Ok(max_bitmap_size as usize);
}
let stop_seg_index = (max_bitmap_size.saturating_sub(1) as usize) / BITMAP_SEGMENT_BYTES;
let mut frag_res = [0u8; BITMAP_SEGMENT_BYTES];
for frag_idx in 0..=stop_seg_index {
let mut fragments: Vec<fjall::Slice> = Vec::with_capacity(src_metas.len());
let mut frag_maxlen = 0usize;
for (sk_str, _) in &src_metas {
let sub_k = kc.bm_segment(sk_str, frag_idx as u32);
let frag = self.data.get(sub_k.as_bytes())?.unwrap_or_default();
if frag.is_empty() {
if op == BitOp::And {
frag_maxlen = 0;
break;
}
} else {
frag_maxlen = frag_maxlen.max(frag.len());
}
fragments.push(frag);
}
let dest_sub_k = kc.bm_segment(dest_str, frag_idx as u32);
if frag_maxlen != 0 || op == BitOp::Not {
let write_len = if op == BitOp::Not {
if frag_idx == stop_seg_index {
if max_bitmap_size.is_multiple_of(BITMAP_SEGMENT_BYTES as u64) {
BITMAP_SEGMENT_BYTES
} else {
(max_bitmap_size % (BITMAP_SEGMENT_BYTES as u64)) as usize
}
} else {
BITMAP_SEGMENT_BYTES
}
} else {
frag_maxlen
};
let frag_slices: Vec<&[u8]> = fragments.iter().map(|f| &**f).collect();
bit_op_exec_into(op, &frag_slices, &mut frag_res[..write_len])?;
batch.insert(&self.data, dest_sub_k.as_bytes(), &frag_res[..write_len]);
} else {
batch.remove(&self.data, dest_sub_k.as_bytes());
}
}
if let Some(old_m) = old_dest_meta {
let old_stop_seg = (old_m.base.size.saturating_sub(1) as usize) / BITMAP_SEGMENT_BYTES;
if old_stop_seg > stop_seg_index {
for seg_idx in (stop_seg_index + 1)..=old_stop_seg {
let seg_k = kc.bm_segment(dest_str, seg_idx as u32);
batch.remove(&self.data, seg_k.as_bytes());
}
}
}
let dest_meta = BitmapMeta::new_with_version(0, max_bitmap_size);
batch.insert(&self.meta, dest_meta_k.as_bytes(), dest_meta.encode());
batch.commit()?;
Ok(max_bitmap_size as usize)
}
pub fn bitfield<K: AsRef<[u8]>>(
&self,
key: K,
ops: &[BitfieldOperation],
) -> Result<Vec<Option<BitfieldValue>>> {
self.exec_bitfield(key, ops, false)
}
pub fn bitfield_read_only<K: AsRef<[u8]>>(
&self,
key: K,
ops: &[BitfieldOperation],
) -> Result<Vec<Option<BitfieldValue>>> {
for op in ops {
if op.op_type != BitfieldOpType::Get {
return Err(Error::invalid_data(
"ERR BITFIELD_RO only supports the GET subcmd",
));
}
}
self.exec_bitfield(key, ops, true)
}
fn exec_bitfield<K: AsRef<[u8]>>(
&self,
key: K,
ops: &[BitfieldOperation],
read_only: bool,
) -> Result<Vec<Option<BitfieldValue>>> {
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
let mut str_bytes = val.to_vec();
let mut results = Vec::with_capacity(ops.len());
let mut modified = false;
let mut view = ArrayBitfieldBitmap::default();
for op in ops {
let bit_offset = op.offset;
let bits = op.encoding.bits();
let first_byte = (bit_offset / 8) as usize;
let last_byte = ((bit_offset + bits as u64 - 1) / 8) as usize;
if op.op_type != BitfieldOpType::Get && last_byte >= str_bytes.len() {
str_bytes.resize(last_byte + 1, 0);
modified = true;
}
view.set_byte_offset(first_byte as u32);
view.reset();
let copy_len = if first_byte < str_bytes.len() {
(str_bytes.len() - first_byte).min(ArrayBitfieldBitmap::SIZE)
} else {
0
};
if copy_len > 0 {
view.set(
first_byte as u32,
&str_bytes[first_byte..first_byte + copy_len],
)?;
}
let old_raw = if op.encoding.is_signed() {
view.get_signed_bitfield(bit_offset, bits)? as u64
} else {
view.get_unsigned_bitfield(bit_offset, bits)?
};
let (ret, new_raw, _) = bitfield_op_calc(op, old_raw);
results.push(ret);
if op.op_type != BitfieldOpType::Get && !read_only && ret.is_some() {
view.set_bitfield(bit_offset, bits, new_raw)?;
let write_bytes =
(last_byte - first_byte + 1).min(ArrayBitfieldBitmap::SIZE);
view.get(
first_byte as u32,
&mut str_bytes[first_byte..first_byte + write_bytes],
)?;
modified = true;
}
}
if modified && !read_only {
let enc_val = encode_string_value(&str_bytes, expire_at);
self.data.insert(&*raw_k, enc_val)?;
}
return Ok(results);
}
}
let bm_meta_k = kc.bm_meta(k_str);
let mut meta = match self.meta.get(bm_meta_k.as_bytes())? {
Some(b) => match BitmapMeta::decode(&b) {
Some(m) if !m.is_expired(now_ms) => m,
_ => BitmapMeta::new_with_version(0, 0),
},
None => BitmapMeta::new_with_version(0, 0),
};
let mut store = SegmentCacheStore::new(self, kc, k_str);
let mut results = Vec::with_capacity(ops.len());
let mut max_bytes = meta.base.size;
let mut has_changes = false;
let mut view = ArrayBitfieldBitmap::default();
for op in ops {
let bit_offset = op.offset;
let bits = op.encoding.bits();
let first_byte = (bit_offset / 8) as u32;
let last_byte = ((bit_offset + bits as u64 - 1) / 8) as u32;
let req_bytes = (last_byte + 1) as u64;
if op.op_type != BitfieldOpType::Get {
max_bytes = max_bytes.max(req_bytes);
}
let first_seg = first_byte / (BITMAP_SEGMENT_BYTES as u32);
let last_seg = last_byte / (BITMAP_SEGMENT_BYTES as u32);
view.set_byte_offset(first_byte);
view.reset();
for s_idx in first_seg..=last_seg {
let seg_slice = store.get(s_idx)?;
let seg_base_byte = s_idx * (BITMAP_SEGMENT_BYTES as u32);
let seg_end_byte = seg_base_byte + (BITMAP_SEGMENT_BYTES as u32);
let inter_start = first_byte.max(seg_base_byte);
let inter_end = (last_byte + 1).min(seg_end_byte);
if inter_start < inter_end {
let seg_rel_start = (inter_start - seg_base_byte) as usize;
let seg_rel_end = (inter_end - seg_base_byte) as usize;
let slice_len = seg_rel_end - seg_rel_start;
let mut msb_slice = [0u8; ArrayBitfieldBitmap::SIZE];
if seg_rel_start < seg_slice.len() {
let avail_end = seg_rel_end.min(seg_slice.len());
let copy_cnt = avail_end - seg_rel_start;
for (dst, &src) in msb_slice[..copy_cnt]
.iter_mut()
.zip(&seg_slice[seg_rel_start..avail_end])
{
*dst = src.reverse_bits();
}
}
view.set(inter_start, &msb_slice[..slice_len])?;
}
}
let old_raw = if op.encoding.is_signed() {
view.get_signed_bitfield(bit_offset, bits)? as u64
} else {
view.get_unsigned_bitfield(bit_offset, bits)?
};
let (ret, new_raw, _) = bitfield_op_calc(op, old_raw);
results.push(ret);
if op.op_type != BitfieldOpType::Get && !read_only && ret.is_some() {
view.set_bitfield(bit_offset, bits, new_raw)?;
for s_idx in first_seg..=last_seg {
let seg_base_byte = s_idx * (BITMAP_SEGMENT_BYTES as u32);
let seg_end_byte = seg_base_byte + (BITMAP_SEGMENT_BYTES as u32);
let inter_start = first_byte.max(seg_base_byte);
let inter_end = (last_byte + 1).min(seg_end_byte);
if inter_start < inter_end {
let seg_rel_start = (inter_start - seg_base_byte) as usize;
let seg_rel_end = (inter_end - seg_base_byte) as usize;
let slice_len = seg_rel_end - seg_rel_start;
let mut msb_slice = [0u8; ArrayBitfieldBitmap::SIZE];
view.get(inter_start, &mut msb_slice[..slice_len])?;
let seg = store.get_mut(s_idx, seg_rel_end)?;
for (dst, &src) in seg[seg_rel_start..seg_rel_end]
.iter_mut()
.zip(&msb_slice[..slice_len])
{
*dst = src.reverse_bits();
}
has_changes = true;
}
}
}
}
if !read_only && (has_changes || max_bytes > meta.base.size) {
let mut batch = self.db.batch();
for (seg_idx, seg_data) in store.dirty_segments() {
let seg_k = kc.bm_segment(k_str, seg_idx);
batch.insert(&self.data, seg_k.as_bytes(), seg_data);
}
meta.base.size = max_bytes;
batch.insert(&self.meta, bm_meta_k.as_bytes(), meta.encode());
batch.commit()?;
}
Ok(results)
}
pub fn get_bitmap_bytes<K: AsRef<[u8]>>(&self, key: K) -> Result<Option<Vec<u8>>> {
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let now_ms = current_now_ms();
let bm_meta_k = kc.bm_meta(k_str);
if let Some(m_bytes) = self.meta.get(bm_meta_k.as_bytes())?
&& let Some(meta) = BitmapMeta::decode(&m_bytes)
{
if meta.is_expired(now_ms) || meta.is_empty() {
return Ok(None);
}
let total_size = meta.base.size as usize;
if total_size > MAX_BITMAP_TO_STRING_BYTES {
return Err(Error::invalid_data(
"The size of the bitmap string exceeds configuration max-bitmap-to-string-mb (512MB)",
));
}
let mut out = vec![0u8; total_size];
let stop_seg = (total_size.saturating_sub(1)) / BITMAP_SEGMENT_BYTES;
for seg_idx in 0..=stop_seg {
let seg_k = kc.bm_segment(k_str, seg_idx as u32);
if let Some(seg_bytes) = self.data.get(seg_k.as_bytes())? {
let seg_start = seg_idx * BITMAP_SEGMENT_BYTES;
let copy_len = seg_bytes.len().min(total_size.saturating_sub(seg_start));
for (dst, &src) in out[seg_start..seg_start + copy_len]
.iter_mut()
.zip(&seg_bytes[..copy_len])
{
*dst = src.reverse_bits();
}
}
}
return Ok(Some(out));
}
let raw_k = kc.raw_key_bytes(key_bytes);
if let Some(raw) = self.data.get(&*raw_k)? {
let (expire_at, val) = decode_string_value(&raw);
if !is_string_expired(expire_at, now_ms) {
return Ok(Some(val.to_vec()));
}
}
Ok(None)
}
pub fn del_bitmap<K: AsRef<[u8]>>(&self, key: K) -> Result<bool> {
let key_bytes = key.as_ref();
let kc = KeyComposer::new("default");
let k_str = str::from_utf8(key_bytes).unwrap_or("");
let bm_meta_k = kc.bm_meta(k_str);
let mut deleted = false;
let mut batch = self.db.batch();
if let Some(m_bytes) = self.meta.get(bm_meta_k.as_bytes())?
&& let Some(meta) = BitmapMeta::decode(&m_bytes)
{
batch.remove(&self.meta, bm_meta_k.as_bytes());
let stop_seg = (meta.base.size.saturating_sub(1) as usize) / BITMAP_SEGMENT_BYTES;
for seg_idx in 0..=stop_seg {
let seg_k = kc.bm_segment(k_str, seg_idx as u32);
batch.remove(&self.data, seg_k.as_bytes());
}
deleted = true;
}
let raw_k = kc.raw_key_bytes(key_bytes);
if self.data.get(&*raw_k)?.is_some() {
batch.remove(&self.data, &*raw_k);
deleted = true;
}
if deleted {
batch.commit()?;
}
Ok(deleted)
}
}