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

Jump to bottom

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 all commits
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
83 changes: 72 additions & 11 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,20 +7,38 @@ use proptest::{
strategy::ValueTree,
test_runner::{basic_result_cache, TestRunner},
};
use reth_trie::{prefix_set::PrefixSetMut, Nibbles};
use reth_trie::{
prefix_set::{PrefixSet, PrefixSetMut},
Nibbles,
};
use std::collections::BTreeSet;

/// Abstraction for aggregating nibbles and freezing it to a type
/// that can be later used for benching.
pub trait PrefixSetMutAbstraction: Default {
Copy link
Collaborator

Choose a reason for hiding this comment

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

needs docs

type Frozen;
fn insert(&mut self, key: Nibbles);
fn freeze(self) -> Self::Frozen;
}

/// Abstractions used for benching
pub trait PrefixSetAbstraction: Default {
fn insert(&mut self, key: Nibbles);
fn contains(&mut self, key: Nibbles) -> bool;
}

impl PrefixSetAbstraction for PrefixSetMut {
impl PrefixSetMutAbstraction for PrefixSetMut {
type Frozen = PrefixSet;

fn insert(&mut self, key: Nibbles) {
Self::insert(self, key)
}

fn freeze(self) -> Self::Frozen {
Self::freeze(self)
}
}

impl PrefixSetAbstraction for PrefixSet {
fn contains(&mut self, key: Nibbles) -> bool {
Self::contains(self, &key)
}
Expand Down Expand Up @@ -56,17 +74,20 @@ pub fn prefix_set_lookups(c: &mut Criterion) {
}
}

fn prefix_set_bench<T: PrefixSetAbstraction>(
fn prefix_set_bench<T>(
group: &mut BenchmarkGroup<'_, WallTime>,
description: &str,
(preload, input, expected): (Vec<Nibbles>, Vec<Nibbles>, Vec<bool>),
) {
) where
T: PrefixSetMutAbstraction,
T::Frozen: PrefixSetAbstraction,
{
let setup = || {
let mut prefix_set = T::default();
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 Expand Up @@ -119,11 +140,19 @@ mod implementations {
keys: BTreeSet<Nibbles>,
}

impl PrefixSetAbstraction for BTreeAnyPrefixSet {
impl PrefixSetMutAbstraction for BTreeAnyPrefixSet {
type Frozen = Self;

fn insert(&mut self, key: Nibbles) {
self.keys.insert(key);
}

fn freeze(self) -> Self::Frozen {
self
}
}

impl PrefixSetAbstraction for BTreeAnyPrefixSet {
fn contains(&mut self, key: Nibbles) -> bool {
self.keys.iter().any(|k| k.has_prefix(&key))
}
Expand All @@ -135,11 +164,19 @@ mod implementations {
last_checked: Option<Nibbles>,
}

impl PrefixSetAbstraction for BTreeRangeLastCheckedPrefixSet {
impl PrefixSetMutAbstraction for BTreeRangeLastCheckedPrefixSet {
type Frozen = Self;

fn insert(&mut self, key: Nibbles) {
self.keys.insert(key);
}

fn freeze(self) -> Self::Frozen {
self
}
}

impl PrefixSetAbstraction for BTreeRangeLastCheckedPrefixSet {
fn contains(&mut self, prefix: Nibbles) -> bool {
let range = match self.last_checked.as_ref() {
// presumably never hit
Expand Down Expand Up @@ -169,12 +206,20 @@ mod implementations {
sorted: bool,
}

impl PrefixSetAbstraction for VecBinarySearchPrefixSet {
impl PrefixSetMutAbstraction for VecBinarySearchPrefixSet {
type Frozen = Self;

fn insert(&mut self, key: Nibbles) {
self.sorted = false;
self.keys.push(key);
}

fn freeze(self) -> Self::Frozen {
self
}
}

impl PrefixSetAbstraction for VecBinarySearchPrefixSet {
fn contains(&mut self, prefix: Nibbles) -> bool {
if !self.sorted {
self.keys.sort();
Expand All @@ -198,12 +243,20 @@ mod implementations {
index: usize,
}

impl PrefixSetAbstraction for VecCursorPrefixSet {
impl PrefixSetMutAbstraction for VecCursorPrefixSet {
type Frozen = Self;

fn insert(&mut self, nibbles: Nibbles) {
self.sorted = false;
self.keys.push(nibbles);
}

fn freeze(self) -> Self::Frozen {
self
}
}

impl PrefixSetAbstraction for VecCursorPrefixSet {
fn contains(&mut self, prefix: Nibbles) -> bool {
if !self.sorted {
self.keys.sort();
Expand Down Expand Up @@ -239,12 +292,20 @@ mod implementations {
sorted: bool,
}

impl PrefixSetAbstraction for VecBinarySearchWithLastFoundPrefixSet {
impl PrefixSetMutAbstraction for VecBinarySearchWithLastFoundPrefixSet {
type Frozen = Self;

fn insert(&mut self, key: Nibbles) {
self.sorted = false;
self.keys.push(key);
}

fn freeze(self) -> Self::Frozen {
self
}
}

impl PrefixSetAbstraction for VecBinarySearchWithLastFoundPrefixSet {
fn contains(&mut self, prefix: Nibbles) -> bool {
if !self.sorted {
self.keys.sort();
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 mut 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 mut 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(), 4); // Length should be 3 (including duplicate)
assert_eq!(prefix_set_mut.keys.capacity(), 4); // Capacity should be 4 (including duplicate)

let mut 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(), 4); // Length should be 3 (including duplicate)
assert_eq!(prefix_set_mut.keys.capacity(), 101); // Capacity should be 101 (including duplicate)

let mut 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