fix: shrink_to_fit PrefixSetMut on freeze (paradigmxyz#7342)

Ruteri · Apr 17, 2024 · a0e542d · a0e542d
1 parent a6f408a
commit a0e542d
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/crates/trie/src/prefix_set/mod.rs b/crates/trie/src/prefix_set/mod.rs
@@ -121,6 +121,9 @@ impl PrefixSetMut {
             self.keys.dedup();
         }
 
+        // we need to shrink in both the sorted and non-sorted cases because deduping may have
+        // occurred either on `freeze`, or during `contains`.
+        self.keys.shrink_to_fit();
         PrefixSet { keys: Arc::new(self.keys), index: self.index }
     }
 }
@@ -186,4 +189,43 @@ mod tests {
         assert!(!prefix_set.contains(&[7, 8]));
         assert_eq!(prefix_set.len(), 3); // Length should be 3 (excluding duplicate)
     }
+
+    #[test]
+    fn test_freeze_shrinks_capacity() {
+        let mut prefix_set = PrefixSetMut::default();
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 3]));
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 4]));
+        prefix_set.insert(Nibbles::from_nibbles([4, 5, 6]));
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate
+
+        assert!(prefix_set.contains(&[1, 2]));
+        assert!(prefix_set.contains(&[4, 5]));
+        assert!(!prefix_set.contains(&[7, 8]));
+        assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
+        assert_eq!(prefix_set.keys.capacity(), 4); // Capacity should be 4 (including duplicate)
+
+        let frozen = prefix_set.freeze();
+        assert_eq!(frozen.keys.len(), 3); // Length should be 3 (excluding duplicate)
+        assert_eq!(frozen.keys.capacity(), 3); // Capacity should be 3 after shrinking
+    }
+
+    #[test]
+    fn test_freeze_shrinks_existing_capacity() {
+        // do the above test but with preallocated capacity
+        let mut prefix_set = PrefixSetMut::with_capacity(101);
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 3]));
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 4]));
+        prefix_set.insert(Nibbles::from_nibbles([4, 5, 6]));
+        prefix_set.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate
+
+        assert!(prefix_set.contains(&[1, 2]));
+        assert!(prefix_set.contains(&[4, 5]));
+        assert!(!prefix_set.contains(&[7, 8]));
+        assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
+        assert_eq!(prefix_set.keys.capacity(), 101); // Capacity should be 101 (including duplicate)
+
+        let frozen = prefix_set.freeze();
+        assert_eq!(frozen.keys.len(), 3); // Length should be 3 (excluding duplicate)
+        assert_eq!(frozen.keys.capacity(), 3); // Capacity should be 3 after shrinking
+    }
 }