refactor: simplify mpt_trie's API

[0xaatif]

trait PartialTrie abstracts over a fully hydrated trie StandardTrie, and a partial trie HashedPartialTrie:

zk_evm/mpt_trie/src/partial_trie.rs

Lines 19 to 30 in 25a8145

	/// Any node in the trie may be replaced by its [hash](Self::hash) in a
	/// [Node::Hash], and the root hash of the trie will remain unchanged.
	///
	/// ```text
	/// R R'
	/// / \ / \
	/// A B H B
	/// / \ \ \
	/// C D E E
	/// ```
	///
	/// That is, if `H` is `A`'s hash, then the roots of `R` and `R'` are the same.

However, our code only uses HashedPartialTrie - we can simplify our API, which will make moving towards a unified backend for SMT and MPT easier (#275)

Changes

• Remove StandardTrie.
• Remove trait PartialTrie, HashedPartialTrie, and replace them with a top level mpt_trie::Node.
• Arc<Box<Node>> -> Arc<Node>.
• Introduce FrozenNode to handle hash caching, instead of RwLock-ing on HashedPartialTrie.
  It turns out that evm_arithmetization only uses FrozenNode, which is a good thing for me to know :)

[Nashtare]

Thanks, overall looking good to me. Mostly nit comments.

[Nashtare]

Thanks Aatif!
We may want another pair of eyes on the refactor (perhaps @BGluth given the library is his work), but otherwise LGTM.

[BGluth]

Hmm... So I get what you're trying to do, but we're throwing away one important optimization with this PR.

With freeze(), we're not really taking advantage of caching hashes at all in trace_decoder. We are "freezing" the trie right when all the mutations stop, then never accessing the cached values again. Then we do some more inserts (mutate the tries some more), and recalculate the same hashes that we previously hashed but have thrown out.

zk_evm/trace_decoder/src/decoding.rs

Lines 227 to 231 in 8b1a1b4

	let mut curr_block_tries = PartialTrieState {
	state: self.tries.state.clone(),
	storage: self.tries.storage.clone(),
	..Default::default()
	};

zk_evm/trace_decoder/src/decoding.rs

Lines 435 to 451 in 8b1a1b4

	let state_trie = create_minimal_state_partial_trie(
	&curr_block_tries.state,
	nodes_used_by_txn.state_accesses.iter().cloned(),
	delta_application_out
	.additional_state_trie_paths_to_not_hash
	.into_iter(),
	)?
	.freeze();
	let txn_k = Nibbles::from_bytes_be(&rlp::encode(&txn_idx)).unwrap();
	let transactions_trie =
	create_trie_subset_wrapped(&curr_block_tries.txn, once(txn_k), TrieType::Txn)?.freeze();
	let receipts_trie =
	create_trie_subset_wrapped(&curr_block_tries.receipt, once(txn_k), TrieType::Receipt)?
	.freeze();

Also if we ever call hash() at some point, mutate the trie a bit, it will still cache any node hashes that were not below any mutated nodes. Essentially, if we're calling hash() in between trie mutations, the current setup is probably going to be more performant.

I personally would not go with the freeze API change and keep the existing setup (maybe without with Arc<_> if possible). With how freeze() is being used, we're not getting any benefits from hashing.

[0xaatif]

Thanks for your review @BGluth :)

I don't quite follow the claim of optimization loss though - could you help me understand?
On the lines you link, we freeze because that's the type expected by evm_arithmetization (which never makes any mutations), I'm not expecting any optimizations there.

AIUI, HashedPartialTrie is strictly equivalent to FrozenNode.

let mut hpt = HashedPartialTrie::default();
hpt.hash();
hpt.insert(...);
hpt.hash();

let frz = FrozenNode::default();
frz.hash();
let thw = frz.thaw();
thw.insert(...);
thaw.frz().hash();

Contain the exact same number of calls to trie_hash - and the internal nodes don't do any cache-ing themselves

Am I missing something?

[0xaatif]

recalculate the same hashes that we previously hashed but have thrown out.

I see that this causes duplicate work, but I think the previous implementation did this too - Nodes never cached their own hashes.

I feel positive about a Cow<Node> with a cached hash, for example, but that should wait for a future PR, no?

[BGluth]

Contain the exact same number of calls to trie_hash - and the internal nodes don't do any cache-ing themselves

I see that this causes duplicate work, but I think the previous implementation did this too - Nodes never cached their own hashes.

So the number of calls to hash() inside trace_decoder is the same between the two interfaces, but the amount of hashing internally is a lot worse with the freeze setup.

With the old setup, every node actually does cache it's own hash. It's actually not obvious when you glance at the code since the definition of Node clearly does not do any caching. However, the non-obvious thing here is the "nodes" in a HashedPartialTrie are actually HashedPartialTries themselves:

zk_evm/mpt_trie/src/partial_trie.rs

Lines 300 to 307 in 3336055

	/// A partial trie that lazily caches hashes for each node as needed.
	/// If you are doing frequent hashing of node, you probably want to use this
	/// `Trie` variant.
	#[derive(Clone, Debug, Default, Deserialize, Serialize)]
	pub struct HashedPartialTrie {
	pub(crate) node: Node<HashedPartialTrie>,
	pub(crate) hash: Arc<RwLock<Option<H256>>>,
	}

The idea here is technically a single Node is a trie, and the trie type (ie. StandardTrie & HashedPartialTrie) potentially wraps the Node with some additional metadata (ie. a cached hash). So this ends up meaning that anytime we calculate a hash, all of the child nodes also calculate and cache their hash. If a mutation ever changes a node, then only the current node & parents up to the root have their cached hash invalidated (note that I actually accidently reversed it in my response above. "it will still cache any node hashes that were not below any mutated nodes" is not correct.). It only affects the node and the nodes upwards.

This means that two different tries can share a common child. Like consider this:

    B
   / \
  L   E

Where the extension node has a very dense sub-trie beneath it. If we perform an insert that converts the extension node into a branch node:

    B
   / \
  L   B

Then both the extension node and branch node will contain the same child node sub-trie (like in terms of the same piece of memory). If the entire child sub-trie was cached, then they will both have access to the already cached hashes.

Also, I'm fine if you want to remove StandardTrie. However, I also wrote this library with the idea that other users outside of us might be using it, so even if we are not using StandardTrie ourselves, there's a chance that someone else is. I don't think there's a huge performance hit from someone using a HashedPartialTrie vs. a StandardTrie and never calling hash(), so if people want to remove it, I'm ok with that (wdyt @Nashtare?). The support is already there, so I'm leaning towards keeping it personally.

[0xaatif]

With the old setup, every node actually does cache it's own hash. It's actually not obvious when you glance at the code since the definition of Node clearly does not do any caching. However, the non-obvious thing here is the "nodes" in a HashedPartialTrie are actually HashedPartialTries themselves

Ah I totally missed that! I'll refactor accordingly :) thanks for the patient explanation :)

As for StandardTrie, I'm keen to keep our codebase lean, and I think our internals are sufficiently exposed, and that there are enough other libraries out there that a user might find it easy to implement their own

[atanmarko]

I think there are few places to update comments from PartialTrie to Node

[atanmarko]

Non-blocking: Could we skip cloning here and work with references?

[BGluth]

This is a bit late, but yeah tries in mpt_trie are all essentially references, so cloning is very cheap.

[atanmarko]

I have skimmed through the PR, seems fine to me. It does simplify navigating the tries implementation.

[BGluth]

@0xaatif Yeah sounds good! Feel free to remove StandardTrie.