Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore(trie): remove PrefixSetMut::contains #10466

Merged
merged 5 commits into from
Aug 23, 2024
Merged

chore(trie): remove PrefixSetMut::contains #10466

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0f7fd5c
chore(trie): remove `PrefixSetMut::contains`
rkrasiuk Aug 22, 2024
93f951b
mut
rkrasiuk Aug 22, 2024
2e01018
fix bench
rkrasiuk Aug 22, 2024
5f0a4a4
fix
rkrasiuk Aug 22, 2024
7cea69c
doc
rkrasiuk Aug 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions crates/trie/trie/benches/prefix_set.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use proptest::{
strategy::ValueTree,
test_runner::{basic_result_cache, TestRunner},
};
use reth_trie::{prefix_set::PrefixSetMut, Nibbles};
use reth_trie::{prefix_set::PrefixSet, Nibbles};
use std::collections::BTreeSet;

/// Abstractions used for benching
Expand All @@ -16,7 +16,7 @@ pub trait PrefixSetAbstraction: Default {
fn contains(&mut self, key: Nibbles) -> bool;
}

impl PrefixSetAbstraction for PrefixSetMut {
impl PrefixSetAbstraction for PrefixSet {
fn insert(&mut self, key: Nibbles) {
Self::insert(self, key)
}
Expand Down Expand Up @@ -66,7 +66,7 @@ fn prefix_set_bench<T: PrefixSetAbstraction>(
for key in &preload {
prefix_set.insert(key.clone());
}
(prefix_set, input.clone(), expected.clone())
(prefix_set.freeze(), input.clone(), expected.clone())
};

let group_id = format!(
Expand Down
105 changes: 36 additions & 69 deletions crates/trie/trie/src/prefix_set.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,65 +75,35 @@ pub struct TriePrefixSets {
/// ```
/// use reth_trie::{prefix_set::PrefixSetMut, Nibbles};
///
/// let mut prefix_set = PrefixSetMut::default();
/// prefix_set.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb]));
/// prefix_set.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb, 0xc]));
/// let mut prefix_set_mut = PrefixSetMut::default();
/// prefix_set_mut.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb]));
/// prefix_set_mut.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb, 0xc]));
/// let prefix_set = prefix_set_mut.freeze();
/// assert!(prefix_set.contains(&[0xa, 0xb]));
/// assert!(prefix_set.contains(&[0xa, 0xb, 0xc]));
/// ```
#[derive(Clone, Default, Debug)]
pub struct PrefixSetMut {
keys: Vec<Nibbles>,
sorted: bool,
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should always be treated as unsorted

index: usize,
}

impl<I> From<I> for PrefixSetMut
where
I: IntoIterator<Item = Nibbles>,
{
fn from(value: I) -> Self {
Self { keys: value.into_iter().collect(), ..Default::default() }
Self { keys: value.into_iter().collect() }
}
}

impl PrefixSetMut {
/// Create [`PrefixSetMut`] with pre-allocated capacity.
pub fn with_capacity(capacity: usize) -> Self {
Self { keys: Vec::with_capacity(capacity), ..Default::default() }
}

/// Returns `true` if any of the keys in the set has the given prefix or
/// if the given prefix is a prefix of any key in the set.
pub fn contains(&mut self, prefix: &[u8]) -> bool {
if !self.sorted {
self.keys.sort();
self.keys.dedup();
self.sorted = true;
}

while self.index > 0 && self.keys[self.index] > *prefix {
self.index -= 1;
}

for (idx, key) in self.keys[self.index..].iter().enumerate() {
if key.has_prefix(prefix) {
self.index += idx;
return true
}

if *key > *prefix {
self.index += idx;
return false
}
}

false
Self { keys: Vec::with_capacity(capacity) }
}

/// Inserts the given `nibbles` into the set.
pub fn insert(&mut self, nibbles: Nibbles) {
self.sorted = false;
self.keys.push(nibbles);
}

Expand All @@ -142,7 +112,6 @@ impl PrefixSetMut {
where
I: IntoIterator<Item = Nibbles>,
{
self.sorted = false;
self.keys.extend(nibbles_iter);
}

Expand All @@ -160,15 +129,12 @@ impl PrefixSetMut {
///
/// If not yet sorted, the elements will be sorted and deduplicated.
pub fn freeze(mut self) -> PrefixSet {
if !self.sorted {
self.keys.sort();
self.keys.dedup();
}

self.keys.sort();
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 }
PrefixSet { keys: Arc::new(self.keys), index: 0 }
}
}

Expand All @@ -185,7 +151,7 @@ impl PrefixSet {
/// Returns `true` if any of the keys in the set has the given prefix or
/// if the given prefix is a prefix of any key in the set.
#[inline]
pub fn contains(&mut self, prefix: &Nibbles) -> bool {
pub fn contains(&mut self, prefix: &[u8]) -> bool {
while self.index > 0 && &self.keys[self.index] > prefix {
self.index -= 1;
}
Expand Down Expand Up @@ -235,12 +201,13 @@ mod tests {

#[test]
fn test_contains_with_multiple_inserts_and_duplicates() {
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
let mut prefix_set_mut = PrefixSetMut::default();
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate

let prefix_set = prefix_set_mut.freeze();
assert!(prefix_set.contains(&[1, 2]));
assert!(prefix_set.contains(&[4, 5]));
assert!(!prefix_set.contains(&[7, 8]));
Expand All @@ -249,40 +216,40 @@ mod tests {

#[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
let mut prefix_set_mut = PrefixSetMut::default();
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate

assert_eq!(prefix_set_mut.keys.len(), 3); // Length should be 3 (excluding duplicate)
assert_eq!(prefix_set_mut.keys.capacity(), 4); // Capacity should be 4 (including duplicate)

let prefix_set = prefix_set_mut.freeze();
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
assert_eq!(prefix_set.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
let mut prefix_set_mut = PrefixSetMut::with_capacity(101);
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate

assert_eq!(prefix_set_mut.keys.len(), 3); // Length should be 3 (excluding duplicate)
assert_eq!(prefix_set_mut.keys.capacity(), 101); // Capacity should be 101 (including duplicate)

let prefix_set = prefix_set_mut.freeze();
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
assert_eq!(prefix_set.keys.capacity(), 3); // Capacity should be 3 after shrinking
}
}
Loading