Skip to content

Commit 487bdba

Browse files
authored
Add for_each_unset_index to BitBuffer (#10866)
Add `for_each_unset_index` to `BitBuffer` ## Summary - Tracking Issue: #10605 `BitBuffer::for_each_set_index` visits every set bit, but nothing does the same for unset bits. The `Unary` encoding in #10609 needs that: each value ends with an unset bit, so a full decode is one walk over the unset bits. ## Changes - `BitBuffer::for_each_unset_index(f)` calls `f(index)` for every unset bit, in ascending order. - It shares one scan with `try_for_each_set_index`, through a private `try_for_each_index` that takes a `const UNSET: bool`. `select` and `select_unset` share their kernels the same way. With `UNSET`, each word is inverted as it is read. - The last word is zero padded past `len`. For unset bits, the padding is set before inverting, so it is never visited. --------- Signed-off-by: rapour <[email protected]>
1 parent a412a53 commit 487bdba

1 file changed

Lines changed: 67 additions & 1 deletion

File tree

‎vortex-buffer/src/bit/buf.rs‎

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -562,12 +562,43 @@ impl BitBuffer {
562562
///
563563
/// Stops and returns the first error from `f`.
564564
#[inline]
565-
pub fn try_for_each_set_index<E, F>(&self, mut f: F) -> Result<(), E>
565+
pub fn try_for_each_set_index<E, F>(&self, f: F) -> Result<(), E>
566+
where
567+
F: FnMut(usize) -> Result<(), E>,
568+
{
569+
self.try_for_each_index::<false, E, F>(f)
570+
}
571+
572+
/// Invoke `f(index)` for every unset bit, in ascending order, processing a `u64` word at a
573+
/// time.
574+
///
575+
/// This is the counterpart of [`for_each_set_index`](Self::for_each_set_index) for unset bits.
576+
#[inline]
577+
pub fn for_each_unset_index<F: FnMut(usize)>(&self, mut f: F) {
578+
let Ok(()) = self.try_for_each_index::<true, _, _>(|index| {
579+
f(index);
580+
Ok::<_, Infallible>(())
581+
});
582+
}
583+
584+
/// The scan behind [`try_for_each_set_index`](Self::try_for_each_set_index). With `UNSET`,
585+
/// each word is inverted as it is read, so the same scan visits the unset bits instead.
586+
#[inline]
587+
fn try_for_each_index<const UNSET: bool, E, F>(&self, mut f: F) -> Result<(), E>
566588
where
567589
F: FnMut(usize) -> Result<(), E>,
568590
{
569591
let mut base = 0usize;
570592
for word in self.chunks().iter_padded() {
593+
let word = if UNSET {
594+
// The last word is zero padded past `len`. Setting the padding before inverting
595+
// keeps it from being visited.
596+
let held = self.len - base;
597+
let padding = if held < 64 { u64::MAX << held } else { 0 };
598+
!(word | padding)
599+
} else {
600+
word
601+
};
571602
if word == u64::MAX {
572603
for k in 0..64 {
573604
f(base + k)?;
@@ -805,7 +836,11 @@ impl<'a> IntoIterator for &'a BitBuffer {
805836

806837
#[cfg(test)]
807838
mod tests {
839+
use rand::RngExt;
840+
use rand::SeedableRng;
841+
use rand::rngs::StdRng;
808842
use rstest::rstest;
843+
use vortex_error::VortexResult;
809844

810845
use crate::ByteBuffer;
811846
use crate::bit::BitBuffer;
@@ -1150,6 +1185,37 @@ mod tests {
11501185
}
11511186
}
11521187

1188+
#[rstest]
1189+
#[case(0)]
1190+
#[case(1)]
1191+
#[case(63)]
1192+
#[case(64)]
1193+
#[case(65)]
1194+
#[case(128)]
1195+
#[case(200)]
1196+
fn test_for_each_unset_index_agrees_with_naive(#[case] len: usize) -> VortexResult<()> {
1197+
let mut rng = StdRng::seed_from_u64(u64::try_from(len)?);
1198+
let random: BitBuffer = (0..len).map(|_| rng.random_bool(0.5)).collect();
1199+
// The same kind of bits, starting 3 bits into their buffer.
1200+
let offset: BitBuffer = (0..len + 3).map(|_| rng.random_bool(0.5)).collect();
1201+
for bits in [
1202+
random,
1203+
offset.slice(3..),
1204+
BitBuffer::new_set(len),
1205+
BitBuffer::new_unset(len),
1206+
] {
1207+
let expected: Vec<usize> = bits
1208+
.iter()
1209+
.enumerate()
1210+
.filter_map(|(index, set)| (!set).then_some(index))
1211+
.collect();
1212+
let mut found = Vec::new();
1213+
bits.for_each_unset_index(|index| found.push(index));
1214+
assert_eq!(found, expected);
1215+
}
1216+
Ok(())
1217+
}
1218+
11531219
#[test]
11541220
fn test_map_cmp_conditional() {
11551221
// map_cmp with conditional logic based on index and bit value

0 commit comments

Comments
 (0)