refactor global rw_counter to chunkctx_table lookup

zkevm-circuits/src/evm_circuit/execution.rs

@@ -20,7 +20,7 @@ use crate::{
         },
         witness::{Block, Call, ExecStep, Transaction},
     },
-    table::LookupTable,
+    table::{chunkctx_table::ChunkCtxFieldTag, LookupTable},
     util::{
         cell_manager::{CMFixedWidthStrategy, CellManager, CellType},
         Challenges, Expr,
@@ -38,7 +38,7 @@ use halo2_proofs::{
     poly::Rotation,
 };
 use std::{
-    collections::{BTreeSet, HashMap},
+    collections::{BTreeSet, HashMap, HashSet},
     iter,
 };
 use strum::IntoEnumIterator;
@@ -383,6 +383,14 @@ impl<F: Field> ExecutionConfig<F> {
 
         let step_curr = Step::new(meta, advices, 0);
         let mut height_map = HashMap::new();
+        let (execute_state_first_step_whitelist, execute_state_last_step_whitelist) = (
+            HashSet::from([
+                ExecutionState::BeginTx,
+                ExecutionState::EndBlock,
+                ExecutionState::BeginChunk,
+            ]),
+            HashSet::from([ExecutionState::EndBlock, ExecutionState::EndChunk]),
+        );
 
         meta.create_gate("Constrain execution state", |meta| {
             let q_usable = meta.query_selector(q_usable);
@@ -392,6 +400,17 @@ impl<F: Field> ExecutionConfig<F> {
 
             let execution_state_selector_constraints = step_curr.state.execution_state.configure();
 
+            let first_step_first_chunk_check = {
+                let exestates = step_curr
+                    .execution_state_selector(execute_state_first_step_whitelist.iter().cloned());
+                iter::once((
+                    "First step first chunk should be BeginTx or EndBlock or BeginChunk",
+                    (1.expr() - is_first_chunk.expr())
+                        * q_step_first.clone()
+                        * (1.expr() - exestates),
+                ))
+            };
+
             let first_step_non_first_chunk_check = {
                 let begin_chunk_selector =
                     step_curr.execution_state_selector([ExecutionState::BeginChunk]);
@@ -426,6 +445,7 @@ impl<F: Field> ExecutionConfig<F> {
             execution_state_selector_constraints
                 .into_iter()
                 .map(move |(name, poly)| (name, q_usable.clone() * q_step.clone() * poly))
+                .chain(first_step_first_chunk_check)
                 .chain(first_step_non_first_chunk_check)
                 .chain(last_step_last_chunk_check)
                 .chain(last_step_non_last_chunk_check)
@@ -507,6 +527,9 @@ impl<F: Field> ExecutionConfig<F> {
                         q_step_first,
                         q_step_last,
                         &step_curr,
+                        chunkctx_table,
+                        &execute_state_first_step_whitelist,
+                        &execute_state_last_step_whitelist,
                         &mut height_map,
                         &mut stored_expressions_map,
                         &mut debug_expressions_map,
@@ -655,6 +678,9 @@ impl<F: Field> ExecutionConfig<F> {
         q_step_first: Selector,
         q_step_last: Selector,
         step_curr: &Step<F>,
+        chunkctx_table: &dyn LookupTable<F>,
+        execute_state_first_step_whitelist: &HashSet<ExecutionState>,
+        execute_state_last_step_whitelist: &HashSet<ExecutionState>,
         height_map: &mut HashMap<ExecutionState, usize>,
         stored_expressions_map: &mut HashMap<ExecutionState, Vec<StoredExpression<F>>>,
         debug_expressions_map: &mut HashMap<ExecutionState, Vec<(String, Expression<F>)>>,
@@ -678,6 +704,7 @@ impl<F: Field> ExecutionConfig<F> {
 
         // Now actually configure the gadget with the correct minimal height
         let step_next = &Step::new(meta, advices, height);
+
         let mut cb = EVMConstraintBuilder::new(
             meta,
             step_curr.clone(),
@@ -699,11 +726,15 @@ impl<F: Field> ExecutionConfig<F> {
             height_map,
             stored_expressions_map,
             debug_expressions_map,
+            execute_state_first_step_whitelist,
+            execute_state_last_step_whitelist,
             instrument,
             G::NAME,
             G::EXECUTION_STATE,
             height,
             cb,
+            chunkctx_table,
+            challenges,
         );
 
         gadget
@@ -721,11 +752,15 @@ impl<F: Field> ExecutionConfig<F> {
         height_map: &mut HashMap<ExecutionState, usize>,
         stored_expressions_map: &mut HashMap<ExecutionState, Vec<StoredExpression<F>>>,
         debug_expressions_map: &mut HashMap<ExecutionState, Vec<(String, Expression<F>)>>,
+        execute_state_first_step_whitelist: &HashSet<ExecutionState>,
+        execute_state_last_step_whitelist: &HashSet<ExecutionState>,
         instrument: &mut Instrument,
         name: &'static str,
         execution_state: ExecutionState,
         height: usize,
         mut cb: EVMConstraintBuilder<F>,
+        chunkctx_table: &dyn LookupTable<F>,
+        challenges: &Challenges<Expression<F>>,
     ) {
         // Enforce the step height for this opcode
         let num_rows_until_next_step_next = cb
@@ -738,6 +773,8 @@ impl<F: Field> ExecutionConfig<F> {
 
         instrument.on_gadget_built(execution_state, &cb);
 
+        let step_curr_rw_counter = cb.curr.state.rw_counter.clone();
+        let step_curr_rw_counter_offset = cb.rw_counter_offset();
         let debug_expressions = cb.debug_expressions.clone();
         let (constraints, stored_expressions, _, meta) = cb.build();
         debug_assert!(
@@ -781,6 +818,53 @@ impl<F: Field> ExecutionConfig<F> {
             }
         }
 
+        // constraint global rw counter value at first/last step via chunkctx_table lookup
+        // we can't do it inside constraint_builder(cb)
+        // because lookup expression in constraint builder DONOT support apply conditional
+        // `step_first/step_last` selector at lookup cell.
+        if execute_state_first_step_whitelist.contains(&execution_state) {
+            meta.lookup_any("first must lookup initial rw_counter", |meta| {
+                let q_usable = meta.query_selector(q_usable);
+                let q_step_first = meta.query_selector(q_step_first);
+                let execute_state_selector = step_curr.execution_state_selector([execution_state]);
+
+                vec![(
+                    q_usable
+                        * q_step_first
+                        * execute_state_selector
+                        * rlc::expr(
+                            &[
+                                ChunkCtxFieldTag::InitialRWC.expr(),
+                                step_curr.state.rw_counter.expr(),
+                            ],
+                            challenges.lookup_input(),
+                        ),
+                    rlc::expr(&chunkctx_table.table_exprs(meta), challenges.lookup_input()),
+                )]
+            });
+        }
+
+        if execute_state_last_step_whitelist.contains(&execution_state) {
+            meta.lookup_any("last step must lookup end rw_counter", |meta| {
+                let q_usable = meta.query_selector(q_usable);
+                let q_step_last = meta.query_selector(q_step_last);
+                let execute_state_selector = step_curr.execution_state_selector([execution_state]);
+                vec![(
+                    q_usable
+                        * q_step_last
+                        * execute_state_selector
+                        * rlc::expr(
+                            &[
+                                ChunkCtxFieldTag::EndRWC.expr(),
+                                step_curr_rw_counter.expr() + step_curr_rw_counter_offset.clone(),
+                            ],
+                            challenges.lookup_input(),
+                        ),
+                    rlc::expr(&chunkctx_table.table_exprs(meta), challenges.lookup_input()),
+                )]
+            });
+        }
+
         // Enforce the state transitions for this opcode
         meta.create_gate("Constrain state machine transitions", |meta| {
             let q_usable = meta.query_selector(q_usable);

zkevm-circuits/src/evm_circuit/execution/begin_chunk.rs

@@ -27,9 +27,6 @@ impl<F: Field> ExecutionGadget<F> for BeginChunkGadget<F> {
     const EXECUTION_STATE: ExecutionState = ExecutionState::BeginChunk;
 
     fn configure(cb: &mut EVMConstraintBuilder<F>) -> Self {
-        // constraint rwc at step first
-        cb.step_first_constraint_rwc();
-
         // state lookup
         cb.step_state_lookup(0.expr());
         let step_state_transition = StepStateTransition::same();

zkevm-circuits/src/evm_circuit/execution/begin_tx.rs

@@ -73,9 +73,6 @@ impl<F: Field> ExecutionGadget<F> for BeginTxGadget<F> {
     const EXECUTION_STATE: ExecutionState = ExecutionState::BeginTx;
 
     fn configure(cb: &mut EVMConstraintBuilder<F>) -> Self {
-        // TODO constraint rwc at step first
-        // cb.step_first_constraint_rwc();
-
         // Use rw_counter of the step which triggers next call as its call_id.
         let call_id = cb.curr.state.rw_counter.clone();
 

zkevm-circuits/src/evm_circuit/execution/end_chunk.rs

@@ -37,8 +37,6 @@ impl<F: Field> ExecutionGadget<F> for EndChunkGadget<F> {
             cb.step_state_lookup(1.expr());
         });
 
-        cb.step_last_constraint_rwc();
-
         // TODO constraint Rw::Padding are append consecutively to avoid malicious insert
 
         Self {

zkevm-circuits/src/evm_circuit/util/constraint_builder.rs

@@ -1540,30 +1540,7 @@ impl<'a, F: Field> EVMConstraintBuilder<'a, F> {
         );
     }
 
-    /// constraint step first rwc should match with chunkctx table value
-    pub(crate) fn step_first_constraint_rwc(&mut self) {
-        // Add first BeginTx step constraint to have tx_id == 1
-        self.step_first(|cb| {
-            cb.chunk_context_lookup(
-                ChunkCtxFieldTag::InitialRWC,
-                cb.curr.state.rw_counter.expr(),
-            );
-        });
-    }
-
-    /// constraint step last rwc should match with chunkctx table value
-    pub(crate) fn step_last_constraint_rwc(&mut self) {
-        // Add first BeginTx step constraint to have tx_id == 1
-        self.step_last(|cb| {
-            cb.chunk_context_lookup(
-                ChunkCtxFieldTag::EndRWC,
-                cb.curr.state.rw_counter.expr() + cb.rw_counter_offset.clone(),
-            );
-        });
-    }
-
     // Validation
-
     pub(crate) fn validate_degree(&self, degree: usize, name: &'static str) {
         // We need to subtract IMPLICIT_DEGREE from MAX_DEGREE because all expressions
         // will be multiplied by state selector and q_step/q_step_first
@@ -1634,6 +1611,8 @@ impl<'a, F: Field> EVMConstraintBuilder<'a, F> {
     }
 
     pub(crate) fn add_lookup(&mut self, name: &str, lookup: Lookup<F>) {
+        // location selector need to be included in lookup expression
+        // which is not covered by cb.condition
         let lookup = match self.condition_expr_opt() {
             Some(condition) => lookup.conditional(condition),
             None => lookup,
@@ -1644,6 +1623,7 @@ impl<'a, F: Field> EVMConstraintBuilder<'a, F> {
             rlc::expr(&lookup.input_exprs(), self.challenges.lookup_input()),
             MAX_DEGREE - IMPLICIT_DEGREE,
         );
+
         self.store_expression(name, compressed_expr, CellType::Lookup(lookup.table()));
     }
 
@@ -1672,6 +1652,11 @@ impl<'a, F: Field> EVMConstraintBuilder<'a, F> {
                 let cell = self.query_cell_with_type(cell_type);
                 self.in_next_step = in_next_step;
 
+                // cb.step_XXXXXXX(|cb| {cb.context_lookup()})
+
+                // gate1: step_first_selector * (lookup_cell.expr() == by_pass_expr()) == 0
+                // lookup_gate = lookup(by_pass_expr())
+
                 // Require the stored value to equal the value of the expression
                 let name = format!("{} (stored expression)", name);
                 self.push_constraint(

zkevm-circuits/src/table/chunkctx_table.rs

@@ -100,6 +100,8 @@ impl ChunkCtxTable {
                         F::from(ChunkCtxFieldTag::EndRWC as u64),
                         F::from(chunkctx.end_rwc as u64),
                     ),
+                    // Empty row for disable lookup
+                    (F::ZERO, F::ZERO),
                 ]
                 .iter()
                 .map(|(tag, value)| {
@@ -122,6 +124,10 @@ impl ChunkCtxTable {
                     Ok(assigned_value)
                 })
                 .collect::<Result<Vec<AssignedCell<F, F>>, Error>>()?;
+
+                // remove last empty cell
+                let assigned_cells = assigned_cells.split_last().unwrap().1;
+
                 Ok(assigned_cells.iter().cloned().collect_tuple().unwrap())
             },
         )