github/akula
1
1
Fork 0
mirror of https://github.com/akula-bft/akula.git synced 2026-04-19 03:00:13 -04:00
Code Issues Packages Projects Releases Wiki Activity

Rustify the trie module (#161)

Browse Source 
* Rustify the trie module

rewrite the line-by-line C++ port in a more idiomatic way,
get rid off interior mutability workarounds, buffer attributes
and clippy exceptions

* fmt
This commit is contained in:
Jochen Müller
2022-02-21 13:08:44 +01:00
committed by GitHub
parent 85298894f9
commit 73fa1f656b
3 changed files with 301 additions and 383 deletions
Download Patch File Download Diff File Expand all files Collapse all files

186
src/trie/hash_builder.rs
View File

@@ -1,4 +1,3 @@
#![allow(clippy::manual_memcpy, clippy::needless_range_loop)]
use crate::{
crypto::keccak256,
models::*,
@@ -7,6 +6,7 @@ use crate::{
util::{assert_subset, prefix_length},
},
};
use bytes::BytesMut;
use ethereum_types::H256;
use rlp::RlpStream;
use std::{boxed::Box, cmp};
@@ -26,8 +26,8 @@ fn encode_path(nibbles: &[u8], terminating: bool) -> Vec<u8> {
i = 1;
}
for j in 1..res.len() {
res[j] = (nibbles[i] << 4) + nibbles[i + 1];
for byte in res.iter_mut().skip(1) {
*byte = (nibbles[i] << 4) + nibbles[i + 1];
i += 2;
}
@@ -35,20 +35,19 @@ fn encode_path(nibbles: &[u8], terminating: bool) -> Vec<u8> {
}
fn wrap_hash(hash: &H256) -> Vec<u8> {
let mut wrapped = vec![0u8; KECCAK_LENGTH + 1];
wrapped[0] = RLP_EMPTY_STRING_CODE + KECCAK_LENGTH as u8;
for i in 0..32 {
wrapped[i + 1] = hash.0[i];
}
wrapped
[
[RLP_EMPTY_STRING_CODE + KECCAK_LENGTH as u8].as_slice(),
hash.0.as_slice(),
]
.concat()
}
fn node_ref(rlp: &[u8]) -> Vec<u8> {
if rlp.len() < KECCAK_LENGTH {
return rlp.to_vec();
rlp.to_vec()
} else {
wrap_hash(&keccak256(rlp))
}
let hash = keccak256(rlp);
wrap_hash(&hash)
}
type NodeCollector<'nc> = Box<dyn FnMut(&[u8], &Node) + Send + Sync + 'nc>;
@@ -59,6 +58,26 @@ enum HashBuilderValue {
Hash(H256),
}
fn leaf_node_rlp(path: &[u8], value: &[u8]) -> BytesMut {
let encoded_path = encode_path(path, true);
let mut stream = RlpStream::new_list(2);
stream.append(&encoded_path);
stream.append(&value);
stream.out()
}
fn extension_node_rlp(path: &[u8], child_ref: &[u8]) -> BytesMut {
let encoded_path = encode_path(path, false);
let mut stream = RlpStream::new_list(2);
stream.append(&encoded_path);
stream.append_raw(child_ref, 1);
stream.out()
}
pub(crate) struct HashBuilder<'nc> {
pub(crate) node_collector: Option<NodeCollector<'nc>>,
key: Vec<u8>,
@@ -68,13 +87,12 @@ pub(crate) struct HashBuilder<'nc> {
tree_masks: Vec<u16>,
hash_masks: Vec<u16>,
stack: Vec<Vec<u8>>,
rlp_buffer: Vec<u8>,
}
impl<'nc> HashBuilder<'nc> {
pub(crate) fn new() -> Self {
pub(crate) fn new(node_collector: Option<NodeCollector<'nc>>) -> Self {
Self {
node_collector: None,
node_collector,
key: vec![],
value: HashBuilderValue::Bytes(vec![]),
is_in_db_trie: false,
@@ -82,15 +100,17 @@ impl<'nc> HashBuilder<'nc> {
tree_masks: vec![],
hash_masks: vec![],
stack: vec![],
rlp_buffer: vec![],
}
}
fn collects_nodes(&self) -> bool {
self.node_collector.is_some()
}
pub(crate) fn add_leaf(&mut self, key: Vec<u8>, value: &[u8]) {
assert!(key > self.key);
if !self.key.is_empty() {
let self_key = self.key.clone();
self.gen_struct_step(self_key.as_slice(), key.as_slice());
self.gen_struct_step(key.as_slice());
}
self.key = key;
self.value = HashBuilderValue::Bytes(value.to_vec());
@@ -99,58 +119,49 @@ impl<'nc> HashBuilder<'nc> {
pub(crate) fn add_branch_node(&mut self, key: Vec<u8>, value: &H256, is_in_db_trie: bool) {
assert!(key > self.key || (self.key.is_empty() && key.is_empty()));
if !self.key.is_empty() {
let self_key = self.key.clone();
self.gen_struct_step(self_key.as_slice(), key.as_slice());
self.gen_struct_step(key.as_slice());
} else if key.is_empty() {
self.stack.push(wrap_hash(value));
}
self.key = key;
self.value = HashBuilderValue::Hash(*value);
self.value = HashBuilderValue::Hash(*value);
self.is_in_db_trie = is_in_db_trie;
}
pub(crate) fn root_hash(&mut self) -> H256 {
self.private_root_hash(true)
pub(crate) fn compute_root_hash(&mut self) -> H256 {
self.finalize();
self.get_root_hash()
}
fn private_root_hash(&mut self, auto_finalize: bool) -> H256 {
if auto_finalize {
self.finalize();
}
if self.stack.is_empty() {
return EMPTY_ROOT;
}
let node_ref = self.stack.last().unwrap();
if node_ref.len() == KECCAK_LENGTH + 1 {
H256::from_slice(&node_ref[1..])
fn get_root_hash(&self) -> H256 {
if let Some(node_ref) = self.stack.last() {
if node_ref.len() == KECCAK_LENGTH + 1 {
H256::from_slice(&node_ref[1..])
} else {
keccak256(node_ref)
}
} else {
keccak256(node_ref)
EMPTY_ROOT
}
}
fn finalize(&mut self) {
if !self.key.is_empty() {
let self_key = self.key.clone();
self.gen_struct_step(self_key.as_slice(), &[]);
self.gen_struct_step(&[]);
self.key.clear();
self.value = HashBuilderValue::Bytes(vec![]);
}
}
fn gen_struct_step(&mut self, current: &[u8], succeeding: &[u8]) {
fn gen_struct_step(&mut self, succeeding: &[u8]) {
let mut build_extensions = false;
let mut current = current.to_vec();
let mut current = self.key.clone();
loop {
let preceding_exists = !self.groups.is_empty();
let preceding_len: usize = self.groups.len().saturating_sub(1);
let preceding_len = if self.groups.is_empty() {
0
} else {
self.groups.len() - 1
};
let common_prefix_len = prefix_length(succeeding, current.as_slice());
let len = cmp::max(preceding_len, common_prefix_len);
assert!(len < current.len());
@@ -174,17 +185,15 @@ impl<'nc> HashBuilder<'nc> {
let short_node_key = current[len_from..].to_vec();
if !build_extensions {
let value = self.value.clone();
match value {
HashBuilderValue::Bytes(ref leaf_value) => {
let x = node_ref(
self.leaf_node_rlp(short_node_key.as_slice(), leaf_value)
.as_slice(),
);
self.stack.push(x);
match &value {
HashBuilderValue::Bytes(leaf_value) => {
let node_ref =
node_ref(leaf_node_rlp(short_node_key.as_slice(), leaf_value).as_ref());
self.stack.push(node_ref);
}
HashBuilderValue::Hash(ref hash) => {
HashBuilderValue::Hash(hash) => {
self.stack.push(wrap_hash(hash));
if self.node_collector.is_some() {
if self.collects_nodes() {
if self.is_in_db_trie {
self.tree_masks[current.len() - 1] |=
1u16 << current.last().unwrap();
@@ -197,7 +206,7 @@ impl<'nc> HashBuilder<'nc> {
}
if build_extensions && !short_node_key.is_empty() {
if self.node_collector.is_some() && len_from > 0 {
if self.collects_nodes() && len_from > 0 {
let flag = 1u16 << current[len_from - 1];
self.hash_masks[len_from - 1] &= !flag;
@@ -209,8 +218,7 @@ impl<'nc> HashBuilder<'nc> {
let stack_last = self.stack.pop().unwrap();
let new_stack_last = node_ref(
self.extension_node_rlp(short_node_key.as_slice(), stack_last.as_slice())
.as_slice(),
extension_node_rlp(short_node_key.as_slice(), stack_last.as_slice()).as_ref(),
);
self.stack.push(new_stack_last);
@@ -225,8 +233,7 @@ impl<'nc> HashBuilder<'nc> {
if !succeeding.is_empty() || preceding_exists {
let child_hashes = self.branch_ref(self.groups[len], self.hash_masks[len]);
let have_node_collector = self.node_collector.is_some();
if have_node_collector {
if self.collects_nodes() {
if len > 0 {
self.hash_masks[len - 1] |= 1u16 << current[len - 1];
}
@@ -236,11 +243,12 @@ impl<'nc> HashBuilder<'nc> {
if len > 0 {
self.tree_masks[len - 1] |= 1u16 << current[len - 1];
}
let mut hashes = Vec::<H256>::with_capacity(child_hashes.len());
for i in 0..child_hashes.len() {
assert_eq!(child_hashes[i].len(), KECCAK_LENGTH + 1);
hashes.push(H256::from_slice(&child_hashes[i][1..]))
}
let hashes = child_hashes
.iter()
.map(|hash| H256::from_slice(&hash[1..]))
.collect();
let mut n = Node::new(
self.groups[len],
self.tree_masks[len],
@@ -248,8 +256,9 @@ impl<'nc> HashBuilder<'nc> {
hashes,
None,
);
if len == 0 {
n.set_root_hash(Some(self.private_root_hash(false)));
n.root_hash = Some(self.get_root_hash());
}
self.node_collector.as_mut().unwrap()(&current[0..len], &n);
@@ -294,34 +303,11 @@ impl<'nc> HashBuilder<'nc> {
}
stream.append_empty_data();
self.rlp_buffer = stream.out().to_vec();
self.stack.truncate(first_child_idx);
self.stack.push(node_ref(self.rlp_buffer.as_slice()));
self.stack.push(node_ref(stream.out().as_ref()));
child_hashes
}
fn leaf_node_rlp(&mut self, path: &[u8], value: &[u8]) -> Vec<u8> {
let encoded_path = encode_path(path, true);
let mut stream = RlpStream::new_list(2);
stream.append(&encoded_path);
stream.append(&value);
self.rlp_buffer = stream.out().to_vec();
self.rlp_buffer.clone()
}
fn extension_node_rlp(&mut self, path: &[u8], child_ref: &[u8]) -> Vec<u8> {
let encoded_path = encode_path(path, false);
let mut stream = RlpStream::new_list(2);
stream.append(&encoded_path);
stream.append_raw(child_ref, 1);
self.rlp_buffer = stream.out().to_vec();
self.rlp_buffer.clone()
}
}
pub(crate) fn pack_nibbles(nibbles: &[u8]) -> Vec<u8> {
@@ -365,8 +351,8 @@ mod tests {
#[test]
fn test_hash_builder_empty_trie() {
let mut hb = HashBuilder::new();
assert_eq!(hb.root_hash(), EMPTY_ROOT);
let mut hb = HashBuilder::new(None);
assert_eq!(hb.compute_root_hash(), EMPTY_ROOT);
}
#[test]
@@ -377,12 +363,12 @@ mod tests {
let val1 = vec![1u8];
let val2 = vec![2u8];
let mut hb = HashBuilder::new();
let mut hb = HashBuilder::new(None);
hb.add_leaf(unpack_nibbles(key1.as_bytes()), val1.as_slice());
hb.add_leaf(unpack_nibbles(key2.as_bytes()), val2.as_slice());
let hash = trie_root(vec![(key1, val1), (key2, val2)]);
let root_hash = hb.root_hash();
let root_hash = hb.compute_root_hash();
assert_eq!(hash, root_hash);
}
@@ -391,16 +377,16 @@ mod tests {
let key0 = hex!("646f").to_vec();
let val0 = hex!("76657262").to_vec();
let mut hb0 = HashBuilder::new();
let mut hb0 = HashBuilder::new(None);
hb0.add_leaf(unpack_nibbles(key0.as_slice()), val0.as_slice());
let hash0 = trie_root(vec![(key0.clone(), val0.clone())]);
assert_eq!(hb0.root_hash(), hash0);
assert_eq!(hb0.compute_root_hash(), hash0);
let key1 = hex!("676f6f64").to_vec();
let val1 = hex!("7075707079").to_vec();
let mut hb1 = HashBuilder::new();
let mut hb1 = HashBuilder::new(None);
hb1.add_leaf(unpack_nibbles(key0.as_slice()), val0.as_slice());
hb1.add_leaf(unpack_nibbles(key1.as_slice()), val1.as_slice());
@@ -409,7 +395,7 @@ mod tests {
(key1.clone(), val1.clone()),
]);
let hash1b = hash1;
assert_eq!(hb1.root_hash(), hash1);
assert_eq!(hb1.compute_root_hash(), hash1);
let mut stream0 = RlpStream::new_list(2);
let path0 = encode_path(unpack_nibbles(&key0[1..]).as_slice(), true);
@@ -439,10 +425,10 @@ mod tests {
let branch_node_rlp = stream.out();
let branch_node_hash = keccak256(branch_node_rlp);
let mut hb2 = HashBuilder::new();
let mut hb2 = HashBuilder::new(None);
hb2.add_branch_node(vec![0x6], &branch_node_hash, false);
assert_eq!(hb2.root_hash(), hash1b);
assert_eq!(hb2.compute_root_hash(), hash1b);
}
#[test]
@@ -450,9 +436,9 @@ mod tests {
let root_hash = H256::from(hex!(
"9fa752911d55c3a1246133fe280785afbdba41f357e9cae1131d5f5b0a078b9c"
));
let mut hb = HashBuilder::new();
let mut hb = HashBuilder::new(None);
hb.add_branch_node(vec![], &root_hash, false);
assert_eq!(hb.root_hash(), root_hash);
assert_eq!(hb.compute_root_hash(), root_hash);
}
#[test]

429
src/trie/intermediate_hashes.rs
View File

@@ -1,4 +1,3 @@
#![allow(clippy::question_mark)]
use crate::{
consensus::{DuoError, ValidationError},
crypto::keccak256,
@@ -12,6 +11,7 @@ use crate::{
util::has_prefix,
},
};
use anyhow::Result;
use parking_lot::Mutex;
use std::marker::PhantomData;
use tempfile::TempDir;
@@ -32,40 +32,43 @@ impl CursorSubNode {
}
fn state_flag(&self) -> bool {
if self.nibble < 0 || self.node.is_none() {
return true;
if let Some(node) = &self.node {
if self.nibble >= 0 {
return node.state_mask & (1u16 << self.nibble) != 0;
}
}
self.node.as_ref().unwrap().state_mask() & (1u16 << self.nibble) != 0
true
}
fn tree_flag(&self) -> bool {
if self.nibble < 0 || self.node.is_none() {
return true;
if let Some(node) = &self.node {
if self.nibble >= 0 {
return node.tree_mask & (1u16 << self.nibble) != 0;
}
}
self.node.as_ref().unwrap().tree_mask() & (1u16 << self.nibble) != 0
true
}
fn hash_flag(&self) -> bool {
if self.node.is_none() {
return false;
} else if self.nibble < 0 {
return self.node.as_ref().unwrap().root_hash().is_some();
match &self.node {
Some(node) => match self.nibble {
-1 => node.root_hash.is_some(),
_ => node.hash_mask & (1u16 << self.nibble) != 0,
},
None => false,
}
self.node.as_ref().unwrap().hash_mask() & (1u16 << self.nibble) != 0
}
fn hash(&self) -> Option<H256> {
if !self.hash_flag() {
return None;
if self.hash_flag() {
let node = self.node.as_ref().unwrap();
match self.nibble {
-1 => node.root_hash,
_ => Some(node.hash_for_nibble(self.nibble)),
}
} else {
None
}
if self.nibble < 0 {
return self.node.as_ref().unwrap().root_hash();
}
let first_nibbles_mask = (1u16 << self.nibble) - 1;
let hash_idx = (self.node.as_ref().unwrap().hash_mask() & first_nibbles_mask).count_ones();
Some(self.node.as_ref().unwrap().hashes()[hash_idx as usize])
}
}
@@ -108,7 +111,7 @@ where
cursor: &'cu mut MdbxCursor<'tx, RW, T>,
changed: &'ps mut PrefixSet,
prefix: &[u8],
) -> anyhow::Result<Cursor<'cu, 'tx, 'ps, T>> {
) -> Result<Cursor<'cu, 'tx, 'ps, T>> {
let mut new_cursor = Self {
cursor: Mutex::new(cursor),
changed,
@@ -121,65 +124,48 @@ where
Ok(new_cursor)
}
fn next(&mut self) -> anyhow::Result<()> {
if self.stack.is_empty() {
return Ok(()); // end-of-tree
}
if !self.can_skip_state && self.children_are_in_trie() {
if self.stack.last().unwrap().nibble < 0 {
self.move_to_next_sibling(true)?;
fn next(&mut self) -> Result<()> {
if let Some(last) = self.stack.last() {
if !self.can_skip_state && self.children_are_in_trie() {
match last.nibble {
-1 => self.move_to_next_sibling(true)?,
_ => self.consume_node(&self.key().unwrap(), false)?,
}
} else {
self.consume_node(&self.key().unwrap(), false)?;
self.move_to_next_sibling(false)?;
}
} else {
self.move_to_next_sibling(false)?;
self.update_skip_state();
}
self.update_skip_state();
Ok(())
}
fn key(&self) -> Option<Vec<u8>> {
if self.stack.is_empty() {
None
} else {
Some(self.stack.last().unwrap().full_key())
}
self.stack.last().map(|n| n.full_key())
}
fn hash(&self) -> Option<H256> {
if self.stack.is_empty() {
return None;
}
self.stack.last().unwrap().hash()
self.stack.last().and_then(|n| n.hash())
}
fn children_are_in_trie(&self) -> bool {
if self.stack.is_empty() {
return false;
}
self.stack.last().unwrap().tree_flag()
}
fn can_skip_state(&self) -> bool {
self.can_skip_state
self.stack.last().map_or(false, |n| n.tree_flag())
}
fn first_uncovered_prefix(&self) -> Option<Vec<u8>> {
let mut k = self.key();
if self.can_skip_state && k.is_some() {
k = increment_key(&k.unwrap());
match &self.key() {
Some(key) => {
if self.can_skip_state {
increment_key(key).map(|k| pack_nibbles(k.as_slice()))
} else {
Some(pack_nibbles(key.as_slice()))
}
}
None => None,
}
if k.is_none() {
return None;
}
Some(pack_nibbles(k.as_ref().unwrap()))
}
fn consume_node(&mut self, to: &[u8], exact: bool) -> anyhow::Result<()> {
fn consume_node(&mut self, to: &[u8], exact: bool) -> Result<()> {
let db_key = [self.prefix.as_slice(), to].concat().to_vec();
let entry = if exact {
self.cursor.lock().seek_exact(db_key)?
@@ -205,17 +191,15 @@ where
let mut node: Option<Node> = None;
if entry.is_some() {
node = Some(unmarshal_node(entry.as_ref().unwrap().1.as_slice()).unwrap());
assert_ne!(node.as_ref().unwrap().state_mask(), 0);
assert_ne!(node.as_ref().unwrap().state_mask, 0);
}
let mut nibble = 0i8;
if node.is_none() || node.as_ref().unwrap().root_hash().is_some() {
nibble = -1;
} else {
while node.as_ref().unwrap().state_mask() & (1u16 << nibble) == 0 {
nibble += 1;
let nibble = match &node {
Some(n) if n.root_hash.is_none() => {
(0i8..16).find(|i| n.state_mask & (1u16 << i) != 0).unwrap()
}
}
_ => -1,
};
if !key.is_empty() && !self.stack.is_empty() {
self.stack[0].nibble = key[0] as i8;
@@ -234,52 +218,41 @@ where
fn move_to_next_sibling(
&mut self,
allow_root_to_child_nibble_within_subnode: bool,
) -> anyhow::Result<()> {
if self.stack.is_empty() {
return Ok(());
}
let sn = self.stack.last().unwrap();
if sn.nibble >= 15 || (sn.nibble < 0 && !allow_root_to_child_nibble_within_subnode) {
self.stack.pop();
self.move_to_next_sibling(false)?;
return Ok(());
}
let sn = self.stack.last_mut().unwrap();
sn.nibble += 1;
if sn.node.is_none() {
let key = self.key();
self.consume_node(key.as_ref().unwrap(), false)?;
return Ok(());
}
while sn.nibble < 16 {
if sn.state_flag() {
) -> Result<()> {
if let Some(sn) = self.stack.last_mut() {
if sn.nibble >= 15 || (sn.nibble < 0 && !allow_root_to_child_nibble_within_subnode) {
self.stack.pop();
self.move_to_next_sibling(false)?;
return Ok(());
}
sn.nibble += 1;
}
self.stack.pop();
self.move_to_next_sibling(false)?;
sn.nibble += 1;
if sn.node.is_none() {
let key = self.key();
self.consume_node(key.as_ref().unwrap(), false)?;
return Ok(());
}
while sn.nibble < 16 {
if sn.state_flag() {
return Ok(());
}
sn.nibble += 1;
}
self.stack.pop();
self.move_to_next_sibling(false)?;
}
Ok(())
}
fn update_skip_state(&mut self) {
if self.key().is_none()
|| self.changed.contains(
[self.prefix.as_slice(), self.key().unwrap().as_slice()]
.concat()
.as_slice(),
)
{
self.can_skip_state = false;
self.can_skip_state = if let Some(key) = self.key() {
let s = [self.prefix.as_slice(), key.as_slice()].concat();
!self.changed.contains(s.as_slice()) && self.stack.last().unwrap().hash_flag()
} else {
self.can_skip_state = self.stack.last().unwrap().hash_flag();
false
}
}
@@ -297,7 +270,7 @@ where
{
txn: &'tx MdbxTransaction<'db, RW, E>,
hb: HashBuilder<'nc>,
storage_collector: Mutex<&'co mut TableCollector<'tmp, tables::TrieStorage>>,
storage_collector: &'co mut TableCollector<'tmp, tables::TrieStorage>,
rlp: Vec<u8>,
_marker: PhantomData<&'db ()>,
}
@@ -314,63 +287,58 @@ where
account_collector: &'co mut TableCollector<'tmp, tables::TrieAccount>,
storage_collector: &'co mut TableCollector<'tmp, tables::TrieStorage>,
) -> Self {
let mut instance = Self {
let node_collector = |unpacked_key: &[u8], node: &Node| {
if !unpacked_key.is_empty() {
account_collector.push(unpacked_key.to_vec(), marshal_node(node));
}
};
Self {
txn,
hb: HashBuilder::new(),
storage_collector: Mutex::new(storage_collector),
hb: HashBuilder::new(Some(Box::new(node_collector))),
storage_collector,
rlp: vec![],
_marker: PhantomData,
};
let node_collector = |unpacked_key: &[u8], node: &Node| {
if unpacked_key.is_empty() {
return;
}
account_collector.push(unpacked_key.to_vec(), marshal_node(node));
};
instance.hb.node_collector = Some(Box::new(node_collector));
instance
}
}
fn calculate_root(&mut self, changed: &mut PrefixSet) -> anyhow::Result<H256> {
fn calculate_root(&mut self, changed: &mut PrefixSet) -> Result<H256> {
let mut state = self.txn.cursor(tables::HashedAccount)?;
let mut trie_db_cursor = self.txn.cursor(tables::TrieAccount)?;
let mut trie = Cursor::new(&mut trie_db_cursor, changed, &[])?;
while trie.key().is_some() {
if trie.can_skip_state() {
assert!(trie.hash().is_some());
while let Some(key) = trie.key() {
if trie.can_skip_state {
self.hb.add_branch_node(
trie.key().unwrap(),
key,
trie.hash().as_ref().unwrap(),
trie.children_are_in_trie(),
);
}
let uncovered = trie.first_uncovered_prefix();
if uncovered.is_none() {
break;
}
let seek_key = match trie.first_uncovered_prefix() {
Some(mut uncovered) => {
uncovered.resize(32, 0);
uncovered
}
None => break,
};
trie.next()?;
let mut seek_key = uncovered.unwrap().to_vec();
seek_key.resize(32, 0);
let mut acc = state.seek(H256::from_slice(seek_key.as_slice()))?;
while acc.is_some() {
let unpacked_key = unpack_nibbles(acc.unwrap().0.as_bytes());
if trie.key().is_some() && trie.key().unwrap() < unpacked_key {
break;
while let Some((address, account)) = acc {
let unpacked_key = unpack_nibbles(address.as_bytes());
if let Some(key) = trie.key() {
if key < unpacked_key {
break;
}
}
let account = acc.unwrap().1;
let storage_root =
self.calculate_storage_root(acc.unwrap().0.as_bytes(), trie.changed_mut())?;
self.calculate_storage_root(address.as_bytes(), trie.changed_mut())?;
self.hb.add_leaf(
unpacked_key,
@@ -381,61 +349,61 @@ where
}
}
Ok(self.hb.root_hash())
Ok(self.hb.compute_root_hash())
}
fn calculate_storage_root(
&self,
&mut self,
key_with_inc: &[u8],
changed: &mut PrefixSet,
) -> anyhow::Result<H256> {
) -> Result<H256> {
let mut state = self.txn.cursor(tables::HashedStorage)?;
let mut trie_db_cursor = self.txn.cursor(tables::TrieStorage)?;
let mut hb = HashBuilder::new();
hb.node_collector = Some(Box::new(|unpacked_storage_key: &[u8], node: &Node| {
let key = [key_with_inc, unpacked_storage_key].concat();
self.storage_collector.lock().push(key, marshal_node(node));
}));
let mut hb = HashBuilder::new(Some(Box::new(
|unpacked_storage_key: &[u8], node: &Node| {
let key = [key_with_inc, unpacked_storage_key].concat();
self.storage_collector.push(key, marshal_node(node));
},
)));
let mut trie = Cursor::new(&mut trie_db_cursor, changed, key_with_inc)?;
while trie.key().is_some() {
if trie.can_skip_state() {
assert!(trie.hash().is_some());
while let Some(key) = trie.key() {
if trie.can_skip_state {
hb.add_branch_node(
trie.key().unwrap(),
key,
trie.hash().as_ref().unwrap(),
trie.children_are_in_trie(),
);
}
let uncovered = trie.first_uncovered_prefix();
if uncovered.is_none() {
break;
}
let seek_key = match trie.first_uncovered_prefix() {
Some(mut uncovered) => {
uncovered.resize(32, 0);
uncovered
}
None => break,
};
trie.next()?;
let mut seek_key = uncovered.unwrap().to_vec();
seek_key.resize(32, 0);
let mut storage = state.seek_both_range(
H256::from_slice(key_with_inc),
H256::from_slice(seek_key.as_slice()),
)?;
while storage.is_some() {
let (storage_location, value) = storage.unwrap();
while let Some((storage_location, value)) = storage {
let unpacked_loc = unpack_nibbles(storage_location.as_bytes());
if trie.key().is_some() && trie.key().unwrap() < unpacked_loc {
break;
if let Some(key) = trie.key() {
if key < unpacked_loc {
break;
}
}
let rlp = rlp::encode(&value);
hb.add_leaf(unpacked_loc, rlp.as_ref());
hb.add_leaf(unpacked_loc, rlp::encode(&value).as_ref());
storage = state.next_dup()?.map(|(_, v)| v);
}
}
Ok(hb.root_hash())
Ok(hb.compute_root_hash())
}
}
@@ -444,7 +412,7 @@ fn do_increment_intermediate_hashes<'db, 'tx, E>(
etl_dir: &TempDir,
expected_root: Option<H256>,
changed: &mut PrefixSet,
) -> Result<H256, DuoError>
) -> std::result::Result<H256, DuoError>
where
'db: 'tx,
E: EnvironmentKind,
@@ -454,15 +422,14 @@ where
let root = {
let mut loader = DbTrieLoader::new(txn, &mut account_collector, &mut storage_collector);
loader.calculate_root(changed)?
};
if let Some(expected_root) = expected_root {
if expected_root != root {
if let Some(expected) = expected_root {
if expected != root {
return Err(DuoError::Validation(Box::new(
ValidationError::WrongStateRoot {
expected: expected_root,
expected,
got: root,
},
)));
@@ -481,7 +448,7 @@ where
fn gather_changes<'db, 'tx, K, E>(
txn: &'tx MdbxTransaction<'db, K, E>,
from: BlockNumber,
) -> anyhow::Result<PrefixSet>
) -> Result<PrefixSet>
where
'db: 'tx,
K: TransactionKind,
@@ -493,26 +460,24 @@ where
let mut account_changes = txn.cursor(tables::AccountChangeSet)?;
let mut data = account_changes.seek(starting_key)?;
while data.is_some() {
let address = data.unwrap().1.address;
let hashed_address = keccak256(address);
while let Some((_, account_change)) = data {
let hashed_address = keccak256(account_change.address);
out.insert(unpack_nibbles(hashed_address.as_bytes()).as_slice());
data = account_changes.next()?;
}
let mut storage_changes = txn.cursor(tables::StorageChangeSet)?;
let mut data = storage_changes.seek(starting_key)?;
while data.is_some() {
let address = data.as_ref().unwrap().0.address;
let location = data.as_ref().unwrap().1.location;
let hashed_address = keccak256(address);
let hashed_location = keccak256(location);
while let Some((key, storage_change)) = data {
let hashed_address = keccak256(key.address);
let hashed_location = keccak256(storage_change.location);
let hashed_key = [
hashed_address.as_bytes(),
unpack_nibbles(hashed_location.as_bytes()).as_slice(),
]
.concat();
out.insert(hashed_key.as_slice());
data = storage_changes.next()?;
}
@@ -525,7 +490,7 @@ pub fn increment_intermediate_hashes<'db, 'tx, E>(
etl_dir: &TempDir,
from: BlockNumber,
expected_root: Option<H256>,
) -> Result<H256, DuoError>
) -> std::result::Result<H256, DuoError>
where
'db: 'tx,
E: EnvironmentKind,
@@ -538,7 +503,7 @@ pub fn regenerate_intermediate_hashes<'db, 'tx, E>(
txn: &'tx MdbxTransaction<'db, RW, E>,
etl_dir: &TempDir,
expected_root: Option<H256>,
) -> Result<H256, DuoError>
) -> std::result::Result<H256, DuoError>
where
'db: 'tx,
E: EnvironmentKind,
@@ -726,7 +691,7 @@ mod tests {
let mut cursor = Cursor::new(&mut trie, &mut changed, prefix_b.as_slice()).unwrap();
assert!(cursor.key().unwrap().is_empty());
assert!(cursor.can_skip_state());
assert!(cursor.can_skip_state);
cursor.next().unwrap();
assert!(cursor.key().is_none());
@@ -737,7 +702,7 @@ mod tests {
let mut cursor = Cursor::new(&mut trie, &mut changed, prefix_b.as_slice()).unwrap();
assert!(cursor.key().unwrap().is_empty());
assert!(!cursor.can_skip_state());
assert!(!cursor.can_skip_state);
cursor.next().unwrap();
assert_eq!(cursor.key().unwrap(), [0x2]);
@@ -808,7 +773,7 @@ mod tests {
let mut cursor = Cursor::new(&mut trie, &mut changed, &prefix_b).unwrap();
assert_eq!(cursor.key(), Some(vec![])); // root
assert!(cursor.can_skip_state()); // due to root_hash
assert!(cursor.can_skip_state); // due to root_hash
cursor.next().unwrap(); // skips to end of trie
assert_eq!(cursor.key(), None);
@@ -819,7 +784,7 @@ mod tests {
let mut cursor = Cursor::new(&mut trie, &mut changed, &prefix_b).unwrap();
assert_eq!(cursor.key(), Some(vec![])); // root
assert!(!cursor.can_skip_state());
assert!(!cursor.can_skip_state);
cursor.next().unwrap();
assert_eq!(cursor.key(), Some(vec![0x2]));
cursor.next().unwrap();
@@ -837,7 +802,7 @@ mod tests {
{
let mut hashed_storage = tx.cursor(tables::HashedStorage).unwrap();
let mut hb = HashBuilder::new();
let mut hb = HashBuilder::new(None);
for (loc, val) in [
(
@@ -865,7 +830,7 @@ mod tests {
hb.add_leaf(unpack_nibbles(loc.as_bytes()), &rlp::encode(&val));
}
hb.root_hash()
hb.compute_root_hash()
}
fn read_all_nodes<K, T>(cursor: MdbxCursor<'_, K, T>) -> HashMap<Vec<u8>, Node>
@@ -891,7 +856,7 @@ mod tests {
let txn = db.begin_mutable().unwrap();
let mut hashed_accounts = txn.cursor(tables::HashedAccount).unwrap();
let mut hb = HashBuilder::new();
let mut hb = HashBuilder::new(None);
let key1 = hex!("B000000000000000000000000000000000000000000000000000000000000000").into();
let a1 = Account {
@@ -981,7 +946,7 @@ mod tests {
// Populate account & storage trie DB tables
// ----------------------------------------------------------------
regenerate_intermediate_hashes(&txn, &temp_dir, Some(hb.root_hash())).unwrap();
regenerate_intermediate_hashes(&txn, &temp_dir, Some(hb.compute_root_hash())).unwrap();
// ----------------------------------------------------------------
// Check account trie
@@ -992,21 +957,21 @@ mod tests {
let node1a = &node_map[&vec![0xB]];
assert_eq!(node1a.state_mask(), 0b1011);
assert_eq!(node1a.tree_mask(), 0b0001);
assert_eq!(node1a.hash_mask(), 0b1001);
assert_eq!(node1a.state_mask, 0b1011);
assert_eq!(node1a.tree_mask, 0b0001);
assert_eq!(node1a.hash_mask, 0b1001);
assert_eq!(node1a.root_hash(), None);
assert_eq!(node1a.hashes().len(), 2);
assert_eq!(node1a.root_hash, None);
assert_eq!(node1a.hashes.len(), 2);
let node2a = &node_map[&vec![0xB, 0x0]];
assert_eq!(node2a.state_mask(), 0b10001);
assert_eq!(node2a.tree_mask(), 0b00000);
assert_eq!(node2a.hash_mask(), 0b10000);
assert_eq!(node2a.state_mask, 0b10001);
assert_eq!(node2a.tree_mask, 0b00000);
assert_eq!(node2a.hash_mask, 0b10000);
assert_eq!(node2a.root_hash(), None);
assert_eq!(node2a.hashes().len(), 1);
assert_eq!(node2a.root_hash, None);
assert_eq!(node2a.hashes.len(), 1);
// ----------------------------------------------------------------
// Check storage trie
@@ -1017,12 +982,12 @@ mod tests {
let node3 = &node_map[&key3.0.to_vec()];
assert_eq!(node3.state_mask(), 0b1010);
assert_eq!(node3.tree_mask(), 0b0000);
assert_eq!(node3.hash_mask(), 0b0010);
assert_eq!(node3.state_mask, 0b1010);
assert_eq!(node3.tree_mask, 0b0000);
assert_eq!(node3.hash_mask, 0b0010);
assert_eq!(node3.root_hash(), Some(storage_root));
assert_eq!(node3.hashes().len(), 1);
assert_eq!(node3.root_hash, Some(storage_root));
assert_eq!(node3.hashes.len(), 1);
// ----------------------------------------------------------------
// Add an account
@@ -1061,15 +1026,15 @@ mod tests {
assert_eq!(node_map.len(), 2);
let node1b = &node_map[&vec![0xB]];
assert_eq!(node1b.state_mask(), 0b1011);
assert_eq!(node1b.tree_mask(), 0b0001);
assert_eq!(node1b.hash_mask(), 0b1011);
assert_eq!(node1b.state_mask, 0b1011);
assert_eq!(node1b.tree_mask, 0b0001);
assert_eq!(node1b.hash_mask, 0b1011);
assert_eq!(node1b.root_hash(), None);
assert_eq!(node1b.root_hash, None);
assert_eq!(node1b.hashes().len(), 3);
assert_eq!(node1a.hashes()[0], node1b.hashes()[0]);
assert_eq!(node1a.hashes()[1], node1b.hashes()[2]);
assert_eq!(node1b.hashes.len(), 3);
assert_eq!(node1a.hashes[0], node1b.hashes[0]);
assert_eq!(node1a.hashes[1], node1b.hashes[2]);
let node2b = &node_map[&vec![0xB, 0x0]];
assert_eq!(node2a, node2b);
@@ -1104,16 +1069,16 @@ mod tests {
assert_eq!(node_map.len(), 1);
let node1c = &node_map[&vec![0xB]];
assert_eq!(node1c.state_mask(), 0b1011);
assert_eq!(node1c.tree_mask(), 0b0000);
assert_eq!(node1c.hash_mask(), 0b1011);
assert_eq!(node1c.state_mask, 0b1011);
assert_eq!(node1c.tree_mask, 0b0000);
assert_eq!(node1c.hash_mask, 0b1011);
assert_eq!(node1c.root_hash(), None);
assert_eq!(node1c.root_hash, None);
assert_eq!(node1c.hashes().len(), 3);
assert_ne!(node1b.hashes()[0], node1c.hashes()[0]);
assert_eq!(node1b.hashes()[1], node1c.hashes()[1]);
assert_eq!(node1b.hashes()[2], node1c.hashes()[2]);
assert_eq!(node1c.hashes.len(), 3);
assert_ne!(node1b.hashes[0], node1c.hashes[0]);
assert_eq!(node1b.hashes[1], node1c.hashes[1]);
assert_eq!(node1b.hashes[2], node1c.hashes[2]);
}
// Delete several accounts
@@ -1141,7 +1106,7 @@ mod tests {
assert_eq!(
read_all_nodes(account_trie),
hashmap! {
vec![0xB] => Node::new(0b1011, 0b0000, 0b1010, vec![node1b.hashes()[1], node1b.hashes()[2]], None)
vec![0xB] => Node::new(0b1011, 0b0000, 0b1010, vec![node1b.hashes[1], node1b.hashes[2]], None)
}
);
}
@@ -1160,7 +1125,7 @@ mod tests {
};
let mut hashed_accounts = txn.cursor(tables::HashedAccount).unwrap();
let mut hb = HashBuilder::new();
let mut hb = HashBuilder::new(None);
for key in [
hex!("30af561000000000000000000000000000000000000000000000000000000000"),
@@ -1177,7 +1142,7 @@ mod tests {
);
}
let expected_root = hb.root_hash();
let expected_root = hb.compute_root_hash();
assert_eq!(
regenerate_intermediate_hashes(&txn, &temp_dir, Some(expected_root)).unwrap(),
expected_root
@@ -1193,12 +1158,12 @@ mod tests {
let node2 = &node_map[&vec![0x3, 0x0, 0xA, 0xF]];
assert_eq!(node2.state_mask(), 0b101100000);
assert_eq!(node2.tree_mask(), 0b000000000);
assert_eq!(node2.hash_mask(), 0b001000000);
assert_eq!(node2.state_mask, 0b101100000);
assert_eq!(node2.tree_mask, 0b000000000);
assert_eq!(node2.hash_mask, 0b001000000);
assert_eq!(node2.root_hash(), None);
assert_eq!(node2.hashes().len(), 1);
assert_eq!(node2.root_hash, None);
assert_eq!(node2.hashes.len(), 1);
}
fn int_to_address(i: u128) -> Address {

69
src/trie/node.rs
View File

@@ -1,14 +1,13 @@
#![allow(clippy::if_same_then_else)]
use crate::{models::KECCAK_LENGTH, trie::util::assert_subset};
use ethereum_types::H256;
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct Node {
state_mask: u16,
tree_mask: u16,
hash_mask: u16,
hashes: Vec<H256>,
root_hash: Option<H256>,
pub(crate) state_mask: u16,
pub(crate) tree_mask: u16,
pub(crate) hash_mask: u16,
pub(crate) hashes: Vec<H256>,
pub(crate) root_hash: Option<H256>,
}
impl Node {
@@ -31,57 +30,27 @@ impl Node {
}
}
fn assign(&mut self, other: &Node) {
self.state_mask = other.state_mask;
self.tree_mask = other.tree_mask;
self.hash_mask = other.hash_mask;
self.hashes = other.hashes.clone();
self.root_hash = other.root_hash;
}
pub(crate) fn state_mask(&self) -> u16 {
self.state_mask
}
pub(crate) fn tree_mask(&self) -> u16 {
self.tree_mask
}
pub(crate) fn hash_mask(&self) -> u16 {
self.hash_mask
}
pub(crate) fn hashes(&self) -> Vec<H256> {
self.hashes.clone()
}
pub(crate) fn root_hash(&self) -> Option<H256> {
self.root_hash
}
pub(crate) fn set_root_hash(&mut self, root_hash: Option<H256>) {
self.root_hash = root_hash;
pub(crate) fn hash_for_nibble(&self, nibble: i8) -> H256 {
let mask = (1u16 << nibble) - 1;
let index = (self.hash_mask & mask).count_ones();
self.hashes[index as usize]
}
}
pub(crate) fn marshal_node(n: &Node) -> Vec<u8> {
let buf_size =
6 + if n.root_hash().is_some() {
KECCAK_LENGTH
} else {
0
} + n.hashes().len() * KECCAK_LENGTH;
let number_of_hashes = n.hashes.len() + if n.root_hash.is_some() { 1 } else { 0 };
let buf_size = number_of_hashes * KECCAK_LENGTH + 6;
let mut buf = Vec::<u8>::with_capacity(buf_size);
buf.extend_from_slice(n.state_mask().to_be_bytes().as_slice());
buf.extend_from_slice(n.tree_mask().to_be_bytes().as_slice());
buf.extend_from_slice(n.hash_mask().to_be_bytes().as_slice());
buf.extend_from_slice(n.state_mask.to_be_bytes().as_slice());
buf.extend_from_slice(n.tree_mask.to_be_bytes().as_slice());
buf.extend_from_slice(n.hash_mask.to_be_bytes().as_slice());
if n.root_hash().is_some() {
buf.extend_from_slice(&*n.root_hash.unwrap().as_bytes());
if let Some(root_hash) = n.root_hash {
buf.extend_from_slice(root_hash.as_bytes());
}
for hash in n.hashes() {
for hash in &n.hashes {
buf.extend_from_slice(hash.as_bytes());
}
@@ -89,9 +58,7 @@ pub(crate) fn marshal_node(n: &Node) -> Vec<u8> {
}
pub(crate) fn unmarshal_node(v: &[u8]) -> Option<Node> {
if v.len() < 6 {
return None;
} else if (v.len() - 6) % KECCAK_LENGTH != 0 {
if v.len() % KECCAK_LENGTH != 6 {
return None;
}