Skip to main content

wedb_embed/api/bitmap/
bitops.rs

1pub use super::r#const::*;
2use crate::{
3  bitmap::opt::{
4    BitOp, BitfieldEncoding, BitfieldOpType, BitfieldOperation, BitfieldOverflow, BitfieldValue,
5  },
6  error::{Error, Result},
7};
8
9/// Computes segment index for the given bit offset (aligned with Kvrocks SegmentSubKeyIndexForBit).
10/// 计算指定位偏移所属的分段索引(对标 Kvrocks SegmentSubKeyIndexForBit / kBitmapSegmentBits)
11#[inline]
12pub const fn segment_index_for_bit(bit_offset: u64) -> u32 {
13  (bit_offset / (BITMAP_SEGMENT_BITS as u64)) as u32
14}
15
16/// Computes starting byte offset of segment for the given bit offset.
17/// 计算指定位偏移所属分段的字节起点偏移
18#[inline]
19pub const fn segment_byte_offset_for_bit(bit_offset: u64) -> u32 {
20  segment_index_for_bit(bit_offset) * (BITMAP_SEGMENT_BYTES as u32)
21}
22
23/// Expands bitmap segment capacity up to min_bytes (aligned with Kvrocks ExpandBitmapSegment).
24/// 扩展分段字节容量至 min_bytes(按需倍增至 1024 字节,对标 Kvrocks ExpandBitmapSegment)
25#[inline]
26pub fn expand_bitmap_segment(segment: &mut Vec<u8>, min_bytes: usize) {
27  debug_assert!(min_bytes <= BITMAP_SEGMENT_BYTES);
28  let old_size = segment.len();
29  if min_bytes > old_size {
30    let new_size = (old_size * 2).clamp(min_bytes, BITMAP_SEGMENT_BYTES);
31    segment.resize(new_size, 0);
32  }
33}
34
35/// Gets bit value in segment using LSB order (aligned with Kvrocks util::lsb::GetBit).
36/// 获取分段中指定位的值(LSB 顺序,对标 Apache Kvrocks util::lsb::GetBit)
37#[inline]
38pub fn get_bit_lsb(segment: &[u8], bit_offset_in_segment: usize) -> u8 {
39  let byte_idx = bit_offset_in_segment >> 3;
40  if byte_idx < segment.len() {
41    (segment[byte_idx] >> (bit_offset_in_segment & 7)) & 1
42  } else {
43    0
44  }
45}
46
47/// Sets bit value in segment using LSB order, returning old bit (aligned with Kvrocks SetBitTo).
48/// 设置分段中指定位的值(LSB 顺序,返回原位值,对标 Apache Kvrocks util::lsb::SetBitTo)
49#[inline]
50pub fn set_bit_lsb(segment: &mut [u8], bit_offset_in_segment: usize, bit: u8) -> u8 {
51  let byte_idx = bit_offset_in_segment >> 3;
52  let shift = bit_offset_in_segment & 7;
53  let old = (segment[byte_idx] >> shift) & 1;
54  if bit != 0 {
55    segment[byte_idx] |= 1 << shift;
56  } else {
57    segment[byte_idx] &= !(1 << shift);
58  }
59  old
60}
61
62/// Gets bit value in byte slice using MSB order.
63/// 获取连续字节中指定位的值(MSB 顺序,用于标准 Redis 兼容格式)
64#[inline]
65pub fn get_bit_from_bytes(bytes: &[u8], bit_offset: usize) -> u8 {
66  let byte_idx = bit_offset >> 3;
67  if byte_idx < bytes.len() {
68    (bytes[byte_idx] >> (7 - (bit_offset & 7))) & 1
69  } else {
70    0
71  }
72}
73
74/// Operation definition.
75/// 设置连续字节中指定位的值(MSB 顺序,用于标准 Redis 兼容格式)
76#[inline]
77pub fn set_bit_in_bytes(bytes: &mut Vec<u8>, bit_offset: usize, bit: u8) -> u8 {
78  let byte_idx = bit_offset >> 3;
79  if byte_idx >= bytes.len() {
80    bytes.resize(byte_idx + 1, 0);
81  }
82  let shift = 7 - (bit_offset & 7);
83  let old = (bytes[byte_idx] >> shift) & 1;
84  if bit != 0 {
85    bytes[byte_idx] |= 1 << shift;
86  } else {
87    bytes[byte_idx] &= !(1 << shift);
88  }
89  old
90}
91
92/// Domain operation (aligned with Apache Kvrocks util::msb::RawBitpos).
93/// 高性能大端序 MSB 位检索(对标 Apache Kvrocks util::msb::RawBitpos)
94#[inline]
95pub fn raw_bitpos(bytes: &[u8], bit: u8) -> Option<usize> {
96  let mut offset = 0usize;
97  let (chunks, remainder) = bytes.as_chunks::<8>();
98  for chunk in chunks {
99    let word = u64::from_be_bytes(*chunk);
100    if bit == 1 {
101      if word != 0 {
102        return Some(offset + word.leading_zeros() as usize);
103      }
104    } else if word != u64::MAX {
105      return Some(offset + (!word).leading_zeros() as usize);
106    }
107    offset += 64;
108  }
109  for &b in remainder {
110    if bit == 1 {
111      if b != 0 {
112        return Some(offset + b.leading_zeros() as usize);
113      }
114    } else if b != 0xFF {
115      return Some(offset + (!b).leading_zeros() as usize);
116    }
117    offset += 8;
118  }
119  None
120}
121
122/// Domain operation (aligned with util::lsb).
123/// 高性能小端序 LSB 位检索(用于 Kvrocks Bitmap 分段原生存储加速,对标 util::lsb)
124#[inline]
125pub fn raw_bitpos_lsb(bytes: &[u8], bit: u8) -> Option<usize> {
126  let mut offset = 0usize;
127  let (chunks, remainder) = bytes.as_chunks::<8>();
128  for chunk in chunks {
129    let word = u64::from_le_bytes(*chunk);
130    if bit == 1 {
131      if word != 0 {
132        return Some(offset + word.trailing_zeros() as usize);
133      }
134    } else if word != u64::MAX {
135      return Some(offset + (!word).trailing_zeros() as usize);
136    }
137    offset += 64;
138  }
139  for &b in remainder {
140    if bit == 1 {
141      if b != 0 {
142        return Some(offset + b.trailing_zeros() as usize);
143      }
144    } else if b != 0xFF {
145      return Some(offset + (!b).trailing_zeros() as usize);
146    }
147    offset += 8;
148  }
149  None
150}
151
152/// Operation definition.
153/// 在单字节中按 LSB 顺序查找指定区间 [start_bit, stop_bit] 内首个目标位(O(1) 零分支快速位运算)
154#[inline]
155pub const fn find_bit_in_byte_lsb(
156  b: u8,
157  bit: u8,
158  start_bit: usize,
159  stop_bit: usize,
160) -> Option<usize> {
161  debug_assert!(start_bit <= stop_bit && stop_bit < 8);
162  let mask = (((1u16 << (stop_bit - start_bit + 1)) - 1) as u8) << start_bit;
163  let target = if bit == 1 { b } else { !b };
164  let masked = target & mask;
165  if masked != 0 {
166    Some(masked.trailing_zeros() as usize)
167  } else {
168    None
169  }
170}
171
172/// Operation definition.
173/// 在单字节中按 MSB 顺序查找指定区间 [start_bit, stop_bit] 内首个目标位(O(1) 零分支快速位运算)
174#[inline]
175pub const fn find_bit_in_byte_msb(
176  b: u8,
177  bit: u8,
178  start_bit: usize,
179  stop_bit: usize,
180) -> Option<usize> {
181  debug_assert!(start_bit <= stop_bit && stop_bit < 8);
182  let mask = (((1u16 << (stop_bit - start_bit + 1)) - 1) as u8) << (7 - stop_bit);
183  let target = if bit == 1 { b } else { !b };
184  let masked = target & mask;
185  if masked != 0 {
186    Some(masked.leading_zeros() as usize)
187  } else {
188    None
189  }
190}
191
192/// Domain operation (aligned with Apache Kvrocks util::RawPopcount).
193/// 高性能 64 位原生 CPU POPCNT 位统计(对标 Apache Kvrocks util::RawPopcount)
194#[inline]
195pub fn raw_popcount(bytes: &[u8]) -> u64 {
196  let mut count = 0u64;
197  let (chunks, remainder) = bytes.as_chunks::<8>();
198  for chunk in chunks {
199    let word = u64::from_ne_bytes(*chunk);
200    count += word.count_ones() as u64;
201  }
202  for &b in remainder {
203    count += b.count_ones() as u64;
204  }
205  count
206}
207
208pub use crate::meta::normalize_bitmap_range as normalize_range;
209
210/// Domain operation (aligned with Kvrocks NormalizeToByteRangeWithPaddingMask).
211/// 位图位索引标准化为字节范围与位掩码(对标 Kvrocks NormalizeToByteRangeWithPaddingMask)
212#[inline]
213pub const fn normalize_bit_range_to_byte_mask(
214  start_bit: i64,
215  end_bit: i64,
216) -> (usize, usize, u8, u8) {
217  debug_assert!(start_bit <= end_bit);
218  let first_byte_neg_mask = (!((1u16 << (8 - (start_bit & 7))) - 1)) as u8;
219  let last_byte_neg_mask = ((1u16 << (7 - (end_bit & 7))) - 1) as u8;
220  let start_byte = (start_bit >> 3) as usize;
221  let end_byte = (end_bit >> 3) as usize;
222  (
223    start_byte,
224    end_byte,
225    first_byte_neg_mask,
226    last_byte_neg_mask,
227  )
228}
229
230/// Domain operation (aligned with Kvrocks).
231/// 支持字节与位索引的范围与掩码归一化(对标 Kvrocks)
232#[inline]
233pub const fn normalize_to_byte_range_with_padding_mask(
234  is_bit_index: bool,
235  start: i64,
236  end: i64,
237) -> (usize, usize, u8, u8) {
238  if is_bit_index {
239    normalize_bit_range_to_byte_mask(start, end)
240  } else {
241    (start as usize, end as usize, 0, 0)
242  }
243}
244
245/// Domain operation (aligned with Kvrocks ArrayBitfieldBitmap).
246/// 9 字节局部小缓冲结构,用于跨分段高精度读取和写入 Bitfield(对标 Kvrocks ArrayBitfieldBitmap)
247#[derive(Debug, Clone)]
248pub struct ArrayBitfieldBitmap {
249  pub buf: [u8; 9],
250  pub byte_offset: u32,
251}
252
253impl Default for ArrayBitfieldBitmap {
254  fn default() -> Self {
255    Self::new(0)
256  }
257}
258
259impl ArrayBitfieldBitmap {
260  pub const SIZE: usize = 9;
261
262  #[inline]
263  pub const fn new(byte_offset: u32) -> Self {
264    Self {
265      buf: [0u8; Self::SIZE],
266      byte_offset,
267    }
268  }
269
270  #[inline]
271  pub fn set_byte_offset(&mut self, byte_offset: u32) {
272    self.byte_offset = byte_offset;
273  }
274
275  #[inline]
276  pub fn reset(&mut self) {
277    self.buf.fill(0);
278  }
279
280  #[inline]
281  pub fn set(&mut self, byte_offset: u32, src: &[u8]) -> Result<()> {
282    let bytes = src.len();
283    if byte_offset < self.byte_offset
284      || (byte_offset + bytes as u32) > (self.byte_offset + Self::SIZE as u32)
285    {
286      return Err(Error::invalid_data(
287        "The range [offset, offset + bytes) is out of bitfield buffer",
288      ));
289    }
290    let rel_offset = (byte_offset - self.byte_offset) as usize;
291    self.buf[rel_offset..rel_offset + bytes].copy_from_slice(src);
292    Ok(())
293  }
294
295  #[inline]
296  pub fn get(&self, byte_offset: u32, dst: &mut [u8]) -> Result<()> {
297    let bytes = dst.len();
298    if byte_offset < self.byte_offset
299      || (byte_offset + bytes as u32) > (self.byte_offset + Self::SIZE as u32)
300    {
301      return Err(Error::invalid_data(
302        "The range [offset, offset + bytes) is out of bitfield buffer",
303      ));
304    }
305    let rel_offset = (byte_offset - self.byte_offset) as usize;
306    dst.copy_from_slice(&self.buf[rel_offset..rel_offset + bytes]);
307    Ok(())
308  }
309
310  #[inline]
311  pub fn get_unsigned_bitfield(&self, bit_offset: u64, bits: u8) -> Result<u64> {
312    if bits == 0 || bits > 63 {
313      return Err(Error::invalid_data("Invalid unsigned bits (1..=63)"));
314    }
315    self.read_raw_bitfield(bit_offset, bits)
316  }
317
318  #[inline]
319  pub fn get_signed_bitfield(&self, bit_offset: u64, bits: u8) -> Result<i64> {
320    if bits == 0 || bits > 64 {
321      return Err(Error::invalid_data("Invalid signed bits (1..=64)"));
322    }
323    let raw = self.read_raw_bitfield(bit_offset, bits)?;
324    let mut val = raw as i64;
325    let msb = 1u64 << (bits - 1);
326    if (raw & msb) != 0 && bits < 64 {
327      let mask = u64::MAX << bits;
328      val |= mask as i64;
329    }
330    Ok(val)
331  }
332
333  #[inline]
334  fn read_raw_bitfield(&self, bit_offset: u64, bits: u8) -> Result<u64> {
335    let first_byte = (bit_offset / 8) as u32;
336    let last_byte = ((bit_offset + bits as u64 - 1) / 8 + 1) as u32;
337    let bytes = (last_byte - first_byte) as usize;
338
339    if first_byte < self.byte_offset
340      || (first_byte + bytes as u32) > (self.byte_offset + Self::SIZE as u32)
341    {
342      return Err(Error::invalid_data("Bitfield range out of buffer"));
343    }
344
345    let rel_bit_offset = (bit_offset - (self.byte_offset as u64 * 8)) as usize;
346    let mut word_bytes = [0u8; 16];
347    word_bytes[7..16].copy_from_slice(&self.buf);
348    let word = u128::from_be_bytes(word_bytes);
349    let shift = 72 - rel_bit_offset - (bits as usize);
350    let mask = if bits == 64 {
351      u64::MAX
352    } else {
353      (1u64 << bits) - 1
354    };
355    Ok(((word >> shift) as u64) & mask)
356  }
357
358  #[inline]
359  pub fn set_bitfield(&mut self, bit_offset: u64, bits: u8, value: u64) -> Result<()> {
360    let first_byte = (bit_offset / 8) as u32;
361    let last_byte = ((bit_offset + bits as u64 - 1) / 8 + 1) as u32;
362    let bytes = (last_byte - first_byte) as usize;
363
364    if first_byte < self.byte_offset
365      || (first_byte + bytes as u32) > (self.byte_offset + Self::SIZE as u32)
366    {
367      return Err(Error::invalid_data("Bitfield range out of buffer"));
368    }
369
370    let rel_bit_offset = (bit_offset - (self.byte_offset as u64 * 8)) as usize;
371    let mut word_bytes = [0u8; 16];
372    word_bytes[7..16].copy_from_slice(&self.buf);
373    let mut word = u128::from_be_bytes(word_bytes);
374    let shift = 72 - rel_bit_offset - (bits as usize);
375    let bit_mask = if bits == 64 {
376      u64::MAX as u128
377    } else {
378      (1u128 << bits) - 1
379    };
380    let mask = bit_mask << shift;
381    let val = ((value as u128) & bit_mask) << shift;
382    word = (word & !mask) | val;
383    let updated_bytes = word.to_be_bytes();
384    self.buf.copy_from_slice(&updated_bytes[7..16]);
385    Ok(())
386  }
387}
388
389/// Domain operation (aligned with Kvrocks detail::SignedBitfieldPlus).
390/// 有符号 BITFIELD 溢出加法运算(对标 Kvrocks detail::SignedBitfieldPlus)
391#[inline]
392pub fn signed_bitfield_plus(
393  value: u64,
394  incr: i64,
395  bits: u8,
396  overflow: BitfieldOverflow,
397) -> (u64, bool) {
398  let max = if bits == 64 {
399    i64::MAX
400  } else {
401    (1i64 << (bits - 1)) - 1
402  };
403  let min = -max - 1;
404
405  let signed_val = value as i64;
406  let max_incr = (max as u64).wrapping_sub(value) as i64;
407  let min_incr = min.wrapping_sub(signed_val);
408
409  if signed_val > max
410    || (bits != 64 && incr > max_incr)
411    || (signed_val >= 0 && incr >= 0 && incr > max_incr)
412  {
413    match overflow {
414      BitfieldOverflow::Wrap => (wrapped_signed_bitfield_plus(value, incr, bits), true),
415      BitfieldOverflow::Sat => (max as u64, true),
416      BitfieldOverflow::Fail => (0, true),
417    }
418  } else if signed_val < min
419    || (bits != 64 && incr < min_incr)
420    || (signed_val < 0 && incr < 0 && incr < min_incr)
421  {
422    match overflow {
423      BitfieldOverflow::Wrap => (wrapped_signed_bitfield_plus(value, incr, bits), true),
424      BitfieldOverflow::Sat => (min as u64, true),
425      BitfieldOverflow::Fail => (0, true),
426    }
427  } else {
428    (signed_val.wrapping_add(incr) as u64, false)
429  }
430}
431
432#[inline]
433const fn wrapped_signed_bitfield_plus(value: u64, incr: i64, bits: u8) -> u64 {
434  let res = value.wrapping_add(incr as u64);
435  if bits < 64 {
436    let mask = u64::MAX << bits;
437    if (res & (1u64 << (bits - 1))) != 0 {
438      res | mask
439    } else {
440      res & !mask
441    }
442  } else {
443    res
444  }
445}
446
447/// Domain operation (aligned with Kvrocks detail::UnsignedBitfieldPlus).
448/// 无符号 BITFIELD 溢出加法运算(对标 Kvrocks detail::UnsignedBitfieldPlus)
449#[inline]
450pub fn unsigned_bitfield_plus(
451  value: u64,
452  incr: i64,
453  bits: u8,
454  overflow: BitfieldOverflow,
455) -> (u64, bool) {
456  let max = if bits == 64 {
457    u64::MAX
458  } else {
459    (1u64 << bits) - 1
460  };
461  let max_incr = max.wrapping_sub(value) as i64;
462  let min_incr = (!value).wrapping_add(1) as i64;
463
464  if value > max || (incr > 0 && incr > max_incr) {
465    match overflow {
466      BitfieldOverflow::Wrap => (wrapped_unsigned_bitfield_plus(value, incr, bits), true),
467      BitfieldOverflow::Sat => (max, true),
468      BitfieldOverflow::Fail => (0, true),
469    }
470  } else if incr < 0 && incr < min_incr {
471    match overflow {
472      BitfieldOverflow::Wrap => (wrapped_unsigned_bitfield_plus(value, incr, bits), true),
473      BitfieldOverflow::Sat => (0, true),
474      BitfieldOverflow::Fail => (0, true),
475    }
476  } else {
477    (value.wrapping_add(incr as u64), false)
478  }
479}
480
481#[inline]
482const fn wrapped_unsigned_bitfield_plus(value: u64, incr: i64, bits: u8) -> u64 {
483  let mask = if bits == 64 { 0 } else { u64::MAX << bits };
484  let res = value.wrapping_add(incr as u64);
485  res & !mask
486}
487
488/// Domain operation (aligned with Kvrocks BitfieldOp).
489/// 执行单步 BITFIELD 逻辑运算(对标 Kvrocks BitfieldOp)
490#[inline]
491pub fn bitfield_op_calc(
492  op: &BitfieldOperation,
493  old_value: u64,
494) -> (Option<BitfieldValue>, u64, bool) {
495  if op.op_type == BitfieldOpType::Get {
496    let val = if op.encoding.is_signed() {
497      BitfieldValue::Signed(old_value as i64)
498    } else {
499      BitfieldValue::Unsigned(old_value)
500    };
501    return (Some(val), old_value, false);
502  }
503
504  let (new_value, is_overflow) = match op.encoding {
505    BitfieldEncoding::Signed(bits) => {
506      let input_val = if op.op_type == BitfieldOpType::Set {
507        op.value as u64
508      } else {
509        old_value
510      };
511      let incr = if op.op_type == BitfieldOpType::Set {
512        0
513      } else {
514        op.value
515      };
516      signed_bitfield_plus(input_val, incr, bits, op.overflow)
517    }
518    BitfieldEncoding::Unsigned(bits) => {
519      let input_val = if op.op_type == BitfieldOpType::Set {
520        op.value as u64
521      } else {
522        old_value
523      };
524      let incr = if op.op_type == BitfieldOpType::Set {
525        0
526      } else {
527        op.value
528      };
529      unsigned_bitfield_plus(input_val, incr, bits, op.overflow)
530    }
531  };
532
533  if op.overflow == BitfieldOverflow::Fail && is_overflow {
534    return (None, old_value, true);
535  }
536
537  let returned_val = if op.op_type == BitfieldOpType::Set {
538    if op.encoding.is_signed() {
539      BitfieldValue::Signed(old_value as i64)
540    } else {
541      BitfieldValue::Unsigned(old_value)
542    }
543  } else if op.encoding.is_signed() {
544    BitfieldValue::Signed(new_value as i64)
545  } else {
546    BitfieldValue::Unsigned(new_value)
547  };
548
549  (Some(returned_val), new_value, false)
550}
551
552/// Operation definition.
553/// 64 位原生字向量化位与操作
554#[inline]
555pub fn bitwise_and(dst: &mut [u8], src: &[u8]) {
556  let common_len = dst.len().min(src.len());
557  let (dst_chunks, dst_rem) = dst[..common_len].as_chunks_mut::<8>();
558  let (src_chunks, src_rem) = src[..common_len].as_chunks::<8>();
559
560  for (d, s) in dst_chunks.iter_mut().zip(src_chunks.iter()) {
561    let dw = u64::from_ne_bytes(*d);
562    let sw = u64::from_ne_bytes(*s);
563    *d = (dw & sw).to_ne_bytes();
564  }
565  for (d, s) in dst_rem.iter_mut().zip(src_rem.iter()) {
566    *d &= *s;
567  }
568  dst[common_len..].fill(0);
569}
570
571/// Operation definition.
572/// 64 位原生字向量化位或操作
573#[inline]
574pub fn bitwise_or(dst: &mut [u8], src: &[u8]) {
575  let len = dst.len().min(src.len());
576  let (dst_chunks, dst_rem) = dst[..len].as_chunks_mut::<8>();
577  let (src_chunks, src_rem) = src[..len].as_chunks::<8>();
578
579  for (d, s) in dst_chunks.iter_mut().zip(src_chunks.iter()) {
580    let dw = u64::from_ne_bytes(*d);
581    let sw = u64::from_ne_bytes(*s);
582    *d = (dw | sw).to_ne_bytes();
583  }
584  for (d, s) in dst_rem.iter_mut().zip(src_rem.iter()) {
585    *d |= *s;
586  }
587}
588
589/// Operation definition.
590/// 64 位原生字向量化位异或操作
591#[inline]
592pub fn bitwise_xor(dst: &mut [u8], src: &[u8]) {
593  let len = dst.len().min(src.len());
594  let (dst_chunks, dst_rem) = dst[..len].as_chunks_mut::<8>();
595  let (src_chunks, src_rem) = src[..len].as_chunks::<8>();
596
597  for (d, s) in dst_chunks.iter_mut().zip(src_chunks.iter()) {
598    let dw = u64::from_ne_bytes(*d);
599    let sw = u64::from_ne_bytes(*s);
600    *d = (dw ^ sw).to_ne_bytes();
601  }
602  for (d, s) in dst_rem.iter_mut().zip(src_rem.iter()) {
603    *d ^= *s;
604  }
605}
606
607/// Operation definition.
608/// 64 位原生字向量化位非操作
609#[inline]
610pub fn bitwise_not(dst: &mut [u8], src: &[u8]) {
611  let len = dst.len().min(src.len());
612  let (dst_chunks, dst_rem) = dst[..len].as_chunks_mut::<8>();
613  let (src_chunks, src_rem) = src[..len].as_chunks::<8>();
614
615  for (d, s) in dst_chunks.iter_mut().zip(src_chunks.iter()) {
616    let sw = u64::from_ne_bytes(*s);
617    *d = (!sw).to_ne_bytes();
618  }
619  for (d, s) in dst_rem.iter_mut().zip(src_rem.iter()) {
620    *d = !*s;
621  }
622}
623
624/// Domain operation (aligned with Kvrocks Bitmap::BitOp).
625/// 高性能 64 位字切片位图操作(零堆分配写入给定缓冲,对标 Kvrocks Bitmap::BitOp)
626#[inline]
627pub fn bit_op_exec_into(op: BitOp, src_slices: &[&[u8]], out: &mut [u8]) -> Result<usize> {
628  let out_len = out.len();
629  match op {
630    BitOp::And => {
631      if let Some((first, rest)) = src_slices.split_first() {
632        let copy_len = first.len().min(out_len);
633        out[..copy_len].copy_from_slice(&first[..copy_len]);
634        out[copy_len..out_len].fill(0);
635        for &src in rest {
636          bitwise_and(out, src);
637        }
638      } else {
639        out.fill(0);
640      }
641    }
642    BitOp::Or => {
643      if let Some((first, rest)) = src_slices.split_first() {
644        let copy_len = first.len().min(out_len);
645        out[..copy_len].copy_from_slice(&first[..copy_len]);
646        out[copy_len..out_len].fill(0);
647        for &src in rest {
648          bitwise_or(out, src);
649        }
650      } else {
651        out.fill(0);
652      }
653    }
654    BitOp::Xor => {
655      if let Some((first, rest)) = src_slices.split_first() {
656        let copy_len = first.len().min(out_len);
657        out[..copy_len].copy_from_slice(&first[..copy_len]);
658        out[copy_len..out_len].fill(0);
659        for &src in rest {
660          bitwise_xor(out, src);
661        }
662      } else {
663        out.fill(0);
664      }
665    }
666    BitOp::Not => {
667      if src_slices.len() != 1 {
668        return Err(Error::invalid_data(
669          "ERR BITOP NOT takes exactly one source key",
670        ));
671      }
672      let src = src_slices[0];
673      let src_len = src.len().min(out_len);
674      bitwise_not(&mut out[..src_len], &src[..src_len]);
675      if out_len > src_len {
676        out[src_len..out_len].fill(0xFF);
677      }
678    }
679  }
680
681  Ok(out_len)
682}
683
684/// Operation definition.
685/// 高性能 64 位字切片位图操作(AND / OR / XOR / NOT)
686pub fn bit_op_exec(op: &str, src_slices: &[&[u8]]) -> Result<Vec<u8>> {
687  let bit_op = op.parse::<BitOp>()?;
688  let max_len = src_slices.iter().map(|s| s.len()).max().unwrap_or(0);
689  let mut out = vec![0u8; max_len];
690  let written = bit_op_exec_into(bit_op, src_slices, &mut out)?;
691  out.truncate(written);
692  Ok(out)
693}
694
695/// Domain operation (aligned with Apache Kvrocks BitmapString::BitCount).
696/// 字符串模式 BITCOUNT(MSB 顺序,对标 Apache Kvrocks BitmapString::BitCount)
697pub fn string_bitcount(
698  val: &[u8],
699  start: Option<i64>,
700  end: Option<i64>,
701  is_bit_index: bool,
702) -> u64 {
703  let strlen = val.len() as i64;
704  let totlen = if is_bit_index { strlen << 3 } else { strlen };
705  let s = start.unwrap_or(0);
706  let e = end.unwrap_or(-1);
707  if s < 0 && e < 0 && s > e {
708    return 0;
709  }
710  let (norm_s, norm_e) = normalize_range(s, e, totlen);
711  if norm_s > norm_e {
712    return 0;
713  }
714
715  let (start_byte, stop_byte, first_mask, last_mask) =
716    normalize_to_byte_range_with_padding_mask(is_bit_index, norm_s, norm_e);
717
718  if start_byte >= val.len() {
719    return 0;
720  }
721  let actual_stop = stop_byte.min(val.len().saturating_sub(1));
722  if start_byte > actual_stop {
723    return 0;
724  }
725
726  let bytes = &val[start_byte..=actual_stop];
727  let cnt = raw_popcount(bytes);
728
729  let mut mask_cnt = 0u64;
730  if first_mask != 0 && start_byte < val.len() {
731    mask_cnt += (val[start_byte] & first_mask).count_ones() as u64;
732  }
733  if last_mask != 0 && actual_stop == stop_byte && actual_stop < val.len() {
734    mask_cnt += (val[actual_stop] & last_mask).count_ones() as u64;
735  }
736  cnt.saturating_sub(mask_cnt)
737}
738
739/// Domain operation (aligned with Apache Kvrocks BitmapString::BitPos).
740/// 字符串模式 BITPOS(MSB 顺序,对标 Apache Kvrocks BitmapString::BitPos)
741pub fn string_bitpos(
742  val: &[u8],
743  bit: u8,
744  start: Option<i64>,
745  end: Option<i64>,
746  stop_given: bool,
747  is_bit_index: bool,
748) -> i64 {
749  let strlen = val.len() as i64;
750  let length = if is_bit_index { strlen * 8 } else { strlen };
751  let s = start.unwrap_or(0);
752  let e = end.unwrap_or(-1);
753  let (norm_s, norm_e) = normalize_range(s, e, length);
754  if norm_s > norm_e {
755    return -1;
756  }
757
758  let mut byte_start = (if is_bit_index { norm_s / 8 } else { norm_s }) as usize;
759  let byte_stop = (if is_bit_index { norm_e / 8 } else { norm_e }) as usize;
760  let bit_in_start_byte = if is_bit_index {
761    (norm_s % 8) as usize
762  } else {
763    0
764  };
765  let bit_in_stop_byte = if is_bit_index {
766    (norm_e % 8) as usize
767  } else {
768    7
769  };
770
771  if is_bit_index && byte_start == byte_stop {
772    if byte_start < val.len()
773      && let Some(bit_idx) =
774        find_bit_in_byte_msb(val[byte_start], bit, bit_in_start_byte, bit_in_stop_byte)
775    {
776      return (byte_start * 8 + bit_idx) as i64;
777    }
778    return -1;
779  }
780
781  if is_bit_index && bit_in_start_byte != 0 {
782    if byte_start < val.len()
783      && let Some(bit_idx) = find_bit_in_byte_msb(val[byte_start], bit, bit_in_start_byte, 7)
784    {
785      return (byte_start * 8 + bit_idx) as i64;
786    }
787    byte_start += 1;
788  }
789
790  if byte_start > byte_stop || byte_start >= val.len() {
791    return if stop_given && bit == 0 {
792      -1
793    } else if bit == 0 {
794      strlen * 8
795    } else {
796      -1
797    };
798  }
799
800  let actual_stop = byte_stop.min(val.len() - 1);
801  let bytes_cnt = actual_stop - byte_start + 1;
802  let pos_opt = raw_bitpos(&val[byte_start..byte_start + bytes_cnt], bit);
803
804  match pos_opt {
805    Some(pos) => {
806      let abs_pos = (pos + byte_start * 8) as i64;
807      if is_bit_index && abs_pos > norm_e {
808        return -1;
809      }
810      abs_pos
811    }
812    None => {
813      if stop_given && bit == 0 {
814        -1
815      } else if bit == 0 {
816        strlen * 8
817      } else {
818        -1
819      }
820    }
821  }
822}