From 1f7c939b827a993226f43b2d38dd64948fd0043e Mon Sep 17 00:00:00 2001 From: Xinding Wei Date: Sun, 27 Aug 2023 19:17:33 -0700 Subject: [PATCH 1/3] PoseidonHasher supports multiple inputs in compact format --- halo2-base/src/poseidon/hasher/mod.rs | 78 ++++++++- .../src/poseidon/hasher/tests/hasher.rs | 152 +++++++++++++++--- 2 files changed, 208 insertions(+), 22 deletions(-) diff --git a/halo2-base/src/poseidon/hasher/mod.rs b/halo2-base/src/poseidon/hasher/mod.rs index 2816c9fa..255dd484 100644 --- a/halo2-base/src/poseidon/hasher/mod.rs +++ b/halo2-base/src/poseidon/hasher/mod.rs @@ -1,7 +1,7 @@ use crate::{ gates::{GateInstructions, RangeInstructions}, poseidon::hasher::{spec::OptimizedPoseidonSpec, state::PoseidonState}, - safe_types::SafeTypeChip, + safe_types::{SafeBool, SafeTypeChip}, utils::BigPrimeField, AssignedValue, Context, QuantumCell::Constant, @@ -49,6 +49,52 @@ impl PoseidonHasherConsts { + // Right padded inputs. No constrains on paddings. + inputs: [AssignedValue; RATE], + // is_final = 1 triggers squeeze. + is_final: SafeBool, + // Length of `inputs`. + len: AssignedValue, +} + +impl PoseidonCompactInput { + /// Create a new PoseidonCompactInput. + pub fn new( + inputs: [AssignedValue; RATE], + is_final: SafeBool, + len: AssignedValue, + ) -> Self { + Self { inputs, is_final, len } + } + + /// Add data validation constraints. + pub fn add_validation_constraints( + &self, + ctx: &mut Context, + range: &impl RangeInstructions, + ) { + range.is_less_than_safe(ctx, self.len, (RATE + 1) as u64); + // Invalid case: (!is_final && len != RATE) ==> !(is_final || len == RATE) + let is_full: AssignedValue = + range.gate().is_equal(ctx, self.len, Constant(F::from(RATE as u64))); + let invalid_cond = range.gate().or(ctx, *self.is_final.as_ref(), is_full); + range.gate().assert_is_const(ctx, &invalid_cond, &F::ZERO); + } +} + +/// 1 logical row of compact output for Poseidon hasher. +#[derive(Getters)] +pub struct PoseidonCompactOutput { + /// hash of 1 logical input. + #[getset(get = "pub")] + hash: AssignedValue, + /// is_final = 1 ==> this is the end of a logical input. + #[getset(get = "pub")] + is_final: SafeBool, +} + impl PoseidonHasher { /// Create a poseidon hasher from an existing spec. pub fn new(spec: OptimizedPoseidonSpec) -> Self { @@ -147,6 +193,36 @@ impl PoseidonHasher, + range: &impl RangeInstructions, + compact_inputs: &[PoseidonCompactInput], + ) -> Vec> + where + F: BigPrimeField, + { + let mut outputs = Vec::with_capacity(compact_inputs.len()); + let mut state = self.init_state().clone(); + for input in compact_inputs { + // Assume this is the last row of a logical input: + // Depending on if len == RATE. + let is_full = range.gate().is_equal(ctx, input.len, Constant(F::from(RATE as u64))); + // Case 1: if len != RATE. + state.permutation(ctx, range.gate(), &input.inputs, Some(input.len), &self.spec); + // Case 2: if len == RATE, an extra permuation is needed for squeeze. + let mut state_2 = state.clone(); + state_2.permutation(ctx, range.gate(), &[], None, &self.spec); + // Select the result of case 1/2 depending on if len == RATE. + let hash = range.gate().select(ctx, state_2.s[1], state.s[1], is_full); + outputs.push(PoseidonCompactOutput { hash, is_final: input.is_final }); + // Reset state to init_state if this is the end of a logical input. + state.select(ctx, range.gate(), input.is_final, self.init_state()); + } + outputs + } } /// Poseidon sponge. This is stateful. diff --git a/halo2-base/src/poseidon/hasher/tests/hasher.rs b/halo2-base/src/poseidon/hasher/tests/hasher.rs index 24a2e18d..081066c6 100644 --- a/halo2-base/src/poseidon/hasher/tests/hasher.rs +++ b/halo2-base/src/poseidon/hasher/tests/hasher.rs @@ -1,9 +1,14 @@ +use std::cmp::max; + use crate::{ - gates::{circuit::builder::RangeCircuitBuilder, range::RangeInstructions}, + gates::{circuit::builder::RangeCircuitBuilder, range::RangeInstructions, RangeChip}, halo2_proofs::halo2curves::bn256::Fr, - poseidon::hasher::{spec::OptimizedPoseidonSpec, PoseidonHasher}, + poseidon::hasher::{spec::OptimizedPoseidonSpec, PoseidonCompactInput, PoseidonHasher}, + safe_types::{SafeBool, SafeTypeChip}, utils::{testing::base_test, BigPrimeField, ScalarField}, + Context, }; +use halo2_proofs_axiom::arithmetic::Field; use pse_poseidon::Poseidon; use rand::Rng; @@ -15,39 +20,96 @@ struct Payload { pub len: usize, } -// check if the results from hasher and native sponge are same. +// check if the results from hasher and native sponge are same for hash_var_len_array. fn hasher_compatiblity_verification< - F: ScalarField, const T: usize, const RATE: usize, const R_F: usize, const R_P: usize, >( - payloads: Vec>, -) where - F: BigPrimeField, -{ - let lookup_bits = 3; + payloads: Vec>, +) { + base_test().k(12).run(|ctx, range| { + // Construct in-circuit Poseidon hasher. Assuming SECURE_MDS = 0. + let spec = OptimizedPoseidonSpec::::new::(); + let mut hasher = PoseidonHasher::::new(spec); + hasher.initialize_consts(ctx, range.gate()); - let mut builder = RangeCircuitBuilder::new(true).use_lookup_bits(lookup_bits); - let range = builder.range_chip(); - let ctx = builder.main(0); + for payload in payloads { + // Construct native Poseidon sponge. + let mut native_sponge = Poseidon::::new(R_F, R_P); + native_sponge.update(&payload.values[..payload.len]); + let native_result = native_sponge.squeeze(); + let inputs = ctx.assign_witnesses(payload.values); + let len = ctx.load_witness(Fr::from(payload.len as u64)); + let hasher_result = hasher.hash_var_len_array(ctx, range, &inputs, len); + assert_eq!(native_result, *hasher_result.value()); + } + }); +} +// check if the results from hasher and native sponge are same for hash_compact_input. +fn hasher_compact_inputs_compatiblity_verification< + const T: usize, + const RATE: usize, + const R_F: usize, + const R_P: usize, +>( + payloads: Vec>, + ctx: &mut Context, + range: &RangeChip, +) { // Construct in-circuit Poseidon hasher. Assuming SECURE_MDS = 0. - let spec = OptimizedPoseidonSpec::::new::(); - let mut hasher = PoseidonHasher::::new(spec); + let spec = OptimizedPoseidonSpec::::new::(); + let mut hasher = PoseidonHasher::::new(spec); hasher.initialize_consts(ctx, range.gate()); + let mut native_results = Vec::with_capacity(payloads.len()); + let mut compact_inputs = Vec::>::new(); + let rate_witness = ctx.load_constant(Fr::from(RATE as u64)); + let true_witness = ctx.load_constant(Fr::ONE); + let false_witness = ctx.load_zero(); for payload in payloads { + assert!(payload.values.len() % RATE == 0); + assert!(payload.values.len() >= payload.len); + assert!(payload.values.len() == RATE || payload.values.len() - payload.len < RATE); + let num_chunk = payload.values.len() / RATE; + let last_chunk_len = RATE - (payload.values.len() - payload.len); + let inputs = ctx.assign_witnesses(payload.values.clone()); + for (chunk_idx, input_chunk) in inputs.chunks(RATE).enumerate() { + let len_witness = if chunk_idx + 1 == num_chunk { + ctx.load_witness(Fr::from(last_chunk_len as u64)) + } else { + rate_witness + }; + let is_final_witness = SafeTypeChip::unsafe_to_bool(if chunk_idx + 1 == num_chunk { + true_witness + } else { + false_witness + }); + compact_inputs.push(PoseidonCompactInput { + inputs: input_chunk.try_into().unwrap(), + len: len_witness, + is_final: is_final_witness, + }); + } // Construct native Poseidon sponge. - let mut native_sponge = Poseidon::::new(R_F, R_P); + let mut native_sponge = Poseidon::::new(R_F, R_P); native_sponge.update(&payload.values[..payload.len]); let native_result = native_sponge.squeeze(); - let inputs = ctx.assign_witnesses(payload.values); - let len = ctx.load_witness(F::from(payload.len as u64)); - let hasher_result = hasher.hash_var_len_array(ctx, &range, &inputs, len); - // 0x1f0db93536afb96e038f897b4fb5548b6aa3144c46893a6459c4b847951a23b4 - assert_eq!(native_result, *hasher_result.value()); + native_results.push(native_result); + } + let compact_outputs = hasher.hash_compact_input(ctx, range, &compact_inputs); + let mut output_offset = 0; + for (compact_output, compact_input) in compact_outputs.iter().zip(compact_inputs) { + // into() doesn't work if ! is in the beginning in the bool expression... + let is_not_final_input: bool = compact_input.is_final.as_ref().value().is_zero().into(); + let is_not_final_output: bool = compact_output.is_final().as_ref().value().is_zero().into(); + assert_eq!(is_not_final_input, is_not_final_output); + if !is_not_final_output { + assert_eq!(native_results[output_offset], *compact_output.hash().value()); + output_offset += 1; + } } } @@ -98,7 +160,7 @@ fn test_poseidon_hasher_compatiblity() { random_payload(RATE * 2 + 1, RATE * 2 + 1, usize::MAX), random_payload(RATE * 5 + 1, RATE * 5 + 1, usize::MAX), ]; - hasher_compatiblity_verification::(payloads); + hasher_compatiblity_verification::(payloads); } } @@ -127,3 +189,51 @@ fn test_poseidon_hasher_with_prover() { } } } + +#[test] +fn test_poseidon_hasher_compact_inputs() { + { + const T: usize = 3; + const RATE: usize = 2; + let payloads = vec![ + // len == 0 + random_payload(RATE, 0, usize::MAX), + // 0 < len < max_len + random_payload(RATE * 2, RATE + 1, usize::MAX), + random_payload(RATE * 5, RATE * 4 + 1, usize::MAX), + // len == max_len + random_payload(RATE * 2, RATE * 2, usize::MAX), + random_payload(RATE * 5, RATE * 5, usize::MAX), + ]; + base_test().k(12).run(|ctx, range| { + hasher_compact_inputs_compatiblity_verification::(payloads, ctx, range); + }); + } +} + +#[test] +fn test_poseidon_hasher_compact_inputs_with_prover() { + { + const T: usize = 3; + const RATE: usize = 2; + let params = vec![ + (RATE, 0), + (RATE * 2, RATE + 1), + (RATE * 5, RATE * 4 + 1), + (RATE * 2, RATE * 2), + (RATE * 5, RATE * 5), + ]; + let init_payloads = params + .iter() + .map(|(max_len, len)| random_payload(*max_len, *len, usize::MAX)) + .collect::>(); + let logic_payloads = params + .iter() + .map(|(max_len, len)| random_payload(*max_len, *len, usize::MAX)) + .collect::>(); + base_test().k(12).bench_builder(init_payloads, logic_payloads, |pool, range, input| { + let ctx = pool.main(); + hasher_compact_inputs_compatiblity_verification::(input, ctx, range); + }); + } +} From 6de76ac89061149ed9fcea2855ba8c6cb7753026 Mon Sep 17 00:00:00 2001 From: Xinding Wei Date: Sun, 27 Aug 2023 19:23:35 -0700 Subject: [PATCH 2/3] Add comments --- halo2-base/src/poseidon/hasher/mod.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/halo2-base/src/poseidon/hasher/mod.rs b/halo2-base/src/poseidon/hasher/mod.rs index 255dd484..07353b1e 100644 --- a/halo2-base/src/poseidon/hasher/mod.rs +++ b/halo2-base/src/poseidon/hasher/mod.rs @@ -128,6 +128,7 @@ impl PoseidonHasher PoseidonHasher Date: Sun, 27 Aug 2023 21:07:09 -0700 Subject: [PATCH 3/3] Remove unnecessary uses --- halo2-base/src/poseidon/hasher/tests/hasher.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/halo2-base/src/poseidon/hasher/tests/hasher.rs b/halo2-base/src/poseidon/hasher/tests/hasher.rs index 081066c6..c72d3c43 100644 --- a/halo2-base/src/poseidon/hasher/tests/hasher.rs +++ b/halo2-base/src/poseidon/hasher/tests/hasher.rs @@ -1,11 +1,9 @@ -use std::cmp::max; - use crate::{ - gates::{circuit::builder::RangeCircuitBuilder, range::RangeInstructions, RangeChip}, + gates::{range::RangeInstructions, RangeChip}, halo2_proofs::halo2curves::bn256::Fr, poseidon::hasher::{spec::OptimizedPoseidonSpec, PoseidonCompactInput, PoseidonHasher}, - safe_types::{SafeBool, SafeTypeChip}, - utils::{testing::base_test, BigPrimeField, ScalarField}, + safe_types::SafeTypeChip, + utils::{testing::base_test, ScalarField}, Context, }; use halo2_proofs_axiom::arithmetic::Field;