From 8e2cca402339f841508ce7f1305676e2b2f62c79 Mon Sep 17 00:00:00 2001
From: Zhang Junyu <junyu92@gmail.com>
Date: Mon, 1 Jul 2024 06:01:03 +0000
Subject: [PATCH] fix: fix memory table and frame table termination

---
 crates/zkwasm/src/circuits/etable/mod.rs    |  1 +
 crates/zkwasm/src/circuits/jtable/assign.rs | 30 +++++++++---------
 crates/zkwasm/src/circuits/mtable/assign.rs | 34 ++++++++++-----------
 crates/zkwasm/src/circuits/mtable/mod.rs    |  1 +
 4 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/crates/zkwasm/src/circuits/etable/mod.rs b/crates/zkwasm/src/circuits/etable/mod.rs
index 1ce6ccccc..10466a92d 100644
--- a/crates/zkwasm/src/circuits/etable/mod.rs
+++ b/crates/zkwasm/src/circuits/etable/mod.rs
@@ -82,6 +82,7 @@ pub struct EventTableCommonConfig<F: FieldExt> {
     enabled_cell: AllocatedBitCell<F>,
     ops: [AllocatedBitCell<F>; OP_CAPABILITY],
 
+    // It MUST NOT be unlimited to ensure termination.
     rest_mops_cell: AllocatedCommonRangeCell<F>,
     rest_call_ops_cell: AllocatedUnlimitedCell<F>,
     rest_return_ops_cell: AllocatedUnlimitedCell<F>,
diff --git a/crates/zkwasm/src/circuits/jtable/assign.rs b/crates/zkwasm/src/circuits/jtable/assign.rs
index 4eba11fbd..2f2e7588a 100644
--- a/crates/zkwasm/src/circuits/jtable/assign.rs
+++ b/crates/zkwasm/src/circuits/jtable/assign.rs
@@ -68,25 +68,23 @@ impl<F: FieldExt> JumpTableChip<F> {
                 )?;
             }
 
-            if i == capability - 1 {
-                ctx.region.assign_advice_from_constant(
-                    || "frame table: entry terminate",
-                    self.config.value,
-                    ctx.offset + FrameTableValueOffset::CallOps as usize,
-                    F::zero(),
-                )?;
-
-                ctx.region.assign_advice_from_constant(
-                    || "frame table: entry terminate",
-                    self.config.value,
-                    ctx.offset + FrameTableValueOffset::ReturnOps as usize,
-                    F::zero(),
-                )?;
-            }
-
             ctx.step(FrameTableValueOffset::Max as usize);
         }
 
+        ctx.region.assign_advice_from_constant(
+            || "frame table: entry terminate",
+            self.config.value,
+            ctx.offset + FrameTableValueOffset::CallOps as usize,
+            F::zero(),
+        )?;
+
+        ctx.region.assign_advice_from_constant(
+            || "frame table: entry terminate",
+            self.config.value,
+            ctx.offset + FrameTableValueOffset::ReturnOps as usize,
+            F::zero(),
+        )?;
+
         ctx.region.assign_fixed(
             || "frame table: inherited",
             self.config.inherited,
diff --git a/crates/zkwasm/src/circuits/mtable/assign.rs b/crates/zkwasm/src/circuits/mtable/assign.rs
index 5ffe3280b..227762b10 100644
--- a/crates/zkwasm/src/circuits/mtable/assign.rs
+++ b/crates/zkwasm/src/circuits/mtable/assign.rs
@@ -31,7 +31,7 @@ impl<F: FieldExt> MemoryTableChip<F> {
     fn assign_fixed(&self, ctx: &mut Context<'_, F>) -> Result<(), Error> {
         let capability = self.maximal_available_rows / MEMORY_TABLE_ENTRY_ROWS as usize;
 
-        for i in 0..capability {
+        for _ in 0..capability {
             ctx.region.assign_fixed(
                 || "mtable: sel",
                 self.config.entry_sel,
@@ -39,26 +39,24 @@ impl<F: FieldExt> MemoryTableChip<F> {
                 || Ok(F::one()),
             )?;
 
-            if i == capability - 1 {
-                ctx.region.assign_advice_from_constant(
-                    || "rest_mops terminate",
-                    self.config.rest_mops_cell.cell.col,
-                    ctx.offset + self.config.rest_mops_cell.cell.rot as usize,
-                    F::zero(),
-                )?;
-
-                #[cfg(feature = "continuation")]
-                ctx.region.assign_advice_from_constant(
-                    || "rest_memory_finalize_ops terminate",
-                    self.config.rest_memory_finalize_ops_cell.cell.col,
-                    ctx.offset + self.config.rest_memory_finalize_ops_cell.cell.rot as usize,
-                    F::zero(),
-                )?;
-            }
-
             ctx.step(MEMORY_TABLE_ENTRY_ROWS as usize);
         }
 
+        ctx.region.assign_advice_from_constant(
+            || "rest_mops terminate",
+            self.config.rest_mops_cell.cell.col,
+            ctx.offset + self.config.rest_mops_cell.cell.rot as usize,
+            F::zero(),
+        )?;
+
+        #[cfg(feature = "continuation")]
+        ctx.region.assign_advice_from_constant(
+            || "rest_memory_finalize_ops terminate",
+            self.config.rest_memory_finalize_ops_cell.cell.col,
+            ctx.offset + self.config.rest_memory_finalize_ops_cell.cell.rot as usize,
+            F::zero(),
+        )?;
+
         Ok(())
     }
 
diff --git a/crates/zkwasm/src/circuits/mtable/mod.rs b/crates/zkwasm/src/circuits/mtable/mod.rs
index 579e1689b..f23f36b65 100644
--- a/crates/zkwasm/src/circuits/mtable/mod.rs
+++ b/crates/zkwasm/src/circuits/mtable/mod.rs
@@ -47,6 +47,7 @@ pub struct MemoryTableConfig<F: FieldExt> {
     start_eid_cell: AllocatedU32StateCell<F>,
     end_eid_cell: AllocatedU32StateCell<F>,
     eid_diff_cell: AllocatedU32StateCell<F>,
+    // It MUST NOT be unlimited to ensure mtable terminates at capability.
     rest_mops_cell: AllocatedCommonRangeCell<F>,
     offset_cell: AllocatedU32Cell<F>,
     offset_diff_cell: AllocatedU32Cell<F>,