From c6fc351b904609da64e96aedb8cd71903027a9d6 Mon Sep 17 00:00:00 2001 From: Jujumba Date: Tue, 16 Jul 2024 12:57:50 +0200 Subject: [PATCH] Fix `insert` --- src/lib.rs | 137 ++++++++++++++++++++++++++++++++++------------------- 1 file changed, 87 insertions(+), 50 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index ea84748..5038422 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1620,7 +1620,7 @@ impl BitVec { self.storage.shrink_to_fit(); } - /// Inserts a given bit at index `at` + /// Inserts a given bit at index `at`, shifting all bits after by one /// /// # Panics /// Panics if `at > nbits` @@ -1638,6 +1638,12 @@ impl BitVec { /// assert_eq!(b.len(), 3); /// assert!(b.eq_vec(&[true, false, true])); ///``` + /// + /// # Time complexity + /// Takes O([`BitVec::len`]) time. All items after the insertion index must be + /// shifted to the right. In the worst case, all elements are shifted when + /// the insertion index is 0. + pub fn insert(&mut self, at: usize, bit: bool) { assert!( at <= self.nbits, @@ -1645,54 +1651,27 @@ impl BitVec { nbits = self.nbits ); - // Initialize if inserting at index zero in a new `BitVec` - if self.nbits == 0 { + let last_block_bits = self.nbits % B::bits(); + let block_at = at / B::bits(); // needed block + let bit_at = at % B::bits(); // index within the block + + if last_block_bits == 0 { self.storage.push(B::zero()); } - let num_blocks = self.nbits / B::bits(); - let w = at / B::bits(); // number of blocks - let b = at % B::bits(); // index within a block - self.nbits += 1; - // is it the last block? - if w == num_blocks { - let num_bits_in_last_block = self.nbits - num_blocks * B::bits(); - // is it the first index in the last block and shifting won't "overflow" this block? - if b == 0 && num_bits_in_last_block < B::bits() { - // if both conditions are met, then we can trivially shift the last block left by one, and set the first bit to `bit` - // though I'm a bit skeptical about this observation, it does require a regirous testing - self.storage[w] = self.storage[w] << 1; - self.set(num_blocks * B::bits(), bit); - return; - } else { - if num_bits_in_last_block == B::bits() { - // if this (the last) block is fully used, then only the last bit in it will be - // moved to the new block - let block = if self.get(self.nbits - 1).unwrap() { - B::one() - } else { - B::zero() - }; - self.storage.push(block); - } - - let mut prev = self.get(at).unwrap(); - let block_start = w * B::bits(); - // shifting all bits at and after `at` left by one - for i in (at + 1)..num_bits_in_last_block { - // I suspect these `get` and `set` calls can be optimized, and some neat - // bit-hacks can be used because we already know with which block we are - // working with - let curr = self.get(block_start + i).unwrap(); // get `i`th bit - self.set(block_start + i, prev); - prev = curr; - } - self.set(block_start + at, bit); - } - } else { - todo!(); + let mut carry = self.storage[block_at] >> (B::bits() - 1); + let lsbits_mask = (B::one() << bit_at) - B::one(); + let set_bit = if bit { B::one() } else { B::zero() } << bit_at; + self.storage[block_at] = (self.storage[block_at] & lsbits_mask) + | ((self.storage[block_at] & !lsbits_mask) << 1) + | set_bit; + + for block_ref in &mut self.storage[block_at + 1..] { + let curr_carry = *block_ref >> (B::bits() - 1); + *block_ref = *block_ref << 1 | carry; + carry = curr_carry; } } } @@ -3131,17 +3110,75 @@ mod tests { } #[test] - fn test_insert() { + fn test_insert_at_zero() { let mut v = BitVec::new(); + v.insert(0, false); + v.insert(0, true); + v.insert(0, false); + v.insert(0, true); + v.insert(0, false); v.insert(0, true); - v.push(true); - v.push(false); - v.insert(1, false); - v.push(true); - v.insert(5, false); assert_eq!(v.len(), 6); + assert_eq!(v.storage().len(), 1); assert!(v.eq_vec(&[true, false, true, false, true, false])); } + + #[test] + fn test_insert_at_end() { + let mut v = BitVec::new(); + + v.insert(v.len(), true); + v.insert(v.len(), false); + v.insert(v.len(), true); + v.insert(v.len(), false); + v.insert(v.len(), true); + v.insert(v.len(), false); + + assert_eq!(v.storage().len(), 1); + assert_eq!(v.len(), 6); + assert!(v.eq_vec(&[true, false, true, false, true, false])); + } + + #[test] + fn test_insert_at_block_boundaries() { + let mut v = BitVec::from_elem(32, false); + + assert_eq!(v.storage().len(), 1); + + v.insert(31, true); + + assert_eq!(v.len(), 33); + + assert!(matches!(v.get(31), Some(true))); + assert!(v.eq_vec(&[ + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, true, false + ])); + + assert_eq!(v.storage().len(), 2); + } + + #[test] + fn test_insert_at_block_boundaries_1() { + let mut v = BitVec::from_elem(64, false); + + assert_eq!(v.storage().len(), 2); + + v.insert(63, true); + + assert_eq!(v.len(), 65); + + assert!(matches!(v.get(63), Some(true))); + assert!(v.eq_vec(&[ + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, false, false, false, false, false, false, false, false, false, + false, false, false, true, false + ])); + } }