Skip to content

Commit

Permalink
Merge pull request #39 from 0xPolygonZero/subtrie_pruning_optimization
Browse files Browse the repository at this point in the history 
Additional tooling for allowing aggresive pruning logic
  • Loading branch information
@BGluth
BGluth authored Feb 21, 2024
2 parents 62b6f5a + 3bd89a1 commit 2d36eba
Show file tree
Hide file tree
Showing 8 changed files with 675 additions and 154 deletions.
99 changes: 1 addition & 98 deletions mpt_trie/src/debug_tools/common.rs
View file
Original file line number Diff line number Diff line change
@@ -1,66 +1,9 @@
//! Common utilities for the debugging tools.
use std::fmt::{self, Display};

use crate::{
nibbles::{Nibble, Nibbles},
nibbles::Nibbles,
partial_trie::{Node, PartialTrie},
utils::TrieNodeType,
};

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(super) enum PathSegment {
Empty,
Hash,
Branch(Nibble),
Extension(Nibbles),
Leaf(Nibbles),
}

impl Display for PathSegment {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
PathSegment::Empty => write!(f, "Empty"),
PathSegment::Hash => write!(f, "Hash"),
PathSegment::Branch(nib) => write!(f, "Branch({})", nib),
PathSegment::Extension(nibs) => write!(f, "Extension({})", nibs),
PathSegment::Leaf(nibs) => write!(f, "Leaf({})", nibs),
}
}
}

impl PathSegment {
pub(super) fn node_type(&self) -> TrieNodeType {
match self {
PathSegment::Empty => TrieNodeType::Empty,
PathSegment::Hash => TrieNodeType::Hash,
PathSegment::Branch(_) => TrieNodeType::Branch,
PathSegment::Extension(_) => TrieNodeType::Extension,
PathSegment::Leaf(_) => TrieNodeType::Leaf,
}
}

pub(super) fn get_key_piece_from_seg_if_present(&self) -> Option<Nibbles> {
match self {
PathSegment::Empty | PathSegment::Hash => None,
PathSegment::Branch(nib) => Some(Nibbles::from_nibble(*nib)),
PathSegment::Extension(nibs) | PathSegment::Leaf(nibs) => Some(*nibs),
}
}
}

pub(super) fn get_segment_from_node_and_key_piece<T: PartialTrie>(
n: &Node<T>,
k_piece: &Nibbles,
) -> PathSegment {
match TrieNodeType::from(n) {
TrieNodeType::Empty => PathSegment::Empty,
TrieNodeType::Hash => PathSegment::Hash,
TrieNodeType::Branch => PathSegment::Branch(k_piece.get_nibble(0)),
TrieNodeType::Extension => PathSegment::Extension(*k_piece),
TrieNodeType::Leaf => PathSegment::Leaf(*k_piece),
}
}

/// Get the key piece from the given node if applicable.
///
/// Note that there is no specific [`Nibble`] associated with a branch like
Expand All @@ -87,43 +30,3 @@ pub(super) fn get_key_piece_from_node<T: PartialTrie>(n: &Node<T>) -> Nibbles {
Node::Extension { nibbles, child: _ } | Node::Leaf { nibbles, value: _ } => *nibbles,
}
}

#[derive(Clone, Debug, Default, Eq, Hash, PartialEq)]
/// A vector of path segments representing a path in the trie.
pub struct NodePath(pub(super) Vec<PathSegment>);

impl NodePath {
pub(super) fn dup_and_append(&self, seg: PathSegment) -> Self {
let mut duped_vec = self.0.clone();
duped_vec.push(seg);

Self(duped_vec)
}

pub(super) fn append(&mut self, seg: PathSegment) {
self.0.push(seg);
}

fn write_elem(f: &mut fmt::Formatter<'_>, seg: &PathSegment) -> fmt::Result {
write!(f, "{}", seg)
}
}

impl Display for NodePath {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let num_elems = self.0.len();

// For everything but the last elem.
for seg in self.0.iter().take(num_elems.saturating_sub(1)) {
Self::write_elem(f, seg)?;
write!(f, " --> ")?;
}

// Avoid the extra `-->` for the last elem.
if let Some(seg) = self.0.last() {
Self::write_elem(f, seg)?;
}

Ok(())
}
}
17 changes: 9 additions & 8 deletions mpt_trie/src/debug_tools/diff.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ use std::{fmt::Display, ops::Deref};

use ethereum_types::H256;

use super::common::{get_key_piece_from_node, get_segment_from_node_and_key_piece, NodePath};
use super::common::get_key_piece_from_node;
use crate::utils::{get_segment_from_node_and_key_piece, TriePath};
use crate::{
nibbles::Nibbles,
partial_trie::{HashedPartialTrie, Node, PartialTrie},
Expand Down Expand Up @@ -83,7 +84,7 @@ pub struct DiffPoint {
/// The depth of the point in both tries.
pub depth: usize,
/// The path of the point in both tries.
pub path: NodePath,
pub path: TriePath,
/// The node key in both tries.
pub key: Nibbles,
/// The node info in the first trie.
Expand All @@ -97,7 +98,7 @@ impl DiffPoint {
child_a: &HashedPartialTrie,
child_b: &HashedPartialTrie,
parent_k: Nibbles,
path: NodePath,
path: TriePath,
) -> Self {
let a_key = parent_k.merge_nibbles(&get_key_piece_from_node(child_a));
let b_key = parent_k.merge_nibbles(&get_key_piece_from_node(child_b));
Expand Down Expand Up @@ -222,7 +223,7 @@ impl DepthNodeDiffState {
parent_k: &Nibbles,
child_a: &HashedPartialTrie,
child_b: &HashedPartialTrie,
path: NodePath,
path: TriePath,
) {
if field
.as_ref()
Expand All @@ -242,7 +243,7 @@ struct DepthDiffPerCallState<'a> {
curr_depth: usize,

// Horribly inefficient, but these are debug tools, so I think we get a pass.
curr_path: NodePath,
curr_path: TriePath,
}

impl<'a> DepthDiffPerCallState<'a> {
Expand All @@ -259,7 +260,7 @@ impl<'a> DepthDiffPerCallState<'a> {
b,
curr_key,
curr_depth,
curr_path: NodePath::default(),
curr_path: TriePath::default(),
}
}

Expand Down Expand Up @@ -411,7 +412,7 @@ fn get_value_from_node<T: PartialTrie>(n: &Node<T>) -> Option<&Vec<u8>> {

#[cfg(test)]
mod tests {
use super::{create_diff_between_tries, DiffPoint, NodeInfo, NodePath};
use super::{create_diff_between_tries, DiffPoint, NodeInfo, TriePath};
use crate::{
nibbles::Nibbles,
partial_trie::{HashedPartialTrie, PartialTrie},
Expand Down Expand Up @@ -447,7 +448,7 @@ mod tests {

let expected = DiffPoint {
depth: 0,
path: NodePath(vec![]),
path: TriePath(vec![]),
key: Nibbles::default(),
a_info: expected_a,
b_info: expected_b,
Expand Down
14 changes: 7 additions & 7 deletions mpt_trie/src/debug_tools/query.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,11 @@ use std::fmt::{self, Display};

use ethereum_types::H256;

use super::common::{
get_key_piece_from_node_pulling_from_key_for_branches, get_segment_from_node_and_key_piece,
NodePath, PathSegment,
};
use super::common::get_key_piece_from_node_pulling_from_key_for_branches;
use crate::{
nibbles::Nibbles,
partial_trie::{Node, PartialTrie, WrappedNode},
utils::{get_segment_from_node_and_key_piece, TriePath, TrieSegment},
};

/// Params controlling how much information is reported in the query output.
Expand Down Expand Up @@ -159,7 +157,9 @@ fn count_non_empty_branch_children_from_mask(mask: u16) -> usize {
/// of the path used for searching for a key in the trie.
pub struct DebugQueryOutput {
k: Nibbles,
node_path: NodePath,

/// The nodes hit during the query.
pub node_path: TriePath,
extra_node_info: Vec<Option<ExtraNodeSegmentInfo>>,
node_found: bool,
params: DebugQueryParams,
Expand Down Expand Up @@ -199,7 +199,7 @@ impl DebugQueryOutput {
fn new(k: Nibbles, params: DebugQueryParams) -> Self {
Self {
k,
node_path: NodePath::default(),
node_path: TriePath::default(),
extra_node_info: Vec::default(),
node_found: false,
params,
Expand All @@ -219,7 +219,7 @@ impl DebugQueryOutput {
// TODO: Make the output easier to read...
fn fmt_node_based_on_debug_params(
f: &mut fmt::Formatter<'_>,
seg: &PathSegment,
seg: &TrieSegment,
extra_seg_info: &Option<ExtraNodeSegmentInfo>,
params: &DebugQueryParams,
) -> fmt::Result {
Expand Down
3 changes: 2 additions & 1 deletion mpt_trie/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@

pub mod nibbles;
pub mod partial_trie;
pub mod special_query;
mod trie_hashing;
pub mod trie_ops;
pub mod trie_subsets;
mod utils;
pub mod utils;

#[cfg(feature = "trie_debug")]
pub mod debug_tools;
Expand Down
91 changes: 88 additions & 3 deletions mpt_trie/src/nibbles.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -363,18 +363,33 @@ impl Nibbles {
/// Appends `Nibbles` to the front.
///
/// # Panics
/// Panics if appending the `Nibble` causes an overflow (total nibbles >
/// Panics if appending the `Nibbles` causes an overflow (total nibbles >
/// 64).
pub fn push_nibbles_front(&mut self, n: &Self) {
let new_count = self.count + n.count;
assert!(new_count <= 64);
self.nibbles_append_safety_asserts(new_count);

let shift_amt = 4 * self.count;

self.count = new_count;
self.packed |= n.packed << shift_amt;
}

/// Appends `Nibbles` to the back.
///
/// # Panics
/// Panics if appending the `Nibbles` causes an overflow (total nibbles >
/// 64).
pub fn push_nibbles_back(&mut self, n: &Self) {
let new_count = self.count + n.count;
self.nibbles_append_safety_asserts(new_count);

let shift_amt = 4 * n.count;

self.count = new_count;
self.packed = (self.packed << shift_amt) | n.packed;
}

/// Gets the nibbles at the range specified, where `0` is the next nibble.
///
/// # Panics
Expand Down Expand Up @@ -765,6 +780,10 @@ impl Nibbles {
assert!(n < 16);
}

fn nibbles_append_safety_asserts(&self, new_count: usize) {
assert!(new_count <= 64);
}

// TODO: REMOVE BEFORE NEXT CRATE VERSION! THIS IS A TEMP HACK!
/// Converts to u256 returning an error if not possible.
pub fn try_into_u256(&self) -> Result<U256, String> {
Expand All @@ -787,6 +806,9 @@ mod tests {
use super::{Nibble, Nibbles, ToNibbles};
use crate::nibbles::FromHexPrefixError;

const LONG_ZERO_NIBS_STR_LEN_63: &str =
"0x000000000000000000000000000000000000000000000000000000000000000";

#[test]
fn get_nibble_works() {
let n = Nibbles::from(0x1234);
Expand Down Expand Up @@ -898,6 +920,69 @@ mod tests {
assert_eq!(res, expected_resulting_nibbles);
}

#[test]
fn push_nibble_front_works() {
test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_front(0x1));
test_and_assert_nib_push_func(0x1, 0x21, |n| n.push_nibble_front(0x2));
test_and_assert_nib_push_func(
Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
Nibbles::from_str("0x1000000000000000000000000000000000000000000000000000000000000000")
.unwrap(),
|n| n.push_nibble_front(0x1),
);
}

#[test]
fn push_nibble_back_works() {
test_and_assert_nib_push_func(Nibbles::default(), 0x1, |n| n.push_nibble_back(0x1));
test_and_assert_nib_push_func(0x1, 0x12, |n| n.push_nibble_back(0x2));
test_and_assert_nib_push_func(
Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
Nibbles::from_str("0x0000000000000000000000000000000000000000000000000000000000000001")
.unwrap(),
|n| n.push_nibble_back(0x1),
);
}

#[test]
fn push_nibbles_front_works() {
test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| {
n.push_nibbles_front(&0x1234.into())
});
test_and_assert_nib_push_func(0x1234, 0x5671234, |n| n.push_nibbles_front(&0x567.into()));
test_and_assert_nib_push_func(
Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
Nibbles::from_str("0x1000000000000000000000000000000000000000000000000000000000000000")
.unwrap(),
|n| n.push_nibbles_front(&0x1.into()),
);
}

#[test]
fn push_nibbles_back_works() {
test_and_assert_nib_push_func(Nibbles::default(), 0x1234, |n| {
n.push_nibbles_back(&0x1234.into())
});
test_and_assert_nib_push_func(0x1234, 0x1234567, |n| n.push_nibbles_back(&0x567.into()));
test_and_assert_nib_push_func(
Nibbles::from_str(LONG_ZERO_NIBS_STR_LEN_63).unwrap(),
Nibbles::from_str("0x0000000000000000000000000000000000000000000000000000000000000001")
.unwrap(),
|n| n.push_nibbles_back(&0x1.into()),
);
}

fn test_and_assert_nib_push_func<F: Fn(&mut Nibbles), S: Into<Nibbles>, E: Into<Nibbles>>(
starting_nibs: S,
expected: E,
f: F,
) {
let mut nibs = starting_nibs.into();
(f)(&mut nibs);

assert_eq!(nibs, expected.into());
}

#[test]
fn get_next_nibbles_works() {
let n: Nibbles = 0x1234.into();
Expand Down Expand Up @@ -1179,7 +1264,7 @@ mod tests {
fn nibbles_from_h256_works() {
assert_eq!(
format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(0))),
"0x0000000000000000000000000000000000000000000000000000000000000000"
"0x0000000000000000000000000000000000000000000000000000000000000000",
);
assert_eq!(
format!("{:x}", Nibbles::from_h256_be(H256::from_low_u64_be(2048))),
Expand Down
Loading

0 comments on commit 2d36eba

Please sign in to comment.