Skip to content

Commit

Permalink
chore: tweaks
Browse files Browse the repository at this point in the history
  • Loading branch information
@0xaatif
0xaatif committed Sep 11, 2024
1 parent c86f69f commit af9251a
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 30 deletions.
31 changes: 17 additions & 14 deletions trace_decoder/src/core.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use crate::{
TxnInfo, TxnMeta, TxnTrace,
};

/// TODO(0xaatif): doc
/// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/275
pub fn entrypoint(
trace: BlockTrace,
other: OtherBlockData,
Expand All @@ -53,9 +53,7 @@ pub fn entrypoint(
checkpoint_state_trie_root,
} = other;

// TODO(0xaatif): docs for the RPC field say this is gwei already:
// https://docs.rs/alloy/0.3.1/alloy/eips/eip4895/struct.Withdrawal.html#structfield.amount
// in any case, this shouldn't be our problem.
// BUG?(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/618
for (_, amt) in &mut withdrawals {
*amt = eth_to_gwei(*amt)
}
Expand Down Expand Up @@ -128,7 +126,8 @@ fn start(
pre_images: BlockTraceTriePreImages,
) -> anyhow::Result<(StateMpt, BTreeMap<H256, StorageTrie>, Hash2Code)> {
Ok(match pre_images {
// TODO(0xaatif): refactor our convoluted input types
// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/issues/401
// refactor our convoluted input types
BlockTraceTriePreImages::Separate(SeparateTriePreImages {
state: SeparateTriePreImage::Direct(state),
storage: SeparateStorageTriesPreImage::MultipleTries(storage),
Expand Down Expand Up @@ -205,7 +204,8 @@ fn batch(txns: Vec<TxnInfo>, batch_size_hint: usize) -> Vec<Vec<Option<TxnInfo>>
.into_iter()
.map(FromIterator::from_iter)
.collect(),
// not enough batches at `hint`, but enough real transactions
// not enough batches at `hint`, but enough real transactions,
// so just split them in half
(2.., ..2) => {
let second = txns.split_off(txns.len() / 2);
vec![txns, second]
Expand Down Expand Up @@ -317,7 +317,7 @@ fn middle<StateTrieT: StateTrie + Clone>(
storage: storage_tries.clone(),
};

// We want to trim the TrieInputs above,
// We want to perform intersections the TrieInputs above,
// but won't know the bounds until after the loop below,
// so store that information here.
let mut storage_masks = BTreeMap::<_, BTreeSet<TrieKey>>::new();
Expand Down Expand Up @@ -464,10 +464,13 @@ fn middle<StateTrieT: StateTrie + Clone>(
..address!("000000000000000000000000000000000000000a");

if !precompiled_addresses.contains(&addr.compat()) {
// TODO(0xaatif): this looks like an optimization,
// but if it's omitted, the tests fail...
// TODO(0xaatif): https://github.com/0xPolygonZero/zk_evm/pull/613
// masking like this SHOULD be a space-saving optimization,
// BUT if it's omitted, we actually get state root mismatches
state_mask.insert(TrieKey::from_address(addr));
}
} // else we don't even need to include them,
// because nodes will only emit a precompiled address if
// the transaction calling them reverted.
}

if do_increment_txn_ix {
Expand Down Expand Up @@ -500,9 +503,9 @@ fn middle<StateTrieT: StateTrie + Clone>(
false => vec![],
},
before: {
before.state.mask(state_mask)?;
before.receipt.mask(batch_first_txn_ix..txn_ix)?;
before.transaction.mask(batch_first_txn_ix..txn_ix)?;
before.state.intersect(state_mask)?;
before.receipt.intersect(batch_first_txn_ix..txn_ix)?;
before.transaction.intersect(batch_first_txn_ix..txn_ix)?;

let keep = storage_masks
.keys()
Expand All @@ -512,7 +515,7 @@ fn middle<StateTrieT: StateTrie + Clone>(

for (addr, mask) in storage_masks {
if let Some(it) = before.storage.get_mut(&keccak_hash::keccak(addr)) {
it.mask(mask)?
it.intersect(mask)?
} // else self_destructed
}
before
Expand Down
7 changes: 4 additions & 3 deletions trace_decoder/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,10 @@
//!
//! **Prover perfomance is a high priority.**
//!
//! The aformentioned trie structures may have subtries _indirected_.
//! The aformentioned trie structures may have subtries _hashed out_.
//! That is, any node (and its children!) may be replaced by its hash,
//! while maintaining provability of its contents:
//!
//! ```text
//! A A
//! / \ / \
Expand All @@ -25,9 +26,9 @@
//! ```
//! (where `H` is the hash of the `D/B\E` subtrie).
//!
//! The principle concern of this library is to step through the transactions,
//! The principle concern of this module is to step through the transactions,
//! and reproduce the _intermediate tries_,
//! while indirecting all possible subtries to minimise prover load
//! while hashing out all possible subtries to minimise prover load
//! (since prover performance is sensitive to the size of the trie).
//! The prover can therefore prove each batch of transactions independently.
//!
Expand Down
28 changes: 15 additions & 13 deletions trace_decoder/src/typed_mpt.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use u4::{AsNibbles, U4};

/// See <https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie>.
///
/// Portions of the trie may be indirected: see [`Self::insert_hash`].
/// Portions of the trie may be _hashed out_: see [`Self::insert_hash`].
#[derive(Debug, Clone, PartialEq, Eq)]
struct TypedMpt<T> {
inner: HashedPartialTrie,
Expand All @@ -28,6 +28,8 @@ impl<T> TypedMpt<T> {
}
}
/// Insert a node which represents an out-of-band sub-trie.
///
/// See [module documentation](super) for more.
fn insert_hash(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> {
self.inner.insert(key.into_nibbles(), hash)?;
Ok(())
Expand Down Expand Up @@ -201,8 +203,8 @@ impl TransactionTrie {
pub const fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie {
&self.untyped
}
/// Indirect (hash) parts of the trie that aren't in `txn_ixs`.
pub fn mask(&mut self, txn_ixs: impl IntoIterator<Item = usize>) -> anyhow::Result<()> {
/// _Hash out_ parts of the trie that aren't in `txn_ixs`.
pub fn intersect(&mut self, txn_ixs: impl IntoIterator<Item = usize>) -> anyhow::Result<()> {
self.untyped = mpt_trie::trie_subsets::create_trie_subset(
&self.untyped,
txn_ixs
Expand Down Expand Up @@ -246,8 +248,8 @@ impl ReceiptTrie {
pub const fn as_hashed_partial_trie(&self) -> &mpt_trie::partial_trie::HashedPartialTrie {
&self.untyped
}
/// Indirect (hash) parts of the trie that aren't in `txn_ixs`.
pub fn mask(&mut self, txn_ixs: impl IntoIterator<Item = usize>) -> anyhow::Result<()> {
/// _Hash out_ parts of the trie that aren't in `txn_ixs`.
pub fn intersect(&mut self, txn_ixs: impl IntoIterator<Item = usize>) -> anyhow::Result<()> {
self.untyped = mpt_trie::trie_subsets::create_trie_subset(
&self.untyped,
txn_ixs
Expand All @@ -273,7 +275,7 @@ pub trait StateTrie {
fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()>;
fn get_by_address(&self, address: Address) -> Option<AccountRlp>;
fn reporting_remove(&mut self, address: Address) -> anyhow::Result<Option<TrieKey>>;
fn mask(&mut self, address: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()>;
fn intersect(&mut self, address: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()>;
fn iter(&self) -> impl Iterator<Item = (H256, AccountRlp)> + '_;
fn root(&self) -> H256;
}
Expand Down Expand Up @@ -322,7 +324,7 @@ impl StateTrie for StateMpt {
#[expect(deprecated)]
self.insert_by_hashed_address(keccak_hash::keccak(address), account)
}
/// Insert an indirected part of the trie
/// Insert an _hashed out_ part of the trie
fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> {
self.typed.insert_hash(key, hash)
}
Expand All @@ -338,7 +340,7 @@ impl StateTrie for StateMpt {
TrieKey::from_address(address),
)
}
fn mask(&mut self, addresses: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
fn intersect(&mut self, addresses: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
let inner = mpt_trie::trie_subsets::create_trie_subset(
self.typed.as_hashed_partial_trie(),
addresses.into_iter().map(TrieKey::into_nibbles),
Expand Down Expand Up @@ -370,7 +372,7 @@ impl From<StateMpt> for HashedPartialTrie {

pub struct StateSmt {
address2state: BTreeMap<Address, AccountRlp>,
indirected: BTreeMap<TrieKey, H256>,
hashed_out: BTreeMap<TrieKey, H256>,
}

impl StateTrie for StateSmt {
Expand All @@ -382,7 +384,7 @@ impl StateTrie for StateSmt {
Ok(self.address2state.insert(address, account))
}
fn insert_hash_by_key(&mut self, key: TrieKey, hash: H256) -> anyhow::Result<()> {
self.indirected.insert(key, hash);
self.hashed_out.insert(key, hash);
Ok(())
}
fn get_by_address(&self, address: Address) -> Option<AccountRlp> {
Expand All @@ -392,7 +394,7 @@ impl StateTrie for StateSmt {
self.address2state.remove(&address);
Ok(None)
}
fn mask(&mut self, address: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
fn intersect(&mut self, address: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
let _ = address;
Ok(())
}
Expand Down Expand Up @@ -440,8 +442,8 @@ impl StorageTrie {
pub fn as_mut_hashed_partial_trie_unchecked(&mut self) -> &mut HashedPartialTrie {
&mut self.untyped
}
/// Indirect (hash) the parts of the trie that aren't in `paths`.
pub fn mask(&mut self, paths: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
/// _Hash out_ the parts of the trie that aren't in `paths`.
pub fn intersect(&mut self, paths: impl IntoIterator<Item = TrieKey>) -> anyhow::Result<()> {
self.untyped = mpt_trie::trie_subsets::create_trie_subset(
&self.untyped,
paths.into_iter().map(TrieKey::into_nibbles),
Expand Down

0 comments on commit af9251a

Please sign in to comment.