[MPT] Add minimal RLP encoding check (privacy-scaling-explorations#1588)

### Description

Some small improvements to the MPT circuit. Most notably:
- Added a check that the RLP encoding of the RLP values are minimal (no
leading zeros) in the main RLP value decoding unit
- Fixed bug: byte lookups were not generated correctly
- Some other misc small improvements

### Issue Link

[_link issue here_]

### Type of change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update

### Contents

- [_item_]

### Rationale

[_design decisions and extended information_]

### How Has This Been Tested?

- All MPT tests run
- make test-all

Author: Brechtpd

README.md

@@ -12,12 +12,12 @@ To run the same tests as the CI, please use: `make test-all`.
 ## Running benchmarks
 
 There are currently several benchmarks to run in the workspace in regards to the circuits.
-All use the `DEGREE` env var to specify the degree of the `K` parameter that you want 
+All use the `DEGREE` env var to specify the degree of the `K` parameter that you want
 to use for your circuit in the bench process.
 -   Keccak Circuit prover benches. -> `DEGREE=16 make packed_multi_keccak_bench`
 -   EVM Circuit prover benches. -> `DEGREE=18 make evm_bench`.
 -   State Circuit prover benches. -> `DEGREE=18 make state_bench`
--   MPT Circuit prover benches. -> `DEGREE=14 make mpt_bench`
+-   MPT Circuit prover benches. -> `DEGREE=15 make mpt_bench`
 
 You can also run all benchmarks by running: `make circuit_benches DEGREE=18`.
 

zkevm-circuits/src/circuit_tools/cached_region.rs

@@ -10,7 +10,10 @@ use std::{
     hash::{Hash, Hasher},
 };
 
-use super::{cell_manager::CellType, constraint_builder::ConstraintBuilder};
+use super::{
+    cell_manager::{CellColumn, CellType},
+    constraint_builder::ConstraintBuilder,
+};
 
 pub trait ChallengeSet<F: Field> {
     fn indexed(&self) -> Vec<&Value<F>>;
@@ -64,13 +67,28 @@ impl<'r, 'b, F: Field> CachedRegion<'r, 'b, F> {
     ) -> Result<(), Error> {
         for (offset, region_id) in self.regions.clone() {
             for stored_expression in cb.get_stored_expressions(region_id).iter() {
-                // println!("stored expression: {}", stored_expression.name);
                 stored_expression.assign(self, challenges, offset)?;
             }
         }
         Ok(())
     }
 
+    pub(crate) fn annotate_columns<C: CellType>(&mut self, cell_columns: &[CellColumn<F, C>]) {
+        for c in cell_columns {
+            self.region.name_column(
+                || {
+                    format!(
+                        "{:?} {:?}: {:?} queried",
+                        c.cell_type.clone(),
+                        c.index,
+                        c.height
+                    )
+                },
+                c.column,
+            );
+        }
+    }
+
     /// Assign an advice column value (witness).
     pub fn assign_advice<'v, V, VR, A, AR>(
         &'v mut self,

zkevm-circuits/src/circuit_tools/cell_manager.rs

@@ -1,10 +1,11 @@
 //! Cell manager
+use super::constraint_builder::ConstraintBuilder;
 use crate::{
     circuit_tools::cached_region::CachedRegion,
+    evm_circuit::util::rlc,
+    table::LookupTable,
     util::{query_expression, word::Word, Expr},
 };
-
-use crate::table::LookupTable;
 use eth_types::Field;
 use halo2_proofs::{
     circuit::{AssignedCell, Value},
@@ -113,47 +114,39 @@ impl<F: Field> WordCell<F> {
 #[derive(Clone, Debug)]
 pub struct CellConfig<C: CellType> {
     pub cell_type: C,
-    pub num_columns: usize,
     pub phase: u8,
     pub is_permute: bool,
 }
 
-impl<C: CellType> From<(C, usize, u8, bool)> for CellConfig<C> {
-    fn from((cell_type, num_columns, phase, is_permute): (C, usize, u8, bool)) -> Self {
+impl<C: CellType> CellConfig<C> {
+    fn new(cell_type: C, phase: u8, is_permute: bool) -> Self {
         Self {
             cell_type,
-            num_columns,
             phase,
             is_permute,
         }
     }
-}
 
-impl<C: CellType> CellConfig<C> {
-    pub fn init_columns<F: Field>(&self, meta: &mut ConstraintSystem<F>) -> Vec<Column<Advice>> {
-        let mut columns = Vec::with_capacity(self.num_columns);
-        for _ in 0..self.num_columns {
-            let tmp = match self.phase {
-                1 => meta.advice_column_in(FirstPhase),
-                2 => meta.advice_column_in(SecondPhase),
-                3 => meta.advice_column_in(ThirdPhase),
-                _ => unreachable!(),
-            };
-            columns.push(tmp);
-        }
+    pub fn init_column<F: Field>(&self, meta: &mut ConstraintSystem<F>) -> Column<Advice> {
+        let column = match self.phase {
+            0 => meta.advice_column_in(FirstPhase),
+            1 => meta.advice_column_in(SecondPhase),
+            2 => meta.advice_column_in(ThirdPhase),
+            _ => unreachable!(),
+        };
         if self.is_permute {
-            let _ = columns
-                .iter()
-                .map(|c| meta.enable_equality(*c))
-                .collect::<Vec<()>>();
+            meta.enable_equality(column);
         }
-        columns
+        column
     }
 }
 
 pub trait CellType:
     Clone + Copy + Debug + PartialEq + Eq + PartialOrd + Ord + Hash + Default
 {
+    /// This is the table type for lookups
+    type TableType: Clone + Copy + Debug + PartialEq + Eq + PartialOrd + Ord + Hash;
+
     fn byte_type() -> Option<Self>;
 
     // The phase that given `Expression` becomes evaluateable.
@@ -171,6 +164,12 @@ pub trait CellType:
     /// Return the storage phase of phase
     fn storage_for_phase(phase: u8) -> Self;
 
+    /// Creates a type from a unique id
+    fn create_type(id: usize) -> Self;
+
+    /// Returns the table type of the lookup (if it's a lookup)
+    fn lookup_table_type(&self) -> Option<Self::TableType>;
+
     /// Return the storage cell of the expression
     fn storage_for_expr<F: Field>(expr: &Expression<F>) -> Self {
         Self::storage_for_phase(Self::expr_phase::<F>(expr))
@@ -184,36 +183,48 @@ pub enum DefaultCellType {
     StoragePhase3,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum DefaultLookupType {}
+
 impl Default for DefaultCellType {
     fn default() -> Self {
         Self::StoragePhase1
     }
 }
 
 impl CellType for DefaultCellType {
+    type TableType = DefaultLookupType;
+
     fn byte_type() -> Option<Self> {
         Some(DefaultCellType::StoragePhase1)
     }
 
     fn storage_for_phase(phase: u8) -> Self {
-        // println!("phase: {}", phase);
         match phase {
-            1 => DefaultCellType::StoragePhase1,
-            2 => DefaultCellType::StoragePhase2,
-            3 => DefaultCellType::StoragePhase3,
+            0 => DefaultCellType::StoragePhase1,
+            1 => DefaultCellType::StoragePhase2,
+            2 => DefaultCellType::StoragePhase3,
             _ => unreachable!(),
         }
     }
+
+    fn create_type(_id: usize) -> Self {
+        unreachable!()
+    }
+
+    fn lookup_table_type(&self) -> Option<Self::TableType> {
+        None
+    }
 }
 
 #[derive(Clone, Debug)]
 pub(crate) struct CellColumn<F, C: CellType> {
     pub(crate) column: Column<Advice>,
-    index: usize,
     pub(crate) cell_type: C,
-    height: usize,
-    cells: Vec<Cell<F>>,
+    pub(crate) cells: Vec<Cell<F>>,
     pub(crate) expr: Expression<F>,
+    pub(super) height: usize,
+    pub(super) index: usize,
 }
 
 impl<F: Field, C: CellType> PartialEq for CellColumn<F, C> {
@@ -244,51 +255,74 @@ impl<F: Field, C: CellType> Expr<F> for CellColumn<F, C> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct CellManager<F, C: CellType> {
     configs: Vec<CellConfig<C>>,
     columns: Vec<CellColumn<F, C>>,
     height: usize,
     height_limit: usize,
+    offset: usize,
 }
 
 impl<F: Field, C: CellType> CellManager<F, C> {
-    pub(crate) fn new(
-        meta: &mut ConstraintSystem<F>,
-        configs: Vec<(C, usize, u8, bool)>,
-        offset: usize,
-        max_height: usize,
-    ) -> Self {
-        let configs = configs
-            .into_iter()
-            .map(|c| c.into())
-            .collect::<Vec<CellConfig<C>>>();
+    pub(crate) fn new(max_height: usize, offset: usize) -> Self {
+        Self {
+            configs: Vec::new(),
+            columns: Vec::new(),
+            height: max_height,
+            height_limit: max_height,
+            offset,
+        }
+    }
 
-        let mut columns = Vec::new();
-        for config in configs.iter() {
-            let cols = config.init_columns(meta);
-            for col in cols.iter() {
-                let mut cells = Vec::new();
-                for r in 0..max_height {
-                    query_expression(meta, |meta| {
-                        cells.push(Cell::new(meta, *col, offset + r));
-                    });
-                }
-                columns.push(CellColumn {
-                    column: *col,
-                    index: columns.len(),
-                    cell_type: config.cell_type,
-                    height: 0,
-                    expr: cells[0].expr(),
-                    cells,
+    pub(crate) fn add_columns(
+        &mut self,
+        meta: &mut ConstraintSystem<F>,
+        cb: &mut ConstraintBuilder<F, C>,
+        cell_type: C,
+        phase: u8,
+        permutable: bool,
+        num_columns: usize,
+    ) {
+        for _ in 0..num_columns {
+            // Add a column of the specified type
+            let config = CellConfig::new(cell_type, phase, permutable);
+            let col = config.init_column(meta);
+            let mut cells = Vec::new();
+            for r in 0..self.height_limit {
+                query_expression(meta, |meta| {
+                    cells.push(Cell::new(meta, col, self.offset + r));
                 });
             }
+            let column_expr = cells[0].expr();
+            self.columns.push(CellColumn {
+                column: col,
+                index: self.columns.len(),
+                cell_type: config.cell_type,
+                height: 0,
+                expr: column_expr.expr(),
+                cells,
+            });
+            self.configs.push(config);
+
+            // For cell types that are lookups, generate the lookup here
+            if let Some(table) = cell_type.lookup_table_type() {
+                cb.add_lookup(
+                    format!("{:?}", table),
+                    vec![column_expr.expr()],
+                    vec![rlc::expr(
+                        &cb.table(table),
+                        cb.lookup_challenge.clone().unwrap(),
+                    )],
+                );
+            }
         }
-        Self {
-            configs,
-            columns,
-            height: max_height,
-            height_limit: max_height,
+    }
+
+    pub(crate) fn restart(&mut self) {
+        self.height = self.height_limit;
+        for col in self.columns.iter_mut() {
+            col.height = 0;
         }
     }