zkEVM Cell Manager refactoring

There are currently two cell managers (abbr. CM) in zkEVM-ce, one in the EVM circuit and other in the Keccak circuit.

Issue https://github.com/privacy-scaling-explorations/zkevm-circuits/issues/1203

Commonalities

In this context a Cell is a pair (column: Column<Advice>, rotation: int). The cell manager can be queried for cells, and it never returns the same. This way each Cell can be used as a different "signal" in the circuit, with the actual underlaying layout in columns and rows of the signals being transparent for the user.

Hence, a cell manager has a method that returns a cell that is always different from previously returned cells.

Differeneces

Columns creation

In the EVM implementation columns are given to the constructor of the cell manager and are fixed, in the Keccak are created on demand when needed.

Columns types

In the EVM implementation, different Columns are of a type, and cells are requested to belong to a column of a specific type. This does not happen in the Keccak implementation.

pub(crate) enum CellType {
    StoragePhase1,
    StoragePhase2,
    StoragePermutation,
    LookupByte,
    Lookup(Table),
}

In the EVM cell manager we havefn query_cell(&mut self, cell_type: CellType) -> Cell<F> while in the keccak we have fn query_cell(&mut self, meta: &mut ConstraintSystem<F>).

Placement strategy

The EVM circuit, will place the requested cell in a column of the corresponding type where it will have the lowest rotation (height). Basically it searches for the column with the lowest current lower cell count for a specific column type. There is a maximum height, the manager will panic if a cell is requested but there is no space for it below that height.

fn next_column(&self, cell_type: CellType) -> usize {
        let mut best_index: Option<usize> = None;
        let mut best_height = self.height;
        for column in self.columns.iter() {
            if column.cell_type == cell_type && column.height < best_height {
                best_index = Some(column.index);
                best_height = column.height;
            }
        }
        // Replace a CellType::Storage by CellType::StoragePermutation if the later has
        // better height
        if cell_type == CellType::StoragePhase1 {
            for column in self.columns.iter() {
                if column.cell_type == CellType::StoragePermutation && column.height < best_height {
                    best_index = Some(column.index);
                    best_height = column.height;
                }
            }
        }
        match best_index {
            Some(index) => index,
            None => unreachable!("not enough cells for query: {:?}", cell_type),
        }
    }

The Keccak has a predefined number of rows (a bit like max height in the EVM version, it will try to find the rows with less columns used), and return a cell in that row and the next column that is unused. If the column does not exist, it will be created. Interestingly, also there is a method that allow the user to get a cell in a particular row, this complicates things for a unification.

Summary

Proposed design

It is clear that internally there should be two strategies to create new cells, but can we provide a common interface and common internal structures?

A common interface could be:

    fn query_top_aligned_cell(&mut self, cell_type: CellType) -> Cell<F>
    fn query_cell(&mut self, cell_type: CellType) -> Cell<F>
    fn query_cells(&mut self, cell_type: CellType, count: usize) -> Vec<Cell<F>>
    fn columns(&self) -> &[CellColumn<F>]
    fn get_width(&self) -> usize
    fn get_height(&self) -> usize

Providing that in the keccak circuit will provide the same cell type always. The important thing is that the interface will be the same to request cells. This is complicated a but by the type of cell request where the row is chosen by the user in the keccak circuit.

Internally, the CM will contain a structure by cell type and inside a vector of the columns in each cell type.

pub struct CellManagerColumns {
    columns: HashMap<CellType, vec<Column>>   
}

impl CellManagerColumns {
    pub add_column(&mut self, CellType, Column<Advice>) { ... }
    pub get_cell_type_width(&self, CellType) -> u32 { ... }
    pub get_column(&self, CellType, column_idx: u32) -> &Column<Advice> { ... }
    fn columns(&self) -> &[CellColumn<F>] { ... }
    pub get_width(&self) -> usize { ... }
}

pub struct CellManager {
    columns: CellManagerColumns
    strategy: CellManagerStrategy
}

impl CellManager {
    new(strategy: CellManagerStrategy) -> Self { ... }
    
    fn query_cell(&mut self, cell_type) -> Cell<F> { ... }
    fn get_width(&self) -> usize { ... }
    fn get_height(&self) -> usize { ... }
}

For each strategy we will have a type to manage it, that implements a trait:

trait CellManagerStrategy {
    on_creation(&mut self, &mut columns: CellManagerColumns)
    query_cell(&mut self, &mut coulmns: CellManagerColumns, cell_type)
    fn get_height(&self) -> usize
}

Note that the Cell Manager can calculate the width independently of the strategy.

We would to implement a FixedWidthStrategy:

struct FixedWidthStrategy {
    height: HashMap<CellType, vec<int>>
}

impl FixedWidthStrategy {
    new(vec<Column<Advice>>, HashMap<CellType, int>, max_height: u32) -> Self
}

impl CellManagerStrategy for FixedWidthStrategy {
    ...
}

And implement a FixedHeightStrategy:

struct FixedHeightStrategy {
    row_width: vec<int>
    cell_type: CellType
}

impl FixedHeightStrategy {
    new(height: u32, cell_type: CellType) -> Self
}

impl CellManagerStrategy for FixedHeightStrategy {
    ...
}

Problem with query by row

We see that the keccak strategy (in our new system FixedHeightStrategy) has a particularity that you can request a particular row cell. We should see in the future why is that, as this seems against the concept of the transparency of the CellManager, and posibly eliminate it.

In the meantime, we can add a method for this in the CellManager and in the CellManagerStrategy, and throw a panic if called in FixedWidthStrategy, as it does not make sense in that case.

General difficulties with this refactoring

The CellManager query method is called in many places in the opcode gadget implementations, and there is many of them.

A good strategy could be to don't change the interface of the methods that are called from this gadgets. A particular point is painful here, the fields on the Cell struct are public currently, ideally we would encapsulate it with methods, but that would require to change all the accesses to them in the opcode gadgets which is impractical.