Skip to content

Commit

Permalink
core: use ValueRef in trie nodes structs (near#8687)
Browse files Browse the repository at this point in the history 
Rather than handling value references as a (size, hash) tuples, use
existing ValueRef type in RawTrieNode and ValueHandle.  This better
documents what the u32 and CryptoHash mean and simplifies some code
by having a single object representing what so far was handled with
two variables.
  • Loading branch information
@mina86
mina86 authored Mar 7, 2023
1 parent 1c77619 commit 615eccf
Show file tree
Hide file tree
Showing 6 changed files with 95 additions and 103 deletions.
19 changes: 18 additions & 1 deletion core/primitives/src/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use borsh::{BorshDeserialize, BorshSerialize};
use near_primitives_core::hash::{hash, CryptoHash};

/// State value reference. Used to charge fees for value length before retrieving the value itself.
#[derive(BorshSerialize, BorshDeserialize, Clone, PartialEq, Eq, Debug)]
#[derive(BorshSerialize, BorshDeserialize, Clone, PartialEq, Eq, Hash)]
pub struct ValueRef {
/// Value length in bytes.
pub length: u32,
Expand All @@ -25,6 +25,23 @@ impl ValueRef {
let length = u32::from_le_bytes(*length);
ValueRef { length, hash: CryptoHash(*hash) }
}

/// Returns length of the referenced value.
pub fn len(&self) -> usize {
usize::try_from(self.length).unwrap()
}
}

impl std::cmp::PartialEq<[u8]> for ValueRef {
fn eq(&self, rhs: &[u8]) -> bool {
self.len() == rhs.len() && self.hash == CryptoHash::hash_bytes(rhs)
}
}

impl std::fmt::Debug for ValueRef {
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(fmt, "({}, {})", self.length, self.hash)
}
}

#[cfg(test)]
Expand Down
11 changes: 6 additions & 5 deletions core/store/src/trie/insert_delete.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use std::collections::HashMap;

use near_primitives::hash::{hash, CryptoHash};
use near_primitives::state::ValueRef;

use crate::trie::nibble_slice::NibbleSlice;
use crate::trie::{
Expand Down Expand Up @@ -615,8 +616,8 @@ impl Trie {
TrieNode::Leaf(key, value) => {
let key = key.clone();
let value = value.clone();
let (value_length, value_hash) = Trie::flatten_value(&mut memory, value);
RawTrieNode::Leaf(key, value_length, value_hash)
let value = Trie::flatten_value(&mut memory, value);
RawTrieNode::Leaf(key, value)
}
};
let raw_node_with_size = RawTrieNodeWithSize { node: raw_node, memory_usage };
Expand All @@ -634,7 +635,7 @@ impl Trie {
Ok(TrieChanges { old_root: *old_root, new_root: last_hash, insertions, deletions })
}

fn flatten_value(memory: &mut NodesStorage, value: ValueHandle) -> (u32, CryptoHash) {
fn flatten_value(memory: &mut NodesStorage, value: ValueHandle) -> ValueRef {
match value {
ValueHandle::InMemory(value_handle) => {
let value = memory.value_ref(value_handle).to_vec();
Expand All @@ -643,9 +644,9 @@ impl Trie {
let (_value, rc) =
memory.refcount_changes.entry(value_hash).or_insert_with(|| (value, 0));
*rc += 1;
(value_length, value_hash)
ValueRef { length: value_length, hash: value_hash }
}
ValueHandle::HashAndSize(value_length, value_hash) => (value_length, value_hash),
ValueHandle::HashAndSize(value) => value,
}
}
}
4 changes: 2 additions & 2 deletions core/store/src/trie/iterator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,15 +239,15 @@ impl<'a> TrieIterator<'a> {
}
(CrumbStatus::At, TrieNode::Branch(_, Some(value))) => {
let hash = match value {
ValueHandle::HashAndSize(_, hash) => *hash,
ValueHandle::HashAndSize(value) => value.hash,
ValueHandle::InMemory(_node) => unreachable!(),
};
Some(IterStep::Value(hash))
}
(CrumbStatus::At, TrieNode::Branch(_, None)) => Some(IterStep::Continue),
(CrumbStatus::At, TrieNode::Leaf(key, value)) => {
let hash = match value {
ValueHandle::HashAndSize(_, hash) => *hash,
ValueHandle::HashAndSize(value) => value.hash,
ValueHandle::InMemory(_node) => unreachable!(),
};
let key = NibbleSlice::from_encoded(key).0;
Expand Down
106 changes: 45 additions & 61 deletions core/store/src/trie/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,14 +97,14 @@ impl std::fmt::Debug for NodeHandle {
#[derive(Clone, Hash)]
enum ValueHandle {
InMemory(StorageValueHandle),
HashAndSize(u32, CryptoHash),
HashAndSize(ValueRef),
}

impl std::fmt::Debug for ValueHandle {
fn fmt(&self, fmtr: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::HashAndSize(size, hash) => write!(fmtr, "{hash}, size:{size}"),
Self::InMemory(handle) => write!(fmtr, "@{}", handle.0),
Self::HashAndSize(value) => write!(fmtr, "{value:?}"),
Self::InMemory(StorageValueHandle(num)) => write!(fmtr, "@{num}"),
}
}
}
Expand Down Expand Up @@ -148,20 +148,14 @@ impl TrieNodeWithSize {
impl TrieNode {
fn new(rc_node: RawTrieNode) -> TrieNode {
match rc_node {
RawTrieNode::Leaf(key, value_length, value_hash) => {
TrieNode::Leaf(key, ValueHandle::HashAndSize(value_length, value_hash))
}
RawTrieNode::Leaf(key, value) => TrieNode::Leaf(key, ValueHandle::HashAndSize(value)),
RawTrieNode::Branch(children, value) => {
let mut new_children: Box<[Option<NodeHandle>; 16]> = Default::default();
for i in 0..children.len() {
new_children[i] = children[i].map(NodeHandle::Hash);
}
TrieNode::Branch(
new_children,
value.map(|(value_length, value_hash)| {
ValueHandle::HashAndSize(value_length, value_hash)
}),
)
let value = value.map(ValueHandle::HashAndSize);
TrieNode::Branch(new_children, value)
}
RawTrieNode::Extension(key, child) => TrieNode::Extension(key, NodeHandle::Hash(child)),
}
Expand Down Expand Up @@ -246,7 +240,7 @@ impl TrieNode {
.expect("InMemory nodes exist, but storage is not provided")
.value_ref(*handle)
.len() as u64,
ValueHandle::HashAndSize(value_length, _value_hash) => *value_length as u64,
ValueHandle::HashAndSize(value) => u64::from(value.length),
};
Self::memory_usage_for_value_length(value_length)
}
Expand Down Expand Up @@ -321,9 +315,9 @@ impl std::fmt::Debug for TrieNode {
#[allow(clippy::large_enum_variant)]
pub enum RawTrieNode {
/// Leaf(key, value_length, value_hash)
Leaf(Vec<u8>, u32, CryptoHash),
Leaf(Vec<u8>, ValueRef),
/// Branch(children, (value_length, value_hash))
Branch([Option<CryptoHash>; 16], Option<(u32, CryptoHash)>),
Branch([Option<CryptoHash>; 16], Option<ValueRef>),
/// Extension(key, child)
Extension(Vec<u8>, CryptoHash),
}
Expand Down Expand Up @@ -361,19 +355,19 @@ impl RawTrieNode {
// size in state_parts = size + 8 for RawTrieNodeWithSize + 8 for borsh vector length
match &self {
// size <= 1 + 4 + 4 + 32 + key_length + value_length
RawTrieNode::Leaf(key, value_length, value_hash) => {
RawTrieNode::Leaf(key, value) => {
out.push(LEAF_NODE);
out.extend((key.len() as u32).to_le_bytes());
out.extend(key);
out.extend((*value_length as u32).to_le_bytes());
out.extend(value_hash.as_bytes());
out.extend(value.length.to_le_bytes());
out.extend(value.hash.as_bytes());
}
// size <= 1 + 4 + 32 + value_length + 2 + 32 * num_children
RawTrieNode::Branch(children, value) => {
if let Some((value_length, value_hash)) = value {
if let Some(value) = value {
out.push(BRANCH_NODE_WITH_VALUE);
out.extend((*value_length as u32).to_le_bytes());
out.extend(value_hash.as_bytes());
out.extend(value.length.to_le_bytes());
out.extend(value.hash.as_bytes());
} else {
out.push(BRANCH_NODE_NO_VALUE);
}
Expand Down Expand Up @@ -408,23 +402,23 @@ impl RawTrieNode {
let key_length = bytes.read_u32::<LittleEndian>()?;
let mut key = vec![0; key_length as usize];
bytes.read_exact(&mut key)?;
let value_length = bytes.read_u32::<LittleEndian>()?;
let length = bytes.read_u32::<LittleEndian>()?;
let mut arr = [0; 32];
bytes.read_exact(&mut arr)?;
let value_hash = CryptoHash(arr);
RawTrieNode::Leaf(key, value_length, value_hash)
let hash = CryptoHash(arr);
RawTrieNode::Leaf(key, ValueRef { length, hash })
}
BRANCH_NODE_NO_VALUE => {
let children = decode_children(&mut bytes)?;
RawTrieNode::Branch(children, None)
}
BRANCH_NODE_WITH_VALUE => {
let value_length = bytes.read_u32::<LittleEndian>()?;
let length = bytes.read_u32::<LittleEndian>()?;
let mut arr = [0; 32];
bytes.read_exact(&mut arr)?;
let value_hash = CryptoHash(arr);
let hash = CryptoHash(arr);
let children = decode_children(&mut bytes)?;
RawTrieNode::Branch(children, Some((value_length, value_hash)))
RawTrieNode::Branch(children, Some(ValueRef { length, hash }))
}
EXTENSION_NODE => {
let key_length = bytes.read_u32::<LittleEndian>()?;
Expand Down Expand Up @@ -657,9 +651,13 @@ impl Trie {
value: &ValueHandle,
) -> Result<(), StorageError> {
match value {
ValueHandle::HashAndSize(_, hash) => {
let bytes = self.storage.retrieve_raw_bytes(hash)?;
memory.refcount_changes.entry(*hash).or_insert_with(|| (bytes.to_vec(), 0)).1 -= 1;
ValueHandle::HashAndSize(value) => {
let bytes = self.storage.retrieve_raw_bytes(&value.hash)?;
memory
.refcount_changes
.entry(value.hash)
.or_insert_with(|| (bytes.to_vec(), 0))
.1 -= 1;
}
ValueHandle::InMemory(_) => {
// do nothing
Expand Down Expand Up @@ -753,16 +751,12 @@ impl Trie {
match self.retrieve_raw_node(hash) {
Ok(Some((_, raw_node))) => {
match raw_node.node {
RawTrieNode::Leaf(key, value_length, value_hash) => {
RawTrieNode::Leaf(key, value) => {
let (slice, _) = NibbleSlice::from_encoded(key.as_slice());
prefix.extend(slice.iter());
writeln!(
f,
"{}Leaf {:?} size:{} child_hash:{} child_prefix:{}",
spaces,
slice,
value_length,
value_hash,
"{spaces}Leaf {slice:?} {value:?} prefix:{}",
self.nibbles_to_string(prefix)
)?;
prefix.truncate(prefix.len() - slice.len());
Expand Down Expand Up @@ -894,12 +888,12 @@ impl Trie {
Some((_bytes, node)) => node.node,
};
match node {
RawTrieNode::Leaf(existing_key, value_length, value_hash) => {
if NibbleSlice::from_encoded(&existing_key).0 == key {
return Ok(Some(ValueRef { length: value_length, hash: value_hash }));
RawTrieNode::Leaf(existing_key, value) => {
return Ok(if NibbleSlice::from_encoded(&existing_key).0 == key {
Some(value)
} else {
return Ok(None);
}
None
});
}
RawTrieNode::Extension(existing_key, child) => {
let existing_key = NibbleSlice::from_encoded(&existing_key).0;
Expand All @@ -912,23 +906,14 @@ impl Trie {
}
RawTrieNode::Branch(mut children, value) => {
if key.is_empty() {
match value {
Some((value_length, value_hash)) => {
return Ok(Some(ValueRef {
length: value_length,
hash: value_hash,
}));
}
None => return Ok(None),
}
} else {
match children[key.at(0) as usize].take() {
Some(x) => {
hash = x;
key = key.mid(1);
}
None => return Ok(None),
return Ok(value);
}
match children[key.at(0) as usize].take() {
Some(x) => {
hash = x;
key = key.mid(1);
}
None => return Ok(None),
}
}
};
Expand Down Expand Up @@ -1120,9 +1105,8 @@ mod tests {
assert!(got.is_err(), "got: {got:?}");
}

let value_length = 3;
let value_hash = CryptoHash::hash_bytes(&[123, 245, 255]);
let node = RawTrieNode::Leaf(vec![1, 2, 3], value_length, value_hash);
let value = ValueRef { length: 3, hash: CryptoHash::hash_bytes(&[123, 245, 255]) };
let node = RawTrieNode::Leaf(vec![1, 2, 3], value.clone());
let encoded = [
0, 3, 0, 0, 0, 1, 2, 3, 3, 0, 0, 0, 194, 40, 8, 24, 64, 219, 69, 132, 86, 52, 110, 175,
57, 198, 165, 200, 83, 237, 211, 11, 194, 83, 251, 33, 145, 138, 234, 226, 7, 242, 186,
Expand All @@ -1141,7 +1125,7 @@ mod tests {

let mut children: [Option<CryptoHash>; 16] = Default::default();
children[3] = Some(Trie::EMPTY_ROOT);
let node = RawTrieNode::Branch(children, Some((value_length, value_hash)));
let node = RawTrieNode::Branch(children, Some(value));
let encoded = [
2, 3, 0, 0, 0, 194, 40, 8, 24, 64, 219, 69, 132, 86, 52, 110, 175, 57, 198, 165, 200,
83, 237, 211, 11, 194, 83, 251, 33, 145, 138, 234, 226, 7, 242, 186, 73, 8, 0, 0, 0, 0,
Expand Down
38 changes: 16 additions & 22 deletions core/store/src/trie/state_parts.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,28 +328,25 @@ mod tests {
if let CrumbStatus::Entering = position {
on_enter(&hash)?;
}
let mut on_enter_value = |value: &ValueHandle| {
if let ValueHandle::HashAndSize(value) = value {
on_enter(&value.hash)
} else {
unreachable!("only possible while mutating")
}
};
match &node.node {
TrieNode::Empty => {
continue;
}
TrieNode::Leaf(_, value) => {
match value {
ValueHandle::HashAndSize(_, hash) => {
on_enter(hash)?;
}
ValueHandle::InMemory(_) => {
unreachable!("only possible while mutating")
}
}
on_enter_value(value)?;
continue;
}
TrieNode::Branch(children, value) => match position {
CrumbStatus::Entering => {
match value {
Some(ValueHandle::HashAndSize(_, hash)) => {
on_enter(hash)?;
}
_ => {}
if let Some(ref value) = value {
on_enter_value(value)?;
}
stack.push((hash, node, CrumbStatus::AtChild(0)));
continue;
Expand Down Expand Up @@ -382,15 +379,12 @@ mod tests {
},
TrieNode::Extension(_key, child) => {
if let CrumbStatus::Entering = position {
match child.clone() {
NodeHandle::InMemory(_) => {
unreachable!("only possible while mutating")
}
NodeHandle::Hash(h) => {
let child = self.retrieve_node(&h)?.1;
stack.push((hash, node, CrumbStatus::Exiting));
stack.push((h, child, CrumbStatus::Entering));
}
if let NodeHandle::Hash(h) = child.clone() {
let child = self.retrieve_node(&h)?.1;
stack.push((h, node, CrumbStatus::Exiting));
stack.push((h, child, CrumbStatus::Entering));
} else {
unreachable!("only possible while mutating")
}
}
}
Expand Down
Loading

0 comments on commit 615eccf

Please sign in to comment.