diff --git a/zkevm-circuits/src/circuit_tools/cell_manager.rs b/zkevm-circuits/src/circuit_tools/cell_manager.rs index 9e75bcd466..59f3d027ba 100644 --- a/zkevm-circuits/src/circuit_tools/cell_manager.rs +++ b/zkevm-circuits/src/circuit_tools/cell_manager.rs @@ -14,7 +14,8 @@ use halo2_proofs::{ }, poly::Rotation, }; -use std::{cmp::Ordering, collections::BTreeMap, fmt::Debug, hash::Hash}; +use lazy_static::__Deref; +use std::{collections::{BTreeMap, HashMap}, fmt::Debug, hash::Hash, cmp::{max, Ordering}}; #[derive(Clone, Debug, Default)] pub(crate) struct Cell { @@ -208,11 +209,10 @@ pub(crate) struct CellColumn { pub(crate) expr: Expression, } + impl PartialEq for CellColumn { fn eq(&self, other: &Self) -> bool { - self.index == other.index - && self.cell_type == other.cell_type - && self.height == other.height + self.index == other.index && self.cell_type == other.cell_type && self.height == other.height } } @@ -236,12 +236,25 @@ impl Expr for CellColumn { } } + #[derive(Clone, Debug)] pub struct CellManager { configs: Vec>, columns: Vec>, height: usize, + width: usize, height_limit: usize, + + // branch ctxs + branch_ctxs: HashMap>, + parent_ctx: Option>, +} + + +#[derive(Default, Clone, Debug)] +struct CmContext{ + parent: Box>>, + columns: Vec>, } impl CellManager { @@ -255,7 +268,8 @@ impl CellManager { .into_iter() .map(|c| c.into()) .collect::>>(); - + + let mut width = 0; let mut columns = Vec::new(); for config in configs.iter() { let cols = config.init_columns(meta); @@ -274,16 +288,87 @@ impl CellManager { expr: cells[0].expr(), cells, }); + width += 1; } } Self { configs, columns, height: max_height, + width, height_limit: max_height, + branch_ctxs: HashMap::new(), + parent_ctx: None, } } + pub(crate) fn cur_to_parent(&mut self) { + let new_parent = match self.parent_ctx.clone() { + // if parent context exists, meaning we are deep in a callstack + // we set it as the parent of new parent + Some(ctx) => CmContext { + parent: Box::new(Some(ctx.clone())), + columns: self.columns.clone(), + }, + // otherwise, this is the fist level of callstack + // the parent of new parent is None + None => CmContext { + parent: Box::new(None), + columns: self.columns.clone(), + } + }; + self.parent_ctx = Some(new_parent); + self.reset(self.height_limit); + } + + pub(crate) fn cur_to_branch(&mut self, name: &str) { + let new_branch = match self.parent_ctx.clone() { + // if parent context exists, meaning we are deep in a callstack + // we set it as the parent of new branch + Some(ctx) => CmContext { + parent: Box::new(Some(ctx.clone())), + columns: self.columns.clone(), + }, + // otherwise, this is the fist level of callstack + // the parent of new branch is None + None => CmContext { + parent: Box::new(None), + columns: self.columns.clone(), + } + }; + self.branch_ctxs.insert(name.to_string(), new_branch); + self.reset(self.height_limit); + } + + pub(crate) fn recover_max_branch(&mut self) { + let mut new_cols = self.columns.clone(); + let parent = self.parent_ctx.clone().expect("Retruning context needs parent"); + self.branch_ctxs + .iter() + .for_each(|(name, ctx)| { + for c in 0..self.width { + new_cols[c] = max(&new_cols[c], &ctx.columns[c]).clone(); + new_cols[c] = max(&new_cols[c], &parent.columns[c]).clone(); + } + }); + self.columns = new_cols; + self.branch_ctxs.clear(); + self.parent_ctx = self.parent_ctx + .clone() + .map(|ctx| ctx.parent.deref().clone()) + .unwrap(); + } + + pub(crate) fn recover_parent(&mut self) { + assert!(self.parent_ctx.is_some(), "No parent context to recover"); + self.columns = self.parent_ctx.clone().unwrap().columns.clone(); + self.parent_ctx + .clone() + .map(|ctx| self.parent_ctx = ctx.parent.deref().clone()) + .unwrap(); + self.branch_ctxs.clear(); + } + pub(crate) fn query_cells(&mut self, cell_type: C, count: usize) -> Vec> { let mut cells = Vec::with_capacity(count); while cells.len() < count {