Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(trie): bincode compatibility for trie updates #11409

Merged
merged 2 commits into from
Oct 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 {
shekhirin marked this conversation as resolved.
Show resolved Hide resolved
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
Loading