reth_trie_parallel/
targets_v2.rs

1//! V2 multiproof targets and chunking.
2
3use alloy_primitives::{map::B256Map, B256};
4use reth_trie::proof_v2;
5
6/// A set of account and storage V2 proof targets. The account and storage targets do not need to
7/// necessarily overlap.
8#[derive(Debug, Default)]
9pub struct MultiProofTargetsV2 {
10    /// The set of account proof targets to generate proofs for.
11    pub account_targets: Vec<proof_v2::Target>,
12    /// The sets of storage proof targets to generate proofs for.
13    pub storage_targets: B256Map<Vec<proof_v2::Target>>,
14}
15
16impl MultiProofTargetsV2 {
17    /// Returns true is there are no account or storage targets.
18    pub fn is_empty(&self) -> bool {
19        self.account_targets.is_empty() && self.storage_targets.is_empty()
20    }
21}
22
23/// An iterator that yields chunks of V2 proof targets of at most `size` account and storage
24/// targets.
25///
26/// Unlike legacy chunking, V2 preserves account targets exactly as they were (with their `min_len`
27/// metadata). Account targets must appear in a chunk. Storage targets for those accounts are
28/// chunked together, but if they exceed the chunk size, subsequent chunks contain only the
29/// remaining storage targets without repeating the account target.
30#[derive(Debug)]
31pub struct ChunkedMultiProofTargetsV2 {
32    /// Remaining account targets to process
33    account_targets: std::vec::IntoIter<proof_v2::Target>,
34    /// Storage targets by account address
35    storage_targets: B256Map<Vec<proof_v2::Target>>,
36    /// Current account being processed (if any storage slots remain)
37    current_account_storage: Option<(B256, std::vec::IntoIter<proof_v2::Target>)>,
38    /// Chunk size
39    size: usize,
40}
41
42impl ChunkedMultiProofTargetsV2 {
43    /// Creates a new chunked iterator for the given targets.
44    pub fn new(targets: MultiProofTargetsV2, size: usize) -> Self {
45        Self {
46            account_targets: targets.account_targets.into_iter(),
47            storage_targets: targets.storage_targets,
48            current_account_storage: None,
49            size,
50        }
51    }
52}
53
54impl Iterator for ChunkedMultiProofTargetsV2 {
55    type Item = MultiProofTargetsV2;
56
57    fn next(&mut self) -> Option<Self::Item> {
58        let mut chunk = MultiProofTargetsV2::default();
59        let mut count = 0;
60
61        // First, finish any remaining storage slots from previous account
62        if let Some((account_addr, ref mut storage_iter)) = self.current_account_storage {
63            let remaining_capacity = self.size - count;
64            let slots: Vec<_> = storage_iter.by_ref().take(remaining_capacity).collect();
65
66            count += slots.len();
67            chunk.storage_targets.insert(account_addr, slots);
68
69            // If iterator is exhausted, clear current_account_storage
70            if storage_iter.len() == 0 {
71                self.current_account_storage = None;
72            }
73        }
74
75        // Process account targets and their storage
76        while count < self.size {
77            let Some(account_target) = self.account_targets.next() else {
78                break;
79            };
80
81            // Add the account target
82            chunk.account_targets.push(account_target);
83            count += 1;
84
85            // Check if this account has storage targets
86            let account_addr = account_target.key();
87            if let Some(storage_slots) = self.storage_targets.remove(&account_addr) {
88                let remaining_capacity = self.size - count;
89
90                if storage_slots.len() <= remaining_capacity {
91                    // Optimization: We can take all slots, just move the vec
92                    count += storage_slots.len();
93                    chunk.storage_targets.insert(account_addr, storage_slots);
94                } else {
95                    // We need to split the storage slots
96                    let mut storage_iter = storage_slots.into_iter();
97                    let slots_in_chunk: Vec<_> =
98                        storage_iter.by_ref().take(remaining_capacity).collect();
99                    count += slots_in_chunk.len();
100
101                    chunk.storage_targets.insert(account_addr, slots_in_chunk);
102
103                    // Save remaining storage slots for next chunk
104                    self.current_account_storage = Some((account_addr, storage_iter));
105                    break;
106                }
107            }
108        }
109
110        // Process any remaining storage-only entries (accounts not in account_targets)
111        while let Some((account_addr, storage_slots)) = self.storage_targets.iter_mut().next() &&
112            count < self.size
113        {
114            let account_addr = *account_addr;
115            let storage_slots = std::mem::take(storage_slots);
116            let remaining_capacity = self.size - count;
117
118            // Always remove from the map - if there are remaining slots they go to
119            // current_account_storage
120            self.storage_targets.remove(&account_addr);
121
122            if storage_slots.len() <= remaining_capacity {
123                // Optimization: We can take all slots, just move the vec
124                count += storage_slots.len();
125                chunk.storage_targets.insert(account_addr, storage_slots);
126            } else {
127                // We need to split the storage slots
128                let mut storage_iter = storage_slots.into_iter();
129                let slots_in_chunk: Vec<_> =
130                    storage_iter.by_ref().take(remaining_capacity).collect();
131
132                chunk.storage_targets.insert(account_addr, slots_in_chunk);
133
134                // Save remaining storage slots for next chunk
135                if storage_iter.len() > 0 {
136                    self.current_account_storage = Some((account_addr, storage_iter));
137                }
138                break;
139            }
140        }
141
142        if chunk.account_targets.is_empty() && chunk.storage_targets.is_empty() {
143            None
144        } else {
145            Some(chunk)
146        }
147    }
148}
reth_trie_parallel/targets_v2.rs

reth_trie_parallel/
targets_v2.rs
