From 2addd74ed478271bfbf248dc9baa91beb62e22dc Mon Sep 17 00:00:00 2001
From: Cecilia Zhang <zhangyixin319@gmail.com>
Date: Thu, 29 Jun 2023 09:56:44 +0800
Subject: [PATCH 1/2] cm with ctx

---
 .../src/circuit_tools/cell_manager.rs         | 95 ++++++++++++++++++-
 1 file changed, 90 insertions(+), 5 deletions(-)

diff --git a/zkevm-circuits/src/circuit_tools/cell_manager.rs b/zkevm-circuits/src/circuit_tools/cell_manager.rs
index 9e75bcd466..68d8372e58 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::{cmp::Ordering, collections::{BTreeMap, HashMap}, fmt::Debug, hash::Hash, cmp::{max, Ordering}};
 
 #[derive(Clone, Debug, Default)]
 pub(crate) struct Cell<F> {
@@ -208,11 +209,10 @@ pub(crate) struct CellColumn<F, C: CellType> {
     pub(crate) expr: Expression<F>,
 }
 
+
 impl<F: Field, C: CellType> PartialEq for CellColumn<F, C> {
     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<F: Field, C: CellType> Expr<F> for CellColumn<F, C> {
     }
 }
 
+
 #[derive(Clone, Debug)]
 pub struct CellManager<F, C: CellType> {
     configs: Vec<CellConfig<C>>,
     columns: Vec<CellColumn<F, C>>,
     height: usize,
+    width: usize,
     height_limit: usize,
+
+    // branch ctxs
+    branch_ctxs: HashMap<String, CmContext<F, C>>,
+    parent_ctx: Option<CmContext<F, C>>,
+}
+
+
+#[derive(Default, Clone, Debug)]
+struct CmContext<F, C: CellType>{
+    parent: Box<Option<CmContext<F, C>>>,
+    columns: Vec<CellColumn<F, C>>,
 }
 
 impl<F: Field, C: CellType> CellManager<F, C> {
@@ -255,7 +268,8 @@ impl<F: Field, C: CellType> CellManager<F, C> {
             .into_iter()
             .map(|c| c.into())
             .collect::<Vec<CellConfig<C>>>();
-
+        
+        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<F: Field, C: CellType> CellManager<F, C> {
                     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<Cell<F>> {
         let mut cells = Vec::with_capacity(count);
         while cells.len() < count {

From 7cbb5b95043cc626cd72d726a41ac250c6753ead Mon Sep 17 00:00:00 2001
From: Cecilia Zhang <zhangyixin319@gmail.com>
Date: Wed, 5 Jul 2023 14:23:05 +0800
Subject: [PATCH 2/2] update

---
 zkevm-circuits/src/circuit_tools/cell_manager.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/zkevm-circuits/src/circuit_tools/cell_manager.rs b/zkevm-circuits/src/circuit_tools/cell_manager.rs
index 68d8372e58..59f3d027ba 100644
--- a/zkevm-circuits/src/circuit_tools/cell_manager.rs
+++ b/zkevm-circuits/src/circuit_tools/cell_manager.rs
@@ -15,7 +15,7 @@ use halo2_proofs::{
     poly::Rotation,
 };
 use lazy_static::__Deref;
-use std::{cmp::Ordering, collections::{BTreeMap, HashMap}, fmt::Debug, hash::Hash, cmp::{max, Ordering}};
+use std::{collections::{BTreeMap, HashMap}, fmt::Debug, hash::Hash, cmp::{max, Ordering}};
 
 #[derive(Clone, Debug, Default)]
 pub(crate) struct Cell<F> {