diff --git a/crates/trie/sparse-parallel/src/lower.rs b/crates/trie/sparse-parallel/src/lower.rs index b7eceb133b..c2438ae926 100644 --- a/crates/trie/sparse-parallel/src/lower.rs +++ b/crates/trie/sparse-parallel/src/lower.rs @@ -128,4 +128,12 @@ impl LowerSparseSubtrie { Self::Blind(None) => {} } } + + /// Returns a heuristic for the in-memory size of this subtrie in bytes. + pub(crate) fn memory_size(&self) -> usize { + match self { + Self::Revealed(subtrie) | Self::Blind(Some(subtrie)) => subtrie.memory_size(), + Self::Blind(None) => 0, + } + } } diff --git a/crates/trie/sparse-parallel/src/trie.rs b/crates/trie/sparse-parallel/src/trie.rs index 4f5eccf8b2..6d12f8562f 100644 --- a/crates/trie/sparse-parallel/src/trie.rs +++ b/crates/trie/sparse-parallel/src/trie.rs @@ -2131,6 +2131,50 @@ impl ParallelSparseTrie { self.subtrie_heat.mark_modified(index); } } + + /// Returns a heuristic for the in-memory size of this trie in bytes. + /// + /// This is an approximation that accounts for: + /// - The upper subtrie nodes and values + /// - All revealed lower subtries nodes and values + /// - The prefix set keys + /// - The branch node masks map + /// - Updates if retained + /// - Update action buffers + /// + /// Note: Heap allocations for hash maps may be larger due to load factor overhead. + pub fn memory_size(&self) -> usize { + let mut size = core::mem::size_of::(); + + // Upper subtrie + size += self.upper_subtrie.memory_size(); + + // Lower subtries (both Revealed and Blind with allocation) + for subtrie in &self.lower_subtries { + size += subtrie.memory_size(); + } + + // Prefix set keys + size += self.prefix_set.len() * core::mem::size_of::(); + + // Branch node masks map + size += self.branch_node_masks.len() * + (core::mem::size_of::() + core::mem::size_of::()); + + // Updates if present + if let Some(updates) = &self.updates { + size += updates.updated_nodes.len() * + (core::mem::size_of::() + core::mem::size_of::()); + size += updates.removed_nodes.len() * core::mem::size_of::(); + } + + // Update actions buffers + for buf in &self.update_actions_buffers { + size += buf.capacity() * core::mem::size_of::(); + } + + size + } } /// Bitset tracking which of the 256 lower subtries were modified in the current cycle. @@ -2825,6 +2869,30 @@ impl SparseSubtrie { pub(crate) fn shrink_values_to(&mut self, size: usize) { self.inner.values.shrink_to(size); } + + /// Returns a heuristic for the in-memory size of this subtrie in bytes. + pub(crate) fn memory_size(&self) -> usize { + let mut size = core::mem::size_of::(); + + // Nodes map: key (Nibbles) + value (SparseNode) + for (path, node) in &self.nodes { + size += core::mem::size_of::(); + size += path.len(); // Nibbles heap allocation + size += node.memory_size(); + } + + // Values map: key (Nibbles) + value (Vec) + for (path, value) in &self.inner.values { + size += core::mem::size_of::(); + size += path.len(); // Nibbles heap allocation + size += core::mem::size_of::>() + value.capacity(); + } + + // Buffers + size += self.inner.buffers.memory_size(); + + size + } } /// Helper type for [`SparseSubtrie`] to mutably access only a subset of fields from the original @@ -3298,6 +3366,19 @@ impl SparseSubtrieBuffers { self.branch_value_stack_buf.clear(); self.rlp_buf.clear(); } + + /// Returns a heuristic for the in-memory size of these buffers in bytes. + const fn memory_size(&self) -> usize { + let mut size = core::mem::size_of::(); + + size += self.path_stack.capacity() * core::mem::size_of::(); + size += self.rlp_node_stack.capacity() * core::mem::size_of::(); + size += self.branch_child_buf.capacity() * core::mem::size_of::(); + size += self.branch_value_stack_buf.capacity() * core::mem::size_of::(); + size += self.rlp_buf.capacity(); + + size + } } /// RLP node path stack item. @@ -8908,4 +8989,41 @@ mod tests { b256!("f000000000000000000000000000000000000000000000000000000000000000"); assert_eq!(ParallelSparseTrie::nibbles_to_padded_b256(&single), expected_single); } + + #[test] + fn test_memory_size() { + // Test that memory_size returns a reasonable value for an empty trie + let trie = ParallelSparseTrie::default(); + let empty_size = trie.memory_size(); + + // Should at least be the size of the struct itself + assert!(empty_size >= core::mem::size_of::()); + + // Create a trie with some data + let mut trie = ParallelSparseTrie::default(); + let nodes = vec![ + ProofTrieNode { + path: Nibbles::from_nibbles_unchecked([0x1, 0x2]), + node: TrieNode::Leaf(LeafNode { + key: Nibbles::from_nibbles_unchecked([0x3, 0x4]), + value: vec![1, 2, 3], + }), + masks: None, + }, + ProofTrieNode { + path: Nibbles::from_nibbles_unchecked([0x5, 0x6]), + node: TrieNode::Leaf(LeafNode { + key: Nibbles::from_nibbles_unchecked([0x7, 0x8]), + value: vec![4, 5, 6], + }), + masks: None, + }, + ]; + trie.reveal_nodes(nodes).unwrap(); + + let populated_size = trie.memory_size(); + + // Populated trie should use more memory than an empty one + assert!(populated_size > empty_size); + } } diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index 8aa6a1ab5b..3387a749df 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -1984,6 +1984,16 @@ impl SparseNode { } } } + + /// Returns the memory size of this node in bytes. + pub const fn memory_size(&self) -> usize { + match self { + Self::Empty | Self::Hash(_) | Self::Branch { .. } => core::mem::size_of::(), + Self::Leaf { key, .. } | Self::Extension { key, .. } => { + core::mem::size_of::() + key.len() + } + } + } } /// A helper struct used to store information about a node that has been removed