From b2c1df36d8fc015b0c2494305f3be70451e4f763 Mon Sep 17 00:00:00 2001
From: Imbris <imbrisf@gmail.com>
Date: Tue, 29 Mar 2022 12:44:50 -0400
Subject: [PATCH] Fix set method to use the new length required (rather than
 index being set) when calculating expansion size. This used to panic
 internally in the set method due to not expanding the storage size. Add test
 for this case.

---
 src/lib.rs   | 9 +++++----
 src/tests.rs | 9 +++++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index cc6592d..0b7c5e7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -158,9 +158,9 @@ where
     // input: Number of bits
     // output:
     //
-    // 1. the number of Vector used
-    // 2. after filling 1, the remaining bytes should be filled
-    // 3. after filling 2, the remaining bits should be filled
+    // 1. the number of vectors completely used
+    // 2. after filling 1, the remaining elements which should be filled
+    // 3. after filling 2, the remaining bits which should be filled
     //
     // notice that this result represents the length of vector
     // so if 3. is 0, it means no extra bits after filling bytes
@@ -564,13 +564,14 @@ where
     /// assert_eq!(bitvec.get(14), Some(false));
     /// ```
     pub fn set(&mut self, index: usize, flag: bool) {
-        let (i, bytes, bits) = Self::bit_to_len(index);
         if self.nbits <= index {
+            let (i, bytes, bits) = Self::bit_to_len(index + 1);
             let new_len = if bytes > 0 || bits > 0 { i + 1 } else { i };
             self.storage
                 .extend((0..new_len - self.storage.len()).map(move |_| A::Item::ZERO));
             self.nbits = index + 1;
         }
+        let (i, bytes, bits) = Self::bit_to_len(index);
         let mut arr = self.storage[i].to_array();
         arr[bytes] = Self::set_bit(flag, arr[bytes], bits as u32);
         self.storage[i] = A::Item::from(arr);
diff --git a/src/tests.rs b/src/tests.rs
index ead18f7..74a78c4 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -48,6 +48,15 @@ fn test_bit_vec_leading_zeros() {
     assert_eq!(bitvec.leading_zeros(), 6);
 }
 
+#[test]
+fn test_set_expands() {
+    let mut bitvec = BitVec::zeros(0);
+    for index in 0..2049 {
+        bitvec.set(index, true);
+        assert_eq!(bitvec.len(), index + 1);
+    }
+}
+
 #[test]
 fn test_bit_vec_resize() {
     for i in (0..3333).filter(|x| x % 13 == 0) {