[word lo/hi] docs and more word utilities (#1441)

### Description

Follow
#1435
design duplicate field design.

### Type of change

- [x] New feature (non-breaking change which adds functionality)

### Contents
- more utilities on `Word` type, and further refactor few to be more
generics

zkevm-circuits/src/evm_circuit/execution/addmod.rs

@@ -226,7 +226,7 @@ impl<F: Field> ExecutionGadget<F> for AddModGadget<F> {
         self.cmp_areduced_n.assign(region, offset, a_reduced, n)?;
 
         self.n_is_zero
-            .assign_value(region, offset, Value::known(Word::from_u256(n)))?;
+            .assign_value(region, offset, Value::known(Word::from(n)))?;
 
         Ok(())
     }

zkevm-circuits/src/evm_circuit/util/common_gadget.rs

@@ -499,7 +499,7 @@ impl<F: Field> TransferWithGasFeeGadget<F> {
             receiver_balance,
         )?;
         self.value_is_zero
-            .assign_value(region, offset, Value::known(Word::from_u256(value)))?;
+            .assign_value(region, offset, Value::known(Word::from(value)))?;
         Ok(())
     }
 }
@@ -593,7 +593,7 @@ impl<F: Field> TransferGadget<F> {
             receiver_balance,
         )?;
         self.value_is_zero
-            .assign_value(region, offset, Value::known(Word::from_u256(value)))?;
+            .assign_value(region, offset, Value::known(Word::from(value)))?;
         Ok(())
     }
 }
@@ -782,19 +782,19 @@ impl<F: Field, const IS_SUCCESS_CALL: bool> CommonCallGadget<F, IS_SUCCESS_CALL>
         )?;
 
         self.value_is_zero
-            .assign(region, offset, Word::from_u256(value))?;
+            .assign(region, offset, Word::from(value))?;
         self.callee_code_hash
             .assign(region, offset, Some(callee_code_hash.to_le_bytes()))?;
         self.is_empty_code_hash.assign_value(
             region,
             offset,
-            Value::known(Word::from_u256(callee_code_hash)),
-            Value::known(Word::from_u256(CodeDB::empty_code_hash().to_word())),
+            Value::known(Word::from(callee_code_hash)),
+            Value::known(Word::from(CodeDB::empty_code_hash().to_word())),
         )?;
         self.callee_not_exists.assign_value(
             region,
             offset,
-            Value::known(Word::from_u256(callee_code_hash)),
+            Value::known(Word::from(callee_code_hash)),
         )?;
         Ok(memory_expansion_gas_cost)
     }
@@ -939,19 +939,19 @@ impl<F: Field> SstoreGasGadget<F> {
         self.value_eq_prev.assign_value(
             region,
             offset,
-            Value::known(Word::from_u256(value)),
-            Value::known(Word::from_u256(value_prev)),
+            Value::known(Word::from(value)),
+            Value::known(Word::from(value_prev)),
         )?;
         self.original_eq_prev.assign_value(
             region,
             offset,
-            Value::known(Word::from_u256(original_value)),
-            Value::known(Word::from_u256(value_prev)),
+            Value::known(Word::from(original_value)),
+            Value::known(Word::from(value_prev)),
         )?;
         self.original_is_zero.assign_value(
             region,
             offset,
-            Value::known(Word::from_u256(original_value)),
+            Value::known(Word::from(original_value)),
         )?;
         Ok(())
     }

zkevm-circuits/src/evm_circuit/util/math_gadget/modulo.rs

@@ -95,17 +95,17 @@ impl<F: Field> ModGadget<F> {
             .assign(region, offset, Some(a_or_zero.to_le_bytes()))?;
         let n_sum = (0..32).fold(0, |acc, idx| acc + n.byte(idx) as u64);
         self.n_is_zero
-            .assign(region, offset, word::Word::from_u64(n_sum))?;
+            .assign(region, offset, word::Word::from(n_sum))?;
         self.a_or_is_zero
-            .assign(region, offset, word::Word::from_u256(a_or_zero))?;
+            .assign(region, offset, word::Word::from(a_or_zero))?;
         self.mul_add_words
             .assign(region, offset, [k, n, r, a_or_zero])?;
         self.lt.assign(region, offset, r, n)?;
         self.eq.assign_value(
             region,
             offset,
-            Value::known(word::Word::from_u256(a)),
-            Value::known(word::Word::from_u256(a_or_zero)),
+            Value::known(word::Word::from(a)),
+            Value::known(word::Word::from(a_or_zero)),
         )?;
 
         Ok(())

zkevm-circuits/src/evm_circuit/util/memory_gadget.rs

@@ -130,7 +130,7 @@ impl<F: Field> MemoryAddressGadget<F> {
             .assign(region, offset, Some(memory_length.to_le_bytes()))?;
 
         self.memory_length_is_zero
-            .assign(region, offset, Word::from_u256(memory_length))?;
+            .assign(region, offset, Word::from(memory_length))?;
         Ok(if memory_length_is_zero {
             0
         } else {

zkevm-circuits/src/util/word.rs

@@ -3,7 +3,7 @@
 // - 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::{Field, ToLittleEndian};
+use eth_types::{Field, ToLittleEndian, H160};
 use gadgets::util::{not, or, Expr};
 use halo2_proofs::{
     circuit::{AssignedCell, Value},
@@ -13,6 +13,9 @@ use itertools::Itertools;
 
 use crate::evm_circuit::util::{from_bytes, CachedRegion, Cell, RandomLinearCombination};
 
+/// evm word 32 bytes, half word 16 bytes
+const N_BYTES_HALF_WORD: usize = 16;
+
 /// The EVM word for witness
 #[derive(Clone, Debug, Copy)]
 pub struct WordLimbs<T, const N: usize> {
@@ -83,6 +86,13 @@ pub trait WordExpr<F> {
 
 impl<F: Field, const N: usize> WordLimbs<Cell<F>, N> {
     /// assign limbs
+    /// N1 is number of bytes to assign, while N is number of limbs.
+    /// N1 % N = 0 (also implies N1 >= N, assuming N1 and N are not 0)
+    /// If N1 > N, then N1 will be chunk into N1 / N size then aggregate to single expression
+    /// then assign to N limbs respectively.
+    /// e.g. N1 = 4 bytes, [b1, b2, b3, b4], and N = 2 limbs [l1, l2]
+    /// It equivalent `l1.assign(b1.expr() + b2.expr * F(256))`, `l2.assign(b3.expr() + b4.expr *
+    /// F(256))`
     pub fn assign<const N1: usize>(
         &self,
         region: &mut CachedRegion<'_, '_, F>,
@@ -100,6 +110,97 @@ impl<F: Field, const N: usize> WordLimbs<Cell<F>, N> {
         })
     }
 
+    /// assign bytes to wordlimbs first half/second half respectively
+    // N_LO, N_HI are number of bytes to assign to first half and second half of size N limbs,
+    // respectively N_LO and N_HI can be different size, the only requirement is N_LO % (N/2)
+    // and N_HI % (N/2) [N/2] limbs will be assigned separately.
+    // E.g. N_LO = 4 => [nl1, nl2, nl3, nl4]
+    // N_HI = 2 => [nh1, nh2]
+    // N = 2 => [l1, l2]
+    // it equivalent l1.assign(nl1.expr() + nl2.expr() * 256 + nl3.expr() * 256^2 +  nl3.expr() *
+    // 256^3) and l2.assign(nh1.expr() + nh2.expr() * 256)
+    fn assign_lo_hi<const N_LO: usize, const N_HI: usize>(
+        &self,
+        region: &mut CachedRegion<'_, '_, F>,
+        offset: usize,
+        bytes_lo_le: [u8; N_LO],
+        bytes_hi_le: Option<[u8; N_HI]>,
+    ) -> Result<Vec<AssignedCell<F, F>>, Error> {
+        assert_eq!(N % 2, 0); // TODO use static_assertion instead
+        assert_eq!(N_LO % (N / 2), 0);
+        assert_eq!(N_HI % (N / 2), 0);
+        let half_limb_size = N / 2;
+
+        // assign lo
+        let bytes_lo_assigned = bytes_lo_le
+            .chunks(N_LO / half_limb_size) // chunk in little endian
+            .map(|chunk| from_bytes::value(chunk))
+            .zip(self.limbs[0..half_limb_size].iter())
+            .map(|(value, cell)| cell.assign(region, offset, Value::known(value)))
+            .collect::<Result<Vec<AssignedCell<F, F>>, _>>()?;
+
+        // assign hi
+        let bytes_hi_assigned = bytes_hi_le.map(|bytes| {
+            bytes
+                .chunks(N_HI / half_limb_size) // chunk in little endian
+                .map(|chunk| from_bytes::value(chunk))
+                .zip(self.limbs[half_limb_size..].iter())
+                .map(|(value, cell)| cell.assign(region, offset, Value::known(value)))
+                .collect::<Result<Vec<AssignedCell<F, F>>, _>>()
+        });
+
+        Ok([
+            bytes_lo_assigned.to_vec(),
+            match bytes_hi_assigned {
+                Some(hi_assigned) => hi_assigned?.to_vec(),
+                None => vec![],
+            },
+        ]
+        .concat())
+    }
+
+    /// assign u256 to wordlimbs
+    pub fn assign_u256(
+        &self,
+        region: &mut CachedRegion<'_, '_, F>,
+        offset: usize,
+        word: eth_types::Word,
+    ) -> Result<Vec<AssignedCell<F, F>>, Error> {
+        self.assign_lo_hi::<N_BYTES_HALF_WORD, N_BYTES_HALF_WORD>(
+            region,
+            offset,
+            word.to_le_bytes()[0..N_BYTES_HALF_WORD].try_into().unwrap(),
+            word.to_le_bytes()[N_BYTES_HALF_WORD..].try_into().ok(),
+        )
+    }
+
+    /// assign h160 to wordlimbs
+    pub fn assign_h160(
+        &self,
+        region: &mut CachedRegion<'_, '_, F>,
+        offset: usize,
+        h160: H160,
+    ) -> Result<Vec<AssignedCell<F, F>>, Error> {
+        let mut bytes = *h160.as_fixed_bytes();
+        bytes.reverse();
+        self.assign_lo_hi::<N_BYTES_HALF_WORD, 4>(
+            region,
+            offset,
+            bytes[0..N_BYTES_HALF_WORD].try_into().unwrap(),
+            bytes[N_BYTES_HALF_WORD..].try_into().ok(),
+        )
+    }
+
+    /// assign u64 to wordlimbs
+    pub fn assign_u64(
+        &self,
+        region: &mut CachedRegion<'_, '_, F>,
+        offset: usize,
+        value: u64,
+    ) -> Result<Vec<AssignedCell<F, F>>, Error> {
+        self.assign_lo_hi(region, offset, value.to_le_bytes(), Option::<[u8; 0]>::None)
+    }
+
     #[deprecated(note = "in fav of to_word trait. Make this private")]
     /// word expr
     pub fn word_expr(&self) -> WordLimbs<Expression<F>, N> {
@@ -122,12 +223,10 @@ impl<T: Clone> Word<T> {
     pub fn new(limbs: [T; 2]) -> Self {
         Self(WordLimbs::<T, 2>::new(limbs))
     }
-
     /// The high 128 bits limb
     pub fn hi(&self) -> T {
         self.0.limbs[1].clone()
     }
-
     /// the low 128 bits limb
     pub fn lo(&self) -> T {
         self.0.limbs[0].clone()
@@ -136,11 +235,14 @@ impl<T: Clone> Word<T> {
     pub fn n() -> usize {
         2
     }
-
     /// word to low and high 128 bits
     pub fn to_lo_hi(&self) -> (T, T) {
         (self.0.limbs[0].clone(), self.0.limbs[1].clone())
     }
+    /// Map the word to other types
+    pub fn map<T2: Clone>(&self, mut func: impl FnMut(T) -> T2) -> Word<T2> {
+        Word(WordLimbs::<T2, 2>::new([func(self.lo()), func(self.hi())]))
+    }
 }
 
 impl<T> std::ops::Deref for Word<T> {
@@ -151,22 +253,56 @@ impl<T> std::ops::Deref for Word<T> {
     }
 }
 
-impl<F: Field> Word<F> {
-    /// Constrct the word from u256
-    pub fn from_u256(value: eth_types::Word) -> Word<F> {
+impl<T: Clone + PartialEq> PartialEq for Word<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.lo() == other.lo() && self.hi() == other.hi()
+    }
+}
+
+impl<F: Field> From<eth_types::Word> for Word<F> {
+    /// Construct the word from u256
+    fn from(value: eth_types::Word) -> Self {
         let bytes = value.to_le_bytes();
         Word::new([
-            from_bytes::value(&bytes[..16]),
-            from_bytes::value(&bytes[16..]),
+            from_bytes::value(&bytes[..N_BYTES_HALF_WORD]),
+            from_bytes::value(&bytes[N_BYTES_HALF_WORD..]),
         ])
     }
-    /// Constrct the word from u64
-    pub fn from_u64(value: u64) -> Word<F> {
+}
+
+impl<F: Field> From<u64> for Word<F> {
+    /// Construct the word from u64
+    fn from(value: u64) -> Self {
         let bytes = value.to_le_bytes();
         Word::new([from_bytes::value(&bytes), F::from(0)])
     }
 }
 
+impl<F: Field> From<u8> for Word<F> {
+    /// Construct the word from u8
+    fn from(value: u8) -> Self {
+        Word::new([F::from(value as u64), F::from(0)])
+    }
+}
+
+impl<F: Field> From<bool> for Word<F> {
+    fn from(value: bool) -> Self {
+        Word::new([F::from(value as u64), F::from(0)])
+    }
+}
+
+impl<F: Field> From<H160> for Word<F> {
+    /// Construct the word from h160
+    fn from(value: H160) -> Self {
+        let mut bytes = *value.as_fixed_bytes();
+        bytes.reverse();
+        Word::new([
+            from_bytes::value(&bytes[..N_BYTES_HALF_WORD]),
+            from_bytes::value(&bytes[N_BYTES_HALF_WORD..]),
+        ])
+    }
+}
+
 impl<F: Field> Word<Cell<F>> {
     /// Assign low 128 bits for the word
     pub fn assign_lo(
@@ -193,7 +329,6 @@ impl<F: Field> Word<Expression<F>> {
     pub fn from_lo_unchecked(lo: Expression<F>) -> Self {
         Self(WordLimbs::<Expression<F>, 2>::new([lo, 0.expr()]))
     }
-
     /// zero word
     pub fn zero() -> Self {
         Self(WordLimbs::<Expression<F>, 2>::new([0.expr(), 0.expr()]))
@@ -231,6 +366,14 @@ impl<F: Field> Word<Expression<F>> {
     pub fn sub_unchecked(self, rhs: Self) -> Self {
         Word::new([self.lo() - rhs.lo(), self.hi() - rhs.hi()])
     }
+
+    /// Compress the lo and hi limbs into an expression without checking the overflow.
+    /// So far only use it for address.
+    /// TODO We should remove it before merging to the main branch.
+    #[deprecated(note = "no overflow check and unsafe. please consider keep word type")]
+    pub fn expr_unchecked(&self) -> Expression<F> {
+        self.lo() + self.hi() * (1 << (N_BYTES_HALF_WORD * 8)).expr()
+    }
 }
 
 impl<F: Field> WordExpr<F> for Word<Expression<F>> {
@@ -284,7 +427,7 @@ pub type WordLegacy<F> = RandomLinearCombination<F, 32>;
 
 impl<F: Field> WordExpr<F> for WordLegacy<F> {
     fn to_word(&self) -> Word<Expression<F>> {
-        Word::from_lo_unchecked(self.expr())
+        Word::new([self.expr(), 0.expr()])
     }
 }