Skip to content

Commit

Permalink
Fix insert
Browse files Browse the repository at this point in the history
  • Loading branch information
@Jujumba
Jujumba committed Jul 16, 2024
1 parent b65a26d commit c6fc351
Showing 1 changed file with 87 additions and 50 deletions.
137 changes: 87 additions & 50 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1620,7 +1620,7 @@ impl<B: BitBlock> BitVec<B> {
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`
Expand All @@ -1638,61 +1638,40 @@ impl<B: BitBlock> BitVec<B> {
/// 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,
"insertion index (is {at}) should be <= nbits (is {nbits})",
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;
}
}
}
Expand Down Expand Up @@ -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
]));
}
}

0 comments on commit c6fc351

Please sign in to comment.