diff --git a/Cargo.lock b/Cargo.lock index d4d30eff4..d005179a2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -920,6 +920,7 @@ dependencies = [ "mozak-sdk", "plonky2", "plonky2_maybe_rayon", + "poseidon2", "proptest", "rand", "serde", @@ -1017,6 +1018,7 @@ dependencies = [ "mozak-examples", "mozak-sdk", "plonky2", + "poseidon2", "proptest", "serde", "serde_json", @@ -1031,6 +1033,7 @@ dependencies = [ "itertools 0.12.1", "once_cell", "plonky2", + "poseidon2", "rand", "rand_chacha", "rkyv", @@ -1240,6 +1243,13 @@ dependencies = [ "plotters-backend", ] +[[package]] +name = "poseidon2" +version = "0.1.0" +dependencies = [ + "plonky2", +] + [[package]] name = "powerfmt" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index 43d6b63cb..8f91b99ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ members = [ "signatures", "state", "wasm-demo", + "poseidon2", ] resolver = "2" diff --git a/circuits/Cargo.toml b/circuits/Cargo.toml index 534805716..8e8f0d7b5 100644 --- a/circuits/Cargo.toml +++ b/circuits/Cargo.toml @@ -22,6 +22,7 @@ mozak-runner = { path = "../runner" } mozak-sdk = { path = "../sdk" } plonky2 = { version = "0", default-features = false } plonky2_maybe_rayon = { version = "0", default-features = false } +poseidon2 = { path = "../poseidon2" } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" starky = { version = "0", default-features = false, features = ["std"] } diff --git a/circuits/src/generation/mod.rs b/circuits/src/generation/mod.rs index a85485269..cde6afe88 100644 --- a/circuits/src/generation/mod.rs +++ b/circuits/src/generation/mod.rs @@ -11,6 +11,7 @@ pub mod io_memory; pub mod memory; pub mod memory_zeroinit; pub mod memoryinit; +pub mod poseidon2_preimage_pack; pub mod xor; use std::borrow::Borrow; use std::fmt::Display; @@ -49,6 +50,7 @@ use crate::generation::memory_zeroinit::generate_memory_zero_init_trace; use crate::generation::memoryinit::{ generate_elf_memory_init_trace, generate_mozak_memory_init_trace, }; +use crate::generation::poseidon2_preimage_pack::generate_poseidon2_preimage_pack_trace; use crate::poseidon2::generation::generate_poseidon2_trace; use crate::poseidon2_output_bytes::generation::generate_poseidon2_output_bytes_trace; use crate::poseidon2_sponge::generation::generate_poseidon2_sponge_trace; @@ -101,6 +103,8 @@ pub fn generate_traces, const D: usize>( let cast_list_commitment_tape_rows = generate_cast_list_commitment_tape_trace(&record.executed); let poseiden2_sponge_rows = generate_poseidon2_sponge_trace(&record.executed); let poseidon2_output_bytes_rows = generate_poseidon2_output_bytes_trace(&poseiden2_sponge_rows); + let poseidon2_preimage_pack_rows = + generate_poseidon2_preimage_pack_trace(&poseiden2_sponge_rows); let poseidon2_rows = generate_poseidon2_trace(&record.executed); let memory_rows = generate_memory_trace( @@ -166,6 +170,7 @@ pub fn generate_traces, const D: usize>( poseidon2_stark: trace_rows_to_poly_values(poseidon2_rows), poseidon2_sponge_stark: trace_rows_to_poly_values(poseiden2_sponge_rows), poseidon2_output_bytes_stark: trace_rows_to_poly_values(poseidon2_output_bytes_rows), + poseidon2_preimage_pack: trace_rows_to_poly_values(poseidon2_preimage_pack_rows), tape_commitments_stark: trace_rows_to_poly_values(tape_commitments_rows), } .build() diff --git a/circuits/src/generation/poseidon2_preimage_pack.rs b/circuits/src/generation/poseidon2_preimage_pack.rs new file mode 100644 index 000000000..174a4adde --- /dev/null +++ b/circuits/src/generation/poseidon2_preimage_pack.rs @@ -0,0 +1,55 @@ +use plonky2::hash::hash_types::RichField; + +use crate::poseidon2_preimage_pack::columns::Poseidon2PreimagePack; +use crate::poseidon2_sponge::columns::Poseidon2Sponge; +use crate::utils::pad_trace_with_default; + +pub fn generate_poseidon2_preimage_pack_trace( + poseidon2_sponge_rows: &[Poseidon2Sponge], +) -> Vec> { + let trace: Vec> = poseidon2_sponge_rows + .iter() + .flat_map(Into::>>::into) + .collect(); + pad_trace_with_default(trace) +} + +#[cfg(test)] +mod tests { + use mozak_runner::vm::Row; + use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; + + use crate::generation::MIN_TRACE_LENGTH; + use crate::poseidon2_sponge::generation::generate_poseidon2_sponge_trace; + use crate::test_utils::{create_poseidon2_test, Poseidon2Test}; + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + #[test] + fn generate_poseidon2_preimage_pack_trace() { + let data = "πŸ˜‡ Mozak is knowledge arguments based technology".to_string(); + let input_start_addr = 1024; + let output_start_addr = 2048; + let (_program, record) = create_poseidon2_test(&[Poseidon2Test { + data, + input_start_addr, + output_start_addr, + }]); + + let step_rows = record.executed; + + let sponge_trace = generate_poseidon2_sponge_trace(&step_rows); + let trace = super::generate_poseidon2_preimage_pack_trace(&sponge_trace); + // for one sponge construct we have one row with gen_output = 1. + // So we expect other padding data to make trace of len MIN_TRACE_LENGTH. + assert_eq!(trace.len(), MIN_TRACE_LENGTH); + } + + #[test] + fn generate_poseidon2_trace_with_dummy() { + let step_rows: Vec> = vec![]; + let sponge_trace = generate_poseidon2_sponge_trace(&step_rows); + let trace = super::generate_poseidon2_preimage_pack_trace(&sponge_trace); + assert_eq!(trace.len(), MIN_TRACE_LENGTH); + } +} diff --git a/circuits/src/lib.rs b/circuits/src/lib.rs index 6011e3734..e1964b4db 100644 --- a/circuits/src/lib.rs +++ b/circuits/src/lib.rs @@ -22,6 +22,7 @@ pub mod memory_zeroinit; pub mod memoryinit; pub mod poseidon2; pub mod poseidon2_output_bytes; +pub mod poseidon2_preimage_pack; pub mod poseidon2_sponge; pub mod program; pub mod program_multiplicities; diff --git a/circuits/src/memory/columns.rs b/circuits/src/memory/columns.rs index 146417d98..e68155287 100644 --- a/circuits/src/memory/columns.rs +++ b/circuits/src/memory/columns.rs @@ -6,6 +6,7 @@ use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::Poseidon2Permutation; use plonky2::iop::ext_target::ExtensionTarget; use plonky2::plonk::circuit_builder::CircuitBuilder; +use poseidon2::mozak_poseidon2; use crate::columns_view::{columns_view_impl, make_col_map}; use crate::cross_table_lookup::Column; @@ -14,7 +15,7 @@ use crate::memory_halfword::columns::HalfWordMemory; use crate::memory_io::columns::InputOutputMemory; use crate::memory_zeroinit::columns::MemoryZeroInit; use crate::memoryinit::columns::{MemoryInit, MemoryInitCtl}; -use crate::poseidon2_output_bytes::columns::{Poseidon2OutputBytes, BYTES_COUNT}; +use crate::poseidon2_output_bytes::columns::Poseidon2OutputBytes; use crate::poseidon2_sponge::columns::Poseidon2Sponge; use crate::rangecheck::columns::RangeCheckCtl; use crate::stark::mozak_stark::{MemoryTable, TableWithTypedOutput}; @@ -123,16 +124,25 @@ impl From<&Poseidon2Sponge> for Vec> { if (value.ops.is_permute + value.ops.is_init_permute).is_zero() { vec![] } else { - let rate = Poseidon2Permutation::::RATE; - // each Field element in preimage represents a byte. - (0..rate) - .map(|i| Memory { - clk: value.clk, - addr: value.input_addr - + F::from_canonical_u8(u8::try_from(i).expect("i > 255")), - is_load: F::ONE, - value: value.preimage[i], - ..Default::default() + // each Field element in preimage represents packed data (packed bytes) + (0..Poseidon2Permutation::::RATE) + .flat_map(|fe_index_inside_preimage| { + let base_address = value.input_addr + + mozak_poseidon2::data_capacity_fe::() + * F::from_canonical_usize(fe_index_inside_preimage); + let unpacked = mozak_poseidon2::unpack_to_field_elements( + &value.preimage[fe_index_inside_preimage], + ); + + (0..mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT) + .map(|byte_index_inside_fe| Memory { + clk: value.clk, + addr: base_address + F::from_canonical_usize(byte_index_inside_fe), + is_load: F::ONE, + value: unpacked[byte_index_inside_fe], + ..Default::default() + }) + .collect::>() }) .collect() } @@ -140,19 +150,17 @@ impl From<&Poseidon2Sponge> for Vec> { } impl From<&Poseidon2OutputBytes> for Vec> { - fn from(value: &Poseidon2OutputBytes) -> Self { - if value.is_executed.is_zero() { + fn from(output: &Poseidon2OutputBytes) -> Self { + if output.is_executed.is_zero() { vec![] } else { - (0..BYTES_COUNT) - .map(|i| Memory { - clk: value.clk, - addr: value.output_addr - + F::from_canonical_u8(u8::try_from(i).expect( - "BYTES_COUNT of poseidon output should be representable by a u8", - )), + (0..) + .zip(output.output_bytes) + .map(|(i, value)| Memory { + clk: output.clk, + addr: output.output_addr + F::from_canonical_usize(i), is_store: F::ONE, - value: value.output_bytes[i], + value, ..Default::default() }) .collect() diff --git a/circuits/src/poseidon2_preimage_pack/columns.rs b/circuits/src/poseidon2_preimage_pack/columns.rs new file mode 100644 index 000000000..92c5df99d --- /dev/null +++ b/circuits/src/poseidon2_preimage_pack/columns.rs @@ -0,0 +1,96 @@ +use itertools::Itertools; +use plonky2::hash::hash_types::RichField; +use poseidon2::mozak_poseidon2; + +use crate::columns_view::{columns_view_impl, make_col_map, NumberOfColumns}; +use crate::linear_combination::Column; +use crate::linear_combination_typed::ColumnWithTypedInput; +use crate::memory::columns::MemoryCtl; +use crate::poseidon2::columns::STATE_SIZE; +use crate::poseidon2_sponge::columns::Poseidon2Sponge; +use crate::stark::mozak_stark::{Poseidon2PreimagePackTable, TableWithTypedOutput}; + +#[repr(C)] +#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)] +pub struct Poseidon2PreimagePack { + pub clk: F, + pub byte_addr: F, + pub bytes: [F; mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT], + pub is_executed: F, +} + +columns_view_impl!(Poseidon2PreimagePack); +make_col_map!(PACK, Poseidon2PreimagePack); + +pub const NUM_POSEIDON2_PREIMAGE_PACK_COLS: usize = Poseidon2PreimagePack::<()>::NUMBER_OF_COLUMNS; + +impl From<&Poseidon2Sponge> for Vec> { + // To make it safe for user to change constants + #[allow(clippy::assertions_on_constants)] + fn from(value: &Poseidon2Sponge) -> Self { + if (value.ops.is_init_permute + value.ops.is_permute).is_zero() { + vec![] + } else { + assert!( + mozak_poseidon2::FIELD_ELEMENTS_RATE <= STATE_SIZE, + "Packing RATE (FIELD_ELEMENTS_RATE) should be less or equal than STATE_SIZE" + ); + let preimage: [F; mozak_poseidon2::FIELD_ELEMENTS_RATE] = value.preimage + [..mozak_poseidon2::FIELD_ELEMENTS_RATE] + .try_into() + .expect("Should succeed since preimage can't be empty"); + // For each FE of preimage we have PACK_CAP bytes + preimage + .iter() + .enumerate() + .map(|(i, fe)| Poseidon2PreimagePack { + clk: value.clk, + byte_addr: value.input_addr + + F::from_canonical_usize(i) * mozak_poseidon2::data_capacity_fe::(), + bytes: mozak_poseidon2::unpack_to_field_elements(fe), + is_executed: F::ONE, + }) + .collect_vec() + } + } +} + +columns_view_impl!(Poseidon2SpongePreimagePackCtl); +#[repr(C)] +#[derive(Clone, Copy, Eq, PartialEq, Debug, Default)] +pub struct Poseidon2SpongePreimagePackCtl { + pub clk: T, + pub value: T, + pub byte_addr: T, +} +#[must_use] +pub fn lookup_for_poseidon2_sponge() -> TableWithTypedOutput> +{ + Poseidon2PreimagePackTable::new( + Poseidon2SpongePreimagePackCtl { + clk: PACK.clk, + value: ColumnWithTypedInput::reduce_with_powers(PACK.bytes, 1 << 8), + byte_addr: PACK.byte_addr, + }, + PACK.is_executed, + ) +} + +#[must_use] +pub fn lookup_for_input_memory() -> Vec>> { + (0..) + .zip(PACK.bytes) + .map(|(index, value)| { + Poseidon2PreimagePackTable::new( + MemoryCtl { + clk: PACK.clk, + is_store: ColumnWithTypedInput::constant(0), + is_load: ColumnWithTypedInput::constant(1), + value, + addr: PACK.byte_addr + index, + }, + PACK.is_executed, + ) + }) + .collect() +} diff --git a/circuits/src/poseidon2_preimage_pack/mod.rs b/circuits/src/poseidon2_preimage_pack/mod.rs new file mode 100644 index 000000000..f3494ff54 --- /dev/null +++ b/circuits/src/poseidon2_preimage_pack/mod.rs @@ -0,0 +1,2 @@ +pub mod columns; +pub mod stark; diff --git a/circuits/src/poseidon2_preimage_pack/stark.rs b/circuits/src/poseidon2_preimage_pack/stark.rs new file mode 100644 index 000000000..538213fbc --- /dev/null +++ b/circuits/src/poseidon2_preimage_pack/stark.rs @@ -0,0 +1,86 @@ +use std::marker::PhantomData; + +use mozak_circuits_derive::StarkNameDisplay; +use plonky2::field::extension::{Extendable, FieldExtension}; +use plonky2::field::packed::PackedField; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::ext_target::ExtensionTarget; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; +use starky::evaluation_frame::StarkFrame; +use starky::stark::Stark; + +use super::columns::{Poseidon2PreimagePack, NUM_POSEIDON2_PREIMAGE_PACK_COLS}; +use crate::columns_view::HasNamedColumns; + +#[derive(Copy, Clone, Default, StarkNameDisplay, Debug)] +#[allow(clippy::module_name_repetitions)] +pub struct Poseidon2PreimagePackStark { + pub _f: PhantomData, +} + +impl HasNamedColumns for Poseidon2PreimagePackStark { + type Columns = Poseidon2PreimagePack; +} + +const COLUMNS: usize = NUM_POSEIDON2_PREIMAGE_PACK_COLS; +const PUBLIC_INPUTS: usize = 0; + +impl, const D: usize> Stark + for Poseidon2PreimagePackStark +{ + type EvaluationFrame = StarkFrame + where + FE: FieldExtension, + P: PackedField; + type EvaluationFrameTarget = + StarkFrame, ExtensionTarget, COLUMNS, PUBLIC_INPUTS>; + + fn eval_packed_generic( + &self, + _vars: &Self::EvaluationFrame, + _yield_constr: &mut ConstraintConsumer

, + ) where + FE: FieldExtension, + P: PackedField, { + } + + fn eval_ext_circuit( + &self, + _builder: &mut CircuitBuilder, + _vars: &Self::EvaluationFrameTarget, + _yield_constr: &mut RecursiveConstraintConsumer, + ) { + } + + fn constraint_degree(&self) -> usize { 3 } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use plonky2::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig}; + use starky::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree}; + + use super::Poseidon2PreimagePackStark; + + const D: usize = 2; + type C = Poseidon2GoldilocksConfig; + type F = >::F; + type S = Poseidon2PreimagePackStark; + + #[test] + // TODO(Roman): FIXME + #[ignore] + fn poseidon2_stark_degree() -> Result<()> { + let stark = S::default(); + test_stark_low_degree(stark) + } + #[test] + fn test_circuit() -> anyhow::Result<()> { + let stark = S::default(); + test_stark_circuit_constraints::(stark)?; + + Ok(()) + } +} diff --git a/circuits/src/poseidon2_sponge/columns.rs b/circuits/src/poseidon2_sponge/columns.rs index 1495d9c8a..7a2dab0e2 100644 --- a/circuits/src/poseidon2_sponge/columns.rs +++ b/circuits/src/poseidon2_sponge/columns.rs @@ -2,6 +2,7 @@ use core::ops::Add; use plonky2::hash::hash_types::NUM_HASH_OUT_ELTS; use plonky2::hash::poseidon2::WIDTH; +use poseidon2::mozak_poseidon2; use crate::columns_view::{columns_view_impl, make_col_map, NumberOfColumns}; use crate::cross_table_lookup::ColumnWithTypedInput; @@ -9,6 +10,7 @@ use crate::linear_combination::Column; use crate::memory::columns::MemoryCtl; use crate::poseidon2::columns::Poseidon2StateCtl; use crate::poseidon2_output_bytes::columns::Poseidon2OutputBytesCtl; +use crate::poseidon2_preimage_pack::columns::Poseidon2SpongePreimagePackCtl; use crate::stark::mozak_stark::{Poseidon2SpongeTable, TableWithTypedOutput}; #[repr(C)] @@ -102,3 +104,24 @@ pub fn lookup_for_input_memory(limb_index: u8) -> TableWithTypedOutput Vec>> { + (0..8) + .zip(COL_MAP.preimage) + .map(|(limb_index, value)| { + Poseidon2SpongeTable::new( + Poseidon2SpongePreimagePackCtl { + clk: COL_MAP.clk, + value, + byte_addr: COL_MAP.input_addr + + limb_index + * i64::try_from(mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT) + .expect("Should be < 255"), + }, + COL_MAP.ops.is_init_permute + COL_MAP.ops.is_permute, + ) + }) + .collect() +} diff --git a/circuits/src/poseidon2_sponge/generation.rs b/circuits/src/poseidon2_sponge/generation.rs index 2f975329c..08f50eb31 100644 --- a/circuits/src/poseidon2_sponge/generation.rs +++ b/circuits/src/poseidon2_sponge/generation.rs @@ -3,6 +3,7 @@ use mozak_runner::vm::Row; use plonky2::hash::hash_types::RichField; use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::Poseidon2Permutation; +use poseidon2::mozak_poseidon2; use crate::poseidon2_sponge::columns::{Ops, Poseidon2Sponge}; use crate::utils::pad_trace_with_default; @@ -40,7 +41,7 @@ fn unroll_sponge_data(row: &Row) -> Vec> { output: sponge_datum.output, gen_output: sponge_datum.gen_output, }); - input_addr += rate_size; + input_addr += u32::try_from(mozak_poseidon2::DATA_PADDING).expect("Should succeed"); input_len -= rate_size; } @@ -65,11 +66,11 @@ pub fn generate_poseidon2_sponge_trace( #[cfg(test)] mod test { - use plonky2::field::types::Field; use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::Poseidon2Permutation; use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; + use poseidon2::mozak_poseidon2; use crate::generation::MIN_TRACE_LENGTH; use crate::poseidon2_sponge::columns::Poseidon2Sponge; @@ -81,7 +82,7 @@ mod test { #[test] fn generate_poseidon2_sponge_trace() { let data = "πŸ˜‡ Mozak is knowledge arguments based technology".to_string(); - let data_len_in_bytes = data.as_bytes().len(); + let data_len_in_bytes = mozak_poseidon2::do_padding(data.as_bytes()).len(); let input_start_addr = 1024; let output_start_addr = 2048; let (_program, record) = create_poseidon2_test(&[Poseidon2Test { @@ -94,7 +95,8 @@ mod test { let trace = super::generate_poseidon2_sponge_trace(&step_rows); let rate_size = Poseidon2Permutation::::RATE; - let sponge_count = data_len_in_bytes / rate_size; + let sponge_count = + (data_len_in_bytes / mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT) / rate_size; for (i, value) in trace.iter().enumerate().take(sponge_count) { assert_eq!( value.input_addr, diff --git a/circuits/src/poseidon2_sponge/stark.rs b/circuits/src/poseidon2_sponge/stark.rs index 1e573269a..287b7684e 100644 --- a/circuits/src/poseidon2_sponge/stark.rs +++ b/circuits/src/poseidon2_sponge/stark.rs @@ -9,6 +9,7 @@ use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::Poseidon2Permutation; use plonky2::iop::ext_target::ExtensionTarget; use plonky2::plonk::circuit_builder::CircuitBuilder; +use poseidon2::mozak_poseidon2; use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer}; use starky::evaluation_frame::{StarkEvaluationFrame, StarkFrame}; use starky::stark::Stark; @@ -55,6 +56,9 @@ impl, const D: usize> Stark for Poseidon2Spon let rate = u8::try_from(Poseidon2Permutation::::RATE).expect("rate > 255"); let state_size = u8::try_from(Poseidon2Permutation::::WIDTH).expect("state_size > 255"); let rate_scalar = P::Scalar::from_canonical_u8(rate); + let padding_scalar = P::Scalar::from_canonical_u8( + u8::try_from(mozak_poseidon2::DATA_PADDING).expect("DATA_PADDING > 255"), + ); let lv: &Poseidon2Sponge

= vars.get_local_values().into(); let nv: &Poseidon2Sponge

= vars.get_next_values().into(); @@ -92,9 +96,9 @@ impl, const D: usize> Stark for Poseidon2Spon // input yield_constr .constraint_transition(not_last_sponge * (lv.input_len - (nv.input_len + rate_scalar))); - // and input_addr increases by RATE + // and input_addr increases by DATA_PADDING yield_constr.constraint_transition( - not_last_sponge * (lv.input_addr - (nv.input_addr - rate_scalar)), + not_last_sponge * (lv.input_addr - (nv.input_addr - padding_scalar)), ); // For each init_permute capacity bits are zero. @@ -141,6 +145,9 @@ impl, const D: usize> Stark for Poseidon2Spon // if row generates output then it must be last rate sized // chunk of input. let rate_ext = builder.constant_extension(F::Extension::from_canonical_u8(rate)); + let padding_ext = builder.constant_extension(F::Extension::from_canonical_u8( + u8::try_from(mozak_poseidon2::DATA_PADDING).expect("DATA_PADDING > 255"), + )); let input_len_sub_rate = builder.sub_extension(lv.input_len, rate_ext); let gen_op_len_check = builder.mul_extension(lv.gen_output, input_len_sub_rate); yield_constr.constraint(builder, gen_op_len_check); @@ -174,8 +181,8 @@ impl, const D: usize> Stark for Poseidon2Spon // length decreases by RATE, note that only actual execution row can consume // input yield_constr.constraint_transition(builder, len_check); - // and input_addr increases by RATE - let nv_input_addr_rate = builder.sub_extension(nv.input_addr, rate_ext); + // and input_addr increases by PADDING + let nv_input_addr_rate = builder.sub_extension(nv.input_addr, padding_ext); let input_addr_diff_rate = builder.sub_extension(lv.input_addr, nv_input_addr_rate); let addr_check = builder.mul_extension(not_last_sponge, input_addr_diff_rate); yield_constr.constraint_transition(builder, addr_check); diff --git a/circuits/src/stark/mozak_stark.rs b/circuits/src/stark/mozak_stark.rs index 9880dc646..2dbf67c1a 100644 --- a/circuits/src/stark/mozak_stark.rs +++ b/circuits/src/stark/mozak_stark.rs @@ -33,6 +33,10 @@ use crate::poseidon2::columns::{Poseidon2State, Poseidon2StateCtl}; use crate::poseidon2::stark::Poseidon2_12Stark; use crate::poseidon2_output_bytes::columns::{Poseidon2OutputBytes, Poseidon2OutputBytesCtl}; use crate::poseidon2_output_bytes::stark::Poseidon2OutputBytesStark; +use crate::poseidon2_preimage_pack::columns::{ + Poseidon2PreimagePack, Poseidon2SpongePreimagePackCtl, +}; +use crate::poseidon2_preimage_pack::stark::Poseidon2PreimagePackStark; use crate::poseidon2_sponge::columns::{Poseidon2Sponge, Poseidon2SpongeCtl}; use crate::poseidon2_sponge::stark::Poseidon2SpongeStark; use crate::program::columns::{InstructionRow, ProgramRom}; @@ -59,11 +63,11 @@ use crate::xor::columns::{XorColumnsView, XorView}; use crate::xor::stark::XorStark; use crate::{ bitshift, cpu, memory, memory_fullword, memory_halfword, memory_io, memory_zeroinit, - memoryinit, poseidon2_output_bytes, poseidon2_sponge, program, program_multiplicities, + memoryinit, poseidon2_output_bytes, poseidon2_preimage_pack, program, program_multiplicities, rangecheck, register, xor, }; -const NUM_CROSS_TABLE_LOOKUP: usize = 17; +const NUM_CROSS_TABLE_LOOKUP: usize = 18; const NUM_PUBLIC_SUB_TABLES: usize = 2; /// STARK Gadgets of Mozak-VM @@ -139,6 +143,8 @@ pub struct MozakStark, const D: usize> { pub poseidon2_sponge_stark: Poseidon2SpongeStark, #[StarkSet(stark_kind = "Poseidon2OutputBytes")] pub poseidon2_output_bytes_stark: Poseidon2OutputBytesStark, + #[StarkSet(stark_kind = "Poseidon2PreimagePack")] + pub poseidon2_preimage_pack: Poseidon2PreimagePackStark, #[StarkSet(stark_kind = "TapeCommitments")] pub tape_commitments_stark: TapeCommitmentsStark, pub cross_table_lookups: [CrossTableLookup; NUM_CROSS_TABLE_LOOKUP], @@ -428,6 +434,7 @@ impl, const D: usize> Default for MozakStark poseidon2_sponge_stark: Poseidon2SpongeStark::default(), poseidon2_stark: Poseidon2_12Stark::default(), poseidon2_output_bytes_stark: Poseidon2OutputBytesStark::default(), + poseidon2_preimage_pack: Poseidon2PreimagePackStark::default(), tape_commitments_stark: TapeCommitmentsStark::default(), // These tables contain only descriptions of the tables. @@ -448,6 +455,7 @@ impl, const D: usize> Default for MozakStark Poseidon2SpongeCpuTable::lookups(), Poseidon2Poseidon2SpongeTable::lookups(), Poseidon2OutputBytesPoseidon2SpongeTable::lookups(), + Poseidon2Sponge2Poseidon2PreimagePackTable::lookups(), EventCommitmentTapeIOLookupTable::lookups(), CastlistCommitmentTapeIOLookupTable::lookups(), ], @@ -635,6 +643,11 @@ table_impl!( TableKind::Poseidon2OutputBytes, Poseidon2OutputBytes ); +table_impl!( + Poseidon2PreimagePackTable, + TableKind::Poseidon2PreimagePack, + Poseidon2PreimagePack +); pub trait Lookups { type Row: IntoIterator; @@ -696,7 +709,7 @@ impl Lookups for IntoMemoryTable { memory_io::columns::lookup_for_memory(TableKind::EventsCommitmentTape), memory_io::columns::lookup_for_memory(TableKind::CastListCommitmentTape), ], - (0..8).map(poseidon2_sponge::columns::lookup_for_input_memory), + poseidon2_preimage_pack::columns::lookup_for_input_memory(), (0..32).map(poseidon2_output_bytes::columns::lookup_for_output_memory), ] .collect(); @@ -912,3 +925,16 @@ impl Lookups for CastlistCommitmentTapeIOLookupTable { ) } } + +pub struct Poseidon2Sponge2Poseidon2PreimagePackTable; + +impl Lookups for Poseidon2Sponge2Poseidon2PreimagePackTable { + type Row = Poseidon2SpongePreimagePackCtl; + + fn lookups_with_typed_output() -> CrossTableLookupWithTypedOutput { + CrossTableLookupWithTypedOutput::new( + crate::poseidon2_sponge::columns::lookup_for_preimage_pack(), + vec![poseidon2_preimage_pack::columns::lookup_for_poseidon2_sponge()], + ) + } +} diff --git a/circuits/src/stark/prover.rs b/circuits/src/stark/prover.rs index 35110cca1..6c3dc7ec2 100644 --- a/circuits/src/stark/prover.rs +++ b/circuits/src/stark/prover.rs @@ -379,9 +379,9 @@ mod tests { use mozak_runner::code; use mozak_runner::instruction::{Args, Instruction, Op}; use plonky2::field::goldilocks_field::GoldilocksField; - use plonky2::field::types::Field; use plonky2::hash::poseidon2::Poseidon2Hash; use plonky2::plonk::config::{GenericHashOut, Hasher}; + use poseidon2::mozak_poseidon2; use crate::stark::mozak_stark::MozakStark; use crate::test_utils::{create_poseidon2_test, Poseidon2Test, ProveAndVerify}; @@ -454,14 +454,14 @@ mod tests { .load_u8(test_datum.output_start_addr + u32::from(i)) }) .collect(); - let mut data_bytes = test_datum.data.as_bytes().to_vec(); - // VM expects input len to be multiple of RATE bits - data_bytes.resize(data_bytes.len().next_multiple_of(8), 0_u8); - let data_fields: Vec = data_bytes - .iter() - .map(|x| GoldilocksField::from_canonical_u8(*x)) - .collect(); - assert_eq!(output, Poseidon2Hash::hash_no_pad(&data_fields).to_bytes()); + let data_fields: Vec = mozak_poseidon2::pack_padded_input( + mozak_poseidon2::do_padding(test_datum.data.as_bytes()).as_slice(), + ); + assert_eq!( + output, + Poseidon2Hash::hash_no_pad(&data_fields).to_bytes(), + "Expected vm-computed output, does not equal to plonky2 version" + ); } MozakStark::prove_and_verify(&program, &record).unwrap(); } diff --git a/circuits/src/test_utils.rs b/circuits/src/test_utils.rs index e5d7ba026..8504871d8 100644 --- a/circuits/src/test_utils.rs +++ b/circuits/src/test_utils.rs @@ -18,6 +18,7 @@ use plonky2::plonk::circuit_data::CircuitConfig; use plonky2::plonk::config::{GenericConfig, Hasher, Poseidon2GoldilocksConfig}; use plonky2::util::log2_ceil; use plonky2::util::timing::TimingTree; +use poseidon2::mozak_poseidon2; use starky::config::StarkConfig; use starky::prover::prove as prove_table; use starky::stark::Stark; @@ -494,9 +495,7 @@ pub fn create_poseidon2_test( let mut memory: Vec<(u32, u8)> = vec![]; for test_datum in test_data { - let mut data_bytes = test_datum.data.as_bytes().to_vec(); - // VM expects input len to be multiple of RATE bits - data_bytes.resize(data_bytes.len().next_multiple_of(8), 0_u8); + let data_bytes = mozak_poseidon2::do_padding(test_datum.data.as_bytes()); let data_len = data_bytes.len(); let input_memory: Vec<(u32, u8)> = izip!((test_datum.input_start_addr..), data_bytes).collect(); diff --git a/examples/Cargo.lock b/examples/Cargo.lock index 66c032d25..89071c7d3 100644 --- a/examples/Cargo.lock +++ b/examples/Cargo.lock @@ -318,6 +318,7 @@ dependencies = [ "itertools", "once_cell", "plonky2", + "poseidon2", "rand", "rand_chacha", "rkyv", @@ -479,6 +480,13 @@ name = "plonky2_util" version = "0.2.0" source = "git+https://github.com/0xmozak/plonky2.git#51f540a0e2a9bd9d6fc6234c6e62d167eaa7c707" +[[package]] +name = "poseidon2" +version = "0.1.0" +dependencies = [ + "plonky2", +] + [[package]] name = "ppv-lite86" version = "0.2.17" diff --git a/poseidon2/Cargo.toml b/poseidon2/Cargo.toml new file mode 100644 index 000000000..028f55da4 --- /dev/null +++ b/poseidon2/Cargo.toml @@ -0,0 +1,14 @@ +[package] +categories = ["development-tools"] +description = "Provides poseidon2 utils for writing programs for mozak platform" +edition = "2021" +keywords = ["hash"] +license = "Apache-2.0" +name = "poseidon2" +readme = "README.md" +repository = "https://github.com/0xmozak/mozak-node/poseidon2" +version = "0.1.0" + +[dependencies] +[target.'cfg(not(target_os="mozakvm"))'.dependencies] +plonky2 = { git = "https://github.com/0xmozak/plonky2.git", default-features = false } diff --git a/poseidon2/README.md b/poseidon2/README.md new file mode 100644 index 000000000..99605ca5e --- /dev/null +++ b/poseidon2/README.md @@ -0,0 +1,2 @@ +# Mozak Poseidon2 helpers + diff --git a/poseidon2/src/lib.rs b/poseidon2/src/lib.rs new file mode 100644 index 000000000..ad867eab9 --- /dev/null +++ b/poseidon2/src/lib.rs @@ -0,0 +1,2 @@ +#![feature(restricted_std)] +pub mod mozak_poseidon2; diff --git a/poseidon2/src/mozak_poseidon2.rs b/poseidon2/src/mozak_poseidon2.rs new file mode 100644 index 000000000..66a31536a --- /dev/null +++ b/poseidon2/src/mozak_poseidon2.rs @@ -0,0 +1,102 @@ +#[cfg(not(target_os = "mozakvm"))] +use plonky2::hash::hash_types::RichField; +#[cfg(not(target_os = "mozakvm"))] +use plonky2::hash::hashing::PlonkyPermutation; +#[cfg(not(target_os = "mozakvm"))] +use plonky2::hash::poseidon2::Poseidon2Permutation; +#[allow(clippy::module_name_repetitions)] +pub const DATA_CAPACITY_PER_FIELD_ELEMENT: usize = 7; +pub const DATA_PADDING: usize = DATA_CAPACITY_PER_FIELD_ELEMENT * FIELD_ELEMENTS_RATE; +pub const EMPTY_BYTES: usize = MAX_BYTES_PER_FIELD_ELEMENT - DATA_CAPACITY_PER_FIELD_ELEMENT; +pub const FIELD_ELEMENTS_RATE: usize = 8; +pub const MAX_BYTES_PER_FIELD_ELEMENT: usize = 8; + +#[must_use] +#[cfg(not(target_os = "mozakvm"))] +pub fn data_capacity_fe() -> F { + F::from_canonical_usize(DATA_CAPACITY_PER_FIELD_ELEMENT) +} + +#[must_use] +#[cfg(not(target_os = "mozakvm"))] +pub fn data_padding_fe() -> F { F::from_canonical_usize(DATA_PADDING) } + +/// Byte padding +/// Bit-Padding schema is used to pad input data +/// Case-A - data length % `DATA_PADDING` != 0 +/// --> Make first bit of the first padded byte to be 1 - `0b0000_0001` +/// Case-B - data length % `DATA_PADDING` == 0 +/// --> Extend padding to next-multiple of `DATA_PADDING` while first bit of +/// the first padded byte will be 1 (same as for Case-A) +#[must_use] +pub fn do_padding(data: &[u8]) -> Vec { + let mut padded = data.to_vec(); + padded.push(1); + padded.resize(padded.len().next_multiple_of(DATA_PADDING), 0); + padded +} + +/// # Panics +/// +/// Panics if `Self::DATA_CAPACITY_PER_FIELD_ELEMENT < +/// Self::BYTES_PER_FIELD_ELEMENT` +#[must_use] +// To make it safe for user to change constants +#[allow(clippy::assertions_on_constants)] +#[cfg(not(target_os = "mozakvm"))] +pub fn pack_padded_input(data: &[u8]) -> Vec { + assert_eq!( + Poseidon2Permutation::::RATE, + FIELD_ELEMENTS_RATE, + "Poseidon2Permutation::::RATE: {:?} differs from mozak_poseidon2::FIELD_ELEMENTS_RATE: {:?} - is not supported", + Poseidon2Permutation::::RATE, + FIELD_ELEMENTS_RATE + ); + assert!( + DATA_CAPACITY_PER_FIELD_ELEMENT < MAX_BYTES_PER_FIELD_ELEMENT, + "For 64 bit field maximum supported packing is 7 bytes" + ); + assert_eq!(data.len() % DATA_PADDING, 0, "Allow only padded byte-data"); + data.chunks(DATA_CAPACITY_PER_FIELD_ELEMENT) + .map(pack_to_field_element) + .collect() +} + +/// # Panics +/// +/// Panics if `leading-zeros + data` isn't convertable to u64 (length > +/// eight bytes) +#[must_use] +#[cfg(not(target_os = "mozakvm"))] +pub fn pack_to_field_element(data: &[u8]) -> F { + // Note: postfix with zeros for LE case + let mut data_extended_with_zeros: Vec = data.to_vec(); + data_extended_with_zeros.extend([0_u8; EMPTY_BYTES]); + assert!( + data_extended_with_zeros.len() <= 8, + "data_extended_with_zeros.len {:?} can't be packed to u64::bytes (8)", + data_extended_with_zeros.len() + ); + + F::from_canonical_u64(u64::from_le_bytes( + data_extended_with_zeros + .as_slice() + .try_into() + .expect("pack bytes to single u64 should succeed"), + )) +} + +/// # Panics +/// When `Self::DATA_CAPACITY_PER_FIELD_ELEMENT` is larger than the number +/// of bytes in a u64, ie 8. +#[cfg(not(target_os = "mozakvm"))] +pub fn unpack_to_bytes(fe: &F) -> [u8; DATA_CAPACITY_PER_FIELD_ELEMENT] { + fe.to_canonical_u64().to_le_bytes()[..DATA_CAPACITY_PER_FIELD_ELEMENT] + .try_into() + .unwrap() +} + +#[cfg(not(target_os = "mozakvm"))] +pub fn unpack_to_field_elements(fe: &F) -> [F; DATA_CAPACITY_PER_FIELD_ELEMENT] { + unpack_to_bytes(fe).map(F::from_canonical_u8) +} diff --git a/runner/Cargo.toml b/runner/Cargo.toml index b13358277..01b54857f 100644 --- a/runner/Cargo.toml +++ b/runner/Cargo.toml @@ -20,6 +20,7 @@ itertools = "0.12" log = "0.4" mozak-sdk = { path = "../sdk" } plonky2 = { version = "0", default-features = false } +poseidon2 = { path = "../poseidon2" } proptest = { version = "1.4", optional = true } serde = { version = "1.0", features = ["derive"] } diff --git a/runner/src/poseidon2.rs b/runner/src/poseidon2.rs index c6e9423db..a19a8d99a 100644 --- a/runner/src/poseidon2.rs +++ b/runner/src/poseidon2.rs @@ -1,11 +1,12 @@ use std::iter::repeat; -use itertools::{chain, izip}; +use itertools::{chain, izip, Itertools}; use mozak_sdk::core::reg_abi::{REG_A1, REG_A2, REG_A3}; use plonky2::hash::hash_types::{HashOut, RichField, NUM_HASH_OUT_ELTS}; use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::{Poseidon2Permutation, WIDTH}; use plonky2::plonk::config::GenericHashOut; +use poseidon2::mozak_poseidon2; use crate::state::{Aux, State}; @@ -24,19 +25,26 @@ pub struct Entry { pub sponge_data: Vec>, } +// TODO(Roman): move it to `poseidon2` crate // Based on hash_n_to_m_no_pad() from plonky2/src/hash/hashing.rs -/// This function is sponge function which uses poseidon2 permutation function. +/// This function is sponge function that uses poseidon2 permutation function. /// Input must be multiple of 8 bytes. It absorbs all input and the squeezes /// `NUM_HASH_OUT_ELTS` Field elements to generate `HashOut`. +/// Why do we use only 4 field elements from our Poseidon2 output, but we are +/// computing 8? (I.e. β€˜rate’ is set to 8.) Technically, we could set the rate +/// to 4 (with permuting 8 -> 8). However, we (Vivek) opted for a rate of 8 is +/// because: first, it's more efficient; with each permutation, a rate of 8/12 +/// (rate/width) achieves higher throughput than 4/8. Second, this approach +/// adheres to the sponge logic defined in Plonky2 +/// # Panics /// -/// # Panics -/// -/// Panics if `PlonkyPermutation` is implemented on `STATE_SIZE` different than +/// Panics if `PlonkyPermutation` is implemented on `STATE_SIZE` different from /// 12. pub fn hash_n_to_m_no_pad>( inputs: &[F], ) -> (HashOut, Vec>) { let permute_and_record_data = |perm: &mut P, sponge_data: &mut Vec>| { + // STATE_SIZE is 12 since it's hard-coded in our stark-backend const STATE_SIZE: usize = 12; assert_eq!(STATE_SIZE, P::WIDTH); let preimage: [F; STATE_SIZE] = perm @@ -48,31 +56,49 @@ pub fn hash_n_to_m_no_pad>( .as_ref() .try_into() .expect("length must be equal to poseidon2 STATE_SIZE"); + // `sponge_data` is previous `perm` and `perm` permutation + // `gen_output` is a flag that will be used later sponge_data.push(SpongeData { preimage, output, gen_output: F::from_bool(false), }); }; - + // `perm` defined in such a way that this statement will be ALWAYS true: + // Some(F::ZERO) == perm.next() let mut perm = P::new(repeat(F::ZERO)); - // input length is expected to be multiple of P::RATE - assert_eq!(inputs.len() % P::RATE, 0); + + assert_eq!( + inputs.len() % P::RATE, + 0, + "input_len: {:?} is expected to be multiple of P::RATE {:?}", + inputs.len(), + P::RATE, + ); let mut sponge_data = Vec::new(); // Absorb all input chunks. + // Divide input-FE to chunks of P::RATE (8), so chunk = 8 FE for chunk in inputs.chunks(P::RATE) { + // put `chunk` elements inside `perm` starting from index 0 - it means always + // put at the beginning of the `perm` perm.set_from_slice(chunk, 0); + // run the function that executes the permutation and append a new `sponge_data` + // element permute_and_record_data(&mut perm, &mut sponge_data); } + // `perm.squeeze` return P::RATE (8) elements, from these elements only first + // NUM_HASH_OUT_ELTS (4) is taken let outputs: [F; NUM_HASH_OUT_ELTS] = perm.squeeze()[..NUM_HASH_OUT_ELTS] .try_into() .expect("squeeze must have minimum NUM_HASH_OUT_ELTS length"); + // set the flag for the last `sponge_data` element sponge_data .last_mut() .expect("Can't fail at least one elem must be there") .gen_output = F::from_bool(true); + // `from` function just takes 4 elements array and create HashOut from it (HashOut::from(outputs), sponge_data) } @@ -81,19 +107,52 @@ impl State { /// # Panics /// /// Panics if hash output of `hash_n_to_m_no_pad` has length different - /// then expected value. + /// from expected value. + /// Note: `ecall_poseidon2` works with 3 parameters: + /// 1) Input-Data - The data we want to hash - represented as `input_ptr` + /// 2) Input-Data-Length - represented as `input_len` + /// 3) Output-Hash - the expected hash value - represented as `output_ptr` + /// Output-Hash size is constant and expected to be 32 bytes - 4 FE - 256b pub fn ecall_poseidon2(self) -> (Aux, Self) { + // In this step, we're taking 3 `ecall_poseidon2` arguments let input_ptr = self.get_register_value(REG_A1); // lengths are in bytes let input_len = self.get_register_value(REG_A2); let output_ptr = self.get_register_value(REG_A3); - let input: Vec = (0..input_len) - .map(|i| F::from_canonical_u8(self.load_u8(input_ptr + i))) - .collect(); + // In this step, we're mapping one-to-one, bytes to FE + // So if initial data is 32 bytes -> input-vector will 32 FE + // Pay attention that `self.load8` loads from memory + // Note: I am not sure why we map byte-to-FE and not 7 bytes to FE + assert_eq!( + input_len % u32::try_from(mozak_poseidon2::DATA_PADDING).expect("Cast from usize to u32 for mozak_poseidon2::DATA_CAPACITY_PER_FIELD_ELEMENT should succeed"), + 0, + "Require padded input data length: {:?} in bytes with respect to {:?}", + input_len, + mozak_poseidon2::DATA_PADDING + ); + let input = mozak_poseidon2::pack_padded_input( + (0..input_len) + .map(|i| self.load_u8(input_ptr + i)) + .collect_vec() + .as_slice(), + ); + // This is the most important step, since here the actual `poseidon2` hash + // computation's taken place. This function returns `computed` hash-value and + // the intermediate `sponge_data` let (hash, sponge_data) = hash_n_to_m_no_pad::>(input.as_slice()); + // In this step, HashOut translated to bytes. Nothing special here, since + // internally it is just to_canonical64 and to 8 bytes. So it is just byte + // representation. The problem is that the same preimage can give us 2 different + // hashes, because we can add F::ORDER as in this example: + // let x = x.to_canonical_u64(); + // x.checked_add(F::ORDER).unwrap_or(x).to_le_bytes() + // poseidon constraints don't ensure this let hash = hash.to_bytes(); - assert_eq!(32, hash.len()); + assert_eq!(32, hash.len(), "Supported hash-length in bytes is: 32"); + // In this step, 2 things happen: + // 1) Fill the Aux.poseidon2 entry + // 2) Store the computed hash inside `output_ptr` let mem_addresses_used: Vec = chain!( (0..input_len).map(|i| input_ptr.wrapping_add(i)), @@ -106,9 +165,8 @@ impl State { poseidon2: Some(Entry { addr: input_ptr, output_addr: output_ptr, - len: input_len.next_multiple_of( - u32::try_from(Poseidon2Permutation::::RATE).expect("RATE > 2^32"), - ), + len: u32::try_from(Poseidon2Permutation::::RATE * sponge_data.len()) + .unwrap(), sponge_data, }), ..Default::default() @@ -127,26 +185,16 @@ impl State { #[cfg(test)] mod tests { use plonky2::field::goldilocks_field::GoldilocksField; - use plonky2::field::types::Field; - use plonky2::hash::hashing::PlonkyPermutation; use plonky2::hash::poseidon2::{Poseidon2Hash, Poseidon2Permutation}; use plonky2::plonk::config::{GenericHashOut, Hasher}; + use poseidon2::mozak_poseidon2; #[test] fn test_hash_n_to_m_no_pad() { let data = "πŸ’₯ Mozak-VM Rocks With Poseidon2"; - let mut data_bytes = data.as_bytes().to_vec(); - // VM expects input length to be multiple of RATE - data_bytes.resize( - data_bytes - .len() - .next_multiple_of(Poseidon2Permutation::::RATE), - 0, + let data_fields: Vec = mozak_poseidon2::pack_padded_input( + mozak_poseidon2::do_padding(data.as_bytes()).as_slice(), ); - let data_fields: Vec = data_bytes - .iter() - .map(|x| GoldilocksField::from_canonical_u8(*x)) - .collect(); let (hash, _sponge_data) = super::hash_n_to_m_no_pad::< GoldilocksField, Poseidon2Permutation, diff --git a/sdk/Cargo.lock b/sdk/Cargo.lock index ed02ad87e..ab1239fc0 100644 --- a/sdk/Cargo.lock +++ b/sdk/Cargo.lock @@ -206,6 +206,7 @@ dependencies = [ "itertools", "once_cell", "plonky2", + "poseidon2", "rand", "rand_chacha", "rkyv", @@ -352,6 +353,13 @@ name = "plonky2_util" version = "0.2.0" source = "git+https://github.com/0xmozak/plonky2.git#51f540a0e2a9bd9d6fc6234c6e62d167eaa7c707" +[[package]] +name = "poseidon2" +version = "0.1.0" +dependencies = [ + "plonky2", +] + [[package]] name = "ppv-lite86" version = "0.2.17" diff --git a/sdk/Cargo.toml b/sdk/Cargo.toml index f9b8f41a2..78fd93e0f 100644 --- a/sdk/Cargo.toml +++ b/sdk/Cargo.toml @@ -12,6 +12,7 @@ version = "0.2.0" [dependencies] itertools = { version = "0.12", default-features = false } once_cell = { version = "1.19", default-features = false, features = ["race"] } +poseidon2 = { path = "../poseidon2" } rkyv = { version = "=0.8.0-alpha.1", default-features = false, features = ["pointer_width_32", "alloc"] } rkyv_derive = "=0.8.0-alpha.1" diff --git a/sdk/src/mozakvm/helpers.rs b/sdk/src/mozakvm/helpers.rs index ec2931e5e..741b4ccf5 100644 --- a/sdk/src/mozakvm/helpers.rs +++ b/sdk/src/mozakvm/helpers.rs @@ -2,7 +2,9 @@ use std::ptr::{addr_of, slice_from_raw_parts}; -use crate::common::types::poseidon2hash::{DIGEST_BYTES, RATE}; +use poseidon2::mozak_poseidon2; + +use crate::common::types::poseidon2hash::DIGEST_BYTES; use crate::common::types::{Poseidon2Hash, ProgramIdentifier}; use crate::mozakvm::linker_symbols::_mozak_self_prog_id; @@ -48,10 +50,7 @@ pub fn get_self_prog_id() -> ProgramIdentifier { /// We use the well known "Bit padding scheme". #[allow(dead_code)] pub fn poseidon2_hash_with_pad(input: &[u8]) -> Poseidon2Hash { - let mut padded_input = input.to_vec(); - padded_input.push(1); - - padded_input.resize(padded_input.len().next_multiple_of(RATE), 0); + let padded_input = mozak_poseidon2::do_padding(input); let mut output = [0; DIGEST_BYTES]; crate::core::ecall::poseidon2( @@ -70,7 +69,7 @@ pub fn poseidon2_hash_with_pad(input: &[u8]) -> Poseidon2Hash { /// would fail otherwise. #[allow(dead_code)] pub fn poseidon2_hash_no_pad(input: &[u8]) -> Poseidon2Hash { - assert!(input.len() % RATE == 0); + assert!(input.len() % mozak_poseidon2::DATA_PADDING == 0); let mut output = [0; DIGEST_BYTES]; crate::core::ecall::poseidon2(input.as_ptr(), input.len(), output.as_mut_ptr()); Poseidon2Hash(output) diff --git a/sdk/src/native/helpers.rs b/sdk/src/native/helpers.rs index 16a78a24a..215a1b38e 100644 --- a/sdk/src/native/helpers.rs +++ b/sdk/src/native/helpers.rs @@ -3,11 +3,10 @@ use std::path::PathBuf; // This file contains code snippets used in native execution use plonky2::field::goldilocks_field::GoldilocksField; -use plonky2::field::types::Field; use plonky2::hash::poseidon2::Poseidon2Hash as Plonky2Poseidon2Hash; use plonky2::plonk::config::{GenericHashOut, Hasher}; +use poseidon2::mozak_poseidon2; -use crate::common::types::poseidon2hash::RATE; use crate::common::types::{Poseidon2Hash, ProgramIdentifier}; /// Represents a stack for call contexts during native execution. @@ -76,15 +75,8 @@ pub fn rm_identity() { /// Hashes the input slice to `Poseidon2Hash` after padding. /// We use the well known "Bit padding scheme". pub fn poseidon2_hash_with_pad(input: &[u8]) -> Poseidon2Hash { - let mut padded_input = input.to_vec(); - padded_input.push(1); - - padded_input.resize(padded_input.len().next_multiple_of(RATE), 0); - let data_fields: Vec = padded_input - .iter() - .map(|x| GoldilocksField::from_canonical_u8(*x)) - .collect(); - + let data_fields: Vec = + mozak_poseidon2::pack_padded_input(mozak_poseidon2::do_padding(input).as_slice()); Poseidon2Hash( Plonky2Poseidon2Hash::hash_no_pad(&data_fields) .to_bytes() @@ -101,11 +93,9 @@ pub fn poseidon2_hash_with_pad(input: &[u8]) -> Poseidon2Hash { /// would fail otherwise. #[allow(unused)] pub fn poseidon2_hash_no_pad(input: &[u8]) -> Poseidon2Hash { - assert!(input.len() % RATE == 0); - let data_fields: Vec = input - .iter() - .map(|x| GoldilocksField::from_canonical_u8(*x)) - .collect(); + assert!(input.len() % mozak_poseidon2::DATA_PADDING == 0); + + let data_fields: Vec = mozak_poseidon2::pack_padded_input(input); Poseidon2Hash( Plonky2Poseidon2Hash::hash_no_pad(&data_fields)