keccak hash circuit added, test added(incomplete)

hash/plonky2/src/bench/keccak.rs

@@ -11,7 +11,8 @@ use plonky2::hash::hash_types::RichField;
 use plonky2_u32::gadgets::arithmetic_u32::CircuitBuilderU32;
 use crate::arithmetic::u64_arithmetic::U64Target;
 use crate::arithmetic::u64_arithmetic::CircuitBuilderU64;
-
+use plonky2::iop::witness::Witness;
+use plonky2::field::types::PrimeField64;
 
 fn generate_data(size: usize) -> Vec<GoldilocksField> {
 
@@ -25,6 +26,27 @@ fn generate_data(size: usize) -> Vec<GoldilocksField> {
 
 }
 
+//TODO: remove the dead codes later
+#[allow(dead_code)]
+pub trait WitnessU64<F: PrimeField64>: Witness<F> {
+    fn set_u64_target(&mut self, target: U64Target, value: u64);
+    // fn get_u64_target(&self, target: U64Target) -> (u64, u64);
+}
+
+impl<T: Witness<F>, F: PrimeField64> WitnessU64<F> for T {
+    fn set_u64_target(&mut self, target: U64Target, value: u64) {
+        self.set_target(target.0[0].0, F::from_canonical_u32((value & 0xFFFFFFFF) as u32));
+        self.set_target(target.0[1].0, F::from_canonical_u32(((value >> 32) & 0xFFFFFFFF) as u32));
+    }
+
+    // fn get_u64_target(&self, target: U64Target) -> (u64, u64) {
+    //     let x_u64 = self.get_target(target.0).to_canonical_u64();
+    //     let low = x_u64 as u32;
+    //     let high = (x_u64 >> 32) as u32;
+    //     (low, high)
+    // }
+}
+
 // TODO: Circuit needs to be implemented
 pub fn keccak_bench(_size: usize) {
 
@@ -48,9 +70,9 @@ pub fn keccak_bench(_size: usize) {
 
 
 // const KECCAK_WIDTH: usize = 1600;
-// const KECCAK_RATE: usize = 1088;
-// const KECCAK_CAPACITY: usize = KECCAK_WIDTH - KECCAK_RATE;
-// const KECCAK_LANES: usize = KECCAK_WIDTH / 64;
+const KECCAK_RATE: usize = 1088;
+// const KECCAK_CsAPACITY: usize = KECCAK_WIDTH - KECCAK_RATE;
+// const KECCAKs_LANES: usize = KECCAK_WIDTH / 64;
 const KECCAK_ROUNDS: usize = 24;
 
 //TODO: remove the dead codes later
@@ -176,7 +198,7 @@ fn chi<F: RichField + Extendable<D>, const D: usize>(
 //TODO: remove the dead codes later
 #[allow(dead_code)]
 // permutation
-fn keccak_permutation<F: RichField + Extendable<D>, const D: usize>(
+fn keccak_f<F: RichField + Extendable<D>, const D: usize>(
     builder: &mut CircuitBuilder<F, D>,
     state: &mut [[U64Target; 5]; 5]
 ) {
@@ -185,6 +207,175 @@ fn keccak_permutation<F: RichField + Extendable<D>, const D: usize>(
         rho(builder, state);
         pi(builder, state);
         chi(builder, state);
-        iota(builder, state, ROUND_CONSTANTS[i] as usize)
+        iota(builder, state, i)
     }
-}
+}
+
+//TODO: remove the dead codes later
+#[allow(dead_code)]
+fn absorb<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    state: &mut [[U64Target; 5]; 5],
+    message: &[U64Target],
+    rate: usize
+) {
+    let chunks = message.chunks(rate / 64);
+    for block in chunks {
+        for (i, &word) in block.iter().enumerate() {
+            let x = i % 5;
+            let y = i / 5;
+            state[x][y] = builder.xor_u64(state[x][y], word);
+        }
+        keccak_f(builder, state); // Assume keccak_f is implemented as a circuit
+    }
+}
+
+//TODO: remove the dead codes later
+#[allow(dead_code)]
+fn squeeze<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    state: &mut [[U64Target; 5]; 5],
+    rate: usize,
+    output_length: usize
+) -> Vec<U64Target> {
+    let mut hash = Vec::new();
+    while hash.len() * 8 < output_length {
+        for y in 0..5 {
+            for x in 0..5 {
+                if (x + 5 * y) * 64 < rate {
+                    hash.push(state[x][y]);
+                    if hash.len() * 8 >= output_length {
+                        return hash;
+                    }
+                }
+            }
+        }
+        keccak_f(builder, state); // Assume keccak_f is implemented as a circuit
+    }
+    hash
+}
+
+//TODO: remove the dead codes later
+#[allow(dead_code)]
+fn keccak_pad<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    message: &[U64Target],
+    rate: usize
+) -> Vec<U64Target> {
+    let mut padded_message = message.to_vec();
+    let rate_words = rate / 64;
+    let mut pad_len = rate_words - (message.len() % rate_words);
+    
+    if pad_len == 0 {
+        pad_len = rate_words;
+    }
+
+    if pad_len == 1 {
+        padded_message.push(U64Target([builder.one_u32(), builder.one_u32()]));
+    } else {
+        padded_message.push(U64Target([builder.one_u32(), builder.zero_u32()]));
+        for _ in 1..(pad_len - 1) {
+            padded_message.push(U64Target([builder.zero_u32(), builder.zero_u32()]));
+        }
+        padded_message.push(U64Target([builder.zero_u32(), builder.one_u32()]));
+    }
+
+    padded_message
+}
+
+//TODO: remove the dead codes later
+#[allow(dead_code)]
+fn keccak256<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    message: &[U64Target]
+) -> Vec<U64Target> {
+    let mut state = [[builder.zero_u64(); 5]; 5];
+    let rate = KECCAK_RATE;
+
+    // Padding
+    let padded_message = keccak_pad(builder, message, rate);
+
+    // Absorbing
+    absorb(builder, &mut state, &padded_message, rate);
+
+    // Squeezing
+    let hash = squeeze(builder, &mut state, rate, 256);
+
+    hash
+}
+
+#[test]
+fn test_keccak256() {
+    // use plonky2_u32::gadgets::arithmetic_u32::U32Target;
+    // use plonky2::iop::target::Target;
+    use plonky2::iop::witness::PartialWitness;
+    use plonky2::plonk::config::PoseidonGoldilocksConfig;
+    // use plonky2_u32::witness::WitnessU32;
+    type F = GoldilocksField;  // Choose the field used in your implementation.
+    const D: usize = 2;  // This should match the extension degree used.
+
+    // Create circuit builder.
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+
+    // Example input message
+    let input = b"hello";
+    
+    // Convert input to U64Target format used in your Keccak circuit.
+    let input_u64: Vec<U64Target> = input.chunks(8)
+        .map(|chunk| {
+            let mut chunk_padded = [0u8; 8];
+            chunk_padded[..chunk.len()].copy_from_slice(chunk);
+            let value = u64::from_le_bytes(chunk_padded);
+            U64Target([
+                builder.constant_u32(value as u32),
+                builder.constant_u32((value >> 32) as u32),
+            ])
+        })
+        .collect();
+
+    
+
+    // Build the Keccak-256 circuit.
+    let _ = keccak256(&mut builder, &input_u64);
+
+    // Generate the circuit and witness.
+    let data = builder.build::<PoseidonGoldilocksConfig>();
+    let mut pw = PartialWitness::new();
+
+    let input_u64_plain: Vec<u64> = input.chunks(8)
+        .map(|chunk| {
+            let mut chunk_padded = [0u8; 8];
+            chunk_padded[..chunk.len()].copy_from_slice(chunk);
+            u64::from_le_bytes(chunk_padded)
+        })
+        .collect();
+
+    for (i, &byte) in input_u64_plain.iter().enumerate() {
+        pw.set_u64_target(input_u64[i], byte as u64);
+    }
+
+    // Run the circuit.
+    let hash_result = data.prove(pw).unwrap();
+
+    let _ = data.verify(hash_result);
+
+    // // Extract the hash result from the circuit output.
+    // let mut output_bytes = Vec::new();
+    // for target in hash_targets {
+    //     let lo = hash_result.get_u32(target.0[0]) as u64;
+    //     let hi = (hash_result.get_u32(target.0[1]) as u64) << 32;
+    //     let combined = lo | hi;
+    //     output_bytes.extend_from_slice(&combined.to_le_bytes());
+    // }
+
+    // // Truncate to 256 bits (32 bytes).
+    // output_bytes.truncate(32);
+
+    // // Compute the expected hash using a reference implementation.
+    // let expected_hash = keccak256_reference(input);
+
+    // // Compare the circuit output with the expected hash.
+    // assert_eq!(output_bytes, expected_hash, "Keccak-256 hash mismatch");
+    }
+