reth_trie_sparse/

traits.rs

//! Traits for sparse trie implementations.

use core::fmt::Debug;

use alloc::{borrow::Cow, vec::Vec};
use alloy_primitives::{
    map::{B256Map, HashMap, HashSet},
    B256,
};
use alloy_trie::BranchNodeCompact;
use reth_execution_errors::SparseTrieResult;
use reth_trie_common::{
    BranchNodeMasks, Nibbles, ProofTrieNodeV2, ProofV2TargetParent, TrieNodeV2,
};

/// Modification epoch assigned to cached sparse trie nodes.
///
/// Epochs must increase monotonically. Nodes materialized from the parent state without being
/// modified use [`Self::UNMODIFIED`].
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct TrieNodeEpoch(u64);

impl TrieNodeEpoch {
    /// Epoch assigned to nodes materialized from the parent state without being modified.
    pub const UNMODIFIED: Self = Self(0);

    /// Creates a new node modification epoch.
    pub const fn new(epoch: u64) -> Self {
        Self(epoch)
    }

    /// Returns the inner epoch.
    pub const fn get(self) -> u64 {
        self.0
    }

    /// Returns whether a node with this epoch should be pruned at the provided cutoff.
    pub const fn should_prune(self, prune_before: Self) -> bool {
        self.0 < prune_before.0
    }
}

/// Describes an update to a leaf in the sparse trie.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeafUpdate {
    /// The leaf value has been changed to the given RLP-encoded value.
    /// Empty Vec indicates the leaf has been removed.
    Changed(Vec<u8>),
    /// The leaf value may have changed, but the new value is not yet known.
    /// Used for optimistic prewarming when the actual value is unavailable.
    Touched,
}

impl LeafUpdate {
    /// Returns true if the leaf update is a change.
    pub const fn is_changed(&self) -> bool {
        matches!(self, Self::Changed(_))
    }

    /// Returns true if the leaf update is a touched update.
    pub const fn is_touched(&self) -> bool {
        matches!(self, Self::Touched)
    }
}

/// Trait defining common operations for revealed sparse trie implementations.
///
/// This trait provides a unified interface for the core trie operations needed by
/// `RevealableSparseTrie`.
pub trait SparseTrie: Sized + Debug + Send + Sync {
    /// Configures the trie to have the given root node revealed.
    ///
    /// # Arguments
    ///
    /// * `root` - The root node to reveal
    /// * `masks` - Trie masks for root branch node
    /// * `retain_updates` - Whether to track updates
    ///
    /// # Returns
    ///
    /// `Ok(())` if successful, or an error if revealing fails.
    ///
    /// # Panics
    ///
    /// May panic if the trie is not new/cleared, and has already revealed nodes.
    fn set_root(
        &mut self,
        root: TrieNodeV2,
        masks: Option<BranchNodeMasks>,
        retain_updates: bool,
    ) -> SparseTrieResult<()>;

    /// Configures the trie to retain information about updates.
    ///
    /// If `retain_updates` is true, the trie will record branch node updates
    /// and deletions. This information can be used to efficiently update
    /// an external database.
    ///
    /// # Arguments
    ///
    /// * `retain_updates` - Whether to track updates
    fn set_updates(&mut self, retain_updates: bool);

    /// Reveals one or more trie nodes if they have not been revealed before.
    ///
    /// This function decodes trie nodes and inserts them into the trie structure. It handles
    /// different node types (leaf, extension, branch) by appropriately adding them to the trie and
    /// recursively revealing their children.
    ///
    /// # Arguments
    ///
    /// * `nodes` - The nodes to be revealed, each having a path and optional set of branch node
    ///   masks. The nodes will be unsorted.
    ///
    /// # Returns
    ///
    /// `Ok(())` if successful, or an error if any of the nodes was not revealed.
    ///
    /// # Note
    ///
    /// The implementation may modify the input nodes. A common thing to do is [`std::mem::replace`]
    /// each node with [`TrieNodeV2::EmptyRoot`] to avoid cloning.
    fn reveal_nodes(&mut self, nodes: &mut [ProofTrieNodeV2]) -> SparseTrieResult<()>;

    /// Calculates and returns the root hash of the trie at the provided epoch.
    ///
    /// This processes dirty nodes by updating their RLP encodings and caching their newest
    /// modification at `new_epoch`, then returns the root hash.
    ///
    /// # Returns
    ///
    /// The root hash of the trie.
    fn root(&mut self, new_epoch: TrieNodeEpoch) -> B256;

    /// Returns true if the root node is cached and does not need any recomputation.
    fn is_root_cached(&self) -> bool;

    /// Returns the root's modification epoch when it is clean, or `None` when it is dirty.
    fn root_epoch(&self) -> Option<TrieNodeEpoch>;

    /// Recalculates and updates the RLP hashes of subtries deeper than a certain level. The level
    /// is defined in the implementation.
    ///
    /// The root node is considered to be at level 0. This method is useful for optimizing
    /// hash recalculations after localized changes to the trie structure.
    fn update_subtrie_hashes(&mut self, new_epoch: TrieNodeEpoch);

    /// Retrieves a reference to the leaf value at the specified path.
    ///
    /// # Arguments
    ///
    /// * `full_path` - The full path to the leaf value
    ///
    /// # Returns
    ///
    /// A reference to the leaf value stored at the given full path, if it is revealed.
    ///
    /// Note: a value can exist in the full trie and this function still returns `None`
    /// because the value has not been revealed.
    ///
    /// Hence a `None` indicates two possibilities:
    /// - The value does not exists in the trie, so it cannot be revealed
    /// - The value has not yet been revealed. In order to determine which is true, one would need
    ///   an exclusion proof.
    fn get_leaf_value(&self, full_path: &Nibbles) -> Option<&Vec<u8>>;

    /// Attempts to find a leaf node at the specified path.
    ///
    /// This method traverses the trie from the root down to the given path, checking
    /// if a leaf exists at that path. It can be used to verify the existence of a leaf
    /// or to generate an exclusion proof (proof that a leaf does not exist).
    ///
    /// # Parameters
    ///
    /// - `full_path`: The path to search for.
    /// - `expected_value`: Optional expected value. If provided, will verify the leaf value
    ///   matches.
    ///
    /// # Returns
    ///
    /// - `Ok(LeafLookup::Exists)` if the leaf exists with the expected value.
    /// - `Ok(LeafLookup::NonExistent)` if the leaf definitely does not exist (exclusion proof).
    /// - `Err(LeafLookupError)` if the search encountered a blinded node or found a different
    ///   value.
    fn find_leaf(
        &self,
        full_path: &Nibbles,
        expected_value: Option<&Vec<u8>>,
    ) -> Result<LeafLookup, LeafLookupError>;

    /// Returns a reference to the current sparse trie updates.
    ///
    /// If no updates have been made/recorded, returns an empty update set.
    fn updates_ref(&self) -> Cow<'_, SparseTrieUpdates>;

    /// Consumes and returns the currently accumulated trie updates.
    ///
    /// This is useful when you want to apply the updates to an external database
    /// and then start tracking a new set of updates.
    ///
    /// # Returns
    ///
    /// The accumulated updates, or an empty set if updates weren't being tracked.
    fn take_updates(&mut self) -> SparseTrieUpdates;

    /// This clears all data structures in the sparse trie, keeping the backing data structures
    /// allocated. An empty root node is inserted at the root.
    ///
    /// This is useful for reusing the trie without needing to reallocate memory.
    fn clear(&mut self);

    /// Collapses nodes last modified before `prune_before` into hash stubs.
    ///
    /// # Preconditions
    ///
    /// The trie must not be dirty. An unmodified revealed root may be pruned because
    /// proof-revealed descendants carry cached RLP nodes.
    ///
    /// # Returns
    ///
    /// The number of nodes converted to hash stubs.
    fn prune(&mut self, prune_before: TrieNodeEpoch) -> usize;

    /// Applies leaf updates to the sparse trie.
    ///
    /// When a [`LeafUpdate::Changed`] is successfully applied, it is removed from the
    /// given [`B256Map`]. If it could not be applied due to blinded nodes, it remains
    /// in the map and the callback is invoked with the required proof target.
    ///
    /// Once that proof is calculated and revealed via [`SparseTrie::reveal_nodes`], the same
    /// `updates` map can be reused to retry the update.
    ///
    /// The callback receives `(key, parent)` where `key` is the full 32-byte hashed key
    /// (right-padded with zeros from the blinded path) and `parent` identifies the revealed logical
    /// parent branch. No known parent indicates that the trie is entirely blind and the proof
    /// must include the root.
    ///
    /// The callback may be invoked multiple times for the same target across retry loops.
    /// Callers should deduplicate if needed.
    ///
    /// [`LeafUpdate::Touched`] behaves identically except it does not modify the leaf value.
    fn update_leaves(
        &mut self,
        updates: &mut B256Map<LeafUpdate>,
        proof_required_fn: impl FnMut(B256, ProofV2TargetParent),
    ) -> SparseTrieResult<()>;
}

/// Tracks modifications to the sparse trie structure.
///
/// Maintains references to both modified and pruned/removed branches, enabling
/// one to make batch updates to a persistent database.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SparseTrieUpdates {
    /// Collection of updated intermediate nodes indexed by full path.
    pub updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
    /// Collection of removed intermediate nodes indexed by full path.
    pub removed_nodes: HashSet<Nibbles>,
}

impl SparseTrieUpdates {
    /// Initialize a [`Self`] with given capacities.
    pub fn with_capacity(num_updated_nodes: usize, num_removed_nodes: usize) -> Self {
        Self {
            updated_nodes: HashMap::with_capacity_and_hasher(num_updated_nodes, Default::default()),
            removed_nodes: HashSet::with_capacity_and_hasher(num_removed_nodes, Default::default()),
        }
    }
}

/// Error type for a leaf lookup operation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeafLookupError {
    /// The path leads to a blinded node, cannot determine if leaf exists.
    /// This means the witness is not complete.
    BlindedNode {
        /// Path to the blinded node.
        path: Nibbles,
        /// Hash of the blinded node.
        hash: B256,
    },
    /// The path leads to a leaf with a different value than expected.
    /// This means the witness is malformed.
    ValueMismatch {
        /// Path to the leaf.
        path: Nibbles,
        /// Expected value.
        expected: Option<Vec<u8>>,
        /// Actual value found.
        actual: Vec<u8>,
    },
}

/// Success value for a leaf lookup operation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LeafLookup {
    /// Leaf exists with expected value.
    Exists,
    /// Leaf does not exist (exclusion proof found).
    NonExistent,
}