Skip to content

Commit

Permalink
feat(trie): bincode compatibility for trie updates (#11409)
Browse files Browse the repository at this point in the history
  • Loading branch information
shekhirin authored Oct 2, 2024
1 parent a07efa7 commit 4491b0d
Show file tree
Hide file tree
Showing 6 changed files with 237 additions and 6 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion crates/evm/execution-types/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,5 +35,5 @@ reth-primitives = { workspace = true, features = ["arbitrary", "test-utils"] }
default = ["std"]
optimism = ["reth-primitives/optimism", "revm/optimism"]
serde = ["dep:serde", "reth-trie/serde", "revm/serde"]
serde-bincode-compat = ["reth-primitives/serde-bincode-compat", "serde_with"]
serde-bincode-compat = ["reth-primitives/serde-bincode-compat", "reth-trie/serde-bincode-compat", "serde_with"]
std = []
8 changes: 4 additions & 4 deletions crates/evm/execution-types/src/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -732,7 +732,7 @@ pub(super) mod serde_bincode_compat {
use alloc::borrow::Cow;
use alloy_primitives::BlockNumber;
use reth_primitives::serde_bincode_compat::SealedBlockWithSenders;
use reth_trie::updates::TrieUpdates;
use reth_trie::serde_bincode_compat::updates::TrieUpdates;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_with::{DeserializeAs, SerializeAs};

Expand All @@ -757,7 +757,7 @@ pub(super) mod serde_bincode_compat {
pub struct Chain<'a> {
blocks: BTreeMap<BlockNumber, SealedBlockWithSenders<'a>>,
execution_outcome: Cow<'a, ExecutionOutcome>,
trie_updates: Cow<'a, Option<TrieUpdates>>,
trie_updates: Option<TrieUpdates<'a>>,
}

impl<'a> From<&'a super::Chain> for Chain<'a> {
Expand All @@ -769,7 +769,7 @@ pub(super) mod serde_bincode_compat {
.map(|(block_number, block)| (*block_number, block.into()))
.collect(),
execution_outcome: Cow::Borrowed(&value.execution_outcome),
trie_updates: Cow::Borrowed(&value.trie_updates),
trie_updates: value.trie_updates.as_ref().map(Into::into),
}
}
}
Expand All @@ -783,7 +783,7 @@ pub(super) mod serde_bincode_compat {
.map(|(block_number, block)| (block_number, block.into()))
.collect(),
execution_outcome: value.execution_outcome.into_owned(),
trie_updates: value.trie_updates.into_owned(),
trie_updates: value.trie_updates.map(Into::into),
}
}
}
Expand Down
5 changes: 5 additions & 0 deletions crates/trie/trie/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@ triehash = { version = "0.8", optional = true }
# `serde` feature
serde = { workspace = true, optional = true }

# `serde-bincode-compat` feature
serde_with = { workspace = true, optional = true }

[dev-dependencies]
# reth
reth-chainspec.workspace = true
Expand All @@ -63,10 +66,12 @@ tokio = { workspace = true, default-features = false, features = [
] }
serde_json.workspace = true
criterion.workspace = true
bincode.workspace = true

[features]
metrics = ["reth-metrics", "dep:metrics"]
serde = ["dep:serde"]
serde-bincode-compat = ["serde_with"]
test-utils = ["triehash", "reth-trie-common/test-utils"]

[[bench]]
Expand Down
11 changes: 11 additions & 0 deletions crates/trie/trie/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,17 @@ pub mod stats;
// re-export for convenience
pub use reth_trie_common::*;

/// Bincode-compatible serde implementations for trie types.
///
/// `bincode` crate allows for more efficient serialization of trie types, because it allows
/// non-string map keys.
///
/// Read more: <https://github.com/paradigmxyz/reth/issues/11370>
#[cfg(all(feature = "serde", feature = "serde-bincode-compat"))]
pub mod serde_bincode_compat {
pub use super::updates::serde_bincode_compat as updates;
}

/// Trie calculation metrics.
#[cfg(feature = "metrics")]
pub mod metrics;
Expand Down
215 changes: 214 additions & 1 deletion crates/trie/trie/src/updates.rs
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,220 @@ fn exclude_empty_from_pair<V>(
iter.into_iter().filter(|(n, _)| !n.is_empty())
}

/// Bincode-compatible trie updates type serde implementations.
#[cfg(all(feature = "serde", feature = "serde-bincode-compat"))]
pub mod serde_bincode_compat {
use std::{
borrow::Cow,
collections::{HashMap, HashSet},
};

use alloy_primitives::B256;
use reth_trie_common::{BranchNodeCompact, Nibbles};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_with::{DeserializeAs, SerializeAs};

/// Bincode-compatible [`super::TrieUpdates`] serde implementation.
///
/// Intended to use with the [`serde_with::serde_as`] macro in the following way:
/// ```rust
/// use reth_trie::{serde_bincode_compat, updates::TrieUpdates};
/// use serde::{Deserialize, Serialize};
/// use serde_with::serde_as;
///
/// #[serde_as]
/// #[derive(Serialize, Deserialize)]
/// struct Data {
/// #[serde_as(as = "serde_bincode_compat::updates::TrieUpdates")]
/// trie_updates: TrieUpdates,
/// }
/// ```
#[derive(Debug, Serialize, Deserialize)]
pub struct TrieUpdates<'a> {
account_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
removed_nodes: Cow<'a, HashSet<Nibbles>>,
storage_tries: HashMap<B256, StorageTrieUpdates<'a>>,
}

impl<'a> From<&'a super::TrieUpdates> for TrieUpdates<'a> {
fn from(value: &'a super::TrieUpdates) -> Self {
Self {
account_nodes: Cow::Borrowed(&value.account_nodes),
removed_nodes: Cow::Borrowed(&value.removed_nodes),
storage_tries: value.storage_tries.iter().map(|(k, v)| (*k, v.into())).collect(),
}
}
}

impl<'a> From<TrieUpdates<'a>> for super::TrieUpdates {
fn from(value: TrieUpdates<'a>) -> Self {
Self {
account_nodes: value.account_nodes.into_owned(),
removed_nodes: value.removed_nodes.into_owned(),
storage_tries: value
.storage_tries
.into_iter()
.map(|(k, v)| (k, v.into()))
.collect(),
}
}
}

impl<'a> SerializeAs<super::TrieUpdates> for TrieUpdates<'a> {
fn serialize_as<S>(source: &super::TrieUpdates, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
TrieUpdates::from(source).serialize(serializer)
}
}

impl<'de> DeserializeAs<'de, super::TrieUpdates> for TrieUpdates<'de> {
fn deserialize_as<D>(deserializer: D) -> Result<super::TrieUpdates, D::Error>
where
D: Deserializer<'de>,
{
TrieUpdates::deserialize(deserializer).map(Into::into)
}
}

/// Bincode-compatible [`super::StorageTrieUpdates`] serde implementation.
///
/// Intended to use with the [`serde_with::serde_as`] macro in the following way:
/// ```rust
/// use reth_trie::{serde_bincode_compat, updates::StorageTrieUpdates};
/// use serde::{Deserialize, Serialize};
/// use serde_with::serde_as;
///
/// #[serde_as]
/// #[derive(Serialize, Deserialize)]
/// struct Data {
/// #[serde_as(as = "serde_bincode_compat::updates::StorageTrieUpdates")]
/// trie_updates: StorageTrieUpdates,
/// }
/// ```
#[derive(Debug, Serialize, Deserialize)]
pub struct StorageTrieUpdates<'a> {
is_deleted: bool,
storage_nodes: Cow<'a, HashMap<Nibbles, BranchNodeCompact>>,
removed_nodes: Cow<'a, HashSet<Nibbles>>,
}

impl<'a> From<&'a super::StorageTrieUpdates> for StorageTrieUpdates<'a> {
fn from(value: &'a super::StorageTrieUpdates) -> Self {
Self {
is_deleted: value.is_deleted,
storage_nodes: Cow::Borrowed(&value.storage_nodes),
removed_nodes: Cow::Borrowed(&value.removed_nodes),
}
}
}

impl<'a> From<StorageTrieUpdates<'a>> for super::StorageTrieUpdates {
fn from(value: StorageTrieUpdates<'a>) -> Self {
Self {
is_deleted: value.is_deleted,
storage_nodes: value.storage_nodes.into_owned(),
removed_nodes: value.removed_nodes.into_owned(),
}
}
}

impl<'a> SerializeAs<super::StorageTrieUpdates> for StorageTrieUpdates<'a> {
fn serialize_as<S>(
source: &super::StorageTrieUpdates,
serializer: S,
) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
StorageTrieUpdates::from(source).serialize(serializer)
}
}

impl<'de> DeserializeAs<'de, super::StorageTrieUpdates> for StorageTrieUpdates<'de> {
fn deserialize_as<D>(deserializer: D) -> Result<super::StorageTrieUpdates, D::Error>
where
D: Deserializer<'de>,
{
StorageTrieUpdates::deserialize(deserializer).map(Into::into)
}
}

#[cfg(test)]
mod tests {
use crate::updates::StorageTrieUpdates;

use super::super::{serde_bincode_compat, TrieUpdates};

use alloy_primitives::B256;
use reth_trie_common::{BranchNodeCompact, Nibbles};
use serde::{Deserialize, Serialize};
use serde_with::serde_as;

#[test]
fn test_trie_updates_bincode_roundtrip() {
#[serde_as]
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
struct Data {
#[serde_as(as = "serde_bincode_compat::TrieUpdates")]
trie_updates: TrieUpdates,
}

let mut data = Data { trie_updates: TrieUpdates::default() };
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);

data.trie_updates.removed_nodes.insert(Nibbles::from_vec(vec![0x0b, 0x0e, 0x0e, 0x0f]));
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);

data.trie_updates.account_nodes.insert(
Nibbles::from_vec(vec![0x0d, 0x0e, 0x0a, 0x0d]),
BranchNodeCompact::default(),
);
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);

data.trie_updates.storage_tries.insert(B256::default(), StorageTrieUpdates::default());
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);
}

#[test]
fn test_storage_trie_updates_bincode_roundtrip() {
#[serde_as]
#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
struct Data {
#[serde_as(as = "serde_bincode_compat::StorageTrieUpdates")]
trie_updates: StorageTrieUpdates,
}

let mut data = Data { trie_updates: StorageTrieUpdates::default() };
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);

data.trie_updates.removed_nodes.insert(Nibbles::from_vec(vec![0x0b, 0x0e, 0x0e, 0x0f]));
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);

data.trie_updates.storage_nodes.insert(
Nibbles::from_vec(vec![0x0d, 0x0e, 0x0a, 0x0d]),
BranchNodeCompact::default(),
);
let encoded = bincode::serialize(&data).unwrap();
let decoded: Data = bincode::deserialize(&encoded).unwrap();
assert_eq!(decoded, data);
}
}
}

#[cfg(all(test, feature = "serde"))]
mod tests {
use super::*;
Expand All @@ -416,7 +630,6 @@ mod tests {
.account_nodes
.insert(Nibbles::from_vec(vec![0x0d, 0x0e, 0x0a, 0x0d]), BranchNodeCompact::default());
let updates_serialized = serde_json::to_string(&default_updates).unwrap();
println!("{updates_serialized}");
let updates_deserialized: TrieUpdates = serde_json::from_str(&updates_serialized).unwrap();
assert_eq!(updates_deserialized, default_updates);

Expand Down

0 comments on commit 4491b0d

Please sign in to comment.