From d2cf4f7c81f178610a3e225deb2129434066eb87 Mon Sep 17 00:00:00 2001 From: Manish Kumar Date: Fri, 31 May 2024 20:42:13 +0530 Subject: [PATCH] keccak hash circuit added, test added(incomplete) --- hash/plonky2/src/bench/keccak.rs | 205 +++++++++++++++++++++++++++++-- 1 file changed, 198 insertions(+), 7 deletions(-) diff --git a/hash/plonky2/src/bench/keccak.rs b/hash/plonky2/src/bench/keccak.rs index 4898ba3..78d0bef 100644 --- a/hash/plonky2/src/bench/keccak.rs +++ b/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 { @@ -25,6 +26,27 @@ fn generate_data(size: usize) -> Vec { } +//TODO: remove the dead codes later +#[allow(dead_code)] +pub trait WitnessU64: Witness { + fn set_u64_target(&mut self, target: U64Target, value: u64); + // fn get_u64_target(&self, target: U64Target) -> (u64, u64); +} + +impl, F: PrimeField64> WitnessU64 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, const D: usize>( //TODO: remove the dead codes later #[allow(dead_code)] // permutation -fn keccak_permutation, const D: usize>( +fn keccak_f, const D: usize>( builder: &mut CircuitBuilder, state: &mut [[U64Target; 5]; 5] ) { @@ -185,6 +207,175 @@ fn keccak_permutation, const D: usize>( rho(builder, state); pi(builder, state); chi(builder, state); - iota(builder, state, ROUND_CONSTANTS[i] as usize) + iota(builder, state, i) } -} \ No newline at end of file +} + +//TODO: remove the dead codes later +#[allow(dead_code)] +fn absorb, const D: usize>( + builder: &mut CircuitBuilder, + 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, const D: usize>( + builder: &mut CircuitBuilder, + state: &mut [[U64Target; 5]; 5], + rate: usize, + output_length: usize +) -> Vec { + 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, const D: usize>( + builder: &mut CircuitBuilder, + message: &[U64Target], + rate: usize +) -> Vec { + 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, const D: usize>( + builder: &mut CircuitBuilder, + message: &[U64Target] +) -> Vec { + 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::::new(config); + + // Example input message + let input = b"hello"; + + // Convert input to U64Target format used in your Keccak circuit. + let input_u64: Vec = 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::(); + let mut pw = PartialWitness::new(); + + let input_u64_plain: Vec = 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"); + } +