reth_trie_common/

target_v2.rs

//! V2 proof targets and chunking.

use crate::Nibbles;
use alloc::vec::Vec;
use alloy_primitives::{keccak256, map::B256Map, B256};
use evm2::evm::AccountChanges;

/// Target describes a proof target. For every proof target given, a proof calculator will calculate
/// and return all nodes whose path is a prefix of the target's `key_nibbles`.
#[derive(Debug, Copy, Clone)]
pub struct ProofV2Target {
    /// The key of the proof target, as nibbles.
    pub key_nibbles: Nibbles,
    /// The known-parent context for this target.
    pub parent: ProofV2TargetParent,
}

impl ProofV2Target {
    /// Returns a new [`ProofV2Target`] which matches all trie nodes whose path is a prefix of this
    /// key.
    pub fn new(key: B256) -> Self {
        // SAFETY: key is a B256 and so is exactly 32-bytes.
        let key_nibbles = unsafe { Nibbles::unpack_unchecked(key.as_slice()) };
        Self { key_nibbles, parent: ProofV2TargetParent::NONE }
    }

    /// Returns the key the target was initialized with.
    pub fn key(&self) -> B256 {
        B256::from_slice(&self.key_nibbles.pack())
    }

    /// Sets the already-revealed parent branch of this target.
    pub const fn with_parent(mut self, parent: ProofV2TargetParent) -> Self {
        self.parent = parent;
        self
    }
}

impl From<B256> for ProofV2Target {
    fn from(key: B256) -> Self {
        Self::new(key)
    }
}

/// The already-revealed parent branch of a [`ProofV2Target`].
///
/// [`Self::NONE`] indicates that no parent is known and the proof must include the actual trie
/// root. A known parent at path length `n` makes the proof start at its direct child at length
/// `n + 1`. In particular, a known parent at path length zero means the root branch is already
/// revealed, so the proof starts at one of its direct children. Known parent path lengths are
/// always less than 64.
///
/// Parent contexts are ordered from broadest to narrowest: [`Self::NONE`] precedes every known
/// parent, and known parents are ordered by path length.
#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ProofV2TargetParent(Option<u8>);

impl ProofV2TargetParent {
    /// No parent branch is known, so the proof must include the actual trie root.
    pub const NONE: Self = Self(None);

    /// Returns a known parent branch with the given logical path length.
    ///
    /// # Panics
    ///
    /// Panics if `path_len` is greater than or equal to 64.
    pub const fn new(path_len: usize) -> Self {
        assert!(path_len < 64, "parent path length must be less than 64");
        Self(Some(path_len as u8))
    }

    /// Returns `true` if the parent branch is already known.
    pub const fn is_known(self) -> bool {
        self.0.is_some()
    }

    /// Returns the logical path length of the known parent branch, if any.
    pub const fn path_len(self) -> Option<usize> {
        match self.0 {
            Some(path_len) => Some(path_len as usize),
            None => None,
        }
    }

    /// Returns the path of the known parent branch for the target path, if any.
    pub fn path(self, mut target_path: Nibbles) -> Option<Nibbles> {
        target_path.truncate(self.path_len()?);
        Some(target_path)
    }
}

/// A set of account and storage V2 proof targets. The account and storage targets do not need to
/// necessarily overlap.
#[derive(Debug, Default)]
pub struct MultiProofTargetsV2 {
    /// The set of account proof targets to generate proofs for.
    pub account_targets: Vec<ProofV2Target>,
    /// The sets of storage proof targets to generate proofs for.
    pub storage_targets: B256Map<Vec<ProofV2Target>>,
}

impl MultiProofTargetsV2 {
    /// Returns true is there are no account or storage targets.
    pub fn is_empty(&self) -> bool {
        self.account_targets.is_empty() && self.storage_targets.is_empty()
    }

    /// Returns the number of items that will be considered during chunking.
    pub fn chunking_length(&self) -> usize {
        self.account_targets.len() +
            self.storage_targets.values().map(|slots| slots.len()).sum::<usize>()
    }

    /// Returns an iterator that yields chunks of the specified size.
    ///
    /// # Panics
    ///
    /// Panics if the chunk size is zero.
    pub fn chunks(self, chunk_size: usize) -> impl Iterator<Item = Self> {
        ChunkedMultiProofTargetsV2::new(self, chunk_size)
    }

    /// Returns a set of [`MultiProofTargetsV2`] and the total amount of storage targets, based on
    /// the given account changes.
    ///
    /// Accounts that were not changed, or that were selfdestructed, get no targets. Since
    /// selfdestruct can only happen in the creating transaction, proofs for those accounts are not
    /// worth prefetching.
    ///
    /// See: <https://eips.ethereum.org/EIPS/eip-6780>
    pub fn from_account_changes<'a>(
        changes: impl IntoIterator<Item = AccountChanges<'a>>,
    ) -> (Self, usize) {
        let mut targets = Self::default();
        let mut storage_target_count = 0;
        for changes in changes {
            if changes.selfdestructed {
                continue
            }

            let storage_slots = changes
                .storage
                .changed_slots()
                .map(|(key, _)| ProofV2Target::new(keccak256(B256::from(*key))))
                .collect::<Vec<_>>();
            let info_changed = changes.original != changes.current;
            if !info_changed && storage_slots.is_empty() {
                continue
            }

            let hashed_address = keccak256(changes.address);
            if info_changed {
                targets.account_targets.push(ProofV2Target::new(hashed_address));
            }
            if !storage_slots.is_empty() {
                storage_target_count += storage_slots.len();
                targets.storage_targets.insert(hashed_address, storage_slots);
            }
        }

        (targets, storage_target_count)
    }
}

/// An iterator that yields chunks of V2 proof targets of at most `size` account and storage
/// targets.
///
/// Unlike legacy chunking, V2 preserves account targets exactly as they were (including their
/// parent metadata). Account targets must appear in a chunk. Storage targets for those accounts
/// are chunked together, but if they exceed the chunk size, subsequent chunks contain only the
/// remaining storage targets without repeating the account target.
#[derive(Debug)]
pub struct ChunkedMultiProofTargetsV2 {
    /// Remaining account targets to process
    account_targets: alloc::vec::IntoIter<ProofV2Target>,
    /// Storage targets by account address
    storage_targets: B256Map<Vec<ProofV2Target>>,
    /// Current account being processed (if any storage slots remain)
    current_account_storage: Option<(B256, alloc::vec::IntoIter<ProofV2Target>)>,
    /// Chunk size
    size: usize,
}

impl ChunkedMultiProofTargetsV2 {
    /// Creates a new chunked iterator for the given targets.
    ///
    /// # Panics
    ///
    /// Panics if `size` is zero.
    pub fn new(targets: MultiProofTargetsV2, size: usize) -> Self {
        assert!(size > 0, "chunk size must be non-zero");
        Self {
            account_targets: targets.account_targets.into_iter(),
            storage_targets: targets.storage_targets,
            current_account_storage: None,
            size,
        }
    }
}

impl Iterator for ChunkedMultiProofTargetsV2 {
    type Item = MultiProofTargetsV2;

    fn next(&mut self) -> Option<Self::Item> {
        let mut chunk = MultiProofTargetsV2::default();
        let mut count = 0;

        // First, finish any remaining storage slots from previous account
        if let Some((account_addr, ref mut storage_iter)) = self.current_account_storage {
            let remaining_capacity = self.size - count;
            let slots: Vec<_> = storage_iter.by_ref().take(remaining_capacity).collect();

            count += slots.len();
            chunk.storage_targets.insert(account_addr, slots);

            // If iterator is exhausted, clear current_account_storage
            if storage_iter.len() == 0 {
                self.current_account_storage = None;
            }
        }

        // Process account targets and their storage
        while count < self.size {
            let Some(account_target) = self.account_targets.next() else {
                break;
            };

            // Add the account target
            chunk.account_targets.push(account_target);
            count += 1;

            // Check if this account has storage targets
            let account_addr = account_target.key();
            if let Some(storage_slots) = self.storage_targets.remove(&account_addr) {
                let remaining_capacity = self.size - count;

                if storage_slots.len() <= remaining_capacity {
                    // Optimization: We can take all slots, just move the vec
                    count += storage_slots.len();
                    chunk.storage_targets.insert(account_addr, storage_slots);
                } else {
                    // We need to split the storage slots
                    let mut storage_iter = storage_slots.into_iter();
                    let slots_in_chunk: Vec<_> =
                        storage_iter.by_ref().take(remaining_capacity).collect();
                    count += slots_in_chunk.len();

                    chunk.storage_targets.insert(account_addr, slots_in_chunk);

                    // Save remaining storage slots for next chunk
                    self.current_account_storage = Some((account_addr, storage_iter));
                    break;
                }
            }
        }

        // Process any remaining storage-only entries (accounts not in account_targets)
        while let Some((account_addr, storage_slots)) = self.storage_targets.iter_mut().next() &&
            count < self.size
        {
            let account_addr = *account_addr;
            let storage_slots = core::mem::take(storage_slots);
            let remaining_capacity = self.size - count;

            // Always remove from the map - if there are remaining slots they go to
            // current_account_storage
            self.storage_targets.remove(&account_addr);

            if storage_slots.len() <= remaining_capacity {
                // Optimization: We can take all slots, just move the vec
                count += storage_slots.len();
                chunk.storage_targets.insert(account_addr, storage_slots);
            } else {
                // We need to split the storage slots
                let mut storage_iter = storage_slots.into_iter();
                let slots_in_chunk: Vec<_> =
                    storage_iter.by_ref().take(remaining_capacity).collect();

                chunk.storage_targets.insert(account_addr, slots_in_chunk);

                // Save remaining storage slots for next chunk
                if storage_iter.len() > 0 {
                    self.current_account_storage = Some((account_addr, storage_iter));
                }
                break;
            }
        }

        if chunk.account_targets.is_empty() && chunk.storage_targets.is_empty() {
            None
        } else {
            Some(chunk)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy_primitives::{Address, U256};
    use evm2::evm::{AccountInfo, StorageOverlay, StorageSlot, Tracked};

    fn overlay(slots: &[(u64, u64, u64)]) -> StorageOverlay {
        let mut storage = StorageOverlay::default();
        for &(key, original, current) in slots {
            storage.slots.insert(
                U256::from(key),
                StorageSlot {
                    value: Tracked::from_parts(U256::from(original), U256::from(current)),
                    ..Default::default()
                },
            );
        }
        storage
    }

    fn changes<'a>(
        address: Address,
        original: &'a AccountInfo,
        current: Option<&'a AccountInfo>,
        selfdestructed: bool,
        storage: &'a StorageOverlay,
    ) -> AccountChanges<'a> {
        AccountChanges {
            address,
            original: Some(original),
            current,
            created: true,
            selfdestructed,
            touched: false,
            code: None,
            storage,
        }
    }

    #[test]
    fn from_account_changes_ignores_unchanged_reads() {
        let info = AccountInfo::default();
        let storage = overlay(&[(7, 8, 8)]);
        let (targets, count) = MultiProofTargetsV2::from_account_changes([AccountChanges {
            created: false,
            ..changes(Address::repeat_byte(1), &info, Some(&info), false, &storage)
        }]);
        assert!(targets.is_empty());
        assert_eq!(count, 0);
    }

    #[test]
    fn from_account_changes_includes_changes() {
        let address = Address::repeat_byte(1);
        let original = AccountInfo::default();
        let current = AccountInfo { nonce: 1, ..original.clone() };
        let storage = overlay(&[(7, 0, 8)]);
        let (targets, count) = MultiProofTargetsV2::from_account_changes([changes(
            address,
            &original,
            Some(&current),
            false,
            &storage,
        )]);
        assert_eq!(
            targets.account_targets.iter().map(ProofV2Target::key).collect::<Vec<_>>(),
            vec![keccak256(address)]
        );
        assert_eq!(
            targets.storage_targets[&keccak256(address)]
                .iter()
                .map(ProofV2Target::key)
                .collect::<Vec<_>>(),
            vec![keccak256(B256::from(U256::from(7)))]
        );
        assert_eq!(count, 1);
    }

    #[test]
    fn from_account_changes_skips_selfdestructed() {
        let address = Address::repeat_byte(1);
        let retained = Address::repeat_byte(2);
        let info = AccountInfo::default();
        let storage = overlay(&[(7, 0, 8)]);
        let (targets, count) = MultiProofTargetsV2::from_account_changes([
            changes(address, &info, None, true, &storage),
            // Lifecycle flags alone do not imply an account trie update.
            changes(retained, &info, Some(&info), false, &storage),
        ]);
        assert!(targets.account_targets.is_empty());
        assert_eq!(targets.storage_targets.len(), 1);
        assert!(targets.storage_targets.contains_key(&keccak256(retained)));
        assert_eq!(count, 1);
    }
}