Merge branch 'main' into bignum-modexp

evm/Cargo.toml

@@ -13,7 +13,7 @@ edition = "2021"
 anyhow = "1.0.40"
 blake2 = "0.10.5"
 env_logger = "0.10.0"
-eth_trie_utils = "0.4.1"
+eth_trie_utils = "0.5.0"
 ethereum-types = "0.14.0"
 hex = { version = "0.4.3", optional = true }
 hex-literal = "0.3.4"

evm/src/cpu/kernel/constants/trie_type.rs

@@ -1,4 +1,8 @@
-use eth_trie_utils::partial_trie::PartialTrie;
+use std::ops::Deref;
+
+use eth_trie_utils::partial_trie::HashedPartialTrie;
+
+use crate::Node;
 
 #[derive(Copy, Clone, Debug)]
 pub(crate) enum PartialTrieType {
@@ -12,13 +16,13 @@ pub(crate) enum PartialTrieType {
 impl PartialTrieType {
     pub(crate) const COUNT: usize = 5;
 
-    pub(crate) fn of(trie: &PartialTrie) -> Self {
+    pub(crate) fn of(trie: &HashedPartialTrie) -> Self {
-        match trie {
+        match trie.deref() {
-            PartialTrie::Empty => Self::Empty,
+            Node::Empty => Self::Empty,
-            PartialTrie::Hash(_) => Self::Hash,
+            Node::Hash(_) => Self::Hash,
-            PartialTrie::Branch { .. } => Self::Branch,
+            Node::Branch { .. } => Self::Branch,
-            PartialTrie::Extension { .. } => Self::Extension,
+            Node::Extension { .. } => Self::Extension,
-            PartialTrie::Leaf { .. } => Self::Leaf,
+            Node::Leaf { .. } => Self::Leaf,
         }
     }
 

evm/src/cpu/kernel/tests/account_code.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use anyhow::Result;
-use eth_trie_utils::partial_trie::PartialTrie;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{Address, BigEndianHash, H256, U256};
 use keccak_hash::keccak;
 use rand::{thread_rng, Rng};
@@ -12,13 +12,14 @@ use crate::cpu::kernel::interpreter::Interpreter;
 use crate::cpu::kernel::tests::mpt::nibbles_64;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::memory::segments::Segment;
+use crate::Node;
 
 // Test account with a given code hash.
 fn test_account(code: &[u8]) -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),
         balance: U256::from(2222),
-        storage_root: PartialTrie::Empty.calc_hash(),
+        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
         code_hash: keccak(code),
     }
 }
@@ -39,7 +40,7 @@ fn prepare_interpreter(
     let load_all_mpts = KERNEL.global_labels["load_all_mpts"];
     let mpt_insert_state_trie = KERNEL.global_labels["mpt_insert_state_trie"];
     let mpt_hash_state_trie = KERNEL.global_labels["mpt_hash_state_trie"];
-    let mut state_trie: PartialTrie = Default::default();
+    let mut state_trie: HashedPartialTrie = Default::default();
     let trie_inputs = Default::default();
 
     interpreter.generation_state.registers.program_counter = load_all_mpts;
@@ -96,7 +97,7 @@ fn prepare_interpreter(
     let hash = H256::from_uint(&interpreter.stack()[0]);
 
     state_trie.insert(k, rlp::encode(account).to_vec());
-    let expected_state_trie_hash = state_trie.calc_hash();
+    let expected_state_trie_hash = state_trie.hash();
     assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())

evm/src/cpu/kernel/tests/balance.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use eth_trie_utils::partial_trie::PartialTrie;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{Address, BigEndianHash, H256, U256};
 use keccak_hash::keccak;
 use rand::{thread_rng, Rng};
@@ -9,13 +9,14 @@ use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::cpu::kernel::interpreter::Interpreter;
 use crate::cpu::kernel::tests::mpt::nibbles_64;
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
+use crate::Node;
 
 // Test account with a given code hash.
 fn test_account(balance: U256) -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),
         balance,
-        storage_root: PartialTrie::Empty.calc_hash(),
+        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
         code_hash: H256::from_uint(&U256::from(8888)),
     }
 }
@@ -30,7 +31,7 @@ fn prepare_interpreter(
     let load_all_mpts = KERNEL.global_labels["load_all_mpts"];
     let mpt_insert_state_trie = KERNEL.global_labels["mpt_insert_state_trie"];
     let mpt_hash_state_trie = KERNEL.global_labels["mpt_hash_state_trie"];
-    let mut state_trie: PartialTrie = Default::default();
+    let mut state_trie: HashedPartialTrie = Default::default();
     let trie_inputs = Default::default();
 
     interpreter.generation_state.registers.program_counter = load_all_mpts;
@@ -87,7 +88,7 @@ fn prepare_interpreter(
     let hash = H256::from_uint(&interpreter.stack()[0]);
 
     state_trie.insert(k, rlp::encode(account).to_vec());
-    let expected_state_trie_hash = state_trie.calc_hash();
+    let expected_state_trie_hash = state_trie.hash();
     assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())

evm/src/cpu/kernel/tests/mpt/hash.rs

@@ -7,6 +7,7 @@ use crate::cpu::kernel::interpreter::Interpreter;
 use crate::cpu::kernel::tests::mpt::{extension_to_leaf, test_account_1_rlp, test_account_2_rlp};
 use crate::generation::mpt::all_mpt_prover_inputs_reversed;
 use crate::generation::TrieInputs;
+use crate::Node;
 
 // TODO: Test with short leaf. Might need to be a storage trie.
 
@@ -24,11 +25,12 @@ fn mpt_hash_empty() -> Result<()> {
 
 #[test]
 fn mpt_hash_empty_branch() -> Result<()> {
-    let children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let children = core::array::from_fn(|_| Node::Empty.into());
-    let state_trie = PartialTrie::Branch {
+    let state_trie = Node::Branch {
         children,
         value: vec![],
-    };
+    }
+    .into();
     let trie_inputs = TrieInputs {
         state_trie,
         transactions_trie: Default::default(),
@@ -42,7 +44,7 @@ fn mpt_hash_empty_branch() -> Result<()> {
 fn mpt_hash_hash() -> Result<()> {
     let hash = H256::random();
     let trie_inputs = TrieInputs {
-        state_trie: PartialTrie::Hash(hash),
+        state_trie: Node::Hash(hash).into(),
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
         storage_tries: vec![],
@@ -53,10 +55,11 @@ fn mpt_hash_hash() -> Result<()> {
 
 #[test]
 fn mpt_hash_leaf() -> Result<()> {
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: 0xABC_u64.into(),
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     let trie_inputs = TrieInputs {
         state_trie,
         transactions_trie: Default::default(),
@@ -80,17 +83,19 @@ fn mpt_hash_extension_to_leaf() -> Result<()> {
 
 #[test]
 fn mpt_hash_branch_to_leaf() -> Result<()> {
-    let leaf = PartialTrie::Leaf {
+    let leaf = Node::Leaf {
         nibbles: 0xABC_u64.into(),
         value: test_account_2_rlp(),
     }
     .into();
-    let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+
+    let mut children = core::array::from_fn(|_| Node::Empty.into());
     children[3] = leaf;
-    let state_trie = PartialTrie::Branch {
+    let state_trie = Node::Branch {
         children,
         value: vec![],
-    };
+    }
+    .into();
 
     let trie_inputs = TrieInputs {
         state_trie,
@@ -124,7 +129,7 @@ fn test_state_trie(trie_inputs: TrieInputs) -> Result<()> {
         interpreter.stack()
     );
     let hash = H256::from_uint(&interpreter.stack()[0]);
-    let expected_state_trie_hash = trie_inputs.state_trie.calc_hash();
+    let expected_state_trie_hash = trie_inputs.state_trie.hash();
     assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())

evm/src/cpu/kernel/tests/mpt/insert.rs

@@ -1,5 +1,6 @@
 use anyhow::Result;
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{BigEndianHash, H256};
 
 use crate::cpu::kernel::aggregator::KERNEL;
@@ -10,6 +11,7 @@ use crate::cpu::kernel::tests::mpt::{
 };
 use crate::generation::mpt::{all_mpt_prover_inputs_reversed, AccountRlp};
 use crate::generation::TrieInputs;
+use crate::Node;
 
 #[test]
 fn mpt_insert_empty() -> Result<()> {
@@ -19,58 +21,64 @@ fn mpt_insert_empty() -> Result<()> {
 #[test]
 fn mpt_insert_leaf_identical_keys() -> Result<()> {
     let key = nibbles_64(0xABC);
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: key,
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, key, test_account_2())
 }
 
 #[test]
 fn mpt_insert_leaf_nonoverlapping_keys() -> Result<()> {
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: nibbles_64(0xABC),
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0x123), test_account_2())
 }
 
 #[test]
 fn mpt_insert_leaf_overlapping_keys() -> Result<()> {
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: nibbles_64(0xABC),
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0xADE), test_account_2())
 }
 
 #[test]
 #[ignore] // TODO: Not valid for state trie, all keys have same len.
 fn mpt_insert_leaf_insert_key_extends_leaf_key() -> Result<()> {
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: 0xABC_u64.into(),
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0xABCDE), test_account_2())
 }
 
 #[test]
 #[ignore] // TODO: Not valid for state trie, all keys have same len.
 fn mpt_insert_leaf_leaf_key_extends_insert_key() -> Result<()> {
-    let state_trie = PartialTrie::Leaf {
+    let state_trie = Node::Leaf {
         nibbles: 0xABCDE_u64.into(),
         value: test_account_1_rlp(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0xABC), test_account_2())
 }
 
 #[test]
 fn mpt_insert_branch_replacing_empty_child() -> Result<()> {
-    let children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let children = core::array::from_fn(|_| Node::Empty.into());
-    let state_trie = PartialTrie::Branch {
+    let state_trie = Node::Branch {
         children,
         value: vec![],
-    };
+    }
+    .into();
 
     test_state_trie(state_trie, nibbles_64(0xABC), test_account_2())
 }
@@ -81,20 +89,21 @@ fn mpt_insert_branch_replacing_empty_child() -> Result<()> {
 #[ignore]
 fn mpt_insert_extension_nonoverlapping_keys() -> Result<()> {
     // Existing keys are 0xABC, 0xABCDEF; inserted key is 0x12345.
-    let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let mut children = core::array::from_fn(|_| Node::Empty.into());
-    children[0xD] = PartialTrie::Leaf {
+    children[0xD] = Node::Leaf {
         nibbles: 0xEF_u64.into(),
         value: test_account_1_rlp(),
     }
     .into();
-    let state_trie = PartialTrie::Extension {
+    let state_trie = Node::Extension {
         nibbles: 0xABC_u64.into(),
-        child: PartialTrie::Branch {
+        child: Node::Branch {
             children,
             value: test_account_1_rlp(),
         }
         .into(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0x12345), test_account_2())
 }
 
@@ -104,48 +113,54 @@ fn mpt_insert_extension_nonoverlapping_keys() -> Result<()> {
 #[ignore]
 fn mpt_insert_extension_insert_key_extends_node_key() -> Result<()> {
     // Existing keys are 0xA, 0xABCD; inserted key is 0xABCDEF.
-    let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let mut children = core::array::from_fn(|_| Node::Empty.into());
-    children[0xB] = PartialTrie::Leaf {
+    children[0xB] = Node::Leaf {
         nibbles: 0xCD_u64.into(),
         value: test_account_1_rlp(),
     }
     .into();
-    let state_trie = PartialTrie::Extension {
+    let state_trie = Node::Extension {
         nibbles: 0xA_u64.into(),
-        child: PartialTrie::Branch {
+        child: Node::Branch {
             children,
             value: test_account_1_rlp(),
         }
         .into(),
-    };
+    }
+    .into();
     test_state_trie(state_trie, nibbles_64(0xABCDEF), test_account_2())
 }
 
 #[test]
 fn mpt_insert_branch_to_leaf_same_key() -> Result<()> {
-    let leaf = PartialTrie::Leaf {
+    let leaf = Node::Leaf {
         nibbles: nibbles_count(0xBCD, 63),
         value: test_account_1_rlp(),
     }
     .into();
 
-    let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let mut children = core::array::from_fn(|_| Node::Empty.into());
     children[0] = leaf;
-    let state_trie = PartialTrie::Branch {
+    let state_trie = Node::Branch {
         children,
         value: vec![],
-    };
+    }
+    .into();
 
     test_state_trie(state_trie, nibbles_64(0xABCD), test_account_2())
 }
 
 /// Note: The account's storage_root is ignored, as we can't insert a new storage_root without the
 /// accompanying trie data. An empty trie's storage_root is used instead.
-fn test_state_trie(mut state_trie: PartialTrie, k: Nibbles, mut account: AccountRlp) -> Result<()> {
+fn test_state_trie(
+    mut state_trie: HashedPartialTrie,
+    k: Nibbles,
+    mut account: AccountRlp,
+) -> Result<()> {
     assert_eq!(k.count, 64);
 
     // Ignore any storage_root; see documentation note.
-    account.storage_root = PartialTrie::Empty.calc_hash();
+    account.storage_root = HashedPartialTrie::from(Node::Empty).hash();
 
     let trie_inputs = TrieInputs {
         state_trie: state_trie.clone(),
@@ -207,7 +222,7 @@ fn test_state_trie(mut state_trie: PartialTrie, k: Nibbles, mut account: Account
     let hash = H256::from_uint(&interpreter.stack()[0]);
 
     state_trie.insert(k, rlp::encode(&account).to_vec());
-    let expected_state_trie_hash = state_trie.calc_hash();
+    let expected_state_trie_hash = state_trie.hash();
     assert_eq!(hash, expected_state_trie_hash);
 
     Ok(())

evm/src/cpu/kernel/tests/mpt/load.rs

@@ -1,5 +1,4 @@
 use anyhow::Result;
-use eth_trie_utils::partial_trie::PartialTrie;
 use ethereum_types::{BigEndianHash, H256, U256};
 
 use crate::cpu::kernel::aggregator::KERNEL;
@@ -9,6 +8,7 @@ use crate::cpu::kernel::interpreter::Interpreter;
 use crate::cpu::kernel::tests::mpt::{extension_to_leaf, test_account_1, test_account_1_rlp};
 use crate::generation::mpt::all_mpt_prover_inputs_reversed;
 use crate::generation::TrieInputs;
+use crate::Node;
 
 #[test]
 fn load_all_mpts_empty() -> Result<()> {
@@ -48,10 +48,11 @@ fn load_all_mpts_empty() -> Result<()> {
 #[test]
 fn load_all_mpts_leaf() -> Result<()> {
     let trie_inputs = TrieInputs {
-        state_trie: PartialTrie::Leaf {
+        state_trie: Node::Leaf {
             nibbles: 0xABC_u64.into(),
             value: test_account_1_rlp(),
-        },
+        }
+        .into(),
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
         storage_tries: vec![],
@@ -100,7 +101,7 @@ fn load_all_mpts_leaf() -> Result<()> {
 fn load_all_mpts_hash() -> Result<()> {
     let hash = H256::random();
     let trie_inputs = TrieInputs {
-        state_trie: PartialTrie::Hash(hash),
+        state_trie: Node::Hash(hash).into(),
         transactions_trie: Default::default(),
         receipts_trie: Default::default(),
         storage_tries: vec![],
@@ -134,11 +135,12 @@ fn load_all_mpts_hash() -> Result<()> {
 
 #[test]
 fn load_all_mpts_empty_branch() -> Result<()> {
-    let children = core::array::from_fn(|_| PartialTrie::Empty.into());
+    let children = core::array::from_fn(|_| Node::Empty.into());
-    let state_trie = PartialTrie::Branch {
+    let state_trie = Node::Branch {
         children,
         value: vec![],
-    };
+    }
+    .into();
     let trie_inputs = TrieInputs {
         state_trie,
         transactions_trie: Default::default(),

evm/src/cpu/kernel/tests/mpt/mod.rs

@@ -1,7 +1,9 @@
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::HashedPartialTrie;
 use ethereum_types::{BigEndianHash, H256, U256};
 
 use crate::generation::mpt::AccountRlp;
+use crate::Node;
 
 mod hash;
 mod hex_prefix;
@@ -46,10 +48,10 @@ pub(crate) fn test_account_2_rlp() -> Vec<u8> {
 }
 
 /// A `PartialTrie` where an extension node leads to a leaf node containing an account.
-pub(crate) fn extension_to_leaf(value: Vec<u8>) -> PartialTrie {
+pub(crate) fn extension_to_leaf(value: Vec<u8>) -> HashedPartialTrie {
-    PartialTrie::Extension {
+    Node::Extension {
         nibbles: 0xABC_u64.into(),
-        child: PartialTrie::Leaf {
+        child: Node::Leaf {
             nibbles: Nibbles {
                 count: 3,
                 packed: 0xDEF.into(),
@@ -58,4 +60,5 @@ pub(crate) fn extension_to_leaf(value: Vec<u8>) -> PartialTrie {
         }
         .into(),
     }
+    .into()
 }

evm/src/generation/mod.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use eth_trie_utils::partial_trie::PartialTrie;
+use eth_trie_utils::partial_trie::HashedPartialTrie;
 use ethereum_types::{Address, BigEndianHash, H256, U256};
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
@@ -59,19 +59,19 @@ pub struct GenerationInputs {
 pub struct TrieInputs {
     /// A partial version of the state trie prior to these transactions. It should include all nodes
     /// that will be accessed by these transactions.
-    pub state_trie: PartialTrie,
+    pub state_trie: HashedPartialTrie,
 
     /// A partial version of the transaction trie prior to these transactions. It should include all
     /// nodes that will be accessed by these transactions.
-    pub transactions_trie: PartialTrie,
+    pub transactions_trie: HashedPartialTrie,
 
     /// A partial version of the receipt trie prior to these transactions. It should include all nodes
     /// that will be accessed by these transactions.
-    pub receipts_trie: PartialTrie,
+    pub receipts_trie: HashedPartialTrie,
 
     /// A partial version of each storage trie prior to these transactions. It should include all
     /// storage tries, and nodes therein, that will be accessed by these transactions.
-    pub storage_tries: Vec<(Address, PartialTrie)>,
+    pub storage_tries: Vec<(Address, HashedPartialTrie)>,
 }
 
 fn apply_metadata_memops<F: RichField + Extendable<D>, const D: usize>(

evm/src/generation/mpt.rs

@@ -1,12 +1,15 @@
 use std::collections::HashMap;
+use std::ops::Deref;
 
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{BigEndianHash, H256, U256};
 use keccak_hash::keccak;
 use rlp_derive::{RlpDecodable, RlpEncodable};
 
 use crate::cpu::kernel::constants::trie_type::PartialTrieType;
 use crate::generation::TrieInputs;
+use crate::Node;
 
 #[derive(RlpEncodable, RlpDecodable, Debug)]
 pub struct AccountRlp {
@@ -21,7 +24,7 @@ impl Default for AccountRlp {
         Self {
             nonce: U256::zero(),
             balance: U256::zero(),
-            storage_root: PartialTrie::Empty.calc_hash(),
+            storage_root: HashedPartialTrie::from(Node::Empty).hash(),
             code_hash: keccak([]),
         }
     }
@@ -70,17 +73,18 @@ pub(crate) fn all_mpt_prover_inputs(trie_inputs: &TrieInputs) -> Vec<U256> {
 /// is serialized as `(TYPE_LEAF, key, value)`, where key is a `(nibbles, depth)` pair and `value`
 /// is a variable-length structure which depends on which trie we're dealing with.
 pub(crate) fn mpt_prover_inputs<F>(
-    trie: &PartialTrie,
+    trie: &HashedPartialTrie,
     prover_inputs: &mut Vec<U256>,
     parse_value: &F,
 ) where
     F: Fn(&[u8]) -> Vec<U256>,
 {
     prover_inputs.push((PartialTrieType::of(trie) as u32).into());
-    match trie {
+
-        PartialTrie::Empty => {}
+    match trie.deref() {
-        PartialTrie::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())),
+        Node::Empty => {}
-        PartialTrie::Branch { children, value } => {
+        Node::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())),
+        Node::Branch { children, value } => {
             if value.is_empty() {
                 prover_inputs.push(U256::zero()); // value_present = 0
             } else {
@@ -92,12 +96,12 @@ pub(crate) fn mpt_prover_inputs<F>(
                 mpt_prover_inputs(child, prover_inputs, parse_value);
             }
         }
-        PartialTrie::Extension { nibbles, child } => {
+        Node::Extension { nibbles, child } => {
             prover_inputs.push(nibbles.count.into());
             prover_inputs.push(nibbles.packed);
             mpt_prover_inputs(child, prover_inputs, parse_value);
         }
-        PartialTrie::Leaf { nibbles, value } => {
+        Node::Leaf { nibbles, value } => {
             prover_inputs.push(nibbles.count.into());
             prover_inputs.push(nibbles.packed);
             let leaf = parse_value(value);
@@ -109,16 +113,16 @@ pub(crate) fn mpt_prover_inputs<F>(
 /// Like `mpt_prover_inputs`, but for the state trie, which is a bit unique since each value
 /// leads to a storage trie which we recursively traverse.
 pub(crate) fn mpt_prover_inputs_state_trie(
-    trie: &PartialTrie,
+    trie: &HashedPartialTrie,
     key: Nibbles,
     prover_inputs: &mut Vec<U256>,
-    storage_tries_by_state_key: &HashMap<Nibbles, &PartialTrie>,
+    storage_tries_by_state_key: &HashMap<Nibbles, &HashedPartialTrie>,
 ) {
     prover_inputs.push((PartialTrieType::of(trie) as u32).into());
-    match trie {
+    match trie.deref() {
-        PartialTrie::Empty => {}
+        Node::Empty => {}
-        PartialTrie::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())),
+        Node::Hash(h) => prover_inputs.push(U256::from_big_endian(h.as_bytes())),
-        PartialTrie::Branch { children, value } => {
+        Node::Branch { children, value } => {
             assert!(value.is_empty(), "State trie should not have branch values");
             prover_inputs.push(U256::zero()); // value_present = 0
 
@@ -135,7 +139,7 @@ pub(crate) fn mpt_prover_inputs_state_trie(
                 );
             }
         }
-        PartialTrie::Extension { nibbles, child } => {
+        Node::Extension { nibbles, child } => {
             prover_inputs.push(nibbles.count.into());
             prover_inputs.push(nibbles.packed);
             let extended_key = key.merge_nibbles(nibbles);
@@ -146,7 +150,7 @@ pub(crate) fn mpt_prover_inputs_state_trie(
                 storage_tries_by_state_key,
             );
         }
-        PartialTrie::Leaf { nibbles, value } => {
+        Node::Leaf { nibbles, value } => {
             let account: AccountRlp = rlp::decode(value).expect("Decoding failed");
             let AccountRlp {
                 nonce,
@@ -155,14 +159,14 @@ pub(crate) fn mpt_prover_inputs_state_trie(
                 code_hash,
             } = account;
 
-            let storage_hash_only = PartialTrie::Hash(storage_root);
+            let storage_hash_only = HashedPartialTrie::new(Node::Hash(storage_root));
             let merged_key = key.merge_nibbles(nibbles);
-            let storage_trie: &PartialTrie = storage_tries_by_state_key
+            let storage_trie: &HashedPartialTrie = storage_tries_by_state_key
                 .get(&merged_key)
                 .copied()
                 .unwrap_or(&storage_hash_only);
 
-            assert_eq!(storage_trie.calc_hash(), storage_root,
+            assert_eq!(storage_trie.hash(), storage_root,
                        "In TrieInputs, an account's storage_root didn't match the associated storage trie hash");
 
             prover_inputs.push(nibbles.count.into());

evm/src/generation/trie_extractor.rs

@@ -2,7 +2,7 @@
 
 use std::collections::HashMap;
 
-use eth_trie_utils::partial_trie::Nibbles;
+use eth_trie_utils::nibbles::Nibbles;
 use ethereum_types::{BigEndianHash, H256, U256};
 
 use crate::cpu::kernel::constants::trie_type::PartialTrieType;

evm/src/lib.rs

@@ -34,6 +34,7 @@ pub mod vars;
 pub mod verifier;
 pub mod witness;
 
+use eth_trie_utils::partial_trie::HashedPartialTrie;
 // Set up Jemalloc
 #[cfg(not(target_env = "msvc"))]
 use jemallocator::Jemalloc;
@@ -41,3 +42,5 @@ use jemallocator::Jemalloc;
 #[cfg(not(target_env = "msvc"))]
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
+
+pub type Node = eth_trie_utils::partial_trie::Node<HashedPartialTrie>;

evm/tests/add11_yml.rs

@@ -2,7 +2,8 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::Address;
 use hex_literal::hex;
 use keccak_hash::keccak;
@@ -16,6 +17,7 @@ use plonky2_evm::generation::{GenerationInputs, TrieInputs};
 use plonky2_evm::proof::BlockMetadata;
 use plonky2_evm::prover::prove;
 use plonky2_evm::verifier::verify_proof;
+use plonky2_evm::Node;
 
 type F = GoldilocksField;
 const D: usize = 2;
@@ -58,7 +60,7 @@ fn add11_yml() -> anyhow::Result<()> {
         ..AccountRlp::default()
     };
 
-    let mut state_trie_before = PartialTrie::Empty;
+    let mut state_trie_before = HashedPartialTrie::from(Node::Empty);
     state_trie_before.insert(
         beneficiary_nibbles,
         rlp::encode(&beneficiary_account_before).to_vec(),
@@ -68,9 +70,9 @@ fn add11_yml() -> anyhow::Result<()> {
 
     let tries_before = TrieInputs {
         state_trie: state_trie_before,
-        transactions_trie: PartialTrie::Empty,
+        transactions_trie: Node::Empty.into(),
-        receipts_trie: PartialTrie::Empty,
+        receipts_trie: Node::Empty.into(),
-        storage_tries: vec![(Address::from_slice(&to), PartialTrie::Empty)],
+        storage_tries: vec![(Address::from_slice(&to), Node::Empty.into())],
     };
 
     let txn = hex!("f863800a83061a8094095e7baea6a6c7c4c2dfeb977efac326af552d87830186a0801ba0ffb600e63115a7362e7811894a91d8ba4330e526f22121c994c4692035dfdfd5a06198379fcac8de3dbfac48b165df4bf88e2088f294b61efb9a65fe2281c76e16");
@@ -110,15 +112,15 @@ fn add11_yml() -> anyhow::Result<()> {
         balance: 0xde0b6b3a76586a0u64.into(),
         code_hash,
         // Storage map: { 0 => 2 }
-        storage_root: PartialTrie::Leaf {
+        storage_root: HashedPartialTrie::from(Node::Leaf {
             nibbles: Nibbles::from_h256_be(keccak([0u8; 32])),
             value: vec![2],
-        }
+        })
-        .calc_hash(),
+        .hash(),
         ..AccountRlp::default()
     };
 
-    let mut expected_state_trie_after = PartialTrie::Empty;
+    let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
     expected_state_trie_after.insert(
         beneficiary_nibbles,
         rlp::encode(&beneficiary_account_after).to_vec(),
@@ -128,7 +130,7 @@ fn add11_yml() -> anyhow::Result<()> {
 
     assert_eq!(
         proof.public_values.trie_roots_after.state_root,
-        expected_state_trie_after.calc_hash()
+        expected_state_trie_after.hash()
     );
 
     verify_proof(&all_stark, proof, &config)

evm/tests/basic_smart_contract.rs

@@ -2,7 +2,8 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{Address, U256};
 use hex_literal::hex;
 use keccak_hash::keccak;
@@ -17,6 +18,7 @@ use plonky2_evm::generation::{GenerationInputs, TrieInputs};
 use plonky2_evm::proof::BlockMetadata;
 use plonky2_evm::prover::prove;
 use plonky2_evm::verifier::verify_proof;
+use plonky2_evm::Node;
 
 type F = GoldilocksField;
 const D: usize = 2;
@@ -61,27 +63,28 @@ fn test_basic_smart_contract() -> anyhow::Result<()> {
     };
 
     let state_trie_before = {
-        let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+        let mut children = core::array::from_fn(|_| Node::Empty.into());
-        children[sender_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[sender_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: sender_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&sender_account_before).to_vec(),
         }
         .into();
-        children[to_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[to_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: to_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&to_account_before).to_vec(),
         }
         .into();
-        PartialTrie::Branch {
+        Node::Branch {
             children,
             value: vec![],
         }
-    };
+    }
+    .into();
 
     let tries_before = TrieInputs {
         state_trie: state_trie_before,
-        transactions_trie: PartialTrie::Empty,
+        transactions_trie: Node::Empty.into(),
-        receipts_trie: PartialTrie::Empty,
+        receipts_trie: Node::Empty.into(),
         storage_tries: vec![],
     };
 
@@ -110,7 +113,7 @@ fn test_basic_smart_contract() -> anyhow::Result<()> {
     let proof = prove::<F, C, D>(&all_stark, &config, inputs, &mut timing)?;
     timing.filter(Duration::from_millis(100)).print();
 
-    let expected_state_trie_after = {
+    let expected_state_trie_after: HashedPartialTrie = {
         let txdata_gas = 2 * 16;
         let gas_used = 21_000 + code_gas + txdata_gas;
 
@@ -128,31 +131,32 @@ fn test_basic_smart_contract() -> anyhow::Result<()> {
             ..to_account_before
         };
 
-        let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+        let mut children = core::array::from_fn(|_| Node::Empty.into());
-        children[beneficiary_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[beneficiary_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: beneficiary_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&beneficiary_account_after).to_vec(),
         }
         .into();
-        children[sender_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[sender_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: sender_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&sender_account_after).to_vec(),
         }
         .into();
-        children[to_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[to_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: to_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&to_account_after).to_vec(),
         }
         .into();
-        PartialTrie::Branch {
+        Node::Branch {
             children,
             value: vec![],
         }
-    };
+    }
+    .into();
 
     assert_eq!(
         proof.public_values.trie_roots_after.state_root,
-        expected_state_trie_after.calc_hash()
+        expected_state_trie_after.hash()
     );
 
     verify_proof(&all_stark, proof, &config)

evm/tests/empty_txn_list.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
-use eth_trie_utils::partial_trie::PartialTrie;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use keccak_hash::keccak;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::plonk::config::PoseidonGoldilocksConfig;
@@ -14,6 +14,7 @@ use plonky2_evm::generation::{GenerationInputs, TrieInputs};
 use plonky2_evm::proof::BlockMetadata;
 use plonky2_evm::prover::prove;
 use plonky2_evm::verifier::verify_proof;
+use plonky2_evm::Node;
 
 type F = GoldilocksField;
 const D: usize = 2;
@@ -30,14 +31,14 @@ fn test_empty_txn_list() -> anyhow::Result<()> {
 
     let block_metadata = BlockMetadata::default();
 
-    let state_trie = PartialTrie::Empty;
+    let state_trie = HashedPartialTrie::from(Node::Empty);
-    let transactions_trie = PartialTrie::Empty;
+    let transactions_trie = HashedPartialTrie::from(Node::Empty);
-    let receipts_trie = PartialTrie::Empty;
+    let receipts_trie = HashedPartialTrie::from(Node::Empty);
     let storage_tries = vec![];
 
-    let state_trie_root = state_trie.calc_hash();
+    let state_trie_root = state_trie.hash();
-    let txns_trie_root = transactions_trie.calc_hash();
+    let txns_trie_root = transactions_trie.hash();
-    let receipts_trie_root = receipts_trie.calc_hash();
+    let receipts_trie_root = receipts_trie.hash();
 
     let mut contract_code = HashMap::new();
     contract_code.insert(keccak(vec![]), vec![]);

evm/tests/simple_transfer.rs

@@ -2,7 +2,8 @@ use std::collections::HashMap;
 use std::time::Duration;
 
 use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
-use eth_trie_utils::partial_trie::{Nibbles, PartialTrie};
+use eth_trie_utils::nibbles::Nibbles;
+use eth_trie_utils::partial_trie::{HashedPartialTrie, PartialTrie};
 use ethereum_types::{Address, U256};
 use hex_literal::hex;
 use keccak_hash::keccak;
@@ -16,6 +17,7 @@ use plonky2_evm::generation::{GenerationInputs, TrieInputs};
 use plonky2_evm::proof::BlockMetadata;
 use plonky2_evm::prover::prove;
 use plonky2_evm::verifier::verify_proof;
+use plonky2_evm::Node;
 
 type F = GoldilocksField;
 const D: usize = 2;
@@ -45,19 +47,20 @@ fn test_simple_transfer() -> anyhow::Result<()> {
     let sender_account_before = AccountRlp {
         nonce: 5.into(),
         balance: eth_to_wei(100_000.into()),
-        storage_root: PartialTrie::Empty.calc_hash(),
+        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
         code_hash: keccak([]),
     };
     let to_account_before = AccountRlp::default();
 
-    let state_trie_before = PartialTrie::Leaf {
+    let state_trie_before = Node::Leaf {
         nibbles: sender_nibbles,
         value: rlp::encode(&sender_account_before).to_vec(),
-    };
+    }
+    .into();
     let tries_before = TrieInputs {
         state_trie: state_trie_before,
-        transactions_trie: PartialTrie::Empty,
+        transactions_trie: HashedPartialTrie::from(Node::Empty),
-        receipts_trie: PartialTrie::Empty,
+        receipts_trie: HashedPartialTrie::from(Node::Empty),
         storage_tries: vec![],
     };
 
@@ -85,7 +88,7 @@ fn test_simple_transfer() -> anyhow::Result<()> {
     let proof = prove::<F, C, D>(&all_stark, &config, inputs, &mut timing)?;
     timing.filter(Duration::from_millis(100)).print();
 
-    let expected_state_trie_after = {
+    let expected_state_trie_after: HashedPartialTrie = {
         let txdata_gas = 2 * 16;
         let gas_used = 21_000 + txdata_gas;
 
@@ -103,31 +106,32 @@ fn test_simple_transfer() -> anyhow::Result<()> {
             ..to_account_before
         };
 
-        let mut children = core::array::from_fn(|_| PartialTrie::Empty.into());
+        let mut children = core::array::from_fn(|_| Node::Empty.into());
-        children[beneficiary_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[beneficiary_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: beneficiary_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&beneficiary_account_after).to_vec(),
         }
         .into();
-        children[sender_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[sender_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: sender_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&sender_account_after).to_vec(),
         }
         .into();
-        children[to_nibbles.get_nibble(0) as usize] = PartialTrie::Leaf {
+        children[to_nibbles.get_nibble(0) as usize] = Node::Leaf {
             nibbles: to_nibbles.truncate_n_nibbles_front(1),
             value: rlp::encode(&to_account_after).to_vec(),
         }
         .into();
-        PartialTrie::Branch {
+        Node::Branch {
             children,
             value: vec![],
         }
+        .into()
     };
 
     assert_eq!(
         proof.public_values.trie_roots_after.state_root,
-        expected_state_trie_after.calc_hash()
+        expected_state_trie_after.hash()
     );
 
     verify_proof(&all_stark, proof, &config)