use crate::bit_chunk_iterator::BitChunks;
use crate::bit_iterator::{BitIndexIterator, BitIndexU32Iterator, BitIterator, BitSliceIterator};
use crate::bit_util::read_u64;
use crate::{
BooleanBufferBuilder, Buffer, MutableBuffer, bit_util, buffer_bin_and, buffer_bin_or,
buffer_bin_xor,
};
use std::ops::{BitAnd, BitOr, BitXor, Not};
#[derive(Debug, Clone, Eq)]
pub struct BooleanBuffer {
buffer: Buffer,
bit_offset: usize,
bit_len: usize,
}
impl PartialEq for BooleanBuffer {
fn eq(&self, other: &Self) -> bool {
if self.bit_len != other.bit_len {
return false;
}
let lhs = self.bit_chunks().iter_padded();
let rhs = other.bit_chunks().iter_padded();
lhs.zip(rhs).all(|(a, b)| a == b)
}
}
impl BooleanBuffer {
pub fn new(buffer: Buffer, bit_offset: usize, bit_len: usize) -> Self {
let total_len = bit_offset.saturating_add(bit_len);
let buffer_len = buffer.len();
let buffer_bit_len = buffer_len.saturating_mul(8);
assert!(
total_len <= buffer_bit_len,
"buffer not large enough (bit_offset: {bit_offset}, bit_len: {bit_len}, buffer_len: {buffer_len})"
);
Self {
buffer,
bit_offset,
bit_len,
}
}
pub fn new_set(length: usize) -> Self {
let mut builder = BooleanBufferBuilder::new(length);
builder.append_n(length, true);
builder.finish()
}
pub fn new_unset(length: usize) -> Self {
let buffer = MutableBuffer::new_null(length).into_buffer();
Self {
buffer,
bit_offset: 0,
bit_len: length,
}
}
pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
let buffer = MutableBuffer::collect_bool(len, f);
Self::new(buffer.into(), 0, len)
}
pub fn from_bits(src: impl AsRef<[u8]>, offset_in_bits: usize, len_in_bits: usize) -> Self {
Self::from_bitwise_unary_op(src, offset_in_bits, len_in_bits, |a| a)
}
pub fn from_bitwise_unary_op<F>(
src: impl AsRef<[u8]>,
offset_in_bits: usize,
len_in_bits: usize,
mut op: F,
) -> Self
where
F: FnMut(u64) -> u64,
{
let end = offset_in_bits + len_in_bits;
let aligned_offset = offset_in_bits & !63;
let aligned_end_bytes = bit_util::ceil(end, 64) * 8;
let src_len = src.as_ref().len();
let slice_end = aligned_end_bytes.min(src_len);
let aligned_start = &src.as_ref()[aligned_offset / 8..slice_end];
let (prefix, aligned_u64s, suffix) = unsafe { aligned_start.as_ref().align_to::<u64>() };
match (prefix, suffix) {
([], []) => {
let result_u64s: Vec<u64> = aligned_u64s.iter().map(|l| op(*l)).collect();
return BooleanBuffer::new(result_u64s.into(), offset_in_bits % 64, len_in_bits);
}
([], suffix) => {
let suffix = read_u64(suffix);
let result_u64s: Vec<u64> = aligned_u64s
.iter()
.cloned()
.chain(std::iter::once(suffix))
.map(&mut op)
.collect();
return BooleanBuffer::new(result_u64s.into(), offset_in_bits % 64, len_in_bits);
}
_ => {}
}
let chunks = aligned_start.chunks_exact(8);
let remainder = chunks.remainder();
let iter = chunks.map(|c| u64::from_le_bytes(c.try_into().unwrap()));
let vec_u64s: Vec<u64> = if remainder.is_empty() {
iter.map(&mut op).collect()
} else {
iter.chain(Some(read_u64(remainder))).map(&mut op).collect()
};
BooleanBuffer::new(vec_u64s.into(), offset_in_bits % 64, len_in_bits)
}
pub fn from_bitwise_binary_op<F>(
left: impl AsRef<[u8]>,
left_offset_in_bits: usize,
right: impl AsRef<[u8]>,
right_offset_in_bits: usize,
len_in_bits: usize,
mut op: F,
) -> Self
where
F: FnMut(u64, u64) -> u64,
{
let left = left.as_ref();
let right = right.as_ref();
if left_offset_in_bits & 0x7 == 0 && right_offset_in_bits & 0x7 == 0 {
let left = &left[left_offset_in_bits / 8..];
let right = &right[right_offset_in_bits / 8..];
unsafe {
let (left_prefix, left_u64s, left_suffix) = left.align_to::<u64>();
let (right_prefix, right_u64s, right_suffix) = right.align_to::<u64>();
if left_prefix.is_empty()
&& right_prefix.is_empty()
&& left_suffix.is_empty()
&& right_suffix.is_empty()
{
let result_u64s = left_u64s
.iter()
.zip(right_u64s.iter())
.map(|(l, r)| op(*l, *r))
.collect::<Vec<u64>>();
return BooleanBuffer {
buffer: Buffer::from(result_u64s),
bit_offset: 0,
bit_len: len_in_bits,
};
}
}
}
let left_chunks = BitChunks::new(left, left_offset_in_bits, len_in_bits);
let right_chunks = BitChunks::new(right, right_offset_in_bits, len_in_bits);
let chunks = left_chunks
.iter()
.zip(right_chunks.iter())
.map(|(left, right)| op(left, right));
let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks) };
let remainder_bytes = bit_util::ceil(left_chunks.remainder_len(), 8);
let rem = op(left_chunks.remainder_bits(), right_chunks.remainder_bits());
let rem = &rem.to_le_bytes()[0..remainder_bytes];
buffer.extend_from_slice(rem);
BooleanBuffer {
buffer: Buffer::from(buffer),
bit_offset: 0,
bit_len: len_in_bits,
}
}
pub fn count_set_bits(&self) -> usize {
self.buffer
.count_set_bits_offset(self.bit_offset, self.bit_len)
}
pub fn find_nth_set_bit_position(&self, start: usize, n: usize) -> usize {
if n == 0 {
return start;
}
self.slice(start, self.bit_len - start)
.set_indices()
.nth(n - 1)
.map(|idx| start + idx + 1)
.unwrap_or(self.bit_len)
}
#[inline]
pub fn bit_chunks(&self) -> BitChunks<'_> {
BitChunks::new(self.values(), self.bit_offset, self.bit_len)
}
#[inline]
pub fn offset(&self) -> usize {
self.bit_offset
}
#[inline]
pub fn len(&self) -> usize {
self.bit_len
}
#[inline]
pub fn is_empty(&self) -> bool {
self.bit_len == 0
}
pub fn shrink_to_fit(&mut self) {
self.buffer.shrink_to_fit();
}
#[inline]
pub fn value(&self, idx: usize) -> bool {
assert!(idx < self.bit_len);
unsafe { self.value_unchecked(idx) }
}
#[inline]
pub unsafe fn value_unchecked(&self, i: usize) -> bool {
unsafe { bit_util::get_bit_raw(self.buffer.as_ptr(), i + self.bit_offset) }
}
#[inline]
pub fn values(&self) -> &[u8] {
&self.buffer
}
pub fn slice(&self, offset: usize, len: usize) -> Self {
assert!(
offset.saturating_add(len) <= self.bit_len,
"the length + offset of the sliced BooleanBuffer cannot exceed the existing length"
);
Self {
buffer: self.buffer.clone(),
bit_offset: self.bit_offset + offset,
bit_len: len,
}
}
pub fn sliced(&self) -> Buffer {
self.buffer.bit_slice(self.bit_offset, self.bit_len)
}
pub fn ptr_eq(&self, other: &Self) -> bool {
self.buffer.as_ptr() == other.buffer.as_ptr()
&& self.bit_offset == other.bit_offset
&& self.bit_len == other.bit_len
}
#[inline]
pub fn inner(&self) -> &Buffer {
&self.buffer
}
pub fn into_inner(self) -> Buffer {
self.buffer
}
pub fn iter(&self) -> BitIterator<'_> {
self.into_iter()
}
pub fn set_indices(&self) -> BitIndexIterator<'_> {
BitIndexIterator::new(self.values(), self.bit_offset, self.bit_len)
}
pub fn set_indices_u32(&self) -> BitIndexU32Iterator<'_> {
BitIndexU32Iterator::new(self.values(), self.bit_offset, self.bit_len)
}
pub fn set_slices(&self) -> BitSliceIterator<'_> {
BitSliceIterator::new(self.values(), self.bit_offset, self.bit_len)
}
}
impl Not for &BooleanBuffer {
type Output = BooleanBuffer;
fn not(self) -> Self::Output {
BooleanBuffer::from_bitwise_unary_op(&self.buffer, self.bit_offset, self.bit_len, |a| !a)
}
}
impl BitAnd<&BooleanBuffer> for &BooleanBuffer {
type Output = BooleanBuffer;
fn bitand(self, rhs: &BooleanBuffer) -> Self::Output {
assert_eq!(self.bit_len, rhs.bit_len);
BooleanBuffer {
buffer: buffer_bin_and(
&self.buffer,
self.bit_offset,
&rhs.buffer,
rhs.bit_offset,
self.bit_len,
),
bit_offset: 0,
bit_len: self.bit_len,
}
}
}
impl BitOr<&BooleanBuffer> for &BooleanBuffer {
type Output = BooleanBuffer;
fn bitor(self, rhs: &BooleanBuffer) -> Self::Output {
assert_eq!(self.bit_len, rhs.bit_len);
BooleanBuffer {
buffer: buffer_bin_or(
&self.buffer,
self.bit_offset,
&rhs.buffer,
rhs.bit_offset,
self.bit_len,
),
bit_offset: 0,
bit_len: self.bit_len,
}
}
}
impl BitXor<&BooleanBuffer> for &BooleanBuffer {
type Output = BooleanBuffer;
fn bitxor(self, rhs: &BooleanBuffer) -> Self::Output {
assert_eq!(self.bit_len, rhs.bit_len);
BooleanBuffer {
buffer: buffer_bin_xor(
&self.buffer,
self.bit_offset,
&rhs.buffer,
rhs.bit_offset,
self.bit_len,
),
bit_offset: 0,
bit_len: self.bit_len,
}
}
}
impl<'a> IntoIterator for &'a BooleanBuffer {
type Item = bool;
type IntoIter = BitIterator<'a>;
fn into_iter(self) -> Self::IntoIter {
BitIterator::new(self.values(), self.bit_offset, self.bit_len)
}
}
impl From<&[bool]> for BooleanBuffer {
fn from(value: &[bool]) -> Self {
let mut builder = BooleanBufferBuilder::new(value.len());
builder.append_slice(value);
builder.finish()
}
}
impl From<Vec<bool>> for BooleanBuffer {
fn from(value: Vec<bool>) -> Self {
value.as_slice().into()
}
}
impl FromIterator<bool> for BooleanBuffer {
fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
let iter = iter.into_iter();
let (hint, _) = iter.size_hint();
let mut builder = BooleanBufferBuilder::new(hint);
iter.for_each(|b| builder.append(b));
builder.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_boolean_new() {
let bytes = &[0, 1, 2, 3, 4];
let buf = Buffer::from(bytes);
let offset = 0;
let len = 24;
let boolean_buf = BooleanBuffer::new(buf.clone(), offset, len);
assert_eq!(bytes, boolean_buf.values());
assert_eq!(offset, boolean_buf.offset());
assert_eq!(len, boolean_buf.len());
assert_eq!(2, boolean_buf.count_set_bits());
assert_eq!(&buf, boolean_buf.inner());
assert_eq!(buf, boolean_buf.clone().into_inner());
assert!(!boolean_buf.is_empty())
}
#[test]
fn test_boolean_data_equality() {
let boolean_buf1 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 32);
let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 32);
assert_eq!(boolean_buf1, boolean_buf2);
let boolean_buf3 = boolean_buf1.slice(8, 16);
assert_ne!(boolean_buf1, boolean_buf3);
let boolean_buf4 = boolean_buf1.slice(0, 32);
assert_eq!(boolean_buf1, boolean_buf4);
let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 0, 2, 3, 4]), 0, 32);
assert_ne!(boolean_buf1, boolean_buf2);
let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 24);
assert_ne!(boolean_buf1, boolean_buf2);
assert!(boolean_buf1.ptr_eq(&boolean_buf1));
assert!(boolean_buf2.ptr_eq(&boolean_buf2));
assert!(!boolean_buf1.ptr_eq(&boolean_buf2));
}
#[test]
fn test_boolean_slice() {
let bytes = &[0, 3, 2, 6, 2];
let boolean_buf1 = BooleanBuffer::new(Buffer::from(bytes), 0, 32);
let boolean_buf2 = BooleanBuffer::new(Buffer::from(bytes), 0, 32);
let boolean_slice1 = boolean_buf1.slice(16, 16);
let boolean_slice2 = boolean_buf2.slice(0, 16);
assert_eq!(boolean_slice1.values(), boolean_slice2.values());
assert_eq!(bytes, boolean_slice1.values());
assert_eq!(16, boolean_slice1.bit_offset);
assert_eq!(16, boolean_slice1.bit_len);
assert_eq!(bytes, boolean_slice2.values());
assert_eq!(0, boolean_slice2.bit_offset);
assert_eq!(16, boolean_slice2.bit_len);
}
#[test]
fn test_boolean_bitand() {
let offset = 0;
let len = 40;
let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
let expected = BooleanBuffer::new(Buffer::from(&[0, 1, 1, 0, 0]), offset, len);
assert_eq!(boolean_buf1 & boolean_buf2, expected);
}
#[test]
fn test_boolean_bitor() {
let offset = 0;
let len = 40;
let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
let expected = BooleanBuffer::new(Buffer::from(&[0, 1, 1, 1, 0]), offset, len);
assert_eq!(boolean_buf1 | boolean_buf2, expected);
}
#[test]
fn test_boolean_bitxor() {
let offset = 0;
let len = 40;
let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
let expected = BooleanBuffer::new(Buffer::from(&[0, 0, 0, 1, 0]), offset, len);
assert_eq!(boolean_buf1 ^ boolean_buf2, expected);
}
#[test]
fn test_boolean_not() {
let offset = 0;
let len = 40;
let buf = Buffer::from(&[0, 1, 1, 0, 0]);
let boolean_buf = &BooleanBuffer::new(buf, offset, len);
let expected = BooleanBuffer::new(Buffer::from(&[255, 254, 254, 255, 255]), offset, len);
assert_eq!(!boolean_buf, expected);
let sliced = boolean_buf.slice(3, 20);
let result = !&sliced;
assert_eq!(result.offset(), 3);
assert_eq!(result.len(), sliced.len());
for i in 0..sliced.len() {
assert_eq!(result.value(i), !sliced.value(i));
}
}
#[test]
fn test_boolean_from_slice_bool() {
let v = [true, false, false];
let buf = BooleanBuffer::from(&v[..]);
assert_eq!(buf.offset(), 0);
assert_eq!(buf.len(), 3);
assert_eq!(buf.values().len(), 1);
assert!(buf.value(0));
}
#[test]
fn test_from_bitwise_unary_op() {
let input_bools = (0..1024)
.map(|_| rand::random::<bool>())
.collect::<Vec<bool>>();
let input_buffer = BooleanBuffer::from(&input_bools[..]);
for offset in 0..1024 {
let result = BooleanBuffer::from_bitwise_unary_op(
input_buffer.values(),
offset,
input_buffer.len() - offset,
|a| !a,
);
let expected = input_bools[offset..]
.iter()
.map(|b| !*b)
.collect::<BooleanBuffer>();
assert_eq!(result, expected);
}
for offset in 0..512 {
let len = 512 - offset; let result =
BooleanBuffer::from_bitwise_unary_op(input_buffer.values(), offset, len, |a| !a);
let expected = input_bools[offset..]
.iter()
.take(len)
.map(|b| !*b)
.collect::<BooleanBuffer>();
assert_eq!(result, expected);
}
}
#[test]
fn test_from_bitwise_unary_op_unaligned_fallback() {
let bytes = (0..80)
.map(|i| (i as u8).wrapping_mul(37).wrapping_add(11))
.collect::<Vec<_>>();
let base = bytes.as_ptr() as usize;
let shift = (0..8).find(|s| (base + s) % 8 != 0).unwrap();
let misaligned = &bytes[shift..];
let src = &misaligned[..24];
let offset = 7;
let len = 96;
let result = BooleanBuffer::from_bitwise_unary_op(src, offset, len, |a| !a);
let expected = (0..len)
.map(|i| !bit_util::get_bit(src, offset + i))
.collect::<BooleanBuffer>();
assert_eq!(result, expected);
assert_eq!(result.offset(), offset % 64);
let src = &misaligned[..13];
let offset = 3;
let len = 100;
let result = BooleanBuffer::from_bitwise_unary_op(src, offset, len, |a| !a);
let expected = (0..len)
.map(|i| !bit_util::get_bit(src, offset + i))
.collect::<BooleanBuffer>();
assert_eq!(result, expected);
assert_eq!(result.offset(), offset % 64);
}
#[test]
fn test_from_bitwise_binary_op() {
let input_bools_left = (0..1024)
.map(|_| rand::random::<bool>())
.collect::<Vec<bool>>();
let input_bools_right = (0..1024)
.map(|_| rand::random::<bool>())
.collect::<Vec<bool>>();
let input_buffer_left = BooleanBuffer::from(&input_bools_left[..]);
let input_buffer_right = BooleanBuffer::from(&input_bools_right[..]);
for left_offset in 0..200 {
for right_offset in [0, 4, 5, 17, 33, 24, 45, 64, 65, 100, 200] {
for len_offset in [0, 1, 44, 100, 256, 300, 512] {
let len = 1024 - len_offset - left_offset.max(right_offset); let result = BooleanBuffer::from_bitwise_binary_op(
input_buffer_left.values(),
left_offset,
input_buffer_right.values(),
right_offset,
len,
|a, b| a & b,
);
let expected = input_bools_left[left_offset..]
.iter()
.zip(&input_bools_right[right_offset..])
.take(len)
.map(|(a, b)| *a & *b)
.collect::<BooleanBuffer>();
assert_eq!(result, expected);
}
}
}
}
#[test]
fn test_extend_trusted_len_sets_byte_len() {
let mut builder = BooleanBufferBuilder::new(0);
let bools: Vec<_> = (0..10).map(|i| i % 2 == 0).collect();
unsafe { builder.extend_trusted_len(bools.into_iter()) };
assert_eq!(builder.as_slice().len(), bit_util::ceil(builder.len(), 8));
}
#[test]
fn test_extend_trusted_len_then_append() {
let mut builder = BooleanBufferBuilder::new(0);
let bools: Vec<_> = (0..9).map(|i| i % 3 == 0).collect();
unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
builder.append(true);
assert_eq!(builder.as_slice().len(), bit_util::ceil(builder.len(), 8));
let finished = builder.finish();
for (i, v) in bools.into_iter().chain(std::iter::once(true)).enumerate() {
assert_eq!(finished.value(i), v, "at index {}", i);
}
}
#[test]
fn test_find_nth_set_bit_position() {
let bools = vec![true, false, true, true, false, true];
let buffer = BooleanBuffer::from(bools);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 1), 1);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 2), 3);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 3), 4);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 4), 6);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 5), 6);
assert_eq!(buffer.clone().find_nth_set_bit_position(1, 1), 3);
assert_eq!(buffer.clone().find_nth_set_bit_position(3, 1), 4);
assert_eq!(buffer.clone().find_nth_set_bit_position(3, 2), 6);
}
#[test]
fn test_find_nth_set_bit_position_large() {
let mut bools = vec![false; 1000];
bools[100] = true;
bools[500] = true;
bools[999] = true;
let buffer = BooleanBuffer::from(bools);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 1), 101);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 2), 501);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 3), 1000);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 4), 1000);
assert_eq!(buffer.clone().find_nth_set_bit_position(101, 1), 501);
}
#[test]
fn test_find_nth_set_bit_position_sliced() {
let bools = vec![false, true, false, true, true, false, true]; let buffer = BooleanBuffer::from(bools);
let slice = buffer.slice(1, 6);
assert_eq!(slice.len(), 6);
assert_eq!(slice.clone().find_nth_set_bit_position(0, 1), 1);
assert_eq!(slice.clone().find_nth_set_bit_position(0, 2), 3);
assert_eq!(slice.clone().find_nth_set_bit_position(0, 3), 4);
assert_eq!(slice.clone().find_nth_set_bit_position(0, 4), 6);
}
#[test]
fn test_find_nth_set_bit_position_all_set() {
let buffer = BooleanBuffer::new_set(100);
for i in 1..=100 {
assert_eq!(buffer.clone().find_nth_set_bit_position(0, i), i);
}
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 101), 100);
}
#[test]
fn test_find_nth_set_bit_position_none_set() {
let buffer = BooleanBuffer::new_unset(100);
assert_eq!(buffer.clone().find_nth_set_bit_position(0, 1), 100);
}
}