diff --git a/ts_bitset/proptest-regressions/lib.txt b/ts_bitset/proptest-regressions/lib.txt index 7abcdd08..428008f3 100644 --- a/ts_bitset/proptest-regressions/lib.txt +++ b/ts_bitset/proptest-regressions/lib.txt @@ -7,3 +7,4 @@ cc 9912420036cb666096da51a46393e3a666c543c9262e210125e71cc5eeb6b081 # shrinks to i = 256 cc 2641a9595a01636de4d21c749d19834a88a316290fe57b6f76acee71a4c82b16 # shrinks to i = 256 cc 99f787fcda492d2e3dc645b67bf52f9f84efd07d57a89bebd3bbe4ef3a8484d6 # shrinks to i = 256 +cc c8662104bad04e8f72621644c6e6a3069600249a7863172feeb2a30d7a8be025 # shrinks to shift_amt = 256, bits = [108] diff --git a/ts_bitset/src/lib.rs b/ts_bitset/src/lib.rs index 13bd2eba..8931ca13 100644 --- a/ts_bitset/src/lib.rs +++ b/ts_bitset/src/lib.rs @@ -88,6 +88,8 @@ impl Bitset { pub const EMPTY: Self = Bitset([0u64; N_WORDS]); /// The bitset with all bits set. pub const FULL: Self = Bitset([u64::MAX; N_WORDS]); + /// The number of bits in the bitset. + pub const N_BITS: usize = N_WORDS * 64; /// Resize this bitset to a new one with `M` 64-bit words. /// @@ -540,6 +542,191 @@ impl core::ops::Not for Bitset { } } +impl Bitset { + /// In-place unbounded shift left. Updates `self` to be `self << rhs`, without bounding + /// the value of `rhs`. + /// + /// If `rhs` is larger or equal to `Self::N_BITS`, the entire value is shifted out, and + /// `Self::EMPTY` is returned. + /// + /// # Panics + /// + /// If the shift amount cannot fit in a `usize`. This should only be possible on targets where + /// `usize` is 16 bits, which the Rust specification allows but is extremely uncommon in + /// practice. + pub fn unbounded_shl_inplace(&mut self, rhs: u32) { + // We need a usize to index the word array. + // Per https://doc.rust-lang.org/reference/types/numeric.html, usize may be as small as + // 16 bits, even though it's more commonly 32 or 64 bits. + // For huge bitsets (hundreds of MiB), shift_words may be up to 26 bits. So, this + // conversion may fail in extreme edge cases. + let rhs: usize = rhs.try_into().unwrap(); + + if rhs >= Self::N_BITS { + *self = Self::EMPTY; + return; + } + + let shift_words: usize = rhs / 64; + let shift_bits = rhs % 64; + let shift_bits_inv = (64 - shift_bits) as u32; + + for tgt_idx in (shift_words..self.0.len()).rev() { + let src_idx = tgt_idx - shift_words; + let src_word1 = self.0[src_idx]; + let src_word2 = if src_idx == 0 { 0 } else { self.0[src_idx - 1] }; + self.0[tgt_idx] = (src_word1 << shift_bits) | (src_word2.unbounded_shr(shift_bits_inv)); + } + self.0[..shift_words].fill(0); + } + + /// Unbounded shift left. Computes `self << rhs`, without bounding the value of `rhs`. + /// + /// If `rhs` is larger or equal to `Self::N_BITS`, the entire value is shifted out, and + /// `Self::EMPTY` is returned. + /// + /// # Panics + /// + /// If the shift amount cannot fit in a `usize`. This should only be possible on targets where + /// `usize` is 16 bits, which the Rust specification allows but is extremely uncommon in + /// practice. + #[inline] + pub fn unbounded_shl(mut self, rhs: u32) -> Self { + self.unbounded_shl_inplace(rhs); + self + } + + /// In-place unbounded shift right. Updates `self` to be `self >> rhs`, without bounding + /// the value of `rhs`. + /// + /// If `rhs` is larger or equal to `Self::N_BITS`, the entire value is shifted out, and + /// `Self::EMPTY` is returned. + /// + /// # Panics + /// + /// If the shift amount cannot fit in a `usize`. This should only be possible on targets where + /// `usize` is 16 bits, which the Rust specification allows but is extremely uncommon in + /// practice. + pub fn unbounded_shr_inplace(&mut self, rhs: u32) { + // We need a usize to index the word array. + // Per https://doc.rust-lang.org/reference/types/numeric.html, usize may be as small as + // 16 bits, even though it's more commonly 32 or 64 bits. + // For huge bitsets (hundreds of MiB), shift_words may be up to 26 bits. So, this + // conversion may fail in extreme edge cases. + let rhs: usize = rhs.try_into().unwrap(); + + if rhs >= Self::N_BITS { + *self = Self::EMPTY; + return; + } + + let shift_words: usize = rhs / 64; + let shift_bits = rhs % 64; + let shift_bits_inv = (64 - shift_bits) as u32; + + for tgt_idx in 0..N_WORDS - shift_words { + let src_idx = tgt_idx + shift_words; + let src_word1 = self.0[src_idx]; + let src_word2 = self.0.get(src_idx + 1).map_or(0, |v| *v); + self.0[tgt_idx] = (src_word1 >> shift_bits) | (src_word2.unbounded_shl(shift_bits_inv)); + } + self.0[N_WORDS - shift_words..].fill(0); + } + + /// Unbounded shift right. Computes `self >> rhs`, without bounding the value of `rhs`. + /// + /// If `rhs` is larger or equal to `Self::N_BITS`, the entire value is shifted out, and + /// `Self::EMPTY` is returned. + /// + /// # Panics + /// + /// If the shift amount cannot fit in a `usize`. This should only be possible on targets where + /// `usize` is 16 bits, which the Rust specification allows but is extremely uncommon in + /// practice. + #[inline] + pub fn unbounded_shr(mut self, rhs: u32) -> Self { + self.unbounded_shr_inplace(rhs); + self + } +} + +macro_rules! shift_impl { + ($t:ty) => { + impl core::ops::Shl<$t> for Bitset { + type Output = Self; + + #[inline] + fn shl(mut self, rhs: $t) -> Self::Output { + self <<= rhs; + self + } + } + + impl core::ops::ShlAssign<$t> for Bitset { + #[inline] + fn shl_assign(&mut self, rhs: $t) { + // This comparison is a no-op for unsigned $t, but required for signed $t. + #[allow(unused_comparisons)] + if rhs < 0 { + panic!("negative shift"); + } + + // Cast may fail on 16b usize, see comment in unbounded_shl. + let rhs: usize = rhs.try_into().unwrap(); + if rhs >= Self::N_BITS { + panic!("attempt to shift left with overflow"); + } + + self.unbounded_shl_inplace(rhs.try_into().unwrap()); + } + } + + impl core::ops::Shr<$t> for Bitset { + type Output = Self; + + #[inline] + fn shr(mut self, rhs: $t) -> Self::Output { + self >>= rhs; + self + } + } + + impl core::ops::ShrAssign<$t> for Bitset { + #[inline] + fn shr_assign(&mut self, rhs: $t) { + // This comparison is a no-op for unsigned $t, but required for signed $t. + #[allow(unused_comparisons)] + if rhs < 0 { + panic!("negative shift"); + } + + // Cast may fail on 16b usize, see comment in unbounded_shl. + let rhs: usize = rhs.try_into().unwrap(); + if rhs >= Self::N_BITS { + panic!("attempt to shift left with overflow"); + } + + self.unbounded_shr_inplace(rhs.try_into().unwrap()); + } + } + }; +} + +shift_impl!(u8); +shift_impl!(u16); +shift_impl!(u32); +shift_impl!(u64); +shift_impl!(u128); + +shift_impl!(i8); +shift_impl!(i16); +shift_impl!(i32); +shift_impl!(i64); +shift_impl!(i128); + +shift_impl!(usize); +shift_impl!(isize); + impl From<[u64; N_WORDS]> for Bitset { #[inline] fn from(value: [u64; N_WORDS]) -> Self { @@ -689,6 +876,60 @@ mod test { Bitset256::with_bits_upto(321); } + #[test] + fn shl() { + // Exhaustive checking of every non-destructive left shift of a one-hot bitset, + // as well as a zero shift and the shift amount that loses the hot bit. + for i in 0..256 { + let destructive_shift_amt = 256 - i; + let bs = Bitset256::EMPTY.with_bit(i); + assert_eq!(bs << 0usize, bs); + for shift in 1..destructive_shift_amt { + let shifted = bs << shift; + assert_eq!(shifted.count_ones(), bs.count_ones()); + assert_eq!(shifted.first_set(), Some(i + shift)); + } + assert_eq!( + bs.unbounded_shl(destructive_shift_amt as u32), + Bitset256::EMPTY + ); + } + } + + #[test] + fn shr() { + // Exhaustive checking of every non-destructive right shift of a one-hot bitset, + // as well as a zero shift and the shift amount that loses the hot bit. + for i in 0..256 { + let destructive_shift_amt = 256 - i; + let bs = Bitset256::EMPTY.with_bit(255 - i); + assert_eq!(bs >> 0usize, bs); + for shift in 1..destructive_shift_amt { + let shifted = bs >> shift; + assert_eq!(shifted.count_ones(), bs.count_ones()); + assert_eq!(shifted.first_set(), Some(255 - i - shift)); + } + assert_eq!( + bs.unbounded_shr(destructive_shift_amt as u32), + Bitset256::EMPTY + ); + } + } + + #[test] + #[should_panic] + fn shl_overflow() { + let bs = Bitset256::EMPTY.with_bit(0); + let _ = bs << 256; + } + + #[test] + #[should_panic] + fn shr_overflow() { + let bs = Bitset256::EMPTY.with_bit(0); + let _ = bs >> 256; + } + proptest::prop_compose! { fn bitset()(bs: [u64; 4]) -> Bitset256 { Bitset(bs) @@ -852,5 +1093,22 @@ mod test { proptest::prop_assert_eq!(Some(i as usize), bits.intersection_top(&other)); } + + #[test] + fn shifts(shift_amt in 0u32..257, bits in nonempty_bitvec()) { + let bs = Bitset256::from_iter(&bits); + + let left = bs.unbounded_shl(shift_amt); + proptest::prop_assert!(left.count_ones() <= bs.count_ones()); + let right = left.unbounded_shr(shift_amt); + proptest::prop_assert_eq!(right.count_ones(), left.count_ones()); + proptest::prop_assert_eq!(bs | right, bs); + + let right = bs.unbounded_shr(shift_amt); + proptest::prop_assert!(right.count_ones() <= bs.count_ones()); + let left = right.unbounded_shl(shift_amt); + proptest::prop_assert_eq!(left.count_ones(), right.count_ones()); + proptest::prop_assert_eq!(bs | left, bs); + } } }