refactor: simplify mpt_trie's API

evm_arithmetization/benches/fibonacci_25m_gas.rs

@@ -25,7 +25,6 @@ use evm_arithmetization::Node;
 use hex_literal::hex;
 use keccak_hash::keccak;
 use mpt_trie::nibbles::Nibbles;
-use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
 use plonky2::field::goldilocks_field::GoldilocksField;
 
 type F = GoldilocksField;
@@ -73,7 +72,7 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
     ];
     let code_hash = keccak(code);
 
-    let empty_trie_root = HashedPartialTrie::from(Node::Empty).hash();
+    let empty_trie_root = Node::from(Node::Empty).hash();
 
     let sender_account_before = AccountRlp {
         nonce: 169.into(),
@@ -93,14 +92,17 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
     state_trie_before.insert(sender_nibbles, rlp::encode(&sender_account_before).to_vec())?;
     state_trie_before.insert(to_nibbles, rlp::encode(&to_account_before).to_vec())?;
 
-    storage_tries.push((sender_state_key, Node::Empty.into()));
-    storage_tries.push((to_state_key, Node::Empty.into()));
+    storage_tries.push((sender_state_key, Node::Empty));
+    storage_tries.push((to_state_key, Node::Empty));
 
     let tries_before = TrieInputs {
-        state_trie: state_trie_before,
-        transactions_trie: Node::Empty.into(),
-        receipts_trie: Node::Empty.into(),
-        storage_tries,
+        state_trie: state_trie_before.freeze(),
+        transactions_trie: Node::Empty.freeze(),
+        receipts_trie: Node::Empty.freeze(),
+        storage_tries: storage_tries
+            .into_iter()
+            .map(|(k, v)| (k, v.freeze()))
+            .collect(),
     };
 
     let gas_used = U256::from(0x17d7840_u32);
@@ -133,7 +135,7 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
     };
     let to_account_after = to_account_before;
 
-    let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
+    let mut expected_state_trie_after = Node::from(Node::Empty);
     expected_state_trie_after
         .insert(sender_nibbles, rlp::encode(&sender_account_after).to_vec())?;
     expected_state_trie_after.insert(to_nibbles, rlp::encode(&to_account_after).to_vec())?;
@@ -159,16 +161,15 @@ fn prepare_setup() -> anyhow::Result<GenerationInputs> {
         bloom: vec![0; 256].into(),
         logs: vec![],
     };
-    let mut receipts_trie = HashedPartialTrie::from(Node::Empty);
+    let mut receipts_trie = Node::from(Node::Empty);
     receipts_trie.insert(
         Nibbles::from_str("0x80").unwrap(),
         rlp::encode(&receipt_0).to_vec(),
     )?;
-    let transactions_trie: HashedPartialTrie = Node::Leaf {
+    let transactions_trie: Node = Node::Leaf {
         nibbles: Nibbles::from_str("0x80").unwrap(),
         value: txn.to_vec(),
-    }
-    .into();
+    };
 
     let trie_roots_after = TrieRoots {
         state_root: expected_state_trie_after.hash(),

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

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

evm_arithmetization/src/cpu/kernel/interpreter.rs

@@ -10,7 +10,6 @@ use std::collections::{BTreeSet, HashMap};
 use anyhow::anyhow;
 use ethereum_types::{BigEndianHash, U256};
 use log::Level;
-use mpt_trie::partial_trie::PartialTrie;
 use plonky2::field::types::Field;
 
 use crate::byte_packing::byte_packing_stark::BytePackingOp;

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

@@ -5,7 +5,6 @@ use ethereum_types::{Address, BigEndianHash, H256, U256};
 use hex_literal::hex;
 use keccak_hash::keccak;
 use mpt_trie::nibbles::Nibbles;
-use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
 use plonky2::field::goldilocks_field::GoldilocksField as F;
 use plonky2::field::types::Field;
 use rand::{thread_rng, Rng};
@@ -62,7 +61,7 @@ fn test_account(code: &[u8]) -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),
         balance: U256::from(2222),
-        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
+        storage_root: Node::from(Node::Empty).hash(),
         code_hash: keccak(code),
     }
 }
@@ -82,7 +81,7 @@ fn prepare_interpreter<F: Field>(
 ) -> Result<()> {
     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: HashedPartialTrie = Default::default();
+    let mut state_trie: Node = Default::default();
     let trie_inputs = Default::default();
 
     initialize_mpts(interpreter, &trie_inputs);
@@ -321,15 +320,15 @@ fn sstore() -> Result<()> {
         ..AccountRlp::default()
     };
 
-    let mut state_trie_before = HashedPartialTrie::from(Node::Empty);
+    let mut state_trie_before = Node::from(Node::Empty);
 
     state_trie_before.insert(addr_nibbles, rlp::encode(&account_before).to_vec())?;
 
     let trie_inputs = TrieInputs {
-        state_trie: state_trie_before.clone(),
-        transactions_trie: Node::Empty.into(),
-        receipts_trie: Node::Empty.into(),
-        storage_tries: vec![(addr_hashed, Node::Empty.into())],
+        state_trie: state_trie_before.clone().freeze(),
+        transactions_trie: Node::Empty.freeze(),
+        receipts_trie: Node::Empty.freeze(),
+        storage_tries: vec![(addr_hashed, Node::Empty.freeze())],
     };
 
     let initial_stack = vec![];
@@ -358,7 +357,7 @@ fn sstore() -> Result<()> {
     let account_after = AccountRlp {
         balance: 0x0de0b6b3a7640000u64.into(),
         code_hash,
-        storage_root: HashedPartialTrie::from(Node::Leaf {
+        storage_root: Node::from(Node::Leaf {
             nibbles: Nibbles::from_h256_be(keccak([0u8; 32])),
             value: vec![2],
         })
@@ -387,7 +386,7 @@ fn sstore() -> Result<()> {
 
     let hash = H256::from_uint(&interpreter.stack()[1]);
 
-    let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
+    let mut expected_state_trie_after = Node::from(Node::Empty);
     expected_state_trie_after.insert(addr_nibbles, rlp::encode(&account_after).to_vec())?;
 
     let expected_state_trie_hash = expected_state_trie_after.hash();
@@ -419,15 +418,15 @@ fn sload() -> Result<()> {
         ..AccountRlp::default()
     };
 
-    let mut state_trie_before = HashedPartialTrie::from(Node::Empty);
+    let mut state_trie_before = Node::from(Node::Empty);
 
     state_trie_before.insert(addr_nibbles, rlp::encode(&account_before).to_vec())?;
 
     let trie_inputs = TrieInputs {
-        state_trie: state_trie_before.clone(),
-        transactions_trie: Node::Empty.into(),
-        receipts_trie: Node::Empty.into(),
-        storage_tries: vec![(addr_hashed, Node::Empty.into())],
+        state_trie: state_trie_before.clone().freeze(),
+        transactions_trie: Node::Empty.freeze(),
+        receipts_trie: Node::Empty.freeze(),
+        storage_tries: vec![(addr_hashed, Node::Empty.freeze())],
     };
 
     let initial_stack = vec![];

evm_arithmetization/src/cpu/kernel/tests/add11.rs

@@ -5,7 +5,7 @@ use ethereum_types::{Address, BigEndianHash, H256};
 use hex_literal::hex;
 use keccak_hash::keccak;
 use mpt_trie::nibbles::Nibbles;
-use mpt_trie::partial_trie::{HashedPartialTrie, Node, PartialTrie};
+use mpt_trie::Node;
 use plonky2::field::goldilocks_field::GoldilocksField as F;
 
 use crate::cpu::kernel::aggregator::KERNEL;
@@ -71,13 +71,16 @@ fn test_add11_yml() {
         .insert(to_nibbles, rlp::encode(&to_account_before).to_vec())
         .unwrap();
 
-    storage_tries.push((to_hashed, Node::Empty.into()));
+    storage_tries.push((to_hashed, Node::Empty));
 
     let tries_before = TrieInputs {
-        state_trie: state_trie_before,
-        transactions_trie: Node::Empty.into(),
-        receipts_trie: Node::Empty.into(),
-        storage_tries,
+        state_trie: state_trie_before.freeze(),
+        transactions_trie: Node::Empty.freeze(),
+        receipts_trie: Node::Empty.freeze(),
+        storage_tries: storage_tries
+            .into_iter()
+            .map(|(k, v)| (k, v.freeze()))
+            .collect(),
     };
 
     let txn = hex!("f863800a83061a8094095e7baea6a6c7c4c2dfeb977efac326af552d87830186a0801ba0ffb600e63115a7362e7811894a91d8ba4330e526f22121c994c4692035dfdfd5a06198379fcac8de3dbfac48b165df4bf88e2088f294b61efb9a65fe2281c76e16");
@@ -111,10 +114,10 @@ fn test_add11_yml() {
             balance: 0xde0b6b3a76586a0u64.into(),
             code_hash,
             // Storage map: { 0 => 2 }
-            storage_root: HashedPartialTrie::from(Node::Leaf {
+            storage_root: Node::Leaf {
                 nibbles: Nibbles::from_h256_be(keccak([0u8; 32])),
                 value: vec![2],
-            })
+            }
             .hash(),
             ..AccountRlp::default()
         };
@@ -127,7 +130,7 @@ fn test_add11_yml() {
         let beacon_roots_account =
             beacon_roots_contract_from_storage(&beacon_roots_account_storage);
 
-        let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
+        let mut expected_state_trie_after = Node::Empty;
         expected_state_trie_after
             .insert(
                 beneficiary_nibbles,
@@ -160,18 +163,17 @@ fn test_add11_yml() {
         bloom: vec![0; 256].into(),
         logs: vec![],
     };
-    let mut receipts_trie = HashedPartialTrie::from(Node::Empty);
+    let mut receipts_trie = Node::Empty;
     receipts_trie
         .insert(
             Nibbles::from_str("0x80").unwrap(),
             rlp::encode(&receipt_0).to_vec(),
         )
         .unwrap();
-    let transactions_trie: HashedPartialTrie = Node::Leaf {
+    let transactions_trie: Node = Node::Leaf {
         nibbles: Nibbles::from_str("0x80").unwrap(),
         value: txn.to_vec(),
-    }
-    .into();
+    };
 
     let trie_roots_after = TrieRoots {
         state_root: expected_state_trie_after.hash(),
@@ -187,7 +189,7 @@ fn test_add11_yml() {
         trie_roots_after,
         contract_code: contract_code.clone(),
         block_metadata,
-        checkpoint_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),
+        checkpoint_state_trie_root: Node::Empty.hash(),
         txn_number_before: 0.into(),
         gas_used_before: 0.into(),
         gas_used_after: gas_used,
@@ -259,13 +261,16 @@ fn test_add11_yml_with_exception() {
         .insert(to_nibbles, rlp::encode(&to_account_before).to_vec())
         .unwrap();
 
-    storage_tries.push((to_hashed, Node::Empty.into()));
+    storage_tries.push((to_hashed, Node::Empty));
 
     let tries_before = TrieInputs {
-        state_trie: state_trie_before,
-        transactions_trie: Node::Empty.into(),
-        receipts_trie: Node::Empty.into(),
-        storage_tries,
+        state_trie: state_trie_before.freeze(),
+        transactions_trie: Node::Empty.freeze(),
+        receipts_trie: Node::Empty.freeze(),
+        storage_tries: storage_tries
+            .into_iter()
+            .map(|(k, v)| (k, v.freeze()))
+            .collect(),
     };
 
     let txn = hex!("f863800a83061a8094095e7baea6a6c7c4c2dfeb977efac326af552d87830186a0801ba0ffb600e63115a7362e7811894a91d8ba4330e526f22121c994c4692035dfdfd5a06198379fcac8de3dbfac48b165df4bf88e2088f294b61efb9a65fe2281c76e16");
@@ -307,7 +312,7 @@ fn test_add11_yml_with_exception() {
         let beacon_roots_account =
             beacon_roots_contract_from_storage(&beacon_roots_account_storage);
 
-        let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
+        let mut expected_state_trie_after = Node::Empty;
         expected_state_trie_after
             .insert(
                 beneficiary_nibbles,
@@ -341,18 +346,17 @@ fn test_add11_yml_with_exception() {
         bloom: vec![0; 256].into(),
         logs: vec![],
     };
-    let mut receipts_trie = HashedPartialTrie::from(Node::Empty);
+    let mut receipts_trie = Node::Empty;
     receipts_trie
         .insert(
             Nibbles::from_str("0x80").unwrap(),
             rlp::encode(&receipt_0).to_vec(),
         )
         .unwrap();
-    let transactions_trie: HashedPartialTrie = Node::Leaf {
+    let transactions_trie: Node = Node::Leaf {
         nibbles: Nibbles::from_str("0x80").unwrap(),
         value: txn.to_vec(),
-    }
-    .into();
+    };
 
     let trie_roots_after = TrieRoots {
         state_root: expected_state_trie_after.hash(),
@@ -368,7 +372,7 @@ fn test_add11_yml_with_exception() {
         trie_roots_after,
         contract_code: contract_code.clone(),
         block_metadata,
-        checkpoint_state_trie_root: HashedPartialTrie::from(Node::Empty).hash(),
+        checkpoint_state_trie_root: Node::Empty.hash(),
         txn_number_before: 0.into(),
         gas_used_before: 0.into(),
         gas_used_after: txn_gas_limit.into(),

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

@@ -1,7 +1,6 @@
 use anyhow::Result;
 use ethereum_types::{Address, BigEndianHash, H256, U256};
 use keccak_hash::keccak;
-use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
 use plonky2::field::goldilocks_field::GoldilocksField as F;
 use plonky2::field::types::Field;
 use rand::{thread_rng, Rng};
@@ -19,7 +18,7 @@ fn test_account(balance: U256) -> AccountRlp {
     AccountRlp {
         nonce: U256::from(1111),
         balance,
-        storage_root: HashedPartialTrie::from(Node::Empty).hash(),
+        storage_root: Node::from(Node::Empty).hash(),
         code_hash: H256::from_uint(&U256::from(8888)),
     }
 }
@@ -33,7 +32,7 @@ fn prepare_interpreter<F: Field>(
 ) -> Result<()> {
     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: HashedPartialTrie = Default::default();
+    let mut state_trie: Node = Default::default();
     let trie_inputs = Default::default();
 
     initialize_mpts(interpreter, &trie_inputs);