From b189aad7e02c22b7eea2a5a9d2b1d9526fb18ac9 Mon Sep 17 00:00:00 2001
From: Wu Sung-Ming <wusm@sea.com>
Date: Mon, 15 May 2023 17:57:53 +0800
Subject: [PATCH] better generic design

---
 keccak256/src/gate_helpers.rs                 |  2 +-
 zkevm-circuits/src/evm_circuit/execution.rs   |  4 +-
 .../src/evm_circuit/execution/add_sub.rs      |  5 +-
 .../src/evm_circuit/execution/addmod.rs       |  4 +-
 zkevm-circuits/src/evm_circuit/util.rs        |  4 +-
 .../src/evm_circuit/util/common_gadget.rs     | 53 ++++++++++---------
 .../evm_circuit/util/constraint_builder.rs    |  4 +-
 .../evm_circuit/util/math_gadget/abs_word.rs  |  5 +-
 .../evm_circuit/util/math_gadget/add_words.rs |  5 +-
 .../evm_circuit/util/math_gadget/cmp_words.rs | 22 ++++----
 .../util/math_gadget/is_equal_word.rs         | 14 ++---
 .../util/math_gadget/is_zero_word.rs          | 18 ++++---
 .../evm_circuit/util/math_gadget/lt_word.rs   |  2 +-
 .../evm_circuit/util/math_gadget/modulo.rs    | 12 ++---
 .../util/math_gadget/mul_add_words.rs         |  2 +-
 .../util/math_gadget/mul_add_words512.rs      |  2 +-
 .../util/math_gadget/mul_word_u64.rs          |  5 +-
 .../src/evm_circuit/util/memory_gadget.rs     |  2 +-
 zkevm-circuits/src/util.rs                    |  2 +
 zkevm-circuits/src/util/word.rs               | 40 ++++++++------
 20 files changed, 116 insertions(+), 91 deletions(-)

diff --git a/keccak256/src/gate_helpers.rs b/keccak256/src/gate_helpers.rs
index 1050f9c3ed..c14015fdc8 100644
--- a/keccak256/src/gate_helpers.rs
+++ b/keccak256/src/gate_helpers.rs
@@ -1,7 +1,7 @@
 use eth_types::Field;
 use num_bigint::BigUint;
 
-/// Convert a bigUint value to FieldExt
+/// Convert a bigUint value to Field
 ///
 /// We assume the input value is smaller than the field size
 pub fn biguint_to_f<F: Field>(x: &BigUint) -> F {
diff --git a/zkevm-circuits/src/evm_circuit/execution.rs b/zkevm-circuits/src/evm_circuit/execution.rs
index a7c889f3aa..310f32c075 100644
--- a/zkevm-circuits/src/evm_circuit/execution.rs
+++ b/zkevm-circuits/src/evm_circuit/execution.rs
@@ -197,7 +197,7 @@ pub(crate) trait ExecutionGadget<F: Field> {
 }
 
 #[derive(Clone, Debug)]
-pub(crate) struct ExecutionConfig<F, T> {
+pub(crate) struct ExecutionConfig<F> {
     // EVM Circuit selector, which enables all usable rows.  The rows where this selector is
     // disabled won't verify any constraint (they can be unused rows or rows with blinding
     // factors).
@@ -312,7 +312,7 @@ pub(crate) struct ExecutionConfig<F, T> {
     error_return_data_out_of_bound: Box<ErrorReturnDataOutOfBoundGadget<F>>,
 }
 
-impl<F: Field, T: WordExpr<F>> ExecutionConfig<F, T> {
+impl<F: Field> ExecutionConfig<F> {
     #[allow(clippy::too_many_arguments)]
     #[allow(clippy::redundant_closure_call)]
     pub(crate) fn configure(
diff --git a/zkevm-circuits/src/evm_circuit/execution/add_sub.rs b/zkevm-circuits/src/evm_circuit/execution/add_sub.rs
index a3dae3fca7..27c918f1fd 100644
--- a/zkevm-circuits/src/evm_circuit/execution/add_sub.rs
+++ b/zkevm-circuits/src/evm_circuit/execution/add_sub.rs
@@ -10,7 +10,10 @@ use crate::{
         },
         witness::{Block, Call, ExecStep, Transaction},
     },
-    util::{word::Word, Expr},
+    util::{
+        word::{Word, WordExpr},
+        Expr,
+    },
 };
 use bus_mapping::evm::OpcodeId;
 use eth_types::Field;
diff --git a/zkevm-circuits/src/evm_circuit/execution/addmod.rs b/zkevm-circuits/src/evm_circuit/execution/addmod.rs
index 75ff6d5679..4e60bca6d3 100644
--- a/zkevm-circuits/src/evm_circuit/execution/addmod.rs
+++ b/zkevm-circuits/src/evm_circuit/execution/addmod.rs
@@ -46,8 +46,8 @@ pub(crate) struct AddModGadget<F> {
     muladd_d_n_r: MulAddWords512Gadget<F>,
 
     n_is_zero: IsZeroWordGadget<F>,
-    cmp_r_n: CmpWordsGadget<F, Word32Cell<F>>,
-    cmp_areduced_n: CmpWordsGadget<F, Word32Cell<F>>,
+    cmp_r_n: CmpWordsGadget<F, Word32Cell<F>, Word32Cell<F>>,
+    cmp_areduced_n: CmpWordsGadget<F, Word32Cell<F>, Word32Cell<F>>,
 }
 
 impl<F: Field> ExecutionGadget<F> for AddModGadget<F> {
diff --git a/zkevm-circuits/src/evm_circuit/util.rs b/zkevm-circuits/src/evm_circuit/util.rs
index c731cbaf6d..ad1f1b309f 100644
--- a/zkevm-circuits/src/evm_circuit/util.rs
+++ b/zkevm-circuits/src/evm_circuit/util.rs
@@ -506,7 +506,7 @@ impl<F: Field, const N: usize> Expr<F> for RandomLinearCombination<F, N> {
 
 pub(crate) type MemoryAddress<F> = RandomLinearCombination<F, N_BYTES_MEMORY_ADDRESS>;
 
-impl<F: FieldExt> WordExpr<F> for MemoryAddress<F> {
+impl<F: Field> WordExpr<F> for MemoryAddress<F> {
     fn to_word(&self) -> Word<Expression<F>> {
         Word::from_lo_unchecked(self.expr())
     }
@@ -514,7 +514,7 @@ impl<F: FieldExt> WordExpr<F> for MemoryAddress<F> {
 
 pub(crate) type AccountAddress<F> = RandomLinearCombination<F, N_BYTES_ACCOUNT_ADDRESS>;
 
-impl<F: FieldExt> WordExpr<F> for AccountAddress<F> {
+impl<F: Field> WordExpr<F> for AccountAddress<F> {
     fn to_word(&self) -> Word<Expression<F>> {
         Word::from_lo_unchecked(self.expr())
     }
diff --git a/zkevm-circuits/src/evm_circuit/util/common_gadget.rs b/zkevm-circuits/src/evm_circuit/util/common_gadget.rs
index 682190caa7..1df3df2020 100644
--- a/zkevm-circuits/src/evm_circuit/util/common_gadget.rs
+++ b/zkevm-circuits/src/evm_circuit/util/common_gadget.rs
@@ -7,7 +7,7 @@ use super::{
 };
 use crate::{
     evm_circuit::{
-        param::{N_BYTES_GAS, N_BYTES_MEMORY_WORD_SIZE},
+        param::{N_BYTES_ACCOUNT_ADDRESS, N_BYTES_GAS, N_BYTES_MEMORY_WORD_SIZE},
         step::ExecutionState,
         table::{FixedTableTag, Lookup},
         util::{
@@ -97,7 +97,7 @@ pub(crate) struct RestoreContextGadget<F> {
     caller_id: Cell<F>,
     caller_is_root: Cell<F>,
     caller_is_create: Cell<F>,
-    caller_code_hash: Cell<F>,
+    caller_code_hash: WordCell<F>,
     caller_program_counter: Cell<F>,
     caller_stack_pointer: Cell<F>,
     caller_gas_left: Cell<F>,
@@ -118,11 +118,10 @@ impl<F: Field> RestoreContextGadget<F> {
     ) -> Self {
         // Read caller's context for restore
         let caller_id = cb.call_context(None, CallContextFieldTag::CallerId);
-        let [caller_is_root, caller_is_create, caller_code_hash, caller_program_counter, caller_stack_pointer, caller_gas_left, caller_memory_word_size, caller_reversible_write_counter] =
+        let [caller_is_root, caller_is_create, caller_program_counter, caller_stack_pointer, caller_gas_left, caller_memory_word_size, caller_reversible_write_counter] =
             [
                 CallContextFieldTag::IsRoot,
                 CallContextFieldTag::IsCreate,
-                CallContextFieldTag::CodeHash,
                 CallContextFieldTag::ProgramCounter,
                 CallContextFieldTag::StackPointer,
                 CallContextFieldTag::GasLeft,
@@ -131,6 +130,9 @@ impl<F: Field> RestoreContextGadget<F> {
             ]
             .map(|field_tag| cb.call_context(Some(caller_id.expr()), field_tag));
 
+        let caller_code_hash =
+            cb.call_context_read_as_word(Some(caller_id.expr()), CallContextFieldTag::CodeHash);
+
         // Update caller's last callee information
         // EIP-211 CREATE/CREATE2 call successful case should set RETURNDATASIZE = 0
         let is_call_create_and_success_expr = cb.curr.state.is_create.expr() * is_success.clone();
@@ -156,8 +158,8 @@ impl<F: Field> RestoreContextGadget<F> {
                 ),
             ),
         ] {
-            cb.call_context_lookup(
-                true.expr(),
+            // TODO review and assure range check
+            cb.call_context_lookup_write_unchecked(
                 Some(caller_id.expr()),
                 field_tag,
                 Word::from_lo_unchecked(value),
@@ -195,7 +197,7 @@ impl<F: Field> RestoreContextGadget<F> {
             call_id: To(caller_id.expr()),
             is_root: To(caller_is_root.expr()),
             is_create: To(caller_is_create.expr()),
-            code_hash: To(caller_code_hash.expr()),
+            code_hash: To(caller_code_hash.to_word()),
             program_counter: To(caller_program_counter.expr()),
             stack_pointer: To(caller_stack_pointer.expr()),
             gas_left: To(gas_left),
@@ -260,7 +262,7 @@ impl<F: Field> RestoreContextGadget<F> {
         }
 
         self.caller_code_hash
-            .assign(region, offset, region.word_rlc(caller_code_hash))?;
+            .assign(region, offset, Word::from_u256(caller_code_hash))?;
 
         Ok(())
     }
@@ -364,7 +366,7 @@ pub(crate) struct TransferWithGasFeeGadget<F> {
     receiver: UpdateBalanceGadget<F, 2, true>,
     receiver_exists: Expression<F>,
     must_create: Expression<F>,
-    pub(crate) value_is_zero: IsZeroWordGadget<F>,
+    pub(crate) value_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
 }
 
 impl<F: Field> TransferWithGasFeeGadget<F> {
@@ -496,7 +498,7 @@ impl<F: Field> TransferWithGasFeeGadget<F> {
 pub(crate) struct TransferGadget<F> {
     sender: UpdateBalanceGadget<F, 2, false>,
     receiver: UpdateBalanceGadget<F, 2, true>,
-    pub(crate) value_is_zero: IsZeroWordGadget<F>,
+    pub(crate) value_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
 }
 
 impl<F: Field> TransferGadget<F> {
@@ -594,12 +596,12 @@ pub(crate) struct CommonCallGadget<F, const IS_SUCCESS_CALL: bool> {
     pub rd_address: MemoryAddressGadget<F>,
     pub memory_expansion: MemoryExpansionGadget<F, 2, N_BYTES_MEMORY_WORD_SIZE>,
 
-    value_is_zero: IsZeroWordGadget<F>,
+    value_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
     pub has_value: Expression<F>,
-    pub callee_code_hash: Word32Cell<F>,
-    pub is_empty_code_hash: IsEqualWordGadget<F, Word<Expression<F>>>,
+    pub callee_code_hash: WordCell<F>,
+    pub is_empty_code_hash: IsEqualWordGadget<F, WordCell<F>, Word<Expression<F>>>,
 
-    pub callee_not_exists: IsZeroWordGadget<F>,
+    pub callee_not_exists: IsZeroWordGadget<F, WordCell<F>>,
 }
 
 impl<F: Field, const IS_SUCCESS_CALL: bool> CommonCallGadget<F, IS_SUCCESS_CALL> {
@@ -664,14 +666,14 @@ impl<F: Field, const IS_SUCCESS_CALL: bool> CommonCallGadget<F, IS_SUCCESS_CALL>
             1.expr() - value_is_zero.expr(),
         );
 
-        let callee_code_hash = cb.query_word32();
+        let callee_code_hash = cb.query_word_unchecked();
         cb.account_read(
             callee_address_word.expr(),
             AccountFieldTag::CodeHash,
             callee_code_hash.to_word(),
         );
         let is_empty_code_hash =
-            IsEqualWordGadget::construct(cb, callee_code_hash.to_word(), cb.empty_code_hash_word());
+            IsEqualWordGadget::construct(cb, callee_code_hash, cb.empty_code_hash_word());
         let callee_not_exists = IsZeroWordGadget::construct(cb, callee_code_hash);
 
         Self {
@@ -732,8 +734,13 @@ impl<F: Field, const IS_SUCCESS_CALL: bool> CommonCallGadget<F, IS_SUCCESS_CALL>
         callee_code_hash: U256,
     ) -> Result<u64, Error> {
         self.gas.assign(region, offset, Some(gas.to_le_bytes()))?;
-        self.callee_address
-            .assign(region, offset, Some(callee_address.to_le_bytes()))?;
+        self.callee_address.assign(
+            region,
+            offset,
+            callee_address.to_le_bytes()[0..N_BYTES_ACCOUNT_ADDRESS]
+                .try_into()
+                .unwrap(),
+        )?;
         self.value
             .assign(region, offset, Some(value.to_le_bytes()))?;
         if IS_SUCCESS_CALL {
@@ -971,15 +978,9 @@ impl<F: Field> CommonErrorGadget<F> {
         let rw_counter_end_of_reversion = cb.query_word_unchecked(); // rw_counter_end_of_reversion just used for read lookup, therefore skip range check
 
         // current call must be failed.
-        cb.call_context_lookup(
-            false.expr(),
-            None,
-            CallContextFieldTag::IsSuccess,
-            Word::zero(),
-        );
+        cb.call_context_lookup_read(None, CallContextFieldTag::IsSuccess, Word::zero());
 
-        cb.call_context_lookup(
-            false.expr(),
+        cb.call_context_lookup_read(
             None,
             CallContextFieldTag::RwCounterEndOfReversion,
             rw_counter_end_of_reversion.to_word(),
diff --git a/zkevm-circuits/src/evm_circuit/util/constraint_builder.rs b/zkevm-circuits/src/evm_circuit/util/constraint_builder.rs
index 74ce952ba3..067c331fe8 100644
--- a/zkevm-circuits/src/evm_circuit/util/constraint_builder.rs
+++ b/zkevm-circuits/src/evm_circuit/util/constraint_builder.rs
@@ -11,7 +11,7 @@ use crate::{
     },
     util::{
         build_tx_log_expression,
-        word::{Word, Word16, Word32, Word32Cell, Word4, WordExpr, WordLimbs},
+        word::{Word, Word16, Word32, Word32Cell, Word4, WordCell, WordExpr, WordLimbs},
         Challenges, Expr,
     },
 };
@@ -423,7 +423,7 @@ impl<'a, F: Field> EVMConstraintBuilder<'a, F> {
     // }
 
     // default query_word is 2 limbs. Each limb is not guaranteed to be 128 bits.
-    pub fn query_word_unchecked<const N: usize>(&mut self) -> Word<Cell<F>> {
+    pub fn query_word_unchecked<const N: usize>(&mut self) -> WordCell<F> {
         Word::new(
             self.query_cells(CellType::StoragePhase1, N)
                 .try_into()
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/abs_word.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/abs_word.rs
index 1f7b8137f9..d9ee02b93f 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/abs_word.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/abs_word.rs
@@ -4,7 +4,10 @@ use crate::{
         math_gadget::*,
         CachedRegion,
     },
-    util::{word::Word32Cell, Expr},
+    util::{
+        word::{Word32Cell, WordExpr},
+        Expr,
+    },
 };
 use eth_types::{Field, ToLittleEndian, Word};
 use gadgets::util::sum;
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/add_words.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/add_words.rs
index 0e831ad9b0..872eca30a1 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/add_words.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/add_words.rs
@@ -4,7 +4,10 @@ use crate::{
         constraint_builder::{ConstrainBuilderCommon, EVMConstraintBuilder},
         pow_of_two_expr, split_u256, sum, CachedRegion, Cell,
     },
-    util::{word::Word32Cell, Expr},
+    util::{
+        word::{Word32Cell, WordExpr},
+        Expr,
+    },
 };
 use eth_types::{Field, ToLittleEndian, ToScalar, Word};
 use halo2_proofs::{circuit::Value, plonk::Error};
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/cmp_words.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/cmp_words.rs
index 894a1ef978..e22268dede 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/cmp_words.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/cmp_words.rs
@@ -11,16 +11,16 @@ use halo2_proofs::plonk::{Error, Expression};
 
 #[derive(Clone, Debug)]
 /// CmpWordsGadget compares two words, exposing `eq`  and `lt`
-pub(crate) struct CmpWordsGadget<F, T> {
+pub(crate) struct CmpWordsGadget<F, T1, T2> {
     comparison_lo: ComparisonGadget<F, 16>,
     comparison_hi: ComparisonGadget<F, 16>,
     pub eq: Expression<F>,
     pub lt: Expression<F>,
-    marker: PhantomData<T>,
+    _marker: PhantomData<(T1, T2)>,
 }
 
-impl<F: Field, T: WordExpr<F>> CmpWordsGadget<F, T> {
-    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, a: T, b: T) -> Self {
+impl<F: Field, T1: WordExpr<F>, T2: WordExpr<F>> CmpWordsGadget<F, T1, T2> {
+    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, a: T1, b: T2) -> Self {
         let (a_lo, a_hi) = a.to_word().to_lo_hi();
         let (b_lo, b_hi) = b.to_word().to_lo_hi();
         // `a.lo <= b.lo`
@@ -45,7 +45,7 @@ impl<F: Field, T: WordExpr<F>> CmpWordsGadget<F, T> {
             comparison_hi,
             lt,
             eq,
-            marker: Default::default(),
+            _marker: Default::default(),
         }
     }
 
@@ -85,14 +85,14 @@ mod tests {
 
     #[derive(Clone)]
     /// CmpWordGadgetTestContainer: require(a == b if CHECK_EQ else a < b)
-    struct CmpWordGadgetTestContainer<F, const CHECK_EQ: bool, T: WordExpr<F>> {
-        cmp_gadget: CmpWordsGadget<F, T>,
-        a_word: T,
-        b_word: T,
+    struct CmpWordGadgetTestContainer<F, const CHECK_EQ: bool, T1, T2> {
+        cmp_gadget: CmpWordsGadget<F, T1, T2>,
+        a_word: T1,
+        b_word: T2,
     }
 
-    impl<F: Field, const CHECK_EQ: bool, T: WordExpr<F>> MathGadgetContainer<F>
-        for CmpWordGadgetTestContainer<F, CHECK_EQ, T>
+    impl<F: Field, const CHECK_EQ: bool, T1: WordExpr<F> + Clone, T2: WordExpr<F> + Clone>
+        MathGadgetContainer<F> for CmpWordGadgetTestContainer<F, CHECK_EQ, T1, T2>
     {
         fn configure_gadget_container(cb: &mut EVMConstraintBuilder<F>) -> Self {
             let a_word = cb.query_word_unchecked();
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/is_equal_word.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/is_equal_word.rs
index 14b4a2e5d3..8d33e232c6 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/is_equal_word.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/is_equal_word.rs
@@ -18,23 +18,23 @@ use super::IsZeroGadget;
 
 /// Returns `1` when `lhs == rhs`, and returns `0` otherwise.
 #[derive(Clone, Debug)]
-pub struct IsEqualWordGadget<F, T> {
+pub struct IsEqualWordGadget<F, T1, T2> {
     is_zero_lo: IsZeroGadget<F>,
     is_zero_hi: IsZeroGadget<F>,
-    marker: PhantomData<T>,
+    _marker: PhantomData<(T1, T2)>,
 }
 
-impl<F: Field, T: WordExpr<F>> IsEqualWordGadget<F, T> {
-    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, lhs: T, rhs: T) -> Self {
-        let (lhs_lo, lhs_hi) = lhs.to_lo_hi();
-        let (rhs_lo, rhs_hi) = rhs.to_lo_hi();
+impl<F: Field, T1: WordExpr<F>, T2: WordExpr<F>> IsEqualWordGadget<F, T1, T2> {
+    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, lhs: T1, rhs: T2) -> Self {
+        let (lhs_lo, lhs_hi) = lhs.to_word().to_lo_hi();
+        let (rhs_lo, rhs_hi) = rhs.to_word().to_lo_hi();
         let is_zero_lo = IsZeroGadget::construct(cb, lhs_lo.clone() - rhs_lo.clone());
         let is_zero_hi = IsZeroGadget::construct(cb, lhs_hi.clone() - rhs_hi.clone());
 
         Self {
             is_zero_lo,
             is_zero_hi,
-            marker: Default::default(),
+            _marker: Default::default(),
         }
     }
 
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/is_zero_word.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/is_zero_word.rs
index 6d81d37793..7a0f692c59 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/is_zero_word.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/is_zero_word.rs
@@ -1,3 +1,5 @@
+use std::marker::PhantomData;
+
 use eth_types::Field;
 use gadgets::util::{or, Expr};
 use halo2_proofs::{
@@ -10,19 +12,20 @@ use crate::{
         constraint_builder::{ConstrainBuilderCommon, EVMConstraintBuilder},
         transpose_val_ret, CachedRegion, Cell, CellType,
     },
-    util::word::{Word, Word32Cell},
+    util::word::{Word, WordExpr},
 };
 
 /// Returns `1` when `word == 0`, and returns `0` otherwise.
 #[derive(Clone, Debug)]
-pub struct IsZeroWordGadget<F> {
+pub struct IsZeroWordGadget<F, T> {
     inverse_lo: Cell<F>,
     inverse_hi: Cell<F>,
     is_zero: Expression<F>,
+    _marker: PhantomData<T>,
 }
 
-impl<F: Field> IsZeroWordGadget<F> {
-    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, word: Word32Cell<F>) -> Self {
+impl<F: Field, T: WordExpr<F>> IsZeroWordGadget<F, T> {
+    pub(crate) fn construct(cb: &mut EVMConstraintBuilder<F>, word: T) -> Self {
         let (word_lo, word_hi) = word.to_word().to_lo_hi();
         let inverse_lo = cb.query_cell_with_type(CellType::storage_for_expr(word_lo));
         let inverse_hi = cb.query_cell_with_type(CellType::storage_for_expr(word_hi));
@@ -54,6 +57,7 @@ impl<F: Field> IsZeroWordGadget<F> {
             inverse_lo,
             inverse_hi,
             is_zero: or::expr([is_zero_lo, is_zero_hi]),
+            _marker: Default::default(),
         }
     }
 
@@ -68,16 +72,16 @@ impl<F: Field> IsZeroWordGadget<F> {
         value: Word<F>,
     ) -> Result<F, Error> {
         let (value_lo, value_hi) = value.to_lo_hi();
-        let inverse_lo = value_lo.invert().unwrap_or(F::zero());
+        let inverse_lo = value_lo.invert().unwrap_or(F::from(0));
         self.inverse_lo
             .assign(region, offset, Value::known(inverse_lo))?;
-        let inverse_hi = value_hi.invert().unwrap_or(F::zero());
+        let inverse_hi = value_hi.invert().unwrap_or(F::from(0));
         self.inverse_hi
             .assign(region, offset, Value::known(inverse_hi))?;
         Ok(if value_lo.is_zero().into() && value_hi.is_zero().into() {
             F::from(2)
         } else {
-            F::zero()
+            F::from(0)
         })
     }
 
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/lt_word.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/lt_word.rs
index a62aa28599..132ef6d12d 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/lt_word.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/lt_word.rs
@@ -2,7 +2,7 @@ use crate::{
     evm_circuit::util::{
         constraint_builder::EVMConstraintBuilder, math_gadget::*, split_u256, CachedRegion, Cell,
     },
-    util::word::Word32Cell,
+    util::word::{Word32Cell, WordExpr},
 };
 use eth_types::{Field, Word};
 use halo2_proofs::plonk::{Error, Expression};
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/modulo.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/modulo.rs
index d6274ce1ec..0058ef5986 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/modulo.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/modulo.rs
@@ -5,7 +5,7 @@ use crate::{
         CachedRegion,
     },
     util::{
-        word::{self, Word32Cell, WordExpr},
+        word::{self, Word32Cell},
         Expr,
     },
 };
@@ -26,9 +26,9 @@ pub(crate) struct ModGadget<F> {
     k: Word32Cell<F>,
     a_or_zero: Word32Cell<F>,
     mul_add_words: MulAddWordsGadget<F>,
-    n_is_zero: IsZeroWordGadget<F>,
-    a_or_is_zero: IsZeroWordGadget<F>,
-    eq: IsEqualWordGadget<F, word::Word<Expression<F>>>,
+    n_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
+    a_or_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
+    eq: IsEqualWordGadget<F, Word32Cell<F>, Word32Cell<F>>,
     lt: LtWordGadget<F>,
 }
 impl<F: Field> ModGadget<F> {
@@ -39,7 +39,7 @@ impl<F: Field> ModGadget<F> {
         let n_is_zero = IsZeroWordGadget::construct(cb, n.clone());
         let a_or_is_zero = IsZeroWordGadget::construct(cb, a_or_zero.clone());
         let mul_add_words = MulAddWordsGadget::construct(cb, [&k, n, r, &a_or_zero]);
-        let eq = IsEqualWordGadget::construct(cb, a.to_word(), a_or_zero.to_word());
+        let eq = IsEqualWordGadget::construct(cb, a.clone(), a_or_zero);
         let lt = LtWordGadget::construct(cb, r, n);
         // Constrain the aux variable a_or_zero to be =a or =0 if n==0:
         // (a == a_or_zero) ^ (n == 0 & a_or_zero == 0)
@@ -118,7 +118,7 @@ mod tests {
         r: Word32Cell<F>,
     }
 
-    impl<F: Field, T: WordExpr<F>> MathGadgetContainer<F> for ModGadgetTestContainer<F> {
+    impl<F: Field> MathGadgetContainer<F> for ModGadgetTestContainer<F> {
         fn configure_gadget_container(cb: &mut EVMConstraintBuilder<F>) -> Self {
             let a = cb.query_word32();
             let n = cb.query_word32();
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words.rs
index 98fedad4f0..b3909a08a4 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words.rs
@@ -4,7 +4,7 @@ use crate::{
         from_bytes, pow_of_two_expr, split_u256, split_u256_limb64, CachedRegion, Cell,
     },
     util::{
-        word::{Word, Word32Cell, Word4},
+        word::{Word, Word32Cell, Word4, WordExpr},
         Expr,
     },
 };
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words512.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words512.rs
index 0ac7bd2856..1541843074 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words512.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_add_words512.rs
@@ -4,7 +4,7 @@ use crate::{
         from_bytes, pow_of_two_expr, split_u256, split_u256_limb64, CachedRegion, Cell,
     },
     util::{
-        word::{self, Word4},
+        word::{self, Word4, WordExpr},
         Expr,
     },
 };
diff --git a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_word_u64.rs b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_word_u64.rs
index 1f36cc105e..958179a222 100644
--- a/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_word_u64.rs
+++ b/zkevm-circuits/src/evm_circuit/util/math_gadget/mul_word_u64.rs
@@ -4,7 +4,10 @@ use crate::{
         constraint_builder::{ConstrainBuilderCommon, EVMConstraintBuilder},
         from_bytes, pow_of_two_expr, split_u256, CachedRegion,
     },
-    util::{word::Word32Cell, Expr},
+    util::{
+        word::{Word32Cell, WordExpr},
+        Expr,
+    },
 };
 use eth_types::{Field, ToLittleEndian, Word};
 use halo2_proofs::{
diff --git a/zkevm-circuits/src/evm_circuit/util/memory_gadget.rs b/zkevm-circuits/src/evm_circuit/util/memory_gadget.rs
index 72794a305e..cf5df536ed 100644
--- a/zkevm-circuits/src/evm_circuit/util/memory_gadget.rs
+++ b/zkevm-circuits/src/evm_circuit/util/memory_gadget.rs
@@ -71,7 +71,7 @@ pub(crate) mod address_high {
 pub(crate) struct MemoryAddressGadget<F> {
     memory_offset: Word32Cell<F>,
     memory_length: Word32Cell<F>,
-    memory_length_is_zero: IsZeroWordGadget<F>,
+    memory_length_is_zero: IsZeroWordGadget<F, Word32Cell<F>>,
 }
 
 impl<F: Field> MemoryAddressGadget<F> {
diff --git a/zkevm-circuits/src/util.rs b/zkevm-circuits/src/util.rs
index 07415f2236..e4e35a8d17 100644
--- a/zkevm-circuits/src/util.rs
+++ b/zkevm-circuits/src/util.rs
@@ -1,6 +1,8 @@
 //! Common utility traits and functions.
 pub mod word;
 
+use std::marker::PhantomData;
+
 use bus_mapping::evm::OpcodeId;
 use halo2_proofs::{
     circuit::{Layouter, Value},
diff --git a/zkevm-circuits/src/util/word.rs b/zkevm-circuits/src/util/word.rs
index 328c2bd30a..ffeac5fbf1 100644
--- a/zkevm-circuits/src/util/word.rs
+++ b/zkevm-circuits/src/util/word.rs
@@ -3,18 +3,17 @@
 // - Limbs: An EVN word is 256 bits. Limbs N means split 256 into N limb. For example, N = 4, each
 //   limb is 256/4 = 64 bits
 
-use eth_types::ToLittleEndian;
+use eth_types::{Field, ToLittleEndian};
 use gadgets::util::{not, or, Expr};
 use halo2_proofs::{
     circuit::{AssignedCell, Value},
-    halo2curves::FieldExt,
     plonk::{Error, Expression},
 };
 use itertools::Itertools;
 
 use crate::evm_circuit::util::{from_bytes, CachedRegion, Cell};
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Copy)]
 pub(crate) struct WordLimbs<T, const N: usize> {
     pub limbs: [T; N],
 }
@@ -53,7 +52,7 @@ pub(crate) trait WordExpr<F> {
     fn to_word(&self) -> Word<Expression<F>>;
 }
 
-impl<F: FieldExt, const N: usize> WordLimbs<Cell<F>, N> {
+impl<F: Field, const N: usize> WordLimbs<Cell<F>, N> {
     pub fn assign<const N1: usize>(
         &self,
         region: &mut CachedRegion<'_, '_, F>,
@@ -76,7 +75,7 @@ impl<F: FieldExt, const N: usize> WordLimbs<Cell<F>, N> {
     }
 }
 
-impl<F: FieldExt, const N: usize> WordExpr<F> for WordLimbs<Cell<F>, N> {
+impl<F: Field, const N: usize> WordExpr<F> for WordLimbs<Cell<F>, N> {
     fn to_word(&self) -> Word<Expression<F>> {
         Word(self.word_expr().to_word_n())
     }
@@ -114,7 +113,7 @@ impl<T> std::ops::Deref for Word<T> {
     }
 }
 
-impl<F: FieldExt> Word<F> {
+impl<F: Field> Word<F> {
     pub fn from_u256(value: eth_types::Word) -> Word<F> {
         let bytes = value.to_le_bytes();
         Word::new([
@@ -125,11 +124,11 @@ impl<F: FieldExt> Word<F> {
 
     pub fn from_u64(value: u64) -> Word<F> {
         let bytes = value.to_le_bytes();
-        Word::new([from_bytes::value(&bytes), F::zero()])
+        Word::new([from_bytes::value(&bytes), F::from(0)])
     }
 }
 
-impl<F: FieldExt> Word<Cell<F>> {
+impl<F: Field> Word<Cell<F>> {
     pub fn assign_lo(
         &self,
         region: &mut CachedRegion<'_, '_, F>,
@@ -138,12 +137,18 @@ impl<F: FieldExt> Word<Cell<F>> {
     ) -> Result<Vec<AssignedCell<F, F>>, Error> {
         Ok(vec![
             self.limbs[0].assign(region, offset, value)?,
-            self.limbs[1].assign(region, offset, Value::known(F::zero()))?,
+            self.limbs[1].assign(region, offset, Value::known(F::from(0)))?,
         ])
     }
 }
 
-impl<F: FieldExt> Word<Expression<F>> {
+impl<F: Field> WordExpr<F> for Word<Cell<F>> {
+    fn to_word(&self) -> Word<Expression<F>> {
+        self.word_expr()
+    }
+}
+
+impl<F: Field> Word<Expression<F>> {
     // create word from lo limb with hi limb as 0. caller need to guaranteed to be 128 bits.
     pub fn from_lo_unchecked(lo: Expression<F>) -> Self {
         Self(WordLimbs::<Expression<F>, 2>::new([lo, 0.expr()]))
@@ -172,11 +177,6 @@ impl<F: FieldExt> Word<Expression<F>> {
         Word::new([self.lo().clone() * selector, self.hi().clone() * selector])
     }
 
-    // Assume selector is 1/0 therefore no overflow check
-    pub fn mul_selector(&self, selector: Expression<F>) -> Self {
-        Word::new([self.lo().clone() * selector, self.hi().clone() * selector])
-    }
-
     // No overflow check on lo/hi limbs
     pub fn add_unchecked(self, rhs: Self) -> Self {
         Word::new([
@@ -186,7 +186,13 @@ impl<F: FieldExt> Word<Expression<F>> {
     }
 }
 
-impl<F: FieldExt, const N1: usize> WordLimbs<Expression<F>, N1> {
+impl<F: Field> WordExpr<F> for Word<Expression<F>> {
+    fn to_word(&self) -> Word<Expression<F>> {
+        self.clone()
+    }
+}
+
+impl<F: Field, const N1: usize> WordLimbs<Expression<F>, N1> {
     /// to_wordlimbs will aggregate nested expressions, which implies during expression evaluation
     /// it need more recursive call. if the converted limbs word will be used in many place,
     /// consider create new low limbs word, have equality constrain, then finally use low limbs
@@ -218,7 +224,7 @@ impl<F: FieldExt, const N1: usize> WordLimbs<Expression<F>, N1> {
     }
 }
 
-impl<F: FieldExt, const N1: usize> WordExpr<F> for WordLimbs<Expression<F>, N1> {
+impl<F: Field, const N1: usize> WordExpr<F> for WordLimbs<Expression<F>, N1> {
     fn to_word(&self) -> Word<Expression<F>> {
         Word(self.to_word_n())
     }