From 916f369445aa3eaed009a648f8d7d4a66252b2cf Mon Sep 17 00:00:00 2001 From: KimiWu Date: Mon, 29 Apr 2024 11:42:27 +0800 Subject: [PATCH 1/2] feat: tx related declaration --- geth-utils/src/mpt.rs | 1 + zkevm-circuits/src/mpt_circuit.rs | 34 ++++++++++++- .../src/mpt_circuit/account_leaf.rs | 8 ++- zkevm-circuits/src/mpt_circuit/lib.rs | 1 + .../src/mpt_circuit/storage_leaf.rs | 8 ++- zkevm-circuits/src/mpt_circuit/witness_row.rs | 51 +++++++++++++++++++ zkevm-circuits/src/table/mpt_table.rs | 9 ++++ zkevm-circuits/src/witness/mpt.rs | 22 +++++++- 8 files changed, 127 insertions(+), 7 deletions(-) diff --git a/geth-utils/src/mpt.rs b/geth-utils/src/mpt.rs index 6f212f2f54..d589c3f7e5 100644 --- a/geth-utils/src/mpt.rs +++ b/geth-utils/src/mpt.rs @@ -19,6 +19,7 @@ pub enum ProofType { StorageChanged = 6, StorageDoesNotExist = 7, AccountCreate = 8, + TransactionUpdate = 9, } #[derive(Default, Debug, Clone)] diff --git a/zkevm-circuits/src/mpt_circuit.rs b/zkevm-circuits/src/mpt_circuit.rs index 5a2c79b02b..2ea37b68e7 100644 --- a/zkevm-circuits/src/mpt_circuit.rs +++ b/zkevm-circuits/src/mpt_circuit.rs @@ -23,6 +23,7 @@ mod param; mod rlp_gadgets; mod start; mod storage_leaf; +mod transaction_leaf; /// MPT witness row pub mod witness_row; @@ -31,6 +32,7 @@ use self::{ helpers::RLPItemView, param::RLP_UNIT_NUM_BYTES, rlp_gadgets::decode_rlp, + transaction_leaf::TransactionLeafConfig, witness_row::{ AccountRowType, ExtensionBranchRowType, Node, StartRowType, StorageRowType, NODE_RLP_TYPES_ACCOUNT, NODE_RLP_TYPES_BRANCH, NODE_RLP_TYPES_START, @@ -48,6 +50,7 @@ use crate::{ helpers::{MPTConstraintBuilder, MainRLPGadget, MptCellType, MptTableType}, start::StartConfig, storage_leaf::StorageLeafConfig, + witness_row::{TransactionRowType, NODE_RLP_TYPES_TRANSACTION}, }, table::{KeccakTable, MPTProofType, MptTable}, util::Challenges, @@ -64,6 +67,7 @@ pub(crate) enum MPTRegion { Branch, Account, Storage, + Transaction, Count, } @@ -74,11 +78,13 @@ pub struct StateMachineConfig { is_branch: Column, is_account: Column, is_storage: Column, + is_transaction: Column, start_config: StartConfig, branch_config: ExtensionBranchConfig, storage_config: StorageLeafConfig, account_config: AccountLeafConfig, + transaction_config: TransactionLeafConfig, } impl StateMachineConfig { @@ -89,10 +95,12 @@ impl StateMachineConfig { is_branch: meta.advice_column(), is_account: meta.advice_column(), is_storage: meta.advice_column(), + is_transaction: meta.advice_column(), start_config: StartConfig::default(), branch_config: ExtensionBranchConfig::default(), storage_config: StorageLeafConfig::default(), account_config: AccountLeafConfig::default(), + transaction_config: TransactionLeafConfig::default(), } } @@ -103,6 +111,7 @@ impl StateMachineConfig { self.is_branch, self.is_account, self.is_storage, + self.is_transaction, ] } @@ -339,10 +348,17 @@ impl MPTConfig { ctx.memory.build_constraints(&mut cb.base, f!(q_first)); cb.base.pop_region(); }, + a!(state_machine.is_transaction) => { + state_machine.step_constraints(meta, &mut cb, TransactionRowType::Count as usize); + cb.base.push_region(MPTRegion::Transaction as usize, TransactionRowType::Count as usize); + state_machine.transaction_config = TransactionLeafConfig::configure(meta, &mut cb, &mut ctx); + ctx.memory.build_constraints(&mut cb.base, f!(q_first)); + cb.base.pop_region(); + }, _ => ctx.memory.build_constraints(&mut cb.base, f!(q_first)), )}; // Only account and storage rows can have lookups, disable lookups on all other rows - ifx! {not!(a!(state_machine.is_account) + a!(state_machine.is_storage)) => { + ifx! {not!(a!(state_machine.is_account) + a!(state_machine.is_storage) + a!(state_machine.is_transaction)) => { require!(a!(ctx.mpt_table.proof_type) => MPTProofType::Disabled.expr()); }} }} @@ -413,6 +429,8 @@ impl MPTConfig { NODE_RLP_TYPES_ACCOUNT.to_vec() } else if node.storage.is_some() { NODE_RLP_TYPES_STORAGE.to_vec() + } else if node.transaction.is_some(){ + NODE_RLP_TYPES_TRANSACTION.to_vec() } else { unreachable!() }; @@ -486,6 +504,20 @@ impl MPTConfig { )?; cached_region.pop_region(); } + else if node.transaction.is_some() { + //println!("{}: transaction", offset); + cached_region.push_region(offset, MPTRegion::Transaction as usize); + assign!(cached_region, (self.state_machine.is_transaction, offset) => "is_transaction", true.scalar())?; + self.state_machine.transaction_config.assign( + &mut cached_region, + self, + &mut memory, + offset, + node, + &rlp_values, + )?; + cached_region.pop_region(); + } offset += node.values.len(); diff --git a/zkevm-circuits/src/mpt_circuit/account_leaf.rs b/zkevm-circuits/src/mpt_circuit/account_leaf.rs index 35d27bc4c6..5a7621337c 100644 --- a/zkevm-circuits/src/mpt_circuit/account_leaf.rs +++ b/zkevm-circuits/src/mpt_circuit/account_leaf.rs @@ -466,6 +466,7 @@ impl AccountLeafConfig { meta, &mut cb.base, address.clone(), + false.expr(), proof_type.clone(), WordLoHi::zero(), config.main_data.new_root.expr(), @@ -480,6 +481,7 @@ impl AccountLeafConfig { meta, &mut cb.base, address.clone(), + false.expr(), proof_type.clone(), WordLoHi::zero(), config.main_data.new_root.expr(), @@ -490,14 +492,15 @@ impl AccountLeafConfig { }}; } elsex { // When the value is set to 0, the leaf is deleted, and if there were only two leaves in the branch, - // the neighbour leaf moves one level up and replaces the branch. When the lookup is executed with + // the neighbor leaf moves one level up and replaces the branch. When the lookup is executed with // the new value set to 0, the lookup fails (without the code below), because the leaf that is returned - // is the neighbour node that moved up (because the branch and the old leaf doesn’t exist anymore), + // is the neighbor node that moved up (because the branch and the old leaf doesn’t exist anymore), // but this leaf doesn’t have the zero value. ctx.mpt_table.constrain( meta, &mut cb.base, address, + false.expr(), proof_type, WordLoHi::zero(), config.main_data.new_root.expr(), @@ -759,6 +762,7 @@ impl AccountLeafConfig { &account.address.iter().cloned().rev().collect::>(), )), storage_key: WordLoHi::zero().into_value(), + transaction_index: Value::known(F::ZERO), proof_type: Value::known(proof_type.scalar()), new_root: main_data.new_root.into_value(), old_root: main_data.old_root.into_value(), diff --git a/zkevm-circuits/src/mpt_circuit/lib.rs b/zkevm-circuits/src/mpt_circuit/lib.rs index d06951aaca..8844417595 100644 --- a/zkevm-circuits/src/mpt_circuit/lib.rs +++ b/zkevm-circuits/src/mpt_circuit/lib.rs @@ -8,4 +8,5 @@ pub mod mod_extension; pub mod mpt; pub mod param; pub mod storage_leaf; +pub mod transaction_leaf; pub mod witness_row; diff --git a/zkevm-circuits/src/mpt_circuit/storage_leaf.rs b/zkevm-circuits/src/mpt_circuit/storage_leaf.rs index 6a698a8195..e66f046f1f 100644 --- a/zkevm-circuits/src/mpt_circuit/storage_leaf.rs +++ b/zkevm-circuits/src/mpt_circuit/storage_leaf.rs @@ -359,6 +359,7 @@ impl StorageLeafConfig { meta, &mut cb.base, config.main_data.address.expr(), + false.expr(), proof_type.clone(), address_item.word(), config.main_data.new_root.expr(), @@ -373,6 +374,7 @@ impl StorageLeafConfig { meta, &mut cb.base, config.main_data.address.expr(), + false.expr(), proof_type.clone(), address_item.word(), config.main_data.new_root.expr(), @@ -383,14 +385,15 @@ impl StorageLeafConfig { }}; } elsex { // When the value is set to 0, the leaf is deleted, and if there were only two leaves in the branch, - // the neighbour leaf moves one level up and replaces the branch. When the lookup is executed with + // the neighbor leaf moves one level up and replaces the branch. When the lookup is executed with // the new value set to 0, the lookup fails (without the code below), because the leaf that is returned - // is the neighbour node that moved up (because the branch and the old leaf doesn’t exist anymore), + // is the neighbor node that moved up (because the branch and the old leaf doesn’t exist anymore), // but this leaf doesn’t have the zero value. ctx.mpt_table.constrain( meta, &mut cb.base, config.main_data.address.expr(), + false.expr(), proof_type, address_item.word(), config.main_data.new_root.expr(), @@ -607,6 +610,7 @@ impl StorageLeafConfig { &MptUpdateRow { address: Value::known(main_data.address), storage_key: address_item.word().into_value(), + transaction_index: Value::known(F::ZERO), proof_type: Value::known(proof_type.scalar()), new_root: main_data.new_root.into_value(), old_root: main_data.old_root.into_value(), diff --git a/zkevm-circuits/src/mpt_circuit/witness_row.rs b/zkevm-circuits/src/mpt_circuit/witness_row.rs index 9ce6c47f52..5794484cdc 100644 --- a/zkevm-circuits/src/mpt_circuit/witness_row.rs +++ b/zkevm-circuits/src/mpt_circuit/witness_row.rs @@ -49,6 +49,20 @@ pub(crate) enum AccountRowType { Key, // hashed account address Count, } +#[derive(Debug, Eq, PartialEq)] +pub(crate) enum TransactionRowType { + KeyS, + ValueS, + KeyC, + ValueC, + Drifted, + ExtNodeKey, + ExtNodeNibbles, + ExtNodeValue, + Index, // transaction index + Key, // encoded transaction index + Count, +} #[derive(Debug, Eq, PartialEq)] pub(crate) enum ExtensionBranchRowType { @@ -201,6 +215,27 @@ pub struct StorageNode { /// RLP bytes denoting the length of the RLP of the long and short modified extension node. pub(crate) mod_list_rlp_bytes: [Hex; 2], } +/// MPT transaction stack trie +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TransactionNode { + /// Transaction indexes + pub index: usize, + /// Encoded transaction index + pub key: Hex, + /// RLP bytes denoting the length of the whole tx leaf stream. + pub list_rlp_bytes: [Hex; 2], + /// RLP bytes denoting the length of the value stream. + pub value_rlp_bytes: [Hex; 2], + /// RLP bytes denoting the length of the RLP stream of the drifted leaf (neighbor leaf). + /// This is only needed in the case when a new branch is created which replaces the existing + /// leaf in the trie and this leaf drifts down into newly created branch. + pub drifted_rlp_bytes: Hex, + /// Denotes whether the extension node nibbles have been modified in either `S` or `C` proof. + /// In these special cases, an additional extension node is inserted (deleted). + pub(crate) is_mod_extension: [bool; 2], + /// RLP bytes denoting the length of the RLP of the long and short modified extension node. + pub(crate) mod_list_rlp_bytes: [Hex; 2], +} /// MPT node #[derive(Clone, Debug, Default, Serialize, Deserialize)] @@ -213,6 +248,8 @@ pub struct Node { pub account: Option, /// A storage leaf node. pub storage: Option, + /// An transaction leaf node. + pub transaction: Option, /// RLP substreams of the node (for example for account leaf it contains substreams for key, /// nonce, balance, storage, codehash, drifted key, wrong key...) pub values: Vec, @@ -290,3 +327,17 @@ pub const NODE_RLP_TYPES_STORAGE: [RlpItemType; StorageRowType::Count as usize] RlpItemType::Hash, RlpItemType::Hash, ]; + +/// RLP types transaction +pub const NODE_RLP_TYPES_TRANSACTION: [RlpItemType; TransactionRowType::Count as usize] = [ + RlpItemType::Key, + RlpItemType::Value, + RlpItemType::Key, + RlpItemType::Value, + RlpItemType::Key, + RlpItemType::Key, + RlpItemType::Nibbles, + RlpItemType::Value, + RlpItemType::Value, + RlpItemType::Hash, +]; diff --git a/zkevm-circuits/src/table/mpt_table.rs b/zkevm-circuits/src/table/mpt_table.rs index 08ae579283..157b66c8a1 100644 --- a/zkevm-circuits/src/table/mpt_table.rs +++ b/zkevm-circuits/src/table/mpt_table.rs @@ -26,6 +26,8 @@ pub enum MPTProofType { StorageChanged, /// Storage does not exist StorageDoesNotExist, + /// Transaction updated + TransactionUpdated, } impl_expr!(MPTProofType); @@ -47,6 +49,8 @@ pub struct MptTable { pub address: Column, /// Storage address pub storage_key: WordLoHi>, + /// Transaction index + pub transaction_index: Column, /// Proof type pub proof_type: Column, /// New MPT root @@ -65,6 +69,7 @@ impl LookupTable for MptTable { self.address, self.storage_key.lo(), self.storage_key.hi(), + self.transaction_index, self.proof_type, self.new_root.lo(), self.new_root.hi(), @@ -85,6 +90,7 @@ impl LookupTable for MptTable { String::from("address"), String::from("storage_key_lo"), String::from("storage_key_hi"), + String::from("transaction_index"), String::from("proof_type"), String::from("new_root_lo"), String::from("new_root_hi"), @@ -104,6 +110,7 @@ impl MptTable { Self { address: meta.advice_column(), storage_key: WordLoHi::new([meta.advice_column(), meta.advice_column()]), + transaction_index: meta.advice_column(), proof_type: meta.advice_column(), new_root: WordLoHi::new([meta.advice_column(), meta.advice_column()]), old_root: WordLoHi::new([meta.advice_column(), meta.advice_column()]), @@ -118,6 +125,7 @@ impl MptTable { meta: &mut VirtualCells<'_, F>, cb: &mut ConstraintBuilder, address: Expression, + transaction_index: Expression, proof_type: Expression, storage_key: WordLoHi>, new_root: WordLoHi>, @@ -128,6 +136,7 @@ impl MptTable { circuit!([meta, cb], { require!(a!(self.address) => address); require!([a!(self.storage_key.lo()), a!(self.storage_key.hi())] => storage_key); + require!(a!(self.transaction_index) => transaction_index); require!(a!(self.proof_type) => proof_type); require!([a!(self.new_root.lo()), a!(self.new_root.hi())] => new_root); require!([a!(self.old_root.lo()), a!(self.old_root.hi())] => old_root); diff --git a/zkevm-circuits/src/witness/mpt.rs b/zkevm-circuits/src/witness/mpt.rs index f2005ddb76..7cb20eac16 100644 --- a/zkevm-circuits/src/witness/mpt.rs +++ b/zkevm-circuits/src/witness/mpt.rs @@ -29,6 +29,7 @@ impl MptUpdate { } } Key::Account { field_tag, .. } => field_tag.into(), + Key::Transaction { .. } => MPTProofType::TransactionUpdated, }; F::from(proof_type as u64) } @@ -46,6 +47,7 @@ pub struct MptUpdates { pub struct MptUpdateRow { pub(crate) address: F, pub(crate) storage_key: WordLoHi, + pub(crate) transaction_index: F, pub(crate) proof_type: F, pub(crate) new_root: WordLoHi, pub(crate) old_root: WordLoHi, @@ -102,6 +104,7 @@ impl MptUpdates { MptUpdateRow { address: Value::known(update.key.address().to_scalar().unwrap()), storage_key: WordLoHi::::from(update.key.storage_key()).into_value(), + transaction_index: Value::known(F::from(update.key.transaction_index())), proof_type: Value::known(update.proof_type()), new_root: WordLoHi::::from(new_root).into_value(), old_root: WordLoHi::::from(old_root).into_value(), @@ -134,6 +137,10 @@ enum Key { storage_key: Word, exists: bool, }, + Transaction { + tx_id: u64, + address: Address, + }, } impl Key { @@ -164,6 +171,7 @@ impl Key { storage_key, exists: false, }, + Key::Transaction { .. } => self, } } else { self @@ -171,13 +179,22 @@ impl Key { } fn address(&self) -> Address { match self { - Self::Account { address, .. } | Self::AccountStorage { address, .. } => *address, + Self::Account { address, .. } + | Self::AccountStorage { address, .. } + | Self::Transaction { address, .. } => *address, } } fn storage_key(&self) -> Word { match self { Self::Account { .. } => Word::zero(), Self::AccountStorage { storage_key, .. } => *storage_key, + Self::Transaction { .. } => Word::zero(), + } + } + fn transaction_index(&self) -> u64 { + match self { + Self::Transaction { tx_id, .. } => *tx_id, + Self::Account { .. } | Self::AccountStorage { .. } => 0, } } } @@ -185,11 +202,12 @@ impl Key { impl MptUpdateRow { /// The individual values of the row, in the column order used by the /// MptTable - pub fn values(&self) -> [F; 12] { + pub fn values(&self) -> [F; 13] { [ self.address.clone(), self.storage_key.lo(), self.storage_key.hi(), + self.transaction_index.clone(), self.proof_type.clone(), self.new_root.lo(), self.new_root.hi(), From 8051045a9fd74b71ba14de03041dda0889da9e43 Mon Sep 17 00:00:00 2001 From: KimiWu Date: Mon, 29 Apr 2024 16:12:30 +0800 Subject: [PATCH 2/2] feat: tx_leaf gadget created (copied from storage_leaf) --- .../src/mpt_circuit/account_leaf.rs | 1 + zkevm-circuits/src/mpt_circuit/helpers.rs | 20 +- zkevm-circuits/src/mpt_circuit/start.rs | 3 +- .../src/mpt_circuit/storage_leaf.rs | 1 + .../src/mpt_circuit/transaction_leaf.rs | 512 ++++++++++++++++++ 5 files changed, 530 insertions(+), 7 deletions(-) create mode 100644 zkevm-circuits/src/mpt_circuit/transaction_leaf.rs diff --git a/zkevm-circuits/src/mpt_circuit/account_leaf.rs b/zkevm-circuits/src/mpt_circuit/account_leaf.rs index 5a7621337c..efd2827df5 100644 --- a/zkevm-circuits/src/mpt_circuit/account_leaf.rs +++ b/zkevm-circuits/src/mpt_circuit/account_leaf.rs @@ -716,6 +716,7 @@ impl AccountLeafConfig { main_data.proof_type, true, address_item.word().compress_f(), + F::ZERO, main_data.new_root, main_data.old_root, )?; diff --git a/zkevm-circuits/src/mpt_circuit/helpers.rs b/zkevm-circuits/src/mpt_circuit/helpers.rs index b780b47a63..5770d79d55 100644 --- a/zkevm-circuits/src/mpt_circuit/helpers.rs +++ b/zkevm-circuits/src/mpt_circuit/helpers.rs @@ -686,6 +686,7 @@ pub(crate) struct MainData { pub(crate) proof_type: Cell, pub(crate) is_below_account: Cell, pub(crate) address: Cell, + pub(crate) transaction_index: Cell, pub(crate) new_root: WordLoHiCell, pub(crate) old_root: WordLoHiCell, } @@ -695,6 +696,7 @@ pub(crate) struct MainDataWitness { pub(crate) proof_type: usize, pub(crate) is_below_account: bool, pub(crate) address: F, + pub(crate) transaction_index: F, pub(crate) new_root: WordLoHi, pub(crate) old_root: WordLoHi, } @@ -709,6 +711,7 @@ impl MainData { proof_type: cb.query_cell(), is_below_account: cb.query_cell(), address: cb.query_cell(), + transaction_index: cb.query_cell(), new_root: cb.query_word_unchecked(), old_root: cb.query_word_unchecked(), }; @@ -720,6 +723,7 @@ impl MainData { main_data.proof_type.expr(), main_data.is_below_account.expr(), main_data.address.expr(), + main_data.transaction_index.expr(), main_data.new_root.lo().expr(), main_data.new_root.hi().expr(), main_data.old_root.lo().expr(), @@ -746,6 +750,7 @@ impl MainData { proof_type: usize, is_below_account: bool, address: F, + transaction_index: F, new_root: WordLoHi, old_root: WordLoHi, ) -> Result<(), Error> { @@ -753,6 +758,7 @@ impl MainData { proof_type.scalar(), is_below_account.scalar(), address, + transaction_index, new_root.lo(), new_root.hi(), old_root.lo(), @@ -775,17 +781,19 @@ impl MainData { self.proof_type.assign(region, offset, values[0])?; self.is_below_account.assign(region, offset, values[1])?; self.address.assign(region, offset, values[2])?; - self.new_root.lo().assign(region, offset, values[3])?; - self.new_root.hi().assign(region, offset, values[4])?; - self.old_root.lo().assign(region, offset, values[5])?; - self.old_root.hi().assign(region, offset, values[6])?; + self.transaction_index.assign(region, offset, values[3])?; + self.new_root.lo().assign(region, offset, values[4])?; + self.new_root.hi().assign(region, offset, values[5])?; + self.old_root.lo().assign(region, offset, values[6])?; + self.old_root.hi().assign(region, offset, values[7])?; Ok(MainDataWitness { proof_type: values[0].get_lower_32() as usize, is_below_account: values[1] == 1.scalar(), address: values[2], - new_root: WordLoHi::new([values[3], values[4]]), - old_root: WordLoHi::new([values[5], values[6]]), + transaction_index: values[3], + new_root: WordLoHi::new([values[4], values[5]]), + old_root: WordLoHi::new([values[6], values[7]]), }) } } diff --git a/zkevm-circuits/src/mpt_circuit/start.rs b/zkevm-circuits/src/mpt_circuit/start.rs index 37ee126e36..18da45fb84 100644 --- a/zkevm-circuits/src/mpt_circuit/start.rs +++ b/zkevm-circuits/src/mpt_circuit/start.rs @@ -107,7 +107,8 @@ impl StartConfig { &mut memory[main_memory()], start.proof_type as usize, false, - 0.scalar(), + F::ZERO, + F::ZERO, root[true.idx()], root[false.idx()], )?; diff --git a/zkevm-circuits/src/mpt_circuit/storage_leaf.rs b/zkevm-circuits/src/mpt_circuit/storage_leaf.rs index e66f046f1f..01f3280d1c 100644 --- a/zkevm-circuits/src/mpt_circuit/storage_leaf.rs +++ b/zkevm-circuits/src/mpt_circuit/storage_leaf.rs @@ -574,6 +574,7 @@ impl StorageLeafConfig { MPTProofType::Disabled as usize, false, F::ZERO, + F::ZERO, WordLoHi::new([F::ZERO, F::ZERO]), WordLoHi::new([F::ZERO, F::ZERO]), )?; diff --git a/zkevm-circuits/src/mpt_circuit/transaction_leaf.rs b/zkevm-circuits/src/mpt_circuit/transaction_leaf.rs new file mode 100644 index 0000000000..3c835a545e --- /dev/null +++ b/zkevm-circuits/src/mpt_circuit/transaction_leaf.rs @@ -0,0 +1,512 @@ +use eth_types::{Field, OpsIdentity, U256}; +use ethers_core::types::transaction; +use gadgets::util::Scalar; +use halo2_proofs::{ + circuit::Value, + plonk::{Error, Expression, VirtualCells}, +}; +use itertools::Itertools; + +use crate::{ + circuit, + circuit_tools::{ + cached_region::CachedRegion, + cell_manager::Cell, + constraint_builder::{RLCChainableRev, RLCable}, + gadgets::{IsEqualGadget, LtGadget}, + }, + mpt_circuit::{ + helpers::{ + key_memory, main_memory, num_nibbles, parent_memory, DriftedGadget, + IsPlaceholderLeafGadget, KeyData, MPTConstraintBuilder, MainData, ParentData, + ParentDataWitness, KECCAK, + }, + param::{EMPTY_TRIE_HASH, KEY_LEN_IN_NIBBLES}, + MPTConfig, MPTContext, MptMemory, RlpItemType, + }, + table::MPTProofType, + util::word::WordLoHi, + witness::MptUpdateRow, +}; + +use super::{ + helpers::{Indexable, KeyDataWitness, ListKeyGadget}, + mod_extension::ModExtensionGadget, + rlp_gadgets::{RLPItemWitness, RLPValueGadget}, + witness_row::{Node, TransactionRowType}, +}; + +#[derive(Clone, Debug, Default)] +pub(crate) struct TransactionLeafConfig { + main_data: MainData, + key_data: [KeyData; 2], + parent_data: [ParentData; 2], + + rlp_key: [ListKeyGadget; 2], + value_rlp_bytes: [[Cell; 1]; 2], + rlp_value: [RLPValueGadget; 2], + is_placeholder_leaf: [IsPlaceholderLeafGadget; 2], + drifted: DriftedGadget, + is_mod_extension: [Cell; 2], + mod_extension: ModExtensionGadget, +} + +impl TransactionLeafConfig { + pub fn configure( + meta: &mut VirtualCells<'_, F>, + cb: &mut MPTConstraintBuilder, + ctx: &mut MPTContext, + ) -> Self { + let mut config = TransactionLeafConfig::default(); + + circuit!([meta, cb], { + let key_items = [ + ctx.rlp_item( + meta, + cb, + TransactionRowType::KeyS as usize, + RlpItemType::Key, + ), + ctx.rlp_item( + meta, + cb, + TransactionRowType::KeyC as usize, + RlpItemType::Key, + ), + ]; + config.value_rlp_bytes = [cb.base.query_bytes(), cb.base.query_bytes()]; + let value_item = [ + ctx.rlp_item( + meta, + cb, + TransactionRowType::ValueS as usize, + RlpItemType::Value, + ), + ctx.rlp_item( + meta, + cb, + TransactionRowType::ValueC as usize, + RlpItemType::Value, + ), + ]; + let drifted_item = ctx.rlp_item( + meta, + cb, + TransactionRowType::Drifted as usize, + RlpItemType::Key, + ); + + let key_item = ctx.rlp_item( + meta, + cb, + TransactionRowType::Key as usize, + RlpItemType::Hash, + ); + + for is_proof_s in [true, false] { + config.is_mod_extension[is_proof_s.idx()] = cb.query_bool(); + } + + config.main_data = MainData::load(cb, &mut ctx.memory[main_memory()], 0.expr()); + + // + require!(config.main_data.is_below_account => false); + + let mut key_rlc = vec![0.expr(); 2]; + let mut value_word = vec![WordLoHi::zero(); 2]; + let mut value_rlp_rlc = vec![0.expr(); 2]; + let mut value_rlp_rlc_mult = vec![0.expr(); 2]; + + let parent_data = &mut config.parent_data; + parent_data[0] = ParentData::load(cb, &mut ctx.memory[parent_memory(true)], 0.expr()); + parent_data[1] = ParentData::load(cb, &mut ctx.memory[parent_memory(false)], 0.expr()); + + let key_data = &mut config.key_data; + key_data[0] = KeyData::load(cb, &mut ctx.memory[key_memory(true)], 0.expr()); + key_data[1] = KeyData::load(cb, &mut ctx.memory[key_memory(false)], 0.expr()); + + for is_proof_s in [true, false] { + ifx! {not!(config.is_mod_extension[is_proof_s.idx()].expr()) => { + // Placeholder leaf checks + config.is_placeholder_leaf[is_proof_s.idx()] = + IsPlaceholderLeafGadget::construct(cb, parent_data[is_proof_s.idx()].hash.expr()); + let is_placeholder_leaf = config.is_placeholder_leaf[is_proof_s.idx()].expr(); + + let rlp_key = &mut config.rlp_key[is_proof_s.idx()]; + *rlp_key = ListKeyGadget::construct(cb, &key_items[is_proof_s.idx()]); + config.rlp_value[is_proof_s.idx()] = RLPValueGadget::construct( + cb, + &config.value_rlp_bytes[is_proof_s.idx()] + .iter() + .map(|c| c.expr()) + .collect::>(), + ); + + // Because the storage value is an rlp encoded string inside another rlp encoded + // string (leaves are always encoded as [key, value], with + // `value` here containing a single stored value) the stored + // value is either stored directly in the RLP encoded string if short, or stored + // wrapped inside another RLP encoded string if long. + let rlp_value = config.rlp_value[is_proof_s.idx()].rlc_value(&cb.key_r); + let rlp_value_rlc_mult = + config.rlp_value[is_proof_s.idx()].rlc_rlp_only_rev(&cb.keccak_r); + let value_lo; + let value_hi; + ( + value_lo, + value_hi, + value_rlp_rlc[is_proof_s.idx()], + value_rlp_rlc_mult[is_proof_s.idx()], + ) = ifx! {config.rlp_value[is_proof_s.idx()].is_short() => { + (rlp_value, 0.expr(), rlp_value_rlc_mult.0.expr(), rlp_value_rlc_mult.1.expr()) + } elsex { + let value = value_item[is_proof_s.idx()].word(); + let value_rlp_rlc = rlp_value_rlc_mult.0.rlc_chain_rev(value_item[is_proof_s.idx()].rlc_chain_data()); + require!(config.rlp_value[is_proof_s.idx()].num_bytes() => value_item[is_proof_s.idx()].num_bytes() + 1.expr()); + (value.lo(), value.hi(), value_rlp_rlc, rlp_value_rlc_mult.1 * value_item[is_proof_s.idx()].mult()) + }}; + value_word[is_proof_s.idx()] = WordLoHi::>::new([value_lo, value_hi]); + + let leaf_rlc = rlp_key.rlc2(&cb.keccak_r).rlc_chain_rev(( + value_rlp_rlc[is_proof_s.idx()].expr(), + value_rlp_rlc_mult[is_proof_s.idx()].expr(), + )); + + // Key + key_rlc[is_proof_s.idx()] = key_data[is_proof_s.idx()].rlc.expr() + + rlp_key.key.expr( + cb, + rlp_key.key_value.clone(), + key_data[is_proof_s.idx()].mult.expr(), + key_data[is_proof_s.idx()].is_odd.expr(), + &cb.key_r.expr(), + ); + // Total number of nibbles needs to be KEY_LEN_IN_NIBBLES + let num_nibbles = + num_nibbles::expr(rlp_key.key_value.len(), key_data[is_proof_s.idx()].is_odd.expr()); + require!(key_data[is_proof_s.idx()].num_nibbles.expr() + num_nibbles => KEY_LEN_IN_NIBBLES); + + // Placeholder leaves default to value `0`. + ifx! {is_placeholder_leaf => { + require!(value_word[is_proof_s.idx()] => WordLoHi::>::zero()); + }} + + // Make sure the RLP encoding is correct. + // storage = [key, "value"] + require!(rlp_key.rlp_list.len() => key_items[is_proof_s.idx()].num_bytes() + config.rlp_value[is_proof_s.idx()].num_bytes()); + + // Check if the leaf is in its parent. + // Check is skipped for placeholder leaves which are dummy leaves. + // Note that the constraint works for the case when there is the placeholder branch above + // the leaf too - in this case `parent_data.hash` contains the hash of the node above the placeholder + // branch. + ifx! {not!(is_placeholder_leaf) => { + // config.is_not_hashed[is_proof_s.idx()] = LtGadget::construct(&mut cb.base, rlp_key.rlp_list.num_bytes(), 32.expr()); + // ifx!{or::expr(&[parent_data[is_proof_s.idx()].is_root.expr(), not!(config.is_not_hashed[is_proof_s.idx()])]) => { + ifx!{parent_data[is_proof_s.idx()].is_root.expr() => { + // Hashed leaf in parent branch + let hash = parent_data[is_proof_s.idx()].hash.expr(); + require!((1.expr(), leaf_rlc.expr(), rlp_key.rlp_list.num_bytes(), hash.lo(), hash.hi()) =>> @KECCAK); + } elsex { + // Non-hashed leaf in parent branch + require!(leaf_rlc => parent_data[is_proof_s.idx()].rlc.expr()); + }} + } elsex { + // For NonExistingStorageProof prove there is no leaf. + + // When there is only one leaf in the trie, `getProof` will always return this leaf - so we will have + // either the required leaf or the wrong leaf, so for NonExistingStorageProof we don't handle this + // case here (handled by WrongLeaf gadget). + + ifx! {parent_data[is_proof_s.idx()].is_root.expr() => { + // If leaf is placeholder and the parent is root (no branch above leaf) and the proof is NonExistingStorageProof, + // the trie needs to be empty. + let empty_hash = WordLoHi::::from(U256::from_big_endian(&EMPTY_TRIE_HASH)); + let hash = parent_data[is_proof_s.idx()].hash.expr(); + require!(hash.lo() => Expression::Constant(empty_hash.lo())); + require!(hash.hi() => Expression::Constant(empty_hash.hi())); + } elsex { + // For NonExistingStorageProof we need to prove that there is nil in the parent branch + // at the `modified_pos` position. + // Note that this does not hold when there is NonExistingStorageProof wrong leaf scenario, + // in this case there is a non-nil leaf. However, in this case the leaf is not a placeholder, + // so the check below is not triggered. + require!(parent_data[is_proof_s.idx()].rlc.expr() => 128.expr()); + }} + + }} + }}; + + // Key done, set the default values + KeyData::store_defaults(cb, &mut ctx.memory[key_memory(is_proof_s)]); + // Store the new parent + ParentData::store( + cb, + &mut ctx.memory[parent_memory(is_proof_s)], + WordLoHi::zero(), + 0.expr(), + true.expr(), + false.expr(), + WordLoHi::zero(), + ); + } + + ifx! {or::expr(&[config.is_mod_extension[0].clone(), config.is_mod_extension[1].clone()]) => { + config.mod_extension = ModExtensionGadget::configure( + meta, + cb, + ctx.clone(), + parent_data, + key_data, + ); + }}; + + // Drifted leaf handling + config.drifted = DriftedGadget::construct( + cb, + &config + .rlp_value + .iter() + .map(|value| value.num_bytes()) + .collect_vec(), + &config.parent_data, + &config.key_data, + &key_rlc, + &value_rlp_rlc, + &value_rlp_rlc_mult, + &drifted_item, + &config.is_mod_extension, + &cb.key_r.expr(), + ); + + // Reset the main memory + // This need to be the last node for this proof + MainData::store( + cb, + &mut ctx.memory[main_memory()], + [ + MPTProofType::Disabled.expr(), + false.expr(), + 0.expr(), + 0.expr(), + 0.expr(), + 0.expr(), + 0.expr(), + ], + ); + + // Put the data in the lookup table + let key_rlc = ifx! {not!(config.parent_data[true.idx()].is_placeholder) => { + key_rlc[true.idx()].expr() + } elsex { + key_rlc[false.idx()].expr() + }}; + // Check that the key item contains the correct key for the path that was taken + require!(key_item.hash_rlc() => key_rlc); + + ifx! {not!(config.parent_data[false.idx()].is_placeholder) => { + ctx.mpt_table.constrain( + meta, + &mut cb.base, + config.main_data.address.expr(), + false.expr(), + MPTProofType::TransactionUpdated.expr(), + WordLoHi::zero(), + config.main_data.new_root.expr(), + config.main_data.old_root.expr(), + value_word[false.idx()].clone(), + value_word[true.idx()].clone(), + ); + } elsex { + // When the value is set to 0, the leaf is deleted, and if there were only two leaves in the branch, + // the neighbour leaf moves one level up and replaces the branch. When the lookup is executed with + // the new value set to 0, the lookup fails (without the code below), because the leaf that is returned + // is the neighbour node that moved up (because the branch and the old leaf doesn’t exist anymore), + // but this leaf doesn’t have the zero value. + ctx.mpt_table.constrain( + meta, + &mut cb.base, + config.main_data.address.expr(), + false.expr(), + MPTProofType::TransactionUpdated.expr(), + WordLoHi::zero(), + config.main_data.new_root.expr(), + config.main_data.old_root.expr(), + WordLoHi::zero(), + value_word[true.idx()].clone(), + ); + }}; + }); + + config + } + + #[allow(clippy::too_many_arguments)] + pub fn assign( + &self, + region: &mut CachedRegion<'_, '_, F>, + mpt_config: &MPTConfig, + memory: &mut MptMemory, + offset: usize, + node: &Node, + rlp_values: &[RLPItemWitness], + ) -> Result<(), Error> { + let storage = &node.storage.clone().unwrap(); + + let key_items = [ + rlp_values[TransactionRowType::KeyS as usize].clone(), + rlp_values[TransactionRowType::KeyC as usize].clone(), + ]; + let value_item = [ + rlp_values[TransactionRowType::ValueS as usize].clone(), + rlp_values[TransactionRowType::ValueC as usize].clone(), + ]; + let drifted_item = rlp_values[TransactionRowType::Drifted as usize].clone(); + // let _key_item = rlp_values[TransactionRowType::Key as usize].clone(); + let tx_idx = rlp_values[TransactionRowType::Index as usize].clone(); + + let main_data = + self.main_data + .witness_load(region, offset, &mut memory[main_memory()], 0)?; + + let mut key_data = vec![KeyDataWitness::default(); 2]; + let mut parent_data = vec![ParentDataWitness::default(); 2]; + let mut key_rlc = vec![0.scalar(); 2]; + let mut value_word = [WordLoHi::zero(); 2]; + for is_proof_s in [true, false] { + self.is_mod_extension[is_proof_s.idx()].assign( + region, + offset, + storage.is_mod_extension[is_proof_s.idx()].scalar(), + )?; + + parent_data[is_proof_s.idx()] = self.parent_data[is_proof_s.idx()].witness_load( + region, + offset, + &mut memory[parent_memory(is_proof_s)], + 0, + )?; + + let rlp_key_witness = self.rlp_key[is_proof_s.idx()].assign( + region, + offset, + &storage.list_rlp_bytes[is_proof_s.idx()], + &key_items[is_proof_s.idx()], + )?; + + key_data[is_proof_s.idx()] = self.key_data[is_proof_s.idx()].witness_load( + region, + offset, + &mut memory[key_memory(is_proof_s)], + 0, + )?; + KeyData::witness_store( + region, + offset, + &mut memory[key_memory(is_proof_s)], + F::ZERO, + F::ONE, + 0, + F::ZERO, + F::ONE, + 0, + )?; + + // Key + (key_rlc[is_proof_s.idx()], _) = rlp_key_witness.key.key( + rlp_key_witness.key_item.clone(), + key_data[is_proof_s.idx()].rlc, + key_data[is_proof_s.idx()].mult, + region.key_r, + ); + + // Value + for (cell, byte) in self.value_rlp_bytes[is_proof_s.idx()] + .iter() + .zip(storage.value_rlp_bytes[is_proof_s.idx()].iter()) + { + cell.assign(region, offset, byte.scalar())?; + } + let value_witness = self.rlp_value[is_proof_s.idx()].assign( + region, + offset, + &storage.value_rlp_bytes[is_proof_s.idx()], + )?; + value_word[is_proof_s.idx()] = if value_witness.is_short() { + WordLoHi::::new([value_witness.rlc_value(region.key_r), 0.scalar()]) + } else { + value_item[is_proof_s.idx()].word() + }; + + ParentData::witness_store( + region, + offset, + &mut memory[parent_memory(is_proof_s)], + WordLoHi::::new([F::ZERO, F::ZERO]), + F::ZERO, + true, + false, + WordLoHi::::new([F::ZERO, F::ZERO]), + )?; + + self.is_placeholder_leaf[is_proof_s.idx()].assign( + region, + offset, + parent_data[is_proof_s.idx()].hash, + )?; + } + + // Drifted leaf handling + self.drifted.assign( + region, + offset, + &parent_data, + &storage.drifted_rlp_bytes, + &drifted_item, + region.key_r, + )?; + + MainData::witness_store( + region, + offset, + &mut memory[main_memory()], + main_data.proof_type, + false, + F::ZERO, + tx_idx.word().lo(), + main_data.new_root, + main_data.old_root, + )?; + + if storage.is_mod_extension[0] || storage.is_mod_extension[1] { + let mod_list_rlp_bytes: [&[u8]; 2] = [ + &storage.mod_list_rlp_bytes[0], + &storage.mod_list_rlp_bytes[1], + ]; + self.mod_extension + .assign(region, offset, rlp_values, mod_list_rlp_bytes)?; + } + + let mut new_value = value_word[false.idx()]; + let old_value = value_word[true.idx()]; + if parent_data[false.idx()].is_placeholder { + new_value = WordLoHi::zero(); + } + mpt_config.mpt_table.assign_cached( + region, + offset, + &MptUpdateRow { + address: Value::known(main_data.address), + storage_key: WordLoHi::new([Value::known(F::ZERO), Value::known(F::ZERO)]), + transaction_index: Value::known(F::ZERO), + proof_type: Value::known(MPTProofType::TransactionUpdated.scalar()), + new_root: main_data.new_root.into_value(), + old_root: main_data.old_root.into_value(), + new_value: new_value.into_value(), + old_value: old_value.into_value(), + }, + )?; + + Ok(()) + } +}