Handle base conversion nicer (#281)

* refactor: use more native biguint radix api

* simplify inspect

Author: ChihChengLiang

keccak256/src/arith_helpers.rs

@@ -1,25 +1,25 @@
 use crate::common::State;
 use itertools::Itertools;
 use num_bigint::BigUint;
-use num_traits::{One, Zero};
+use num_traits::Zero;
 use pairing::arithmetic::FieldExt;
 use pairing::bn256::Fr as Fp;
 use std::ops::{Index, IndexMut};
 
-pub const B2: u64 = 2;
-pub const B13: u64 = 13;
-pub const B9: u64 = 9;
+pub const B2: u8 = 2;
+pub const B13: u8 = 13;
+pub const B9: u8 = 9;
 
 /// Base 9 coef mapper scalers
 /// f_logic(x1, x2, x3, x4) = x1 ^ (!x2 & x3) ^ x4
 /// f_arith(x1, x2, x3, x4) = 2*x1 + x2 + 3*x3 + 2*x4
 /// where x1, x2, x3, x4 are binary.
 /// We have the property that `0 <= f_arith(...) < 9` and
 /// the map from `f_arith(...)` to `f_logic(...)` is injective.
-pub const A1: u64 = 2u64;
-pub const A2: u64 = 1u64;
-pub const A3: u64 = 3u64;
-pub const A4: u64 = 2u64;
+pub const A1: u8 = 2;
+pub const A2: u8 = 1;
+pub const A3: u8 = 3;
+pub const A4: u8 = 2;
 
 pub type Lane13 = BigUint;
 pub type Lane9 = BigUint;
@@ -87,10 +87,6 @@ impl Clone for StateBigInt {
     }
 }
 
-pub fn mod_u64(a: &BigUint, b: u64) -> u64 {
-    (a % b).iter_u64_digits().take(1).next().unwrap_or(0)
-}
-
 pub fn convert_b2_to_b13(a: u64) -> Lane13 {
     let mut lane13: BigUint = Zero::zero();
     for i in 0..64 {
@@ -117,7 +113,7 @@ pub fn convert_b2_to_b9(a: u64) -> Lane9 {
 ///
 /// For example, if we have 5 bits set and 7 bits unset, then we have `x` as 5
 /// and the xor result to be 1.
-pub fn convert_b13_coef(x: u64) -> u64 {
+pub fn convert_b13_coef(x: u8) -> u8 {
     assert!(x < 13);
     x & 1
 }
@@ -127,55 +123,49 @@ pub fn convert_b13_coef(x: u64) -> u64 {
 ///
 /// The input `x` is a chunk of a base 9 number and it represents the arithmatic
 /// result of `2*a + b + 3*c + 2*d`, where `a`, `b`, `c`, and `d` each is a bit.
-pub fn convert_b9_coef(x: u64) -> u64 {
+pub fn convert_b9_coef(x: u8) -> u8 {
     assert!(x < 9);
-    let bit_table: [u64; 9] = [0, 0, 1, 1, 0, 0, 1, 1, 0];
+    let bit_table: [u8; 9] = [0, 0, 1, 1, 0, 0, 1, 1, 0];
     bit_table[x as usize]
 }
 
+// We assume the input comes from Theta step and has 65 chunks
 pub fn convert_b13_lane_to_b9(x: Lane13, rot: u32) -> Lane9 {
-    let mut base = BigUint::from(B9).pow(rot);
-    let mut special_chunk = Zero::zero();
-    let mut raw = x;
-    let mut acc: Lane9 = Zero::zero();
-
-    for i in 0..65 {
-        let remainder: u64 = mod_u64(&raw, B13);
-        if i == 0 || i == 64 {
-            special_chunk += remainder;
-        } else {
-            acc += convert_b13_coef(remainder) * base.clone();
-        }
-        raw /= B13;
-        base *= B9;
-        if i == 63 - rot {
-            base = One::one();
-        }
-    }
-    acc += convert_b13_coef(special_chunk) * BigUint::from(B9).pow(rot);
-    acc
+    // 65 chunks
+    let mut chunks = x.to_radix_le(B13.into());
+    chunks.resize(65, 0);
+    // 0 and 64 was separated in Theta, we now combined them together
+    let special = chunks.get(0).unwrap() + chunks.get(64).unwrap();
+    // middle 63 chunks
+    let middle = chunks.get(1..64).unwrap();
+    // split at offset
+    let (left, right) = middle.split_at(63 - rot as usize);
+    // rotated has 64 chunks
+    // left is rotated right, and the right is wrapped over to left
+    // with the special chunk in the middle
+    let rotated: Vec<u8> = right
+        .iter()
+        .chain(vec![special].iter())
+        .chain(left.iter())
+        .map(|&x| convert_b13_coef(x))
+        .collect_vec();
+    BigUint::from_radix_le(&rotated, B9.into()).unwrap_or_default()
 }
 
 pub fn convert_lane<F>(
     lane: BigUint,
-    from_base: u64,
-    to_base: u64,
+    from_base: u8,
+    to_base: u8,
     coef_transform: F,
 ) -> BigUint
 where
-    F: Fn(u64) -> u64,
+    F: Fn(u8) -> u8,
 {
-    let mut base: BigUint = One::one();
-    let mut raw = lane;
-    let mut acc: BigUint = Zero::zero();
-
-    for _ in 0..64 {
-        let remainder: u64 = mod_u64(&raw, B9);
-        acc += coef_transform(remainder) * base.clone();
-        raw /= from_base;
-        base *= to_base;
-    }
-    acc
+    let chunks = lane.to_radix_be(from_base.into());
+    let converted_chunks: Vec<u8> =
+        chunks.iter().map(|&x| coef_transform(x)).collect();
+    BigUint::from_radix_be(&converted_chunks, to_base.into())
+        .unwrap_or_default()
 }
 
 pub fn convert_b9_lane_to_b13(x: Lane9) -> Lane13 {
@@ -215,17 +205,11 @@ pub fn big_uint_to_field<F: FieldExt>(a: &BigUint) -> F {
 
 /// This function allows us to inpect coefficients of big-numbers in different
 /// bases.
-pub fn inspect(x: BigUint, name: &str, base: u64) {
-    let mut raw = x.clone();
-    let mut info: Vec<(u32, u64)> = vec![];
-
-    for i in 0..65 {
-        let remainder: u64 = mod_u64(&raw, base);
-        raw /= base;
-        if remainder != 0 {
-            info.push((i, remainder));
-        }
-    }
+pub fn inspect(x: BigUint, name: &str, base: u8) {
+    let mut chunks = x.to_radix_le(base.into());
+    chunks.resize(65, 0);
+    let info: Vec<(usize, u8)> =
+        (0..65).zip(chunks.iter().copied()).collect_vec();
     println!("inspect {} {} info {:?}", name, x, info);
 }
 
@@ -283,3 +267,34 @@ pub fn state_bigint_to_field<F: FieldExt, const N: usize>(
     arr[0..N].copy_from_slice(&vector[0..N]);
     arr
 }
+
+pub fn f_from_radix_be<F: FieldExt>(buf: &[u8], base: u8) -> F {
+    let base = F::from(base.into());
+    buf.iter()
+        .fold(F::zero(), |acc, &x| acc * base + F::from(x.into()))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use num_bigint::BigUint;
+    #[test]
+    fn test_convert_b13_lane_to_b9() {
+        // the number 1 is chosen that `convert_b13_coef` has no effect
+        let mut a = vec![0, 1, 1, 1];
+        a.resize(65, 0);
+        let lane = BigUint::from_radix_le(&a, B13.into()).unwrap_or_default();
+        assert_eq!(
+            convert_b13_lane_to_b9(lane.clone(), 0),
+            BigUint::from_radix_le(&a, B9.into()).unwrap_or_default()
+        );
+
+        // rotate by 4
+        let mut b = vec![0, 0, 0, 0, 0, 1, 1, 1];
+        b.resize(65, 0);
+        assert_eq!(
+            convert_b13_lane_to_b9(lane, 4),
+            BigUint::from_radix_le(&b, B9.into()).unwrap_or_default()
+        );
+    }
+}

keccak256/src/gates/gate_helpers.rs

@@ -29,9 +29,9 @@ pub fn f_to_biguint<F: FieldExt>(x: F) -> Option<BigUint> {
     Option::from(BigUint::from_bytes_le(&x.to_bytes()[..]))
 }
 
-pub fn biguint_mod(x: &BigUint, modulus: u64) -> u64 {
-    match (x % modulus).to_u64_digits().first() {
-        Some(d) => *d,
-        None => 0u64,
-    }
+pub fn biguint_mod(x: &BigUint, modulus: u8) -> u8 {
+    x.to_radix_le(modulus.into())
+        .first()
+        .copied()
+        .unwrap_or_default()
 }

keccak256/src/gates/rho_checks.rs

@@ -178,16 +178,16 @@ struct RotatingVariables {
     block_count: Option<u32>,
     // step2 acc and step3 acc
     block_count_acc: [u32; 2],
-    low_value: u64,
-    high_value: Option<u64>,
+    low_value: u8,
+    high_value: Option<u8>,
 }
 
 impl RotatingVariables {
     fn from(lane_base_13: BigUint, rotation: u32) -> Result<Self, Error> {
         let chunk_idx = 1;
         let step = get_step_size(chunk_idx, rotation);
-        let input_raw = lane_base_13.clone() / B13;
-        let input_coef = input_raw.clone() % B13.pow(step);
+        let input_raw = lane_base_13.clone() / (B13 as u64);
+        let input_coef = input_raw.clone() % (B13 as u64).pow(step);
         let input_power_of_base = BigUint::from(B13);
         let input_acc = lane_base_13.clone();
         let (block_count, output_coef) =
@@ -204,7 +204,7 @@ impl RotatingVariables {
         if step == 2 || step == 3 {
             block_count_acc[(step - 2) as usize] += block_count;
         }
-        let low_value: u64 = biguint_mod(&lane_base_13, B13);
+        let low_value = biguint_mod(&lane_base_13, B13);
         Ok(Self {
             rotation,
             chunk_idx,
@@ -228,16 +228,18 @@ impl RotatingVariables {
         let mut new = self.clone();
         new.chunk_idx += self.step;
         new.step = get_step_size(new.chunk_idx, self.rotation);
-        new.input_raw /= B13.pow(self.step);
-        new.input_coef = new.input_raw.clone() % B13.pow(new.step);
-        new.input_power_of_base *= B13.pow(self.step);
+        let input_pow_of_step = (B13 as u64).pow(self.step);
+        new.input_raw /= input_pow_of_step;
+        new.input_coef = new.input_raw.clone() % (B13 as u64).pow(new.step);
+        new.input_power_of_base *= input_pow_of_step;
         new.input_acc -=
             self.input_power_of_base.clone() * self.input_coef.clone();
+        let output_pow_of_step = (B9 as u64).pow(self.step);
         new.output_power_of_base =
             if is_at_rotation_offset(new.chunk_idx, self.rotation) {
                 BigUint::one()
             } else {
-                self.output_power_of_base.clone() * B9.pow(self.step)
+                self.output_power_of_base.clone() * output_pow_of_step
             };
         new.output_acc +=
             self.output_power_of_base.clone() * self.output_coef.clone();
@@ -250,11 +252,11 @@ impl RotatingVariables {
             );
             new.block_count = None;
             new.high_value = Some(biguint_mod(&self.input_raw, B13));
-            let high = new.high_value.unwrap();
+            let high: u8 = new.high_value.unwrap();
             new.output_coef =
                 BigUint::from(convert_b13_coef(high + self.low_value));
-            let expect =
-                new.low_value + high * BigUint::from(B13).pow(LANE_SIZE);
+            let expect = (new.low_value as u64)
+                + (high as u64) * BigUint::from(B13).pow(LANE_SIZE);
             assert_eq!(
                 new.input_acc,
                 expect,
@@ -431,7 +433,8 @@ impl<F: FieldExt> ChunkRotateConversionConfig<F> {
             fix_cols,
         );
 
-        let power_of_b13 = F::from(B13).pow(&[chunk_idx.into(), 0, 0, 0]);
+        let power_of_b13 =
+            F::from(B13.into()).pow(&[chunk_idx.into(), 0, 0, 0]);
 
         // | coef | 13**x | acc       |
         // |------|-------|-----------|
@@ -451,7 +454,7 @@ impl<F: FieldExt> ChunkRotateConversionConfig<F> {
                 q_enable * (acc_next - acc + coef * power_of_base),
             )]
         });
-        let power_of_b9 = F::from(B9).pow(&[
+        let power_of_b9 = F::from(B9.into()).pow(&[
             ((rotation + chunk_idx) % LANE_SIZE).into(),
             0,
             0,
@@ -573,8 +576,9 @@ impl<F: FieldExt> SpecialChunkConfig<F> {
                 .query_advice(base_9_acc, Rotation::next())
                 - meta.query_advice(base_9_acc, Rotation::cur());
             let last_b9_coef = meta.query_advice(last_b9_coef, Rotation::cur());
-            let pow_of_9 =
-                Expression::Constant(F::from(B9).pow(&[rotation, 0, 0, 0]));
+            let pow_of_9 = Expression::Constant(
+                F::from(B9.into()).pow(&[rotation, 0, 0, 0]),
+            );
             vec![(
                 "delta_base_9_acc === (high_value + low_value) * 9**rotation",
                 meta.query_selector(q_enable)

keccak256/src/gates/rho_helpers.rs

@@ -1,7 +1,6 @@
 use crate::{
     arith_helpers::{convert_b13_coef, B13, B9},
     common::LANE_SIZE,
-    gates::gate_helpers::biguint_mod,
 };
 
 use std::convert::TryInto;
@@ -72,7 +71,7 @@ pub const STEP3_RANGE: u64 = 169;
 
 /// Get the block count from an input chunks
 ///
-/// The input is chunks of a base 13 number in descending order of signiticance.
+/// The input is chunks of a base 13 number in big endian.
 /// For example, if the input is `[1, 12, 3, 7]`, it represents a coefficient
 /// `1*13^3 + 12*13^2 + 3*13 + 7`. The example only happens when `step = 4`. If
 /// we have a `step = 3`, the first chunk must be 0. It could be the case that
@@ -84,7 +83,7 @@ pub const STEP3_RANGE: u64 = 169;
 /// block count to be 170.
 ///
 /// This would fail the final block count check.
-pub fn get_block_count(b13_chunks: [u64; BASE_NUM_OF_CHUNKS as usize]) -> u32 {
+pub fn get_block_count(b13_chunks: [u8; BASE_NUM_OF_CHUNKS as usize]) -> u32 {
     // could be 0, 1, 2, 3, 4
     let non_zero_chunk_count = BASE_NUM_OF_CHUNKS as usize
         - b13_chunks.iter().take_while(|x| **x == 0).count();
@@ -93,20 +92,13 @@ pub fn get_block_count(b13_chunks: [u64; BASE_NUM_OF_CHUNKS as usize]) -> u32 {
 }
 
 pub fn get_block_count_and_output_coef(input_coef: BigUint) -> (u32, u64) {
-    // b13_chunks is in descending order of more significant chunk
-    let b13_chunks: [u64; BASE_NUM_OF_CHUNKS as usize] = {
-        let mut x = input_coef;
-        let mut b13_chunks: Vec<u64> = vec![];
-        for _ in 0..BASE_NUM_OF_CHUNKS {
-            b13_chunks.push(biguint_mod(&x, B13));
-            x /= B13;
-        }
-        b13_chunks.reverse();
-        b13_chunks.try_into().unwrap()
-    };
-    let output_coef: u64 = b13_chunks.iter().fold(0, |acc, b13_chunk| {
-        let b9_chunk = convert_b13_coef(*b13_chunk);
-        acc * B9 + b9_chunk
+    let mut b13_chunks = input_coef.to_radix_le(B13.into());
+    b13_chunks.resize(BASE_NUM_OF_CHUNKS as usize, 0);
+    b13_chunks.reverse();
+    let b13_chunks: [u8; BASE_NUM_OF_CHUNKS as usize] =
+        b13_chunks.try_into().unwrap();
+    let output_coef: u64 = b13_chunks.iter().fold(0, |acc, &b13_chunk| {
+        acc * B9 as u64 + convert_b13_coef(b13_chunk) as u64
     });
     let block_count = get_block_count(b13_chunks);
     (block_count, output_coef)