add grow_and_insert method (#92)

Co-authored-by: James Liu <contact@jamessliu.com>

Author: shuoli84

benches/benches/benches.rs

@@ -91,6 +91,19 @@ fn insert(c: &mut Criterion) {
     });
 }
 
+fn grow_and_insert(c: &mut Criterion) {
+    const N: usize = 1_000_000;
+    let mut fb = FixedBitSet::with_capacity(N);
+
+    c.bench_function("grow_and_insert", |b| {
+        b.iter(|| {
+            for i in 0..N {
+                fb.grow_and_insert(i);
+            }
+        })
+    });
+}
+
 fn union_with(c: &mut Criterion) {
     const N: usize = 1_000_000;
     let mut fb_a = FixedBitSet::with_capacity(N);
@@ -159,5 +172,6 @@ criterion_group!(
     symmetric_difference_with,
     count_ones,
     clear,
+    grow_and_insert,
 );
 criterion_main!(benches);

src/lib.rs

@@ -123,6 +123,22 @@ impl FixedBitSet {
         }
     }
 
+    /// Grows the internal size of the bitset before inserting a bit
+    ///
+    /// Unlike `insert`, this cannot panic, but may allocate if the bit is outside of the existing buffer's range.
+    ///
+    /// This is faster than calling `grow` then `insert` in succession.
+    #[inline]
+    pub fn grow_and_insert(&mut self, bits: usize) {
+        self.grow(bits + 1);
+
+        let (blocks, rem) = div_rem(bits);
+        // SAFETY: The above grow ensures that the block is inside the Vec's allocation.
+        unsafe {
+            *self.data.get_unchecked_mut(blocks) |= 1 << rem;
+        }
+    }
+
     /// The length of the [`FixedBitSet`] in bits.
     ///
     /// Note: `len` includes both set and unset bits.
@@ -1043,6 +1059,18 @@ mod tests {
         assert!(fb.contains(64));
     }
 
+    #[test]
+    fn grow_and_insert() {
+        let mut fb = FixedBitSet::default();
+        for i in 0..100 {
+            if i % 3 == 0 {
+                fb.grow_and_insert(i);
+            }
+        }
+
+        assert_eq!(fb.count_ones(..), 34);
+    }
+
     #[test]
     fn test_toggle() {
         let mut fb = FixedBitSet::with_capacity(16);