diff --git a/core/store/src/trie/mem/construction.rs b/core/store/src/trie/mem/construction.rs index 3eee0e2ae7f..40478ada131 100644 --- a/core/store/src/trie/mem/construction.rs +++ b/core/store/src/trie/mem/construction.rs @@ -1,7 +1,9 @@ use super::arena::Arena; +use super::freelist::{ReusableVecU8, VecU8Freelist}; use super::node::MemTrieNodeId; use crate::trie::mem::node::InputMemTrieNode; use crate::NibbleSlice; +use elastic_array::ElasticArray4; use near_primitives::state::FlatStateValue; /// Algorithm to construct a trie from a given stream of sorted leaf values. @@ -64,6 +66,7 @@ use near_primitives::state::FlatStateValue; pub struct TrieConstructor<'a> { arena: &'a mut Arena, segments: Vec, + trail_freelist: VecU8Freelist, } /// A segment of the rightmost path of the trie under construction, as @@ -75,79 +78,137 @@ struct TrieConstructionSegment { is_branch: bool, // The trail, an edge below this node. If this is a branch node, // it is the rightmost child edge. It is an encoded NibbleSlice. - trail: Vec, + // We use a freelist here because this would otherwise be frequently + // allocated and deallocated and slow down construction significantly. + trail: ReusableVecU8, // If present, it is either a Leaf node or BranchWithValue. value: Option, // Only used if is_branch is true. The children that are already // constructed. The last child currently being constructed is not in here. - children: Vec<(u8, MemTrieNodeId)>, + children: ElasticArray4<(u8, MemTrieNodeId)>, // Only used for extension nodes; the child that is already constructed. child: Option, } impl TrieConstructionSegment { /// Prepares a segment that represents a branch node, possibly with value. - fn new_branch(initial_trail: Vec, value: Option) -> Self { - Self { is_branch: true, trail: initial_trail, value, children: Vec::new(), child: None } + fn new_branch(initial_trail: ReusableVecU8, value: Option) -> Self { + Self { + is_branch: true, + trail: initial_trail, + value, + children: Default::default(), + child: None, + } } /// Prepares a segment that represents an extension node. - fn new_extension(trail: Vec) -> Self { + fn new_extension(trail: ReusableVecU8) -> Self { let nibbles = NibbleSlice::from_encoded(&trail); assert!(!nibbles.1); // nibble slice is not leaf assert!(!nibbles.0.is_empty()); // extension nibbles cannot be empty - Self { is_branch: false, trail, value: None, children: Vec::new(), child: None } + Self { is_branch: false, trail, value: None, children: Default::default(), child: None } } /// Prepares a segment that represents a leaf node. - fn new_leaf(trail: Vec, value: FlatStateValue) -> Self { + fn new_leaf(trail: ReusableVecU8, value: FlatStateValue) -> Self { let nibbles = NibbleSlice::from_encoded(&trail); assert!(nibbles.1); - Self { is_branch: false, trail, value: Some(value), children: Vec::new(), child: None } + Self { + is_branch: false, + trail, + value: Some(value), + children: Default::default(), + child: None, + } } fn is_leaf(&self) -> bool { self.value.is_some() && !self.is_branch } - fn into_node(self, arena: &mut Arena) -> MemTrieNodeId { + fn to_node(&self, arena: &mut Arena) -> MemTrieNodeId { let input_node = if self.is_branch { assert!(!self.children.is_empty()); assert!(self.child.is_none()); let mut children = [None; 16]; - for (i, child) in self.children.into_iter() { - children[i as usize] = Some(child); + for (i, child) in self.children.iter() { + children[*i as usize] = Some(*child); } - if let Some(value) = self.value { + if let Some(value) = &self.value { InputMemTrieNode::BranchWithValue { children, value } } else { InputMemTrieNode::Branch { children } } - } else if let Some(value) = self.value { + } else if let Some(value) = &self.value { assert!(self.child.is_none()); assert!(self.children.is_empty()); - InputMemTrieNode::Leaf { value, extension: self.trail.into_boxed_slice() } + InputMemTrieNode::Leaf { value, extension: &self.trail } } else { assert!(self.child.is_some()); assert!(self.children.is_empty()); - InputMemTrieNode::Extension { - extension: self.trail.into_boxed_slice(), - child: self.child.unwrap(), - } + InputMemTrieNode::Extension { extension: &self.trail, child: self.child.unwrap() } }; MemTrieNodeId::new(arena, input_node) } } +/// A helper trait to make the construction code more readable. +/// +/// Whenever we encode nibbles to vector, we want to use a vector from the +/// freelist; otherwise allocation is quite slow. +trait NibblesHelper { + fn encode_to_vec(&self, freelist: &mut VecU8Freelist, is_leaf: bool) -> ReusableVecU8; + fn encode_leftmost_to_vec( + &self, + freelist: &mut VecU8Freelist, + len: usize, + is_leaf: bool, + ) -> ReusableVecU8; +} + +impl NibblesHelper for NibbleSlice<'_> { + fn encode_to_vec(&self, freelist: &mut VecU8Freelist, is_leaf: bool) -> ReusableVecU8 { + let mut vec = freelist.alloc(); + self.encode_to(is_leaf, vec.vec_mut()); + vec + } + + fn encode_leftmost_to_vec( + &self, + freelist: &mut VecU8Freelist, + len: usize, + is_leaf: bool, + ) -> ReusableVecU8 { + let mut vec = freelist.alloc(); + self.encode_leftmost_to(len, is_leaf, vec.vec_mut()); + vec + } +} + impl<'a> TrieConstructor<'a> { pub fn new(arena: &'a mut Arena) -> Self { - Self { arena, segments: vec![] } + // We should only have as many allocations as the number of segments + // alive, which is at most the length of keys. We give a generous + // margin on top of that. If this is exceeded in production, an error + // is printed; if exceeded in debug, it panics. + const EXPECTED_FREELIST_MAX_ALLOCATIONS: usize = 4096; + Self { + arena, + segments: vec![], + trail_freelist: VecU8Freelist::new(EXPECTED_FREELIST_MAX_ALLOCATIONS), + } + } + + fn recycle_segment(&mut self, segment: TrieConstructionSegment) { + self.trail_freelist.free(segment.trail); } /// Encodes the bottom-most segment into a node, and pops it off the stack. fn pop_segment(&mut self) { let segment = self.segments.pop().unwrap(); - let node = segment.into_node(self.arena); + let node = segment.to_node(self.arena); + self.recycle_segment(segment); let parent = self.segments.last_mut().unwrap(); if parent.is_branch { parent.children.push((NibbleSlice::from_encoded(&parent.trail).0.at(0), node)); @@ -205,9 +266,16 @@ impl<'a> TrieConstructor<'a> { assert_eq!(was_leaf, segment.child.is_none()); let top_segment = TrieConstructionSegment::new_extension( - extension_nibbles.encoded_leftmost(common_prefix_len, false).to_vec(), + extension_nibbles.encode_leftmost_to_vec( + &mut self.trail_freelist, + common_prefix_len, + false, + ), ); - segment.trail = extension_nibbles.mid(common_prefix_len).encoded(was_leaf).to_vec(); + let new_trail = extension_nibbles + .mid(common_prefix_len) + .encode_to_vec(&mut self.trail_freelist, was_leaf); + self.trail_freelist.free(std::mem::replace(&mut segment.trail, new_trail)); self.segments.push(top_segment); self.segments.push(segment); nibbles = nibbles.mid(common_prefix_len); @@ -218,8 +286,9 @@ impl<'a> TrieConstructor<'a> { if self.segments.last().unwrap().is_branch { // If the existing segment is a branch, we simply add another // case of the branch. - self.segments.last_mut().unwrap().trail = - nibbles.encoded_leftmost(1, false).to_vec(); + let segment = self.segments.last_mut().unwrap(); + let new_trail = nibbles.encode_leftmost_to_vec(&mut self.trail_freelist, 1, false); + self.trail_freelist.free(std::mem::replace(&mut segment.trail, new_trail)); nibbles = nibbles.mid(1); break; } else { @@ -235,7 +304,7 @@ impl<'a> TrieConstructor<'a> { // for this branch node begins with the old trail's first nibble. // We'll insert the new leaf later. let mut top_segment = TrieConstructionSegment::new_branch( - extension_nibbles.encoded_leftmost(1, false).to_vec(), + extension_nibbles.encode_leftmost_to_vec(&mut self.trail_freelist, 1, false), None, ); if extension_nibbles.len() > 1 || was_leaf { @@ -244,7 +313,9 @@ impl<'a> TrieConstructor<'a> { // Similarly, if the old segment had just 1 nibble but was a // leaf, we still need to keep the leaf but now with empty // trail on the leaf. - segment.trail = extension_nibbles.mid(1).encoded(was_leaf).to_vec(); + let new_trail = + extension_nibbles.mid(1).encode_to_vec(&mut self.trail_freelist, was_leaf); + self.trail_freelist.free(std::mem::replace(&mut segment.trail, new_trail)); self.segments.push(top_segment); self.segments.push(segment); // The bottom segment is no longer relevant to our new leaf, @@ -256,12 +327,14 @@ impl<'a> TrieConstructor<'a> { // extension segment's child directly to the branch node. top_segment.children.push((extension_nibbles.at(0), segment.child.unwrap())); self.segments.push(top_segment); + self.recycle_segment(segment); } // At this point we have popped the old case of the branch node, // so we advance the branch node to point to our new leaf // segment that we'll add below. - self.segments.last_mut().unwrap().trail = - nibbles.encoded_leftmost(1, false).to_vec(); + let segment = self.segments.last_mut().unwrap(); + let new_trail = nibbles.encode_leftmost_to_vec(&mut self.trail_freelist, 1, false); + self.trail_freelist.free(std::mem::replace(&mut segment.trail, new_trail)); nibbles = nibbles.mid(1); break; } @@ -276,30 +349,36 @@ impl<'a> TrieConstructor<'a> { assert!(!nibbles.is_empty()); // In order for a leaf node to have another leaf below it, it needs // to be converted to a branch node with value. - let segment = self.segments.pop().unwrap(); + let mut segment = self.segments.pop().unwrap(); let (extension_nibbles, was_leaf) = NibbleSlice::from_encoded(&segment.trail); assert!(was_leaf); if !extension_nibbles.is_empty() { // If the original leaf node had an extension within it, we need // to create an extension above the branch node. let top_segment = TrieConstructionSegment::new_extension( - extension_nibbles.encoded(false).to_vec(), + extension_nibbles.encode_to_vec(&mut self.trail_freelist, false), ); self.segments.push(top_segment); } // Now let's construct our branch node, and add our new leaf node below it. let mid_segment = TrieConstructionSegment::new_branch( - nibbles.encoded_leftmost(1, false).to_vec(), - segment.value, + nibbles.encode_leftmost_to_vec(&mut self.trail_freelist, 1, false), + std::mem::take(&mut segment.value), + ); + let bottom_segment = TrieConstructionSegment::new_leaf( + nibbles.mid(1).encode_to_vec(&mut self.trail_freelist, true), + value, ); - let bottom_segment = - TrieConstructionSegment::new_leaf(nibbles.mid(1).encoded(true).to_vec(), value); self.segments.push(mid_segment); self.segments.push(bottom_segment); + self.recycle_segment(segment); } else { // Otherwise we're at one branch of a branch node (or we're at root), // so just append the leaf. - let segment = TrieConstructionSegment::new_leaf(nibbles.encoded(true).to_vec(), value); + let segment = TrieConstructionSegment::new_leaf( + nibbles.encode_to_vec(&mut self.trail_freelist, true), + value, + ); self.segments.push(segment); } } @@ -313,6 +392,21 @@ impl<'a> TrieConstructor<'a> { while self.segments.len() > 1 { self.pop_segment(); } - self.segments.into_iter().next().map(|segment| segment.into_node(self.arena)) + if self.segments.is_empty() { + None + } else { + let segment = self.segments.pop().unwrap(); + let ret = segment.to_node(self.arena); + self.recycle_segment(segment); + Some(ret) + } + } +} + +impl<'a> Drop for TrieConstructor<'a> { + fn drop(&mut self) { + for segment in std::mem::take(&mut self.segments) { + self.recycle_segment(segment); + } } } diff --git a/core/store/src/trie/mem/flexible_data/children.rs b/core/store/src/trie/mem/flexible_data/children.rs index d1093deec5e..05011c2b94c 100644 --- a/core/store/src/trie/mem/flexible_data/children.rs +++ b/core/store/src/trie/mem/flexible_data/children.rs @@ -40,11 +40,11 @@ impl FlexibleDataHeader for EncodedChildrenHeader { fn encode_flexible_data( &self, - children: [Option; 16], + children: &[Option; 16], target: &mut ArenaSliceMut<'_>, ) { let mut j = 0; - for (i, child) in children.into_iter().enumerate() { + for (i, child) in children.iter().enumerate() { if self.mask & (1 << i) != 0 { target.write_pos_at(j, child.unwrap().pos); j += size_of::(); diff --git a/core/store/src/trie/mem/flexible_data/encoding.rs b/core/store/src/trie/mem/flexible_data/encoding.rs index db85f8a0d04..74c095ffc8f 100644 --- a/core/store/src/trie/mem/flexible_data/encoding.rs +++ b/core/store/src/trie/mem/flexible_data/encoding.rs @@ -51,7 +51,7 @@ impl<'a> RawEncoder<'a> { /// flexibly-sized part, as returned by `header.flexible_data_length()`. /// Note that the header itself is NOT encoded; only the flexible part is. /// The header is expected to have been encoded earlier. - pub fn encode_flexible(&mut self, header: &T, data: T::InputData) { + pub fn encode_flexible(&mut self, header: &T, data: &T::InputData) { let length = header.flexible_data_length(); header.encode_flexible_data(data, &mut self.data.subslice_mut(self.pos, length)); self.pos += length; diff --git a/core/store/src/trie/mem/flexible_data/extension.rs b/core/store/src/trie/mem/flexible_data/extension.rs index 3c67fc926e3..1c173e1d6a1 100644 --- a/core/store/src/trie/mem/flexible_data/extension.rs +++ b/core/store/src/trie/mem/flexible_data/extension.rs @@ -16,9 +16,9 @@ impl BorshFixedSize for EncodedExtensionHeader { } impl FlexibleDataHeader for EncodedExtensionHeader { - type InputData = Box<[u8]>; + type InputData = [u8]; type View<'a> = ArenaSlice<'a>; - fn from_input(extension: &Box<[u8]>) -> EncodedExtensionHeader { + fn from_input(extension: &[u8]) -> EncodedExtensionHeader { EncodedExtensionHeader { length: extension.len() as u16 } } @@ -26,7 +26,7 @@ impl FlexibleDataHeader for EncodedExtensionHeader { self.length as usize } - fn encode_flexible_data(&self, extension: Box<[u8]>, target: &mut ArenaSliceMut<'_>) { + fn encode_flexible_data(&self, extension: &[u8], target: &mut ArenaSliceMut<'_>) { target.raw_slice_mut().copy_from_slice(&extension); } diff --git a/core/store/src/trie/mem/flexible_data/mod.rs b/core/store/src/trie/mem/flexible_data/mod.rs index 58fb2b032b9..40025f49051 100644 --- a/core/store/src/trie/mem/flexible_data/mod.rs +++ b/core/store/src/trie/mem/flexible_data/mod.rs @@ -28,7 +28,7 @@ pub mod value; /// with multiple flexibly-sized parts with relative ease. pub trait FlexibleDataHeader { /// The type of the original form of data to be used for encoding. - type InputData; + type InputData: ?Sized; /// The type of a view of the decoded data, which may reference the memory /// that we are decoding from, and therefore having a lifetime. type View<'a>; @@ -45,7 +45,7 @@ pub trait FlexibleDataHeader { /// slice. This function must be implemented in a way that writes /// exactly `self.flexible_data_length()` bytes to the given memory /// slice. The caller must ensure that the memory slice is large enough. - fn encode_flexible_data(&self, data: Self::InputData, target: &mut ArenaSliceMut<'_>); + fn encode_flexible_data(&self, data: &Self::InputData, target: &mut ArenaSliceMut<'_>); /// Decodes the flexibly-sized part of the data from the given memory /// slice. This function must be implemented in a consistent manner diff --git a/core/store/src/trie/mem/flexible_data/value.rs b/core/store/src/trie/mem/flexible_data/value.rs index 99f9a6f75a5..52d08795c9e 100644 --- a/core/store/src/trie/mem/flexible_data/value.rs +++ b/core/store/src/trie/mem/flexible_data/value.rs @@ -60,7 +60,7 @@ impl FlexibleDataHeader for EncodedValueHeader { } } - fn encode_flexible_data(&self, value: FlatStateValue, target: &mut ArenaSliceMut<'_>) { + fn encode_flexible_data(&self, value: &FlatStateValue, target: &mut ArenaSliceMut<'_>) { let (length, inlined) = self.decode(); match value { FlatStateValue::Ref(value_ref) => { @@ -71,7 +71,7 @@ impl FlexibleDataHeader for EncodedValueHeader { FlatStateValue::Inlined(v) => { assert!(inlined); assert_eq!(length, v.len() as u32); - target.raw_slice_mut().copy_from_slice(&v); + target.raw_slice_mut().copy_from_slice(v); } } } diff --git a/core/store/src/trie/mem/freelist.rs b/core/store/src/trie/mem/freelist.rs new file mode 100644 index 00000000000..cb28d2c390f --- /dev/null +++ b/core/store/src/trie/mem/freelist.rs @@ -0,0 +1,90 @@ +use std::ops::Deref; + +/// A simple freelist for `Vec` that is used to avoid unnecessary +/// re-allocations. +pub struct VecU8Freelist { + free: Vec, + + // There are two ways we detect that the freelist is not being used properly: + // 1. If a ReusableVecU8 is dropped before being returned to the freelist, + // it will panic in debug. + // 2. If the number of allocations exceeds the expected number of allocations, + // it will panic in debug and log an error in production the first time + // it happens. + num_allocs: usize, + expected_allocs: usize, +} + +/// A wrapper around `Vec` that makes it harder to accidentally drop it +/// without returning it to the freelist. +pub struct ReusableVecU8 { + vec: Vec, +} + +impl Deref for ReusableVecU8 { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + &self.vec + } +} + +impl ReusableVecU8 { + /// Used to actually free the vector in the end. + fn internal_free(&mut self) { + std::mem::take(&mut self.vec); + } + + pub fn vec_mut(&mut self) -> &mut Vec { + &mut self.vec + } +} + +impl Drop for ReusableVecU8 { + fn drop(&mut self) { + // Do not drop without returning it to the freelist. + debug_assert_eq!(self.vec.capacity(), 0); + } +} + +impl VecU8Freelist { + /// Create a new freelist with an expected number of non-reused allocations. + /// The expected number is used to detect incorrect usage of the freelist. + pub fn new(expected_allocs: usize) -> Self { + Self { free: Vec::with_capacity(expected_allocs), num_allocs: 0, expected_allocs } + } + + /// Returns a byte array, by either reusing one or allocating a new one + /// if there's none to reuse. + pub fn alloc(&mut self) -> ReusableVecU8 { + if let Some(vec) = self.free.pop() { + vec + } else { + self.num_allocs += 1; + if self.num_allocs == self.expected_allocs { + // If this triggers, it means that we're not using the freelist properly. + if cfg!(debug_assertions) { + panic!("Too many freelist allocations; expected {}", self.expected_allocs); + } else { + tracing::error!(target: "memtrie", "Too many freelist allocations; expected {}", self.expected_allocs); + } + } + ReusableVecU8 { vec: Vec::new() } + } + } + + /// Returns a byte array to the freelist. + pub fn free(&mut self, mut vec: ReusableVecU8) { + vec.vec.clear(); + self.free.push(vec); + } +} + +impl Drop for VecU8Freelist { + fn drop(&mut self) { + // This is where the returned byte arrays are actually freed. + for mut vec in self.free.drain(..) { + vec.internal_free(); + } + } +} diff --git a/core/store/src/trie/mem/mod.rs b/core/store/src/trie/mem/mod.rs index 0a41ed8527f..3d7d604727d 100644 --- a/core/store/src/trie/mem/mod.rs +++ b/core/store/src/trie/mem/mod.rs @@ -11,6 +11,7 @@ use std::collections::{BTreeMap, HashMap}; mod arena; mod construction; pub(crate) mod flexible_data; +mod freelist; pub mod iter; pub mod loading; pub mod lookup; @@ -224,13 +225,10 @@ mod tests { let root = MemTrieNodeId::new( arena, InputMemTrieNode::Leaf { - value: FlatStateValue::Inlined( + value: &FlatStateValue::Inlined( format!("{}", height).into_bytes(), ), - extension: NibbleSlice::new(&[]) - .encoded(true) - .to_vec() - .into_boxed_slice(), + extension: &NibbleSlice::new(&[]).encoded(true), }, ); root.as_ptr_mut(arena.memory_mut()).compute_hash_recursively(); diff --git a/core/store/src/trie/mem/node/encoding.rs b/core/store/src/trie/mem/node/encoding.rs index 218fb8df1c3..caa93a8f298 100644 --- a/core/store/src/trie/mem/node/encoding.rs +++ b/core/store/src/trie/mem/node/encoding.rs @@ -180,7 +180,7 @@ impl MemTrieNodeId { value: value_header, }); data.encode_flexible(&extension_header, extension); - data.encode_flexible(&value_header, value); + data.encode_flexible(&value_header, &value); data.finish() } InputMemTrieNode::Extension { extension, child } => { @@ -209,7 +209,7 @@ impl MemTrieNodeId { nonleaf: NonLeafHeader::new(memory_usage, node_hash), children: children_header, }); - data.encode_flexible(&children_header, children); + data.encode_flexible(&children_header, &children); data.finish() } InputMemTrieNode::BranchWithValue { children, value } => { @@ -227,8 +227,8 @@ impl MemTrieNodeId { children: children_header, value: value_header, }); - data.encode_flexible(&children_header, children); - data.encode_flexible(&value_header, value); + data.encode_flexible(&children_header, &children); + data.encode_flexible(&value_header, &value); data.finish() } }; diff --git a/core/store/src/trie/mem/node/mod.rs b/core/store/src/trie/mem/node/mod.rs index eb1f0bf755f..5f13af2a4cb 100644 --- a/core/store/src/trie/mem/node/mod.rs +++ b/core/store/src/trie/mem/node/mod.rs @@ -43,6 +43,14 @@ impl MemTrieNodeId { } } +/// This is for internal use only, so that we can put `MemTrieNodeId` in an +/// ElasticArray. +impl Default for MemTrieNodeId { + fn default() -> Self { + Self { pos: ArenaPos::invalid() } + } +} + /// Pointer to an in-memory trie node that allows read-only access to the node /// and all its descendants. #[derive(Clone, Copy, PartialEq, Eq)] @@ -79,11 +87,11 @@ impl<'a> MemTrieNodePtr<'a> { /// Used to construct a new in-memory trie node. #[derive(PartialEq, Eq, Debug, Clone)] -pub enum InputMemTrieNode { - Leaf { value: FlatStateValue, extension: Box<[u8]> }, - Extension { extension: Box<[u8]>, child: MemTrieNodeId }, +pub enum InputMemTrieNode<'a> { + Leaf { value: &'a FlatStateValue, extension: &'a [u8] }, + Extension { extension: &'a [u8], child: MemTrieNodeId }, Branch { children: [Option; 16] }, - BranchWithValue { children: [Option; 16], value: FlatStateValue }, + BranchWithValue { children: [Option; 16], value: &'a FlatStateValue }, } /// A view of the encoded data of `MemTrieNode`, obtainable via diff --git a/core/store/src/trie/mem/node/tests.rs b/core/store/src/trie/mem/node/tests.rs index bb9ff973c80..37a3b0db9c0 100644 --- a/core/store/src/trie/mem/node/tests.rs +++ b/core/store/src/trie/mem/node/tests.rs @@ -11,8 +11,8 @@ fn test_basic_leaf_node_inlined() { let node = MemTrieNodeId::new( &mut arena, InputMemTrieNode::Leaf { - extension: vec![0, 1, 2, 3, 4].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![5, 6, 7, 8, 9]), + extension: &[0, 1, 2, 3, 4], + value: &FlatStateValue::Inlined(vec![5, 6, 7, 8, 9]), }, ); let view = node.as_ptr(arena.memory()).view(); @@ -44,8 +44,8 @@ fn test_basic_leaf_node_ref() { let node = MemTrieNodeId::new( &mut arena, InputMemTrieNode::Leaf { - extension: vec![0, 1, 2, 3, 4].into_boxed_slice(), - value: FlatStateValue::Ref(ValueRef { hash: test_hash, length: 5 }), + extension: &[0, 1, 2, 3, 4], + value: &FlatStateValue::Ref(ValueRef { hash: test_hash, length: 5 }), }, ); let view = node.as_ptr(arena.memory()).view(); @@ -75,10 +75,7 @@ fn test_basic_leaf_node_empty_extension_empty_value() { let mut arena = Arena::new("".to_owned()); let node = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Leaf { - extension: vec![].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![]), - }, + InputMemTrieNode::Leaf { extension: &[], value: &FlatStateValue::Inlined(vec![]) }, ); let view = node.as_ptr(arena.memory()).view(); assert_eq!( @@ -105,13 +102,13 @@ fn test_basic_extension_node() { let child = MemTrieNodeId::new( &mut arena, InputMemTrieNode::Leaf { - extension: vec![0, 1, 2, 3, 4].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![5, 6, 7, 8, 9]), + extension: &[0, 1, 2, 3, 4], + value: &FlatStateValue::Inlined(vec![5, 6, 7, 8, 9]), }, ); let node = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Extension { extension: vec![5, 6, 7, 8, 9].into_boxed_slice(), child }, + InputMemTrieNode::Extension { extension: &[5, 6, 7, 8, 9], child }, ); node.as_ptr_mut(arena.memory_mut()).compute_hash_recursively(); let child_ptr = child.as_ptr(arena.memory()); @@ -152,17 +149,11 @@ fn test_basic_branch_node() { let mut arena = Arena::new("".to_owned()); let child1 = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Leaf { - extension: vec![].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![1]), - }, + InputMemTrieNode::Leaf { extension: &[], value: &FlatStateValue::Inlined(vec![1]) }, ); let child2 = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Leaf { - extension: vec![1].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![2]), - }, + InputMemTrieNode::Leaf { extension: &[1], value: &FlatStateValue::Inlined(vec![2]) }, ); let node = MemTrieNodeId::new( &mut arena, @@ -222,23 +213,17 @@ fn test_basic_branch_with_value_node() { let mut arena = Arena::new("".to_owned()); let child1 = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Leaf { - extension: vec![].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![1]), - }, + InputMemTrieNode::Leaf { extension: &[], value: &FlatStateValue::Inlined(vec![1]) }, ); let child2 = MemTrieNodeId::new( &mut arena, - InputMemTrieNode::Leaf { - extension: vec![1].into_boxed_slice(), - value: FlatStateValue::Inlined(vec![2]), - }, + InputMemTrieNode::Leaf { extension: &[1], value: &FlatStateValue::Inlined(vec![2]) }, ); let node = MemTrieNodeId::new( &mut arena, InputMemTrieNode::BranchWithValue { children: branch_array(vec![(0, child1), (15, child2)]), - value: FlatStateValue::Inlined(vec![3, 4, 5]), + value: &FlatStateValue::Inlined(vec![3, 4, 5]), }, ); diff --git a/core/store/src/trie/mem/updating.rs b/core/store/src/trie/mem/updating.rs index 2eb3400a5f4..495987a4e43 100644 --- a/core/store/src/trie/mem/updating.rs +++ b/core/store/src/trie/mem/updating.rs @@ -876,7 +876,7 @@ pub fn apply_memtrie_changes( let node_ids_with_hashes = &changes.node_ids_with_hashes; for (node_id, node_hash) in node_ids_with_hashes.iter() { let node = updated_nodes.get(*node_id).unwrap().clone().unwrap(); - let node = match node { + let node = match &node { UpdatedMemTrieNode::Empty => unreachable!(), UpdatedMemTrieNode::Branch { children, value } => { let mut new_children = [None; 16]; @@ -896,7 +896,7 @@ pub fn apply_memtrie_changes( UpdatedMemTrieNode::Extension { extension, child } => { InputMemTrieNode::Extension { extension, - child: map_to_new_node_id(child, &updated_to_new_map), + child: map_to_new_node_id(*child, &updated_to_new_map), } } UpdatedMemTrieNode::Leaf { extension, value } => { diff --git a/core/store/src/trie/nibble_slice.rs b/core/store/src/trie/nibble_slice.rs index 9e945fe7c2d..471c8eb66f8 100644 --- a/core/store/src/trie/nibble_slice.rs +++ b/core/store/src/trie/nibble_slice.rs @@ -102,11 +102,8 @@ impl<'a> NibbleSlice<'a> { /// Get the nibble at position `i`. #[inline(always)] pub fn at(&self, i: usize) -> u8 { - if (self.offset + i) & 1 == 1 { - self.data[(self.offset + i) / 2] & 15u8 - } else { - self.data[(self.offset + i) / 2] >> 4 - } + let shift = if (self.offset + i) & 1 == 1 { 0 } else { 4 }; + (self.data[(self.offset + i) / 2] >> shift) & 0xf } /// Return object which represents a view on to this slice (further) offset by `i` nibbles. @@ -147,15 +144,38 @@ impl<'a> NibbleSlice<'a> { /// Encode while nibble slice in prefixed hex notation, noting whether it `is_leaf`. #[inline] pub fn encoded(&self, is_leaf: bool) -> ElasticArray36 { + let mut to = ElasticArray36::new(); let l = self.len(); - let mut r = ElasticArray36::new(); - let mut i = l % 2; - r.push(if i == 1 { 0x10 + self.at(0) } else { 0 } + if is_leaf { 0x20 } else { 0 }); - while i < l { - r.push(self.at(i) * 16 + self.at(i + 1)); - i += 2; + let parity = l % 2; + let mut first_byte: u8 = 0; + if parity == 1 { + first_byte = 0x10 + self.at(0); } - r + if is_leaf { + first_byte |= 0x20; + } + to.push(first_byte); + let from_byte = (self.offset + parity) / 2; + to.append_slice(&self.data[from_byte..]); + to + } + + /// Same as `encoded`, but writes the result to the given vector. + #[inline] + pub fn encode_to(&self, is_leaf: bool, to: &mut Vec) { + let l = self.len(); + let parity = l % 2; + let mut first_byte: u8 = 0; + if parity == 1 { + first_byte = 0x10 + self.at(0); + } + if is_leaf { + first_byte |= 0x20; + } + to.clear(); + to.push(first_byte); + let from_byte = (self.offset + parity) / 2; + to.extend_from_slice(&self.data[from_byte..]); } pub fn merge_encoded(&self, other: &Self, is_leaf: bool) -> ElasticArray36 { @@ -195,6 +215,18 @@ impl<'a> NibbleSlice<'a> { r } + /// Same as `encoded_leftmost`, but writes the result to the given vector. + pub fn encode_leftmost_to(&self, n: usize, is_leaf: bool, to: &mut Vec) { + let l = min(self.len(), n); + to.resize(1 + l / 2, 0); + let mut i = l % 2; + to[0] = if i == 1 { 0x10 + self.at(0) } else { 0 } + if is_leaf { 0x20 } else { 0 }; + while i < l { + to[i / 2 + 1] = self.at(i) * 16 + self.at(i + 1); + i += 2; + } + } + // Helper to convert nibbles to bytes. pub fn nibbles_to_bytes(nibbles: &[u8]) -> Vec { assert_eq!(nibbles.len() % 2, 0); @@ -294,6 +326,11 @@ mod tests { let n = NibbleSlice::new(D); assert_eq!(n.encoded(false), ElasticArray36::from_slice(&[0x00, 0x01, 0x23, 0x45])); assert_eq!(n.encoded(true), ElasticArray36::from_slice(&[0x20, 0x01, 0x23, 0x45])); + let mut v = vec![]; + n.encode_to(false, &mut v); + assert_eq!(v, vec![0x00, 0x01, 0x23, 0x45]); + n.encode_to(true, &mut v); + assert_eq!(v, vec![0x20, 0x01, 0x23, 0x45]); assert_eq!(n.mid(1).encoded(false), ElasticArray36::from_slice(&[0x11, 0x23, 0x45])); assert_eq!(n.mid(1).encoded(true), ElasticArray36::from_slice(&[0x31, 0x23, 0x45])); } @@ -308,10 +345,24 @@ mod tests { } fn encode_decode(nibbles: &[u8], is_leaf: bool) { - let n = NibbleSlice::encode_nibbles(nibbles, is_leaf); - let (n, is_leaf_decoded) = NibbleSlice::from_encoded(&n); + let encoded = NibbleSlice::encode_nibbles(nibbles, is_leaf); + let (n, is_leaf_decoded) = NibbleSlice::from_encoded(&encoded); assert_eq!(&n.iter().collect::>(), nibbles); - assert_eq!(is_leaf_decoded, is_leaf) + assert_eq!(is_leaf_decoded, is_leaf); + let reencoded = n.encoded(is_leaf); + assert_eq!(&reencoded, &encoded); + let mut reencoded2 = vec![]; + n.encode_to(is_leaf, &mut reencoded2); + assert_eq!(&reencoded2, &encoded.to_vec()); + + for i in 0..nibbles.len() { + let encoded = NibbleSlice::encode_nibbles(&nibbles[..i], is_leaf); + let leftmost_encoded = n.encoded_leftmost(i, is_leaf); + assert_eq!(&leftmost_encoded, &encoded); + let mut leftmost_encoded2 = vec![]; + n.encode_leftmost_to(i, is_leaf, &mut leftmost_encoded2); + assert_eq!(&leftmost_encoded2, &encoded.to_vec()); + } } #[test]