fix: shrink_to_fit PrefixSetMut on freeze (#7342)

This commit is contained in:
Dan Cline
2024-04-01 12:27:34 -04:00
committed by GitHub
parent ab8a4f1ff6
commit d6933bad99

View File

@ -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
}
}