Add mpt lookups to state circuit

prover/src/compute_proof.rs

@@ -0,0 +1,96 @@
+use bus_mapping::circuit_input_builder::BuilderClient;
+use bus_mapping::rpc::GethClient;
+use ethers_providers::Http;
+use halo2_proofs::{
+    pairing::bn256::{Fr, G1Affine},
+    plonk::*,
+    poly::commitment::Params,
+    transcript::{Blake2bWrite, Challenge255},
+};
+use rand::SeedableRng;
+use rand_xorshift::XorShiftRng;
+
+use std::str::FromStr;
+use std::time::Instant;
+
+use strum::IntoEnumIterator;
+use zkevm_circuits::evm_circuit::{
+    table::FixedTableTag, test::TestCircuit, witness::block_convert,
+};
+use zkevm_circuits::state_circuit::StateCircuit;
+
+use crate::structs::Proofs;
+
+/// Gathers debug trace(s) from `rpc_url` for block `block_num` with `params`
+/// created via the `gen_params` tool.
+/// Expects a go-ethereum node with debug & archive capabilities on `rpc_url`.
+pub async fn compute_proof(
+    params: &Params<G1Affine>,
+    block_num: &u64,
+    rpc_url: &str,
+) -> Result<Proofs, Box<dyn std::error::Error>> {
+    // request & build the inputs for the circuits
+    let time_started = Instant::now();
+    let url = Http::from_str(rpc_url)?;
+    let geth_client = GethClient::new(url);
+    let builder = BuilderClient::new(geth_client).await?;
+    let builder = builder.gen_inputs(*block_num).await?;
+
+    // TODO: only {evm,state}_proof are implemented right now
+    let evm_proof;
+    let state_proof;
+    let block = block_convert(&builder.block, &builder.code_db);
+    {
+        // generate evm_circuit proof
+        let circuit = TestCircuit::<Fr>::new(block.clone(), FixedTableTag::iter().collect());
+
+        // TODO: can this be pre-generated to a file?
+        // related
+        // https://github.com/zcash/halo2/issues/443
+        // https://github.com/zcash/halo2/issues/449
+        let vk = keygen_vk(params, &circuit)?;
+        let pk = keygen_pk(params, vk, &circuit)?;
+
+        // Create randomness
+        let rng = XorShiftRng::from_seed([
+            0x59, 0x62, 0xbe, 0x5d, 0x76, 0x3d, 0x31, 0x8d, 0x17, 0xdb, 0x37, 0x32, 0x54, 0x06,
+            0xbc, 0xe5,
+        ]);
+
+        // create a proof
+        let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+        create_proof(params, &pk, &[circuit], &[], rng, &mut transcript)?;
+        evm_proof = transcript.finalize();
+    }
+
+    {
+        // generate state_circuit proof
+        const N_ROWS: usize = 1 << 16;
+        // TODO: use different values for word_randomness and lookup_randomness.
+        let circuit =
+            StateCircuit::<Fr, N_ROWS>::new(block.rws, block.randomness, block.randomness);
+
+        // TODO: same quest like in the first scope
+        let vk = keygen_vk(params, &circuit)?;
+        let pk = keygen_pk(params, vk, &circuit)?;
+
+        // Create randomness
+        let rng = XorShiftRng::from_seed([
+            0x59, 0x62, 0xbe, 0x5d, 0x76, 0x3d, 0x31, 0x8d, 0x17, 0xdb, 0x37, 0x32, 0x54, 0x06,
+            0xbc, 0xe5,
+        ]);
+
+        // create a proof
+        let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+        create_proof(params, &pk, &[circuit], &[], rng, &mut transcript)?;
+        state_proof = transcript.finalize();
+    }
+
+    let ret = Proofs {
+        evm_proof: evm_proof.into(),
+        state_proof: state_proof.into(),
+        duration: Instant::now().duration_since(time_started).as_millis() as u64,
+    };
+
+    Ok(ret)
+}

zkevm-circuits/src/state_circuit.rs

@@ -2,6 +2,7 @@
 mod constraint_builder;
 mod lexicographic_ordering;
 mod lookups;
+mod mpt_updates;
 mod multiple_precision_integer;
 mod random_linear_combination;
 #[cfg(test)]
@@ -23,6 +24,7 @@ use halo2_proofs::{
 };
 use lexicographic_ordering::Config as LexicographicOrderingConfig;
 use lookups::{Chip as LookupsChip, Config as LookupsConfig, Queries as LookupsQueries};
+use mpt_updates::MptUpdates;
 use multiple_precision_integer::{Chip as MpiChip, Config as MpiConfig, Queries as MpiQueries};
 use random_linear_combination::{Chip as RlcChip, Config as RlcConfig, Queries as RlcQueries};
 #[cfg(test)]
@@ -45,6 +47,7 @@ pub struct StateCircuitConfig<F> {
     initial_value: Column<Advice>, /* Assigned value at the start of the block. For Rw::Account
                                     * and Rw::AccountStorage rows this is the committed value in
                                     * the MPT, for others, it is 0. */
+    state_root: Column<Advice>,
     lexicographic_ordering: LexicographicOrderingConfig,
     lookups: LookupsConfig,
     power_of_randomness: [Expression<F>; N_BYTES_WORD - 1],
@@ -221,6 +224,7 @@ type Lookup<F> = (&'static str, Expression<F>, Expression<F>);
 #[derive(Default, Clone)]
 pub struct StateCircuit<F> {
     pub(crate) rows: Vec<Rw>,
+    updates: MptUpdates,
     pub(crate) n_rows: usize,
     pub(crate) randomness: F,
     #[cfg(test)]
@@ -231,9 +235,11 @@ impl<F: Field> StateCircuit<F> {
     /// make a new state circuit from an RwMap
     pub fn new(randomness: F, rw_map: RwMap, n_rows: usize) -> Self {
         let rows = rw_map.table_assignments();
+        let updates = MptUpdates::mock_from(&rows);
         Self {
             randomness,
             rows,
+            updates,
             n_rows,
             #[cfg(test)]
             overrides: HashMap::new(),
@@ -246,6 +252,11 @@ impl<F: Field> StateCircuit<F> {
             .map(|exp| vec![self.randomness.pow(&[exp, 0, 0, 0]); self.n_rows])
             .collect()
     }
+
+    /// Number of active (i.e. non-Rw::Start) rows in the circuit.
+    pub fn n_active_rows(&self) -> usize {
+        self.rows.len()
+    }
 }
 
 impl<F: Field> Circuit<F> for StateCircuit<F>
@@ -362,11 +373,22 @@ fn queries<F: Field>(meta: &mut VirtualCells<'_, F>, c: &StateCircuitConfig<F>)
             - meta.query_advice(first_different_limb.bits[1], Rotation::cur())
             - meta.query_advice(first_different_limb.bits[2], Rotation::cur())
             - meta.query_advice(first_different_limb.bits[3], Rotation::cur()),
+        last_access: 1.expr()
+            - meta.query_advice(first_different_limb.bits[0], Rotation::next())
+                * meta.query_advice(first_different_limb.bits[1], Rotation::next())
+                * meta.query_advice(first_different_limb.bits[2], Rotation::next())
+                * meta.query_advice(first_different_limb.bits[3], Rotation::next()),
         // 1 if first_different_limb is in the rw counter, 0 otherwise (i.e. any of the 4 most
         // significant bits are 0)
         not_first_access: meta.query_advice(first_different_limb.bits[0], Rotation::cur())
             * meta.query_advice(first_different_limb.bits[1], Rotation::cur())
             * meta.query_advice(first_different_limb.bits[2], Rotation::cur())
             * meta.query_advice(first_different_limb.bits[3], Rotation::cur()),
+        state_root: meta.query_advice(c.state_root, Rotation::cur()),
+        state_root_prev: meta.query_advice(c.state_root, Rotation::prev()),
     }
 }
+
+fn pow<F: Field>(x: F, exponent: usize) -> F {
+    x.pow(&[exponent.try_into().unwrap(), 0, 0, 0])
+}

zkevm-circuits/src/state_circuit/constraint_builder.rs

@@ -11,6 +11,7 @@ use crate::{
 use eth_types::Field;
 use gadgets::binary_number::BinaryNumberConfig;
 use halo2_proofs::plonk::Expression;
+use itertools::Itertools;
 use strum::IntoEnumIterator;
 
 #[derive(Clone)]
@@ -43,10 +44,13 @@ pub struct Queries<F: Field> {
     pub storage_key: RlcQueries<F, N_BYTES_WORD>,
     pub initial_value: Expression<F>,
     pub initial_value_prev: Expression<F>,
+    pub state_root: Expression<F>,
+    pub state_root_prev: Expression<F>,
     pub lookups: LookupsQueries<F>,
-    pub power_of_randomness: [Expression<F>; N_BYTES_WORD - 1],
+    pub powers_of_lookup_randomness: [Expression<F>; 6],
     pub first_access: Expression<F>,
     pub not_first_access: Expression<F>,
+    pub last_access: Expression<F>,
 }
 
 type Constraint<F> = (&'static str, Expression<F>);
@@ -154,6 +158,14 @@ impl<F: Field> ConstraintBuilder<F> {
         );
         self.require_zero("Start value is 0", q.value());
         self.require_zero("Start initial_value is 0", q.initial_value());
+
+        self.condition(q.lexicographic_ordering_selector.clone(), |cb| {
+            cb.require_equal(
+                "state_root is unchanged for Start",
+                q.state_root(),
+                q.state_root_prev(),
+            )
+        });
     }
 
     fn build_memory_constraints(&mut self, q: &Queries<F>) {
@@ -172,6 +184,12 @@ impl<F: Field> ConstraintBuilder<F> {
             (q.rw_table.value.clone(), q.lookups.u8.clone()),
         );
         self.require_zero("initial Memory value is 0", q.initial_value());
+
+        self.require_equal(
+            "state_root is unchanged for Memory",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_stack_constraints(&mut self, q: &Queries<F>) {
@@ -193,14 +211,23 @@ impl<F: Field> ConstraintBuilder<F> {
         });
 
         self.require_zero("initial Stack value is 0", q.initial_value.clone());
+
+        self.require_equal(
+            "state_root is unchanged for Stack",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_account_storage_constraints(&mut self, q: &Queries<F>) {
-        // TODO: cold VS warm
         self.require_zero("field_tag is 0 for AccountStorage", q.field_tag());
 
-        // TODO: add mpt lookup for committed value and final value in an access
-        // group.
+        self.condition(q.last_access(), |cb| {
+            cb.add_lookup(
+                "mpt_update exists in mpt circuit for AccountStorage last access",
+                (q.mpt_lookup(), q.lookups.mpt_update.clone()),
+            )
+        });
     }
     fn build_tx_access_list_account_constraints(&mut self, q: &Queries<F>) {
         self.require_zero("field_tag is 0 for TxAccessListAccount", q.field_tag());
@@ -213,6 +240,12 @@ impl<F: Field> ConstraintBuilder<F> {
             "initial TxAccessListAccount value is false",
             q.initial_value(),
         );
+
+        self.require_equal(
+            "state_root is unchanged for TxAccessListAccount",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_tx_access_list_account_storage_constraints(&mut self, q: &Queries<F>) {
@@ -225,6 +258,12 @@ impl<F: Field> ConstraintBuilder<F> {
             "initial TxAccessListAccountStorage value is false",
             q.initial_value(),
         );
+
+        self.require_equal(
+            "state_root is unchanged for TxAccessListAccountStorage",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_tx_refund_constraints(&mut self, q: &Queries<F>) {
@@ -235,6 +274,12 @@ impl<F: Field> ConstraintBuilder<F> {
             q.rw_table.storage_key.clone(),
         );
         self.require_zero("initial TxRefund value is 0", q.initial_value());
+
+        self.require_equal(
+            "state_root is unchanged for TxAccessListAccountStorage",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_account_constraints(&mut self, q: &Queries<F>) {
@@ -249,8 +294,12 @@ impl<F: Field> ConstraintBuilder<F> {
             set::<F, AccountFieldTag>(),
         );
 
-        // TODO: add mpt lookup for committed value and final value in an access
-        // group.
+        self.condition(q.last_access(), |cb| {
+            cb.add_lookup(
+                "mpt_update exists in mpt circuit for Account last access",
+                (q.mpt_lookup(), q.lookups.mpt_update.clone()),
+            )
+        });
     }
 
     fn build_account_destructed_constraints(&mut self, q: &Queries<F>) {
@@ -273,7 +322,14 @@ impl<F: Field> ConstraintBuilder<F> {
             "field_tag in CallContextFieldTag range",
             (q.field_tag(), q.lookups.call_context_field_tag.clone()),
         );
+
         // TODO: explain why call context doesn't need an initial value.
+
+        self.require_equal(
+            "state_root is unchanged for CallContext",
+            q.state_root(),
+            q.state_root_prev(),
+        );
     }
 
     fn build_tx_log_constraints(&mut self, q: &Queries<F>) {
@@ -284,6 +340,12 @@ impl<F: Field> ConstraintBuilder<F> {
         );
         self.require_zero("initial TxLog value is 0", q.initial_value());
 
+        self.require_equal(
+            "state_root is unchanged for TxLog",
+            q.state_root(),
+            q.state_root_prev(),
+        );
+
         // Comment out the following field_tag-related constraints as it is
         // duplicated between state circuit and evm circuit. For more information, please refer to https://github.com/privacy-scaling-explorations/zkevm-specs/issues/221
         // cb.require_zero(
@@ -443,6 +505,10 @@ impl<F: Field> Queries<F> {
         self.first_access.clone()
     }
 
+    fn last_access(&self) -> Expression<F> {
+        self.last_access.clone()
+    }
+
     fn address_change(&self) -> Expression<F> {
         self.rw_table.address.clone() - self.rw_table.prev_address.clone()
     }
@@ -466,6 +532,30 @@ impl<F: Field> Queries<F> {
     fn tx_log_id_prev(&self) -> Expression<F> {
         from_digits(&self.address.limbs_prev[3..5], (1u64 << 16).expr())
     }
+
+    fn state_root(&self) -> Expression<F> {
+        self.state_root.clone()
+    }
+
+    fn state_root_prev(&self) -> Expression<F> {
+        self.state_root_prev.clone()
+    }
+
+    fn mpt_lookup(&self) -> Expression<F> {
+        [
+            self.storage_key.encoded.clone(),
+            self.field_tag(),
+            self.state_root(),
+            self.state_root_prev(),
+            self.value(),
+            self.initial_value(),
+        ]
+        .iter()
+        .zip_eq(&self.powers_of_lookup_randomness)
+        .fold(self.address.value.clone(), |result, (a, b)| {
+            result + a.clone() * b.clone()
+        })
+    }
 }
 
 fn from_digits<F: Field>(digits: &[Expression<F>], base: Expression<F>) -> Expression<F> {