perf: optimize etable assignment

DelphinusLab · Jun 3, 2024 · 54d30ef · 54d30ef
1 parent 0ab19c4
commit 54d30ef
Show file tree

Hide file tree

Showing 35 changed files with 100 additions and 63 deletions.
diff --git a/crates/specs/src/encode/memory_table.rs b/crates/specs/src/encode/memory_table.rs
@@ -1,17 +1,26 @@
-use num_bigint::ToBigUint;
+use num_bigint::BigUint;
+use num_traits::One;
+use static_assertions::const_assert;
 
 use super::FromBn;
 use crate::encode::COMMON_RANGE_OFFSET;
 
-pub fn encode_memory_table_entry<T: FromBn>(offset: T, location_type: T, is_i32: T) -> T {
-    const END_SHIFT: u32 = OFFSET_SHIFT + COMMON_RANGE_OFFSET;
-    const OFFSET_SHIFT: u32 = LOCATION_TYPE_SHIFT + COMMON_RANGE_OFFSET;
-    const LOCATION_TYPE_SHIFT: u32 = IS_I32_SHIFT + 1;
-    const IS_I32_SHIFT: u32 = 0;
+const _END_SHIFT: u32 = OFFSET_SHIFT + COMMON_RANGE_OFFSET;
+const OFFSET_SHIFT: u32 = LOCATION_TYPE_SHIFT + COMMON_RANGE_OFFSET;
+const LOCATION_TYPE_SHIFT: u32 = IS_I32_SHIFT + 1;
+const IS_I32_SHIFT: u32 = 0;
+
+const_assert!(_END_SHIFT < 240);
 
-    assert!(END_SHIFT < 240);
+lazy_static! {
+    pub static ref MEMORY_TABLE_ENTRY_OFFSET: BigUint = BigUint::one() << OFFSET_SHIFT;
+    pub static ref MEMORY_TABLE_ENTRY_LOCATION_TYPE: BigUint =
+        BigUint::one() << LOCATION_TYPE_SHIFT;
+    pub static ref MEMORY_TABLE_ENTRY_IS_I32: BigUint = BigUint::one() << IS_I32_SHIFT;
+}
 
-    offset * T::from_bn(&(1u64.to_biguint().unwrap() << OFFSET_SHIFT))
-        + location_type * T::from_bn(&(1u64.to_biguint().unwrap() << LOCATION_TYPE_SHIFT))
-        + is_i32 * T::from_bn(&(1u64.to_biguint().unwrap() << IS_I32_SHIFT))
+pub fn encode_memory_table_entry<T: FromBn>(offset: T, location_type: T, is_i32: T) -> T {
+    offset * T::from_bn(&MEMORY_TABLE_ENTRY_OFFSET)
+        + location_type * T::from_bn(&MEMORY_TABLE_ENTRY_LOCATION_TYPE)
+        + is_i32 * T::from_bn(&MEMORY_TABLE_ENTRY_IS_I32)
 }
diff --git a/crates/specs/src/imtable.rs b/crates/specs/src/imtable.rs
@@ -30,10 +30,6 @@ impl InitMemoryTable {
         Self(map)
     }
 
-    pub fn entries(&self) -> &HashMap<(LocationType, u32), InitMemoryTableEntry> {
-        &self.0
-    }
-
     pub fn to_string(&self) -> String {
         serde_json::to_string(&self.0).unwrap()
     }

diff --git a/crates/zkwasm/src/circuits/etable/assign.rs b/crates/zkwasm/src/circuits/etable/assign.rs
@@ -21,6 +21,7 @@ use super::EVENT_TABLE_ENTRY_ROWS;
 use crate::circuits::cell::CellExpression;
 use crate::circuits::jtable::FrameEtablePermutationCells;
 use crate::circuits::utils::bn_to_field;
+use crate::circuits::utils::step_status::FieldHelper;
 use crate::circuits::utils::step_status::Status;
 use crate::circuits::utils::step_status::StepStatus;
 use crate::circuits::utils::table_entry::EventTableWithMemoryInfo;
@@ -395,16 +396,19 @@ impl<F: FieldExt> EventTableChip<F> {
                 let mut ctx = Context::new(region);
                 ctx.offset = (chunk_size * chunk_index) * (EVENT_TABLE_ENTRY_ROWS as usize);
 
+                let mut field_helper = FieldHelper::default();
+
                 for (index, entry) in entries.iter().enumerate() {
                     let index = chunk_index * chunk_size + index;
 
                     let instruction = entry.eentry.get_instruction(itable);
 
-                    let step_status = StepStatus {
+                    let mut step_status = StepStatus {
                         current: &status[index],
                         next: &status[index + 1],
                         configure_table,
                         frame_table_returned_lookup: &frame_table_returned_lookup,
+                        field_helper: &mut field_helper,
                     };
 
                     {
@@ -438,7 +442,10 @@ impl<F: FieldExt> EventTableChip<F> {
 
                     {
                         let op_config = op_configs.get(&((&instruction.opcode).into())).unwrap();
-                        op_config.0.assign(&mut ctx, &step_status, &entry).unwrap();
+                        op_config
+                            .0
+                            .assign(&mut ctx, &mut step_status, &entry)
+                            .unwrap();
                     }
 
                     // Be careful, the function will step context.

diff --git a/crates/zkwasm/src/circuits/etable/mod.rs b/crates/zkwasm/src/circuits/etable/mod.rs
@@ -121,7 +121,7 @@ pub trait EventTableOpcodeConfig<F: FieldExt> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error>;
     fn memory_writing_ops(&self, _: &EventTableEntry) -> u32 {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_bin.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_bin.rs
@@ -380,7 +380,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BinConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         let (class, var_type, shift, left, right, value) = match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_bin_bit.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_bin_bit.rs
@@ -123,7 +123,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BinBitConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         let (class, vtype, left, right, value) = match entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_bin_shift.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_bin_shift.rs
@@ -344,7 +344,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BinShiftConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         let (class, left, right, value, power, is_eight_bytes, _is_sign) =

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_br.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_br.rs
@@ -90,7 +90,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_br_if.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_br_if.rs
@@ -140,7 +140,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrIfConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {
@@ -199,8 +199,10 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrIfConfig<F> {
                 }
 
                 self.cond_cell.assign(ctx, cond)?;
-                self.cond_inv_cell
-                    .assign(ctx, F::from(cond).invert().unwrap_or(F::zero()))?;
+                if cond != 0 {
+                    self.cond_inv_cell
+                        .assign(ctx, step.field_helper.invert(cond))?;
+                }
                 self.cond_is_zero_cell
                     .assign(ctx, if cond == 0 { F::one() } else { F::zero() })?;
                 self.cond_is_not_zero_cell

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_br_if_eqz.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_br_if_eqz.rs
@@ -136,7 +136,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrIfEqzConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {
@@ -193,8 +193,10 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrIfEqzConfig<F> {
                     }
                 }
 
-                self.cond_inv_cell
-                    .assign(ctx, F::from(cond).invert().unwrap_or(F::zero()))?;
+                if cond != 0 {
+                    self.cond_inv_cell
+                        .assign(ctx, step.field_helper.invert(cond))?;
+                }
                 self.cond_is_zero_cell
                     .assign(ctx, if cond == 0 { F::one() } else { F::zero() })?;
                 self.cond_is_not_zero_cell

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_br_table.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_br_table.rs
@@ -177,7 +177,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for BrTableConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_call.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_call.rs
@@ -74,7 +74,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for CallConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_call_host_foreign_circuit.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_call_host_foreign_circuit.rs
@@ -117,7 +117,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for ExternalCallHostCircuitConfig<F>
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_call_indirect.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_call_indirect.rs
@@ -128,7 +128,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for CallIndirectConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_const.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_const.rs
@@ -73,7 +73,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for ConstConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {
@@ -94,7 +94,6 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for ConstConfig<F> {
             }
             StepInfo::I64Const { value } => {
                 self.value.assign(ctx, *value as u64)?;
-                self.is_i32.assign(ctx, F::zero())?;
                 self.memory_table_lookup_stack_write.assign(
                     ctx,
                     step.current.eid,

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_conversion.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_conversion.rs
@@ -219,7 +219,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for ConversionConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         let (is_sign_op, value, value_type, result, result_type, padding, shift) =

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_drop.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_drop.rs
@@ -40,7 +40,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for DropConfig {
     fn assign(
         &self,
         _: &mut Context<'_, F>,
-        _: &StepStatus,
+        _: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_global_get.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_global_get.rs
@@ -82,7 +82,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for GlobalGetConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_global_set.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_global_set.rs
@@ -82,7 +82,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for GlobalSetConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_load.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_load.rs
@@ -426,7 +426,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for LoadConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_local_get.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_local_get.rs
@@ -90,7 +90,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for LocalGetConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_local_set.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_local_set.rs
@@ -90,7 +90,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for LocalSetConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_local_tee.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_local_tee.rs
@@ -90,7 +90,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for LocalTeeConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_memory_grow.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_memory_grow.rs
@@ -119,7 +119,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for MemoryGrowConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_memory_size.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_memory_size.rs
@@ -64,7 +64,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for MemorySizeConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_rel.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_rel.rs
@@ -333,7 +333,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for RelConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         let (class, var_type, lhs, rhs, value, diff) = match entry.eentry.step_info {
@@ -413,8 +413,9 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for RelConfig<F> {
             .assign(ctx, rhs.into(), var_type == VarType::I32, op_is_sign)?;
         self.diff.assign(ctx, diff.into())?;
 
-        self.diff_inv
-            .assign(ctx, F::from(diff).invert().unwrap_or(F::zero()))?;
+        if diff != 0 {
+            self.diff_inv.assign(ctx, step.field_helper.invert(diff))?;
+        }
         {
             self.res_is_eq.assign_bool(ctx, lhs == rhs)?;
             self.res_is_gt.assign_bool(ctx, lhs > rhs)?;

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_return.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_return.rs
@@ -131,7 +131,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for ReturnConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_select.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_select.rs
@@ -144,7 +144,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for SelectConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match &entry.eentry.step_info {
@@ -158,8 +158,9 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for SelectConfig<F> {
                 self.val1.assign(ctx, *val1)?;
                 self.val2.assign(ctx, *val2)?;
                 self.cond.assign(ctx, *cond)?;
-                self.cond_inv
-                    .assign(ctx, F::from(*cond).invert().unwrap_or(F::zero()))?;
+                if *cond != 0 {
+                    self.cond_inv.assign(ctx, step.field_helper.invert(*cond))?;
+                }
                 self.res.assign(ctx, *result)?;
                 self.is_i32.assign_bool(ctx, *vtype == VarType::I32)?;
 

diff --git a/crates/zkwasm/src/circuits/etable/op_configure/op_store.rs b/crates/zkwasm/src/circuits/etable/op_configure/op_store.rs
@@ -450,7 +450,7 @@ impl<F: FieldExt> EventTableOpcodeConfig<F> for StoreConfig<F> {
     fn assign(
         &self,
         ctx: &mut Context<'_, F>,
-        step: &StepStatus,
+        step: &mut StepStatus<F>,
         entry: &EventTableEntryWithMemoryInfo,
     ) -> Result<(), Error> {
         match entry.eentry.step_info {