Merge pull request #693 from mir-protocol/partial_trie

Have witness generation take a partial trie instead of Merkle proofs

evm/src/generation/mod.rs

@@ -1,4 +1,4 @@
-use ethereum_types::U256;
+use ethereum_types::Address;
 use plonky2::field::extension::Extendable;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::field::types::Field;
@@ -7,25 +7,33 @@ use plonky2::hash::hash_types::RichField;
 use crate::all_stark::AllStark;
 use crate::cpu::bootstrap_kernel::generate_bootstrap_kernel;
 use crate::cpu::columns::NUM_CPU_COLUMNS;
+use crate::generation::partial_trie::PartialTrie;
 use crate::generation::state::GenerationState;
 use crate::util::trace_rows_to_poly_values;
 
 pub(crate) mod memory;
+pub mod partial_trie;
 pub(crate) mod state;
 
-/// A piece of data which has been encoded using Recursive Length Prefix (RLP) serialization.
-/// See https://ethereum.org/en/developers/docs/data-structures-and-encoding/rlp/
-pub type RlpBlob = Vec<u8>;
-
-/// Merkle proofs are encoded using an RLP blob for each node in the path.
-pub type RlpMerkleProof = Vec<RlpBlob>;
-
 #[allow(unused)] // TODO: Should be used soon.
 pub struct TransactionData {
     pub signed_txn: Vec<u8>,
 
-    /// A Merkle proof for each interaction with the state trie, ordered chronologically.
-    pub trie_proofs: Vec<RlpMerkleProof>,
+    /// A partial version of the state trie prior to this transaction. It should include all nodes
+    /// that will be accessed by this transaction.
+    pub state_trie: PartialTrie,
+
+    /// A partial version of the transaction trie prior to this transaction. It should include all
+    /// nodes that will be accessed by this transaction.
+    pub transaction_trie: PartialTrie,
+
+    /// A partial version of the receipt trie prior to this transaction. It should include all nodes
+    /// that will be accessed by this transaction.
+    pub receipt_trie: PartialTrie,
+
+    /// A partial version of each storage trie prior to this transaction. It should include all
+    /// storage tries, and nodes therein, that will be accessed by this transaction.
+    pub storage_tries: Vec<(Address, PartialTrie)>,
 }
 
 #[allow(unused)] // TODO: Should be used soon.
@@ -47,11 +55,9 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         memory,
         keccak_inputs,
         logic_ops,
-        prover_inputs,
         ..
     } = state;
     assert_eq!(current_cpu_row, [F::ZERO; NUM_CPU_COLUMNS].into());
-    assert_eq!(prover_inputs, vec![], "Not all prover inputs were consumed");
 
     let cpu_trace = trace_rows_to_poly_values(cpu_rows);
     let keccak_trace = all_stark.keccak_stark.generate_trace(keccak_inputs);
@@ -60,14 +66,6 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     vec![cpu_trace, keccak_trace, logic_trace, memory_trace]
 }
 
-fn generate_txn<F: Field>(state: &mut GenerationState<F>, txn: &TransactionData) {
-    // TODO: Add transaction RLP to prover_input.
-
-    // Supply Merkle trie proofs as prover inputs.
-    for proof in &txn.trie_proofs {
-        let proof = proof
-            .iter()
-            .flat_map(|node_rlp| node_rlp.iter().map(|byte| U256::from(*byte)));
-        state.prover_inputs.extend(proof);
-    }
+fn generate_txn<F: Field>(_state: &mut GenerationState<F>, _txn: &TransactionData) {
+    // TODO
 }

evm/src/generation/partial_trie.rs

@@ -0,0 +1,32 @@
+use ethereum_types::U256;
+
+/// A partial trie, or a sub-trie thereof. This mimics the structure of an Ethereum trie, except
+/// with an additional `Hash` node type, representing a node whose data is not needed to process
+/// our transaction.
+pub enum PartialTrie {
+    /// An empty trie.
+    Empty,
+    /// The digest of trie whose data does not need to be stored.
+    Hash(U256),
+    /// A branch node, which consists of 16 children and an optional value.
+    Branch {
+        children: [Box<PartialTrie>; 16],
+        value: Option<U256>,
+    },
+    /// An extension node, which consists of a list of nibbles and a single child.
+    Extension {
+        nibbles: Nibbles,
+        child: Box<PartialTrie>,
+    },
+    /// A leaf node, which consists of a list of nibbles and a value.
+    Leaf { nibbles: Nibbles, value: Vec<u8> },
+}
+
+/// A sequence of nibbles.
+pub struct Nibbles {
+    /// The number of nibbles in this sequence.
+    pub count: usize,
+    /// A packed encoding of these nibbles. Only the first (least significant) `4 * count` bits are
+    /// used. The rest are unused and should be zero.
+    pub packed: U256,
+}

evm/src/generation/state.rs

@@ -19,9 +19,6 @@ pub(crate) struct GenerationState<F: Field> {
 
     pub(crate) keccak_inputs: Vec<[u64; keccak::keccak_stark::NUM_INPUTS]>,
     pub(crate) logic_ops: Vec<logic::Operation>,
-
-    /// Non-deterministic inputs provided by the prover.
-    pub(crate) prover_inputs: Vec<U256>,
 }
 
 impl<F: Field> GenerationState<F> {
@@ -116,7 +113,6 @@ impl<F: Field> Default for GenerationState<F> {
             memory: MemoryState::default(),
             keccak_inputs: vec![],
             logic_ops: vec![],
-            prover_inputs: vec![],
         }
     }
 }

evm/tests/transfer_to_new_addr.rs

@@ -4,6 +4,7 @@ use plonky2::plonk::config::PoseidonGoldilocksConfig;
 use plonky2::util::timing::TimingTree;
 use plonky2_evm::all_stark::AllStark;
 use plonky2_evm::config::StarkConfig;
+use plonky2_evm::generation::partial_trie::PartialTrie;
 use plonky2_evm::generation::{generate_traces, TransactionData};
 use plonky2_evm::prover::prove;
 use plonky2_evm::verifier::verify_proof;
@@ -20,15 +21,11 @@ fn test_simple_transfer() -> anyhow::Result<()> {
 
     let txn = TransactionData {
         signed_txn: hex!("f85f050a82520894000000000000000000000000000000000000000064801ca0fa56df5d988638fad8798e5ef75a1e1125dc7fb55d2ac4bce25776a63f0c2967a02cb47a5579eb5f83a1cabe4662501c0059f1b58e60ef839a1b0da67af6b9fb38").to_vec(),
-        trie_proofs: vec![
-            vec![
-                hex!("f874a1202f93d0dfb1562c03c825a33eec4438e468c17fff649ae844c004065985ae2945b850f84e058a152d02c7e14af6800000a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").to_vec(),
-            ],
-            vec![
-                hex!("f8518080a0d36b8b6b60021940d5553689fb33e5d45e649dd8f4f211d26566238a83169da58080a0c62aa627943b70321f89a8b2fea274ecd47116e62042077dcdc0bdca7c1f66738080808080808080808080").to_vec(),
-                hex!("f873a03f93d0dfb1562c03c825a33eec4438e468c17fff649ae844c004065985ae2945b850f84e068a152d02c7e14af67ccb4ca056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0c5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470").to_vec(),
-            ],
-        ]
+        // TODO: Add trie with sender account.
+        state_trie: PartialTrie::Empty,
+        transaction_trie: PartialTrie::Empty,
+        receipt_trie: PartialTrie::Empty,
+        storage_tries: vec![],
     };
 
     let traces = generate_traces(&all_stark, &[txn]);