reth_trie_common/

prefix_set.rs

use crate::Nibbles;
use alloc::{sync::Arc, vec::Vec};
use alloy_primitives::{
    map::{B256Map, B256Set},
    B256,
};
use core::ops::Range;

/// Collection of mutable prefix sets.
#[derive(Clone, Default, Debug, PartialEq, Eq)]
pub struct TriePrefixSetsMut {
    /// A set of account prefixes that have changed.
    pub account_prefix_set: PrefixSetMut,
    /// A map containing storage changes with the hashed address as key and a set of storage key
    /// prefixes as the value.
    pub storage_prefix_sets: B256Map<PrefixSetMut>,
    /// A set of hashed addresses of destroyed accounts.
    pub destroyed_accounts: B256Set,
}

impl TriePrefixSetsMut {
    /// Returns `true` if all prefix sets are empty.
    pub fn is_empty(&self) -> bool {
        self.account_prefix_set.is_empty() &&
            self.storage_prefix_sets.is_empty() &&
            self.destroyed_accounts.is_empty()
    }

    /// Extends prefix sets with contents of another prefix set.
    pub fn extend(&mut self, other: Self) {
        self.account_prefix_set.extend(other.account_prefix_set);
        for (hashed_address, prefix_set) in other.storage_prefix_sets {
            self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set);
        }
        self.destroyed_accounts.extend(other.destroyed_accounts);
    }

    /// Extends prefix sets with contents of another prefix set by reference.
    pub fn extend_ref(&mut self, other: &Self) {
        self.account_prefix_set.extend_ref(&other.account_prefix_set);
        for (hashed_address, prefix_set) in &other.storage_prefix_sets {
            self.storage_prefix_sets.entry(*hashed_address).or_default().extend_ref(prefix_set);
        }
        self.destroyed_accounts.extend(other.destroyed_accounts.iter().copied());
    }

    /// Returns a `TriePrefixSets` with the same elements as these sets.
    ///
    /// If not yet sorted, the elements will be sorted and deduplicated.
    pub fn freeze(self) -> TriePrefixSets {
        TriePrefixSets {
            account_prefix_set: self.account_prefix_set.freeze(),
            storage_prefix_sets: self
                .storage_prefix_sets
                .into_iter()
                .map(|(hashed_address, prefix_set)| (hashed_address, prefix_set.freeze()))
                .collect(),
            destroyed_accounts: self.destroyed_accounts,
        }
    }

    /// Clears the prefix sets and destroyed accounts map.
    pub fn clear(&mut self) {
        self.destroyed_accounts.clear();
        self.storage_prefix_sets.clear();
        self.account_prefix_set.clear();
    }
}

/// Collection of trie prefix sets.
#[derive(Default, Debug, Clone)]
pub struct TriePrefixSets {
    /// A set of account prefixes that have changed.
    pub account_prefix_set: PrefixSet,
    /// A map containing storage changes with the hashed address as key and a set of storage key
    /// prefixes as the value.
    pub storage_prefix_sets: B256Map<PrefixSet>,
    /// A set of hashed addresses of destroyed accounts.
    pub destroyed_accounts: B256Set,
}

/// A container for efficiently storing and checking for the presence of key prefixes.
///
/// This data structure stores a set of `Nibbles` and provides methods to insert
/// new elements and check whether any existing element has a given prefix.
///
/// Internally, this implementation stores keys in an unsorted `Vec<Nibbles>` together with an
/// `all` flag. The `all` flag indicates that every entry should be considered changed and that
/// individual keys can be ignored.
///
/// Sorting and deduplication do not happen during insertion or membership checks on this mutable
/// structure. Instead, keys are sorted and deduplicated when converting into the immutable
/// `PrefixSet` via `freeze()`. The immutable `PrefixSet` provides `contains` and relies on the
/// sorted and unique keys produced by `freeze()`; it does not perform additional sorting or
/// deduplication.
///
/// This guarantees that a `PrefixSet` constructed from a `PrefixSetMut` is always sorted and
/// deduplicated.
/// # Examples
///
/// ```
/// use reth_trie_common::{prefix_set::PrefixSetMut, Nibbles};
///
/// 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(&Nibbles::from_nibbles_unchecked([0xa, 0xb])));
/// assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([0xa, 0xb, 0xc])));
/// ```
#[derive(PartialEq, Eq, Clone, Default, Debug)]
pub struct PrefixSetMut {
    /// Flag indicating that any entry should be considered changed.
    /// If set, the keys will be discarded.
    all: bool,
    keys: Vec<Nibbles>,
}

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

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

    /// Create [`PrefixSetMut`] that considers all key changed.
    pub const fn all() -> Self {
        Self { all: true, keys: Vec::new() }
    }

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

    /// Extend prefix set with contents of another prefix set.
    pub fn extend(&mut self, other: Self) {
        self.all |= other.all;
        self.keys.extend(other.keys);
    }

    /// Extend prefix set with contents of another prefix set by reference.
    pub fn extend_ref(&mut self, other: &Self) {
        self.all |= other.all;
        self.keys.extend(other.keys.iter().copied());
    }

    /// Appends prefix set keys from another mutable prefix set, leaving it empty.
    pub fn append(&mut self, other: &mut Self) {
        self.all |= other.all;
        other.all = false;
        self.keys.append(&mut other.keys);
    }

    /// Extend prefix set keys with contents of provided iterator.
    pub fn extend_keys<I>(&mut self, keys: I)
    where
        I: IntoIterator<Item = Nibbles>,
    {
        self.keys.extend(keys);
    }

    /// Returns an iterator over all currently retained keys.
    pub fn iter(&self) -> core::slice::Iter<'_, Nibbles> {
        self.keys.iter()
    }

    /// Returns the number of elements in the set.
    pub const fn len(&self) -> usize {
        self.keys.len()
    }

    /// Returns `true` if the set is empty and `all` flag is not set.
    pub const fn is_empty(&self) -> bool {
        !self.all && self.keys.is_empty()
    }

    /// Clears the inner vec for reuse, setting `all` to `false`.
    pub fn clear(&mut self) {
        self.all = false;
        self.keys.clear();
    }

    /// Returns a `PrefixSet` with the same elements as this set.
    ///
    /// If not yet sorted, the elements will be sorted and deduplicated.
    pub fn freeze(mut self) -> PrefixSet {
        if !self.all {
            self.keys.sort_unstable();
            self.keys.dedup();
            // Shrink after deduplication to release unused capacity.
            self.keys.shrink_to_fit();
        }
        PrefixSet::new(self.all, self.keys)
    }
}

impl<'a> IntoIterator for &'a PrefixSetMut {
    type Item = &'a Nibbles;
    type IntoIter = core::slice::Iter<'a, Nibbles>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

/// A sorted prefix set that has an immutable _sorted_ list of unique keys.
///
/// See also [`PrefixSetMut::freeze`].
#[derive(Debug, Clone)]
pub struct PrefixSet {
    /// Flag indicating that any entry should be considered changed.
    all: bool,
    index: usize,
    /// `None` for an empty set to avoid allocating an empty vector.
    keys: Option<Arc<Vec<Nibbles>>>,
}

impl Default for PrefixSet {
    fn default() -> Self {
        Self::new(false, Vec::new())
    }
}

impl<I> From<I> for PrefixSet
where
    I: Iterator<Item = B256>,
{
    fn from(keys: I) -> Self {
        let (lower_bound, upper_bound) = keys.size_hint();
        let mut unpacked = Vec::with_capacity(upper_bound.unwrap_or(lower_bound));
        let mut previous = None;

        for key in keys {
            debug_assert!(
                previous.is_none_or(|previous| previous <= key),
                "prefix set keys must be sorted"
            );
            if previous != Some(key) {
                unpacked.push(Nibbles::unpack(key));
            }
            previous = Some(key);
        }

        Self::new(false, unpacked)
    }
}

impl PrefixSet {
    fn new(all: bool, keys: Vec<Nibbles>) -> Self {
        Self { index: 0, all, keys: (!all && !keys.is_empty()).then(|| Arc::new(keys)) }
    }

    /// Creates a prefix set that considers every path changed.
    pub fn all_paths() -> Self {
        Self::new(true, Vec::new())
    }

    /// Returns `true` if any of the keys in the set has the given prefix
    ///
    /// # Note on Mutability
    ///
    /// This method requires `&mut self` (unlike typical `contains` methods) because it maintains an
    /// internal position tracker (`self.index`) between calls. This enables significant performance
    /// optimization for sequential lookups in sorted order, which is common during trie traversal.
    ///
    /// The `index` field allows subsequent searches to start where previous ones left off,
    /// avoiding repeated full scans of the prefix array when keys are accessed in nearby ranges.
    ///
    /// This optimization was inspired by Silkworm's implementation and significantly improves
    /// incremental state root calculation performance
    /// ([see PR #2417](https://github.com/paradigmxyz/reth/pull/2417)).
    #[inline]
    pub fn contains(&mut self, prefix: &Nibbles) -> bool {
        if self.all {
            return true
        }

        let keys = self.keys.as_deref().map(Vec::as_slice).unwrap_or_default();
        while self.index > 0 && &keys[self.index] > prefix {
            self.index -= 1;
        }

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

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

        false
    }

    /// Returns `true` if any key in the set falls within the given half-open range
    /// `[start, end)`.
    ///
    /// Like [`Self::contains`], this method maintains the internal index for sequential access
    /// optimization.
    #[inline]
    pub fn contains_range(&mut self, range: Range<&Nibbles>) -> bool {
        if self.all {
            return true
        }

        let keys = self.keys.as_deref().map(Vec::as_slice).unwrap_or_default();
        while self.index > 0 && &keys[self.index] >= range.end {
            self.index -= 1;
        }

        for (idx, key) in keys[self.index..].iter().enumerate() {
            if key >= range.start && key < range.end {
                self.index += idx;
                return true
            }

            if key >= range.end {
                self.index += idx;
                return false
            }
        }

        false
    }

    /// Returns `true` if any key in the set is at or after `start`.
    #[inline]
    pub fn contains_from(&mut self, start: &Nibbles) -> bool {
        if self.all {
            return true
        }

        let keys = self.keys.as_deref().map(Vec::as_slice).unwrap_or_default();
        while self.index > 0 && &keys[self.index] > start {
            self.index -= 1;
        }

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

        false
    }

    /// Returns an iterator over reference to _all_ nibbles regardless of cursor position.
    pub fn iter(&self) -> core::slice::Iter<'_, Nibbles> {
        self.slice().iter()
    }

    /// Returns the underlying sorted prefix slice.
    pub fn slice(&self) -> &[Nibbles] {
        self.keys.as_deref().map(Vec::as_slice).unwrap_or_default()
    }

    /// Returns true if every entry should be considered changed.
    pub const fn all(&self) -> bool {
        self.all
    }

    /// Returns the number of elements in the set.
    pub fn len(&self) -> usize {
        self.slice().len()
    }

    /// Returns `true` if the set is empty and `all` flag is not set.
    pub const fn is_empty(&self) -> bool {
        !self.all && self.keys.is_none()
    }
}

impl<'a> IntoIterator for &'a PrefixSet {
    type Item = &'a Nibbles;
    type IntoIter = core::slice::Iter<'a, Nibbles>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_primitives::B256;

    #[test]
    fn test_contains_with_multiple_inserts_and_duplicates() {
        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(&Nibbles::from_nibbles_unchecked([1, 2])));
        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
        assert_eq!(prefix_set.len(), 3); // Length should be 3 (excluding duplicate)
    }

    #[test]
    fn test_contains_from() {
        let first = Nibbles::from_nibbles([1, 2, 3]);
        let middle = Nibbles::from_nibbles([1, 2, 4]);
        let last = Nibbles::from_nibbles([4, 5, 6]);
        let mut prefix_set = PrefixSetMut::from([first, middle, last]).freeze();

        assert!(prefix_set.contains_range(&first..&middle));
        assert!(prefix_set.contains_from(&middle));
        assert!(prefix_set.contains_from(&last));
        assert!(!prefix_set.contains_from(&Nibbles::from_nibbles([5])));
    }

    #[test]
    fn test_from_sorted_b256_iterator() {
        let first = B256::with_last_byte(1);
        let second = B256::with_last_byte(2);
        let prefix_set = PrefixSet::from([first, second].into_iter());

        assert_eq!(prefix_set.slice(), &[Nibbles::unpack(first), Nibbles::unpack(second)]);
    }

    #[test]
    fn test_from_sorted_b256_iterator_with_duplicates() {
        let first = B256::with_last_byte(1);
        let second = B256::with_last_byte(2);
        let prefix_set = PrefixSet::from([first, first, second].into_iter());

        assert_eq!(prefix_set.slice(), &[Nibbles::unpack(first), Nibbles::unpack(second)]);
    }

    #[cfg(debug_assertions)]
    #[test]
    #[should_panic(expected = "prefix set keys must be sorted")]
    fn test_from_unsorted_b256_iterator() {
        let _: PrefixSet = [B256::with_last_byte(2), B256::with_last_byte(1)].into_iter().into();
    }

    #[test]
    fn test_freeze_shrinks_capacity() {
        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 is 4 (before deduplication)
        assert_eq!(prefix_set_mut.keys.capacity(), 4); // Capacity is 4 (before deduplication)

        let mut prefix_set = prefix_set_mut.freeze();
        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([1, 2])));
        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
        assert_eq!(prefix_set.slice().len(), 3); // Length should be 3 (excluding duplicate)
        assert_eq!(prefix_set.keys.as_ref().unwrap().capacity(), 3); // Capacity after shrinking
    }

    #[test]
    fn test_freeze_shrinks_existing_capacity() {
        // do the above test but with preallocated capacity
        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 is 4 (before deduplication)
        assert_eq!(prefix_set_mut.keys.capacity(), 101); // Capacity is 101 (before deduplication)

        let mut prefix_set = prefix_set_mut.freeze();
        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([1, 2])));
        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
        assert_eq!(prefix_set.slice().len(), 3); // Length should be 3 (excluding duplicate)
        assert_eq!(prefix_set.keys.as_ref().unwrap().capacity(), 3); // Capacity after shrinking
    }

    #[test]
    fn test_empty_prefix_sets_do_not_allocate() {
        assert!(PrefixSet::default().keys.is_none());
        assert!(PrefixSetMut::default().freeze().keys.is_none());
        assert!(PrefixSet::from(core::iter::empty()).keys.is_none());
        assert!(PrefixSet::all_paths().keys.is_none());
    }

    #[test]
    fn test_prefix_set_all_extend() {
        let mut prefix_set_mut = PrefixSetMut::default();
        prefix_set_mut.extend(PrefixSetMut::all());
        assert!(prefix_set_mut.all);
    }

    #[test]
    fn test_prefix_set_slice_returns_frozen_keys() {
        let path_a = Nibbles::from_nibbles([1, 2, 3]);
        let path_b = Nibbles::from_nibbles([4, 5, 6]);
        let mut prefix_set_mut = PrefixSetMut::default();
        prefix_set_mut.insert(path_b);
        prefix_set_mut.insert(path_a);
        prefix_set_mut.insert(path_b);

        let prefix_set = prefix_set_mut.freeze();
        assert_eq!(prefix_set.slice(), &[path_a, path_b]);
    }

    #[test]
    fn test_trie_prefix_sets_mut_extend_ref() {
        let account_path = Nibbles::from_nibbles([1, 2]);
        let storage_path = Nibbles::from_nibbles([3, 4]);
        let storage_account = B256::with_last_byte(1);
        let destroyed_account = B256::with_last_byte(2);
        let other = TriePrefixSetsMut {
            account_prefix_set: PrefixSetMut::from([account_path]),
            storage_prefix_sets: B256Map::from_iter([(
                storage_account,
                PrefixSetMut::from([storage_path]),
            )]),
            destroyed_accounts: B256Set::from_iter([destroyed_account]),
        };

        let mut prefix_sets = TriePrefixSetsMut::default();
        prefix_sets.extend_ref(&other);

        let frozen = prefix_sets.freeze();
        assert_eq!(frozen.account_prefix_set.slice(), &[account_path]);
        assert_eq!(frozen.storage_prefix_sets[&storage_account].slice(), &[storage_path]);
        assert!(frozen.destroyed_accounts.contains(&destroyed_account));
        assert_eq!(other.account_prefix_set.len(), 1);
    }
}