From eaf8647fdf09ec33588b968e4a66401d2927253a Mon Sep 17 00:00:00 2001
From: Chris Fallin <chris@cfallin.org>
Date: Thu, 12 Aug 2021 14:40:18 -0700
Subject: [PATCH] BitVec: remove zero words to avoid expanding when
 unnecessary.

---
 src/bitvec.rs | 55 ++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 48 insertions(+), 7 deletions(-)

diff --git a/src/bitvec.rs b/src/bitvec.rs
index a226c0b..ab45389 100644
--- a/src/bitvec.rs
+++ b/src/bitvec.rs
@@ -43,15 +43,33 @@ impl AdaptiveMap {
     fn expand(&mut self) {
         match self {
             &mut Self::Small {
-                len,
-                ref keys,
-                ref values,
+                ref mut len,
+                ref mut keys,
+                ref mut values,
             } => {
-                let mut map = FxHashMap::default();
-                for i in 0..len {
-                    map.insert(keys[i as usize], values[i as usize]);
+                // Note: we *may* remain as `Small` if there are any
+                // zero elements. Try removing them first, before we
+                // commit to a memory allocation.
+                if values.iter().any(|v| *v == 0) {
+                    let mut out = 0;
+                    for i in 0..(*len as usize) {
+                        if values[i] == 0 {
+                            continue;
+                        }
+                        if out < i {
+                            keys[out] = keys[i];
+                            values[out] = values[i];
+                        }
+                        out += 1;
+                    }
+                    *len = out as u32;
+                } else {
+                    let mut map = FxHashMap::default();
+                    for i in 0..(*len as usize) {
+                        map.insert(keys[i], values[i]);
+                    }
+                    *self = Self::Large(map);
                 }
-                *self = Self::Large(map);
             }
             _ => {}
         }
@@ -256,6 +274,15 @@ impl BitVec {
             set_bits(bits).map(move |i| BITS_PER_WORD * word_idx + i)
         })
     }
+
+    /// Is the adaptive data structure in "small" mode? This is meant
+    /// for testing assertions only.
+    pub(crate) fn is_small(&self) -> bool {
+        match &self.elems {
+            &AdaptiveMap::Small { .. } => true,
+            _ => false,
+        }
+    }
 }
 
 fn set_bits(bits: u64) -> impl Iterator<Item = usize> {
@@ -309,4 +336,18 @@ mod test {
 
         assert_eq!(sum, checksum);
     }
+
+    #[test]
+    fn test_expand_remove_zero_elems() {
+        let mut vec = BitVec::new();
+        // Set 12 different words (this is the max small-mode size).
+        for i in 0..12 {
+            vec.set(64 * i, true);
+        }
+        // Now clear a bit, and set a bit in a different word. We
+        // should still be in small mode.
+        vec.set(64 * 5, false);
+        vec.set(64 * 100, true);
+        assert!(vec.is_small());
+    }
 }